Deduce captures(none) for a return place and parameters

Extend attribute deduction to determine whether parameters using
indirect pass mode might have their address captured. Similarly to
the deduction of `readonly` attribute this information facilitates
memcpy optimizations.

Author: tmiasko

compiler/rustc_codegen_llvm/src/abi.rs

@@ -39,13 +39,11 @@ trait ArgAttributesExt {
 const ABI_AFFECTING_ATTRIBUTES: [(ArgAttribute, llvm::AttributeKind); 1] =
     [(ArgAttribute::InReg, llvm::AttributeKind::InReg)];
 
-const OPTIMIZATION_ATTRIBUTES: [(ArgAttribute, llvm::AttributeKind); 6] = [
+const OPTIMIZATION_ATTRIBUTES: [(ArgAttribute, llvm::AttributeKind); 4] = [
     (ArgAttribute::NoAlias, llvm::AttributeKind::NoAlias),
-    (ArgAttribute::CapturesAddress, llvm::AttributeKind::CapturesAddress),
     (ArgAttribute::NonNull, llvm::AttributeKind::NonNull),
     (ArgAttribute::ReadOnly, llvm::AttributeKind::ReadOnly),
     (ArgAttribute::NoUndef, llvm::AttributeKind::NoUndef),
-    (ArgAttribute::CapturesReadOnly, llvm::AttributeKind::CapturesReadOnly),
 ];
 
 fn get_attrs<'ll>(this: &ArgAttributes, cx: &CodegenCx<'ll, '_>) -> SmallVec<[&'ll Attribute; 8]> {
@@ -81,15 +79,23 @@ fn get_attrs<'ll>(this: &ArgAttributes, cx: &CodegenCx<'ll, '_>) -> SmallVec<[&'
         }
         for (attr, llattr) in OPTIMIZATION_ATTRIBUTES {
             if regular.contains(attr) {
-                // captures(...) is only available since LLVM 21.
-                if (attr == ArgAttribute::CapturesReadOnly || attr == ArgAttribute::CapturesAddress)
-                    && llvm_util::get_version() < (21, 0, 0)
-                {
-                    continue;
-                }
                 attrs.push(llattr.create_attr(cx.llcx));
             }
         }
+        // captures(...) is only available since LLVM 21.
+        if (21, 0, 0) <= llvm_util::get_version() {
+            const CAPTURES_ATTRIBUTES: [(ArgAttribute, llvm::AttributeKind); 3] = [
+                (ArgAttribute::CapturesNone, llvm::AttributeKind::CapturesNone),
+                (ArgAttribute::CapturesAddress, llvm::AttributeKind::CapturesAddress),
+                (ArgAttribute::CapturesReadOnly, llvm::AttributeKind::CapturesReadOnly),
+            ];
+            for (attr, llattr) in CAPTURES_ATTRIBUTES {
+                if regular.contains(attr) {
+                    attrs.push(llattr.create_attr(cx.llcx));
+                    break;
+                }
+            }
+        }
     } else if cx.tcx.sess.opts.unstable_opts.sanitizer.contains(SanitizerSet::MEMORY) {
         // If we're not optimising, *but* memory sanitizer is on, emit noundef, since it affects
         // memory sanitizer's behavior.

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -289,6 +289,7 @@ pub(crate) enum AttributeKind {
     DeadOnUnwind = 43,
     DeadOnReturn = 44,
     CapturesReadOnly = 45,
+    CapturesNone = 46,
 }
 
 /// LLVMIntPredicate

compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp

@@ -245,6 +245,7 @@ enum class LLVMRustAttributeKind {
   DeadOnUnwind = 43,
   DeadOnReturn = 44,
   CapturesReadOnly = 45,
+  CapturesNone = 46,
 };
 
 static Attribute::AttrKind fromRust(LLVMRustAttributeKind Kind) {
@@ -339,6 +340,7 @@ static Attribute::AttrKind fromRust(LLVMRustAttributeKind Kind) {
 #endif
   case LLVMRustAttributeKind::CapturesAddress:
   case LLVMRustAttributeKind::CapturesReadOnly:
+  case LLVMRustAttributeKind::CapturesNone:
     report_fatal_error("Should be handled separately");
   }
   report_fatal_error("bad LLVMRustAttributeKind");
@@ -390,6 +392,9 @@ extern "C" void LLVMRustEraseInstFromParent(LLVMValueRef Instr) {
 extern "C" LLVMAttributeRef
 LLVMRustCreateAttrNoValue(LLVMContextRef C, LLVMRustAttributeKind RustAttr) {
 #if LLVM_VERSION_GE(21, 0)
+  if (RustAttr == LLVMRustAttributeKind::CapturesNone) {
+    return wrap(Attribute::getWithCaptureInfo(*unwrap(C), CaptureInfo::none()));
+  }
   if (RustAttr == LLVMRustAttributeKind::CapturesAddress) {
     return wrap(Attribute::getWithCaptureInfo(
         *unwrap(C), CaptureInfo(CaptureComponents::Address)));

compiler/rustc_middle/src/middle/deduced_param_attrs.rs

@@ -2,19 +2,21 @@ use rustc_macros::{Decodable, Encodable, HashStable};
 
 use crate::ty::{Ty, TyCtxt, TypingEnv};
 
-/// Flags that dictate how a parameter is mutated. If the flags are empty, the param is
-/// read-only. If non-empty, it is read-only if *all* flags' conditions are met.
+/// Summarizes how a parameter (a return place or an argument) is used inside a MIR body.
 #[derive(Clone, Copy, PartialEq, Debug, Decodable, Encodable, HashStable)]
-pub struct DeducedReadOnlyParam(u8);
+pub struct UsageSummary(u8);
 
 bitflags::bitflags! {
-    impl DeducedReadOnlyParam: u8 {
-        /// This parameter is dropped. It is read-only if `!needs_drop`.
-        const IF_NO_DROP = 1 << 0;
-        /// This parameter is borrowed. It is read-only if `Freeze`.
-        const IF_FREEZE   = 1 << 1;
-        /// This parameter is mutated. It is never read-only.
-        const MUTATED     = 1 << 2;
+    impl UsageSummary: u8 {
+        /// This parameter is dropped when it `needs_drop`.
+        const DROP = 1 << 0;
+        /// There is a shared borrow to this parameter.
+        /// It allows for mutation unless parameter is `Freeze`.
+        const SHARED_BORROW = 1 << 1;
+        /// This parameter is mutated (excluding through a drop or a shared borrow).
+        const MUTATE = 1 << 2;
+        /// This parameter is captured (excluding through a drop).
+        const CAPTURE = 1 << 3;
     }
 }
 
@@ -24,43 +26,53 @@ bitflags::bitflags! {
 /// These can be useful for optimization purposes when a function is directly called. We compute
 /// them and store them into the crate metadata so that downstream crates can make use of them.
 ///
-/// Right now, we only have `read_only`, but `no_capture` and `no_alias` might be useful in the
+/// Right now, we have `readonly` and `captures(none)`, but `no_alias` might be useful in the
 /// future.
 #[derive(Clone, Copy, PartialEq, Debug, Decodable, Encodable, HashStable)]
 pub struct DeducedParamAttrs {
-    /// The parameter is marked immutable in the function.
-    pub read_only: DeducedReadOnlyParam,
-}
-
-// By default, consider the parameters to be mutated.
-impl Default for DeducedParamAttrs {
-    #[inline]
-    fn default() -> DeducedParamAttrs {
-        DeducedParamAttrs { read_only: DeducedReadOnlyParam::MUTATED }
-    }
+    pub usage: UsageSummary,
 }
 
 impl DeducedParamAttrs {
+    /// Returns true if no attributes have been deduced.
     #[inline]
     pub fn is_default(self) -> bool {
-        self.read_only.contains(DeducedReadOnlyParam::MUTATED)
+        self.usage.contains(UsageSummary::MUTATE | UsageSummary::CAPTURE)
     }
 
+    /// For parameters passed indirectly, returns true if pointer is never written through.
     pub fn read_only<'tcx>(
         &self,
         tcx: TyCtxt<'tcx>,
         typing_env: TypingEnv<'tcx>,
         ty: Ty<'tcx>,
     ) -> bool {
-        let read_only = self.read_only;
-        // We have to check *all* set bits; only if all checks pass is this truly read-only.
-        if read_only.contains(DeducedReadOnlyParam::MUTATED) {
+        // Only if all checks pass is this truly read-only.
+        if self.usage.contains(UsageSummary::MUTATE) {
+            return false;
+        }
+        if self.usage.contains(UsageSummary::DROP) && ty.needs_drop(tcx, typing_env) {
+            return false;
+        }
+        if self.usage.contains(UsageSummary::SHARED_BORROW) && !ty.is_freeze(tcx, typing_env) {
             return false;
         }
-        if read_only.contains(DeducedReadOnlyParam::IF_NO_DROP) && ty.needs_drop(tcx, typing_env) {
+        true
+    }
+
+    /// For parameters passed indirectly, returns true if pointer is not captured, i.e., its
+    /// address is not captured, and pointer is used neither for reads nor writes after function
+    /// returns.
+    pub fn captures_none<'tcx>(
+        &self,
+        tcx: TyCtxt<'tcx>,
+        typing_env: TypingEnv<'tcx>,
+        ty: Ty<'tcx>,
+    ) -> bool {
+        if self.usage.contains(UsageSummary::CAPTURE) {
             return false;
         }
-        if read_only.contains(DeducedReadOnlyParam::IF_FREEZE) && !ty.is_freeze(tcx, typing_env) {
+        if self.usage.contains(UsageSummary::DROP) && ty.needs_drop(tcx, typing_env) {
             return false;
         }
         true

compiler/rustc_mir_transform/src/deduce_param_attrs.rs

@@ -11,65 +11,91 @@
 
 use rustc_hir::def_id::LocalDefId;
 use rustc_index::IndexVec;
-use rustc_middle::middle::deduced_param_attrs::{DeducedParamAttrs, DeducedReadOnlyParam};
+use rustc_middle::middle::deduced_param_attrs::{DeducedParamAttrs, UsageSummary};
 use rustc_middle::mir::visit::{MutatingUseContext, NonMutatingUseContext, PlaceContext, Visitor};
 use rustc_middle::mir::*;
 use rustc_middle::ty::{self, Ty, TyCtxt};
 use rustc_session::config::OptLevel;
 
-/// A visitor that determines which arguments have been mutated. We can't use the mutability field
-/// on LocalDecl for this because it has no meaning post-optimization.
-struct DeduceReadOnly {
-    /// Each bit is indexed by argument number, starting at zero (so 0 corresponds to local decl
-    /// 1). The bit is false if the argument may have been mutated or true if we know it hasn't
-    /// been up to the point we're at.
-    read_only: IndexVec<usize, DeducedReadOnlyParam>,
+/// A visitor that determines how a return place and arguments are used inside MIR body.
+/// To determine whether a local is mutated we can't use the mutability field on LocalDecl
+/// because it has no meaning post-optimization.
+struct DeduceParamAttrs {
+    /// Summarizes how a return place and arguments are used inside MIR body.
+    usage: IndexVec<Local, UsageSummary>,
 }
 
-impl DeduceReadOnly {
-    /// Returns a new DeduceReadOnly instance.
-    fn new(arg_count: usize) -> Self {
-        Self { read_only: IndexVec::from_elem_n(DeducedReadOnlyParam::empty(), arg_count) }
+impl DeduceParamAttrs {
+    /// Returns a new DeduceParamAttrs instance.
+    fn new(body: &Body<'_>) -> Self {
+        let mut this =
+            Self { usage: IndexVec::from_elem_n(UsageSummary::empty(), body.arg_count + 1) };
+        // Code generation indicates that a return place is writable. To avoid setting both
+        // `readonly` and `writable` attributes, when return place is never written to, mark it as
+        // mutated.
+        this.usage[RETURN_PLACE] |= UsageSummary::MUTATE;
+        this
     }
 
-    /// Returns whether the given local is a parameter and its index.
-    fn as_param(&self, local: Local) -> Option<usize> {
-        // Locals and parameters are shifted by `RETURN_PLACE`.
-        let param_index = local.as_usize().checked_sub(1)?;
-        if param_index < self.read_only.len() { Some(param_index) } else { None }
+    /// Returns whether a local is the return place or an argument and returns its index.
+    fn as_param(&self, local: Local) -> Option<Local> {
+        if local.index() < self.usage.len() { Some(local) } else { None }
     }
 }
 
-impl<'tcx> Visitor<'tcx> for DeduceReadOnly {
+impl<'tcx> Visitor<'tcx> for DeduceParamAttrs {
     fn visit_place(&mut self, place: &Place<'tcx>, context: PlaceContext, _location: Location) {
-        // We're only interested in arguments.
-        let Some(param_index) = self.as_param(place.local) else { return };
+        // We're only interested in the return place or an argument.
+        let Some(i) = self.as_param(place.local) else { return };
 
         match context {
-            // Not mutating, so it's fine.
+            // Not actually using the local.
             PlaceContext::NonUse(..) => {}
-            // Dereference is not a mutation.
+            // Neither mutated nor captured.
             _ if place.is_indirect_first_projection() => {}
             // This is a `Drop`. It could disappear at monomorphization, so mark it specially.
             PlaceContext::MutatingUse(MutatingUseContext::Drop)
                 // Projection changes the place's type, so `needs_drop(local.ty)` is not
                 // `needs_drop(place.ty)`.
                 if place.projection.is_empty() => {
-                self.read_only[param_index] |= DeducedReadOnlyParam::IF_NO_DROP;
+                    self.usage[i] |= UsageSummary::DROP;
+            }
+            PlaceContext::MutatingUse(
+                  MutatingUseContext::Call
+                | MutatingUseContext::Yield
+                | MutatingUseContext::Drop
+                | MutatingUseContext::Borrow
+                | MutatingUseContext::RawBorrow) => {
+                self.usage[i] |= UsageSummary::MUTATE;
+                self.usage[i] |= UsageSummary::CAPTURE;
+            }
+            PlaceContext::MutatingUse(
+                  MutatingUseContext::Store
+                | MutatingUseContext::SetDiscriminant
+                | MutatingUseContext::AsmOutput
+                | MutatingUseContext::Projection
+                | MutatingUseContext::Retag) => {
+                self.usage[i] |= UsageSummary::MUTATE;
             }
-            // This is a mutation, so mark it as such.
-            PlaceContext::MutatingUse(..)
-            // Whether mutating though a `&raw const` is allowed is still undecided, so we
-            // disable any sketchy `readonly` optimizations for now.
             | PlaceContext::NonMutatingUse(NonMutatingUseContext::RawBorrow) => {
-                self.read_only[param_index] |= DeducedReadOnlyParam::MUTATED;
+                // Whether mutating though a `&raw const` is allowed is still undecided, so we
+                // disable any sketchy `readonly` optimizations for now.
+                self.usage[i] |= UsageSummary::MUTATE;
+                self.usage[i] |= UsageSummary::CAPTURE;
             }
-            // Not mutating if the parameter is `Freeze`.
             PlaceContext::NonMutatingUse(NonMutatingUseContext::SharedBorrow) => {
-                self.read_only[param_index] |= DeducedReadOnlyParam::IF_FREEZE;
+                // Not mutating if the parameter is `Freeze`.
+                self.usage[i] |= UsageSummary::SHARED_BORROW;
+                self.usage[i] |= UsageSummary::CAPTURE;
             }
             // Not mutating, so it's fine.
-            PlaceContext::NonMutatingUse(..) => {}
+            PlaceContext::NonMutatingUse(
+                  NonMutatingUseContext::Inspect
+                | NonMutatingUseContext::Copy
+                | NonMutatingUseContext::Move
+                | NonMutatingUseContext::FakeBorrow
+                | NonMutatingUseContext::PlaceMention
+                | NonMutatingUseContext::Projection) => {}
         }
     }
 
@@ -98,11 +124,11 @@ impl<'tcx> Visitor<'tcx> for DeduceReadOnly {
         if let TerminatorKind::Call { ref args, .. } = terminator.kind {
             for arg in args {
                 if let Operand::Move(place) = arg.node
-                    // We're only interested in arguments.
-                    && let Some(param_index) = self.as_param(place.local)
                     && !place.is_indirect_first_projection()
+                    && let Some(i) = self.as_param(place.local)
                 {
-                    self.read_only[param_index] |= DeducedReadOnlyParam::MUTATED;
+                    self.usage[i] |= UsageSummary::MUTATE;
+                    self.usage[i] |= UsageSummary::CAPTURE;
                 }
             }
         };
@@ -154,10 +180,9 @@ pub(super) fn deduced_param_attrs<'tcx>(
     if matches!(fn_ty.kind(), ty::FnDef(..))
         && fn_ty
             .fn_sig(tcx)
-            .inputs()
+            .inputs_and_output()
             .skip_binder()
             .iter()
-            .cloned()
             .all(type_will_always_be_passed_directly)
     {
         return &[];
@@ -170,13 +195,13 @@ pub(super) fn deduced_param_attrs<'tcx>(
 
     // Grab the optimized MIR. Analyze it to determine which arguments have been mutated.
     let body: &Body<'tcx> = tcx.optimized_mir(def_id);
-    let mut deduce_read_only = DeduceReadOnly::new(body.arg_count);
-    deduce_read_only.visit_body(body);
-    tracing::trace!(?deduce_read_only.read_only);
+    let mut deduce = DeduceParamAttrs::new(body);
+    deduce.visit_body(body);
+    tracing::trace!(?deduce.usage);
 
-    let mut deduced_param_attrs: &[_] = tcx.arena.alloc_from_iter(
-        deduce_read_only.read_only.into_iter().map(|read_only| DeducedParamAttrs { read_only }),
-    );
+    let mut deduced_param_attrs: &[_] = tcx
+        .arena
+        .alloc_from_iter(deduce.usage.into_iter().map(|usage| DeducedParamAttrs { usage }));
 
     // Trailing parameters past the size of the `deduced_param_attrs` array are assumed to have the
     // default set of attributes, so we don't have to store them explicitly. Pop them off to save a

compiler/rustc_target/src/callconv/mod.rs

@@ -113,13 +113,14 @@ mod attr_impl {
     pub struct ArgAttribute(u8);
     bitflags::bitflags! {
         impl ArgAttribute: u8 {
-            const NoAlias   = 1 << 1;
-            const CapturesAddress = 1 << 2;
-            const NonNull   = 1 << 3;
-            const ReadOnly  = 1 << 4;
-            const InReg     = 1 << 5;
-            const NoUndef = 1 << 6;
-            const CapturesReadOnly = 1 << 7;
+            const CapturesNone     = 0b111;
+            const CapturesAddress  = 0b110;
+            const CapturesReadOnly = 0b100;
+            const NoAlias  = 1 << 3;
+            const NonNull  = 1 << 4;
+            const ReadOnly = 1 << 5;
+            const InReg    = 1 << 6;
+            const NoUndef  = 1 << 7;
         }
     }
     rustc_data_structures::external_bitflags_debug! { ArgAttribute }