Auto merge of #75255 - davidtwco:polymorphisation-symbol-mangling-v0-upvar-closures, r=lcnr

instance: polymorphize upvar closures/generators

This PR modifies how instances are polymorphized so that closures and generators have any closures or generators captured within their upvars also polymorphized.

With the new symbol mangling, a fully polymorphised closure will produce the same symbol regardless of what it was instantiated with. However, when that polymorphised closure captures another closure as an upvar, then the type of that other closure in the upvar substitution wouldn't have been polymorphised. The other closure will still refer to the initial substitutions. Therefore, the polymorphised closure will end up hashing differently but producing the same symbol - triggering `assert_symbols_are_distinct` in MIR partitioning. The old mangling scheme had a hash at the end that meant this didn't happen (this would still have been an issue, we just didn't have a way to notice).

See [this Zulip discussion for further elaboration](https://rust-lang.zulipchat.com/#narrow/stream/216091-t-compiler.2Fwg-polymorphization/topic/symbol.20mangling.20v0.20.E2.9C.95.20polymorphisation/near/206152008).

r? @eddyb
cc @lcnr

src/librustc_middle/ty/flags.rs

@@ -85,6 +85,8 @@ impl FlagComputation {
             }
 
             &ty::Generator(_, ref substs, _) => {
+                self.add_flags(TypeFlags::MAY_POLYMORPHIZE);
+
                 let substs = substs.as_generator();
                 let should_remove_further_specializable =
                     !self.flags.contains(TypeFlags::STILL_FURTHER_SPECIALIZABLE);
@@ -107,6 +109,8 @@ impl FlagComputation {
             }
 
             &ty::Closure(_, substs) => {
+                self.add_flags(TypeFlags::MAY_POLYMORPHIZE);
+
                 let substs = substs.as_closure();
                 let should_remove_further_specializable =
                     !self.flags.contains(TypeFlags::STILL_FURTHER_SPECIALIZABLE);
@@ -192,6 +196,8 @@ impl FlagComputation {
             }
 
             &ty::FnDef(_, substs) => {
+                self.add_flags(TypeFlags::MAY_POLYMORPHIZE);
+
                 self.add_substs(substs);
             }
 

src/librustc_middle/ty/fold.rs

@@ -150,6 +150,12 @@ pub trait TypeFoldable<'tcx>: fmt::Debug + Clone {
         self.has_type_flags(TypeFlags::STILL_FURTHER_SPECIALIZABLE)
     }
 
+    /// Does this value contain closures, generators or functions such that it may require
+    /// polymorphization?
+    fn may_polymorphize(&self) -> bool {
+        self.has_type_flags(TypeFlags::MAY_POLYMORPHIZE)
+    }
+
     /// A visitor that does not recurse into types, works like `fn walk_shallow` in `Ty`.
     fn visit_tys_shallow(&self, visit: impl FnMut(Ty<'tcx>) -> bool) -> bool {
         pub struct Visitor<F>(F);

src/librustc_middle/ty/instance.rs

@@ -474,26 +474,7 @@ impl<'tcx> Instance<'tcx> {
         }
 
         if let InstanceDef::Item(def) = self.def {
-            let unused = tcx.unused_generic_params(def.did);
-
-            if unused.is_empty() {
-                // Exit early if every parameter was used.
-                return self;
-            }
-
-            debug!("polymorphize: unused={:?}", unused);
-            let polymorphized_substs =
-                InternalSubsts::for_item(tcx, def.did, |param, _| match param.kind {
-                // If parameter is a const or type parameter..
-                ty::GenericParamDefKind::Const | ty::GenericParamDefKind::Type { .. } if
-                    // ..and is within range and unused..
-                    unused.contains(param.index).unwrap_or(false) =>
-                        // ..then use the identity for this parameter.
-                        tcx.mk_param_from_def(param),
-                // Otherwise, use the parameter as before.
-                _ => self.substs[param.index as usize],
-            });
-
+            let polymorphized_substs = polymorphize(tcx, def.did, self.substs);
             debug!("polymorphize: self={:?} polymorphized_substs={:?}", self, polymorphized_substs);
             Self { def: self.def, substs: polymorphized_substs }
         } else {
@@ -502,6 +483,82 @@ impl<'tcx> Instance<'tcx> {
     }
 }
 
+fn polymorphize<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    def_id: DefId,
+    substs: SubstsRef<'tcx>,
+) -> SubstsRef<'tcx> {
+    debug!("polymorphize({:?}, {:?})", def_id, substs);
+    let unused = tcx.unused_generic_params(def_id);
+    debug!("polymorphize: unused={:?}", unused);
+
+    struct PolymorphizationFolder<'tcx> {
+        tcx: TyCtxt<'tcx>,
+    };
+
+    impl ty::TypeFolder<'tcx> for PolymorphizationFolder<'tcx> {
+        fn tcx<'a>(&'a self) -> TyCtxt<'tcx> {
+            self.tcx
+        }
+
+        fn fold_ty(&mut self, ty: Ty<'tcx>) -> Ty<'tcx> {
+            debug!("fold_ty: ty={:?}", ty);
+            match ty.kind {
+                ty::Closure(def_id, substs) => {
+                    let polymorphized_substs = polymorphize(self.tcx, def_id, substs);
+                    if substs == polymorphized_substs {
+                        ty
+                    } else {
+                        self.tcx.mk_closure(def_id, polymorphized_substs)
+                    }
+                }
+                ty::FnDef(def_id, substs) => {
+                    let polymorphized_substs = polymorphize(self.tcx, def_id, substs);
+                    if substs == polymorphized_substs {
+                        ty
+                    } else {
+                        self.tcx.mk_fn_def(def_id, polymorphized_substs)
+                    }
+                }
+                ty::Generator(def_id, substs, movability) => {
+                    let polymorphized_substs = polymorphize(self.tcx, def_id, substs);
+                    if substs == polymorphized_substs {
+                        ty
+                    } else {
+                        self.tcx.mk_generator(def_id, polymorphized_substs, movability)
+                    }
+                }
+                _ => ty.super_fold_with(self),
+            }
+        }
+    }
+
+    InternalSubsts::for_item(tcx, def_id, |param, _| {
+        let is_unused = unused.contains(param.index).unwrap_or(false);
+        debug!("polymorphize: param={:?} is_unused={:?}", param, is_unused);
+        match param.kind {
+            // If parameter is a const or type parameter..
+            ty::GenericParamDefKind::Const | ty::GenericParamDefKind::Type { .. } if
+                // ..and is within range and unused..
+                unused.contains(param.index).unwrap_or(false) =>
+                    // ..then use the identity for this parameter.
+                    tcx.mk_param_from_def(param),
+
+            // If the parameter does not contain any closures or generators, then use the
+            // substitution directly.
+            _ if !substs.may_polymorphize() => substs[param.index as usize],
+
+            // Otherwise, use the substitution after polymorphizing.
+            _ => {
+                let arg = substs[param.index as usize];
+                let polymorphized_arg = arg.fold_with(&mut PolymorphizationFolder { tcx });
+                debug!("polymorphize: arg={:?} polymorphized_arg={:?}", arg, polymorphized_arg);
+                ty::GenericArg::from(polymorphized_arg)
+            }
+        }
+    })
+}
+
 fn needs_fn_once_adapter_shim(
     actual_closure_kind: ty::ClosureKind,
     trait_closure_kind: ty::ClosureKind,

src/librustc_middle/ty/mod.rs

@@ -575,6 +575,10 @@ bitflags! {
         /// Does this value have parameters/placeholders/inference variables which could be
         /// replaced later, in a way that would change the results of `impl` specialization?
         const STILL_FURTHER_SPECIALIZABLE = 1 << 17;
+
+        /// Does this value contain closures, generators or functions such that it may require
+        /// polymorphization?
+        const MAY_POLYMORPHIZE = 1 << 18;
     }
 }
 

src/test/codegen-units/polymorphization/pr-75255.rs

@@ -0,0 +1,52 @@
+// compile-flags:-Zpolymorphize=on -Zprint-mono-items=lazy -Copt-level=1
+// ignore-tidy-linelength
+
+#![crate_type = "rlib"]
+
+// Test that only one copy of `Iter::map` and `iter::repeat` are generated.
+
+fn unused<T>() -> u64 {
+    42
+}
+
+fn foo<T>() {
+    let x = [1, 2, 3, std::mem::size_of::<T>()];
+    x.iter().map(|_| ());
+}
+
+//~ MONO_ITEM fn core::iter[0]::adapters[0]::{{impl}}[29]::new[0]<core::slice[0]::Iter[0]<usize>, pr_75255::foo[0]::{{closure}}[0]<T>> @@ pr_75255-cgu.0[External]
+//~ MONO_ITEM fn core::iter[0]::traits[0]::iterator[0]::Iterator[0]::map[0]<core::slice[0]::Iter[0]<usize>, (), pr_75255::foo[0]::{{closure}}[0]<T>> @@ pr_75255-cgu.1[Internal]
+
+fn bar<T>() {
+    std::iter::repeat(unused::<T>);
+}
+
+//~ MONO_ITEM fn core::iter[0]::sources[0]::repeat[0]<fn() -> u64> @@ pr_75255-cgu.1[Internal]
+
+pub fn dispatch() {
+    foo::<String>();
+    foo::<Vec<String>>();
+
+    bar::<String>();
+    bar::<Vec<String>>();
+}
+
+// These are all the items that aren't relevant to the test.
+//~ MONO_ITEM fn core::mem[0]::size_of[0]<alloc::string[0]::String[0]> @@ pr_75255-cgu.1[Internal]
+//~ MONO_ITEM fn core::mem[0]::size_of[0]<alloc::vec[0]::Vec[0]<alloc::string[0]::String[0]>> @@ pr_75255-cgu.1[Internal]
+//~ MONO_ITEM fn core::mem[0]::size_of[0]<usize> @@ pr_75255-cgu.1[Internal]
+//~ MONO_ITEM fn core::ptr[0]::const_ptr[0]::{{impl}}[0]::add[0]<usize> @@ pr_75255-cgu.1[Internal]
+//~ MONO_ITEM fn core::ptr[0]::const_ptr[0]::{{impl}}[0]::is_null[0]<usize> @@ pr_75255-cgu.1[Internal]
+//~ MONO_ITEM fn core::ptr[0]::const_ptr[0]::{{impl}}[0]::offset[0]<usize> @@ pr_75255-cgu.1[Internal]
+//~ MONO_ITEM fn core::ptr[0]::const_ptr[0]::{{impl}}[0]::wrapping_add[0]<u8> @@ pr_75255-cgu.1[Internal]
+//~ MONO_ITEM fn core::ptr[0]::const_ptr[0]::{{impl}}[0]::wrapping_offset[0]<u8> @@ pr_75255-cgu.1[Internal]
+//~ MONO_ITEM fn core::ptr[0]::non_null[0]::{{impl}}[3]::new_unchecked[0]<usize> @@ pr_75255-cgu.1[Internal]
+//~ MONO_ITEM fn core::ptr[0]::null[0]<u8> @@ pr_75255-cgu.1[Internal]
+//~ MONO_ITEM fn core::slice[0]::{{impl}}[0]::as_ptr[0]<usize> @@ pr_75255-cgu.1[Internal]
+//~ MONO_ITEM fn core::slice[0]::{{impl}}[0]::iter[0]<usize> @@ pr_75255-cgu.1[Internal]
+//~ MONO_ITEM fn core::slice[0]::{{impl}}[0]::len[0]<usize> @@ pr_75255-cgu.1[Internal]
+//~ MONO_ITEM fn pr_75255::dispatch[0] @@ pr_75255-cgu.1[External]
+//~ MONO_ITEM fn pr_75255::foo[0]<alloc::string[0]::String[0]> @@ pr_75255-cgu.1[Internal]
+//~ MONO_ITEM fn pr_75255::foo[0]<alloc::vec[0]::Vec[0]<alloc::string[0]::String[0]>> @@ pr_75255-cgu.1[Internal]
+//~ MONO_ITEM fn pr_75255::bar[0]<alloc::string[0]::String[0]> @@ pr_75255-cgu.1[Internal]
+//~ MONO_ITEM fn pr_75255::bar[0]<alloc::vec[0]::Vec[0]<alloc::string[0]::String[0]>> @@ pr_75255-cgu.1[Internal]

src/test/ui/polymorphization/closure_in_upvar/fn.rs

@@ -0,0 +1,29 @@
+// build-pass
+// compile-flags:-Zpolymorphize=on -Zsymbol-mangling-version=v0
+
+fn foo(f: impl Fn()) {
+    let x = |_: ()| ();
+
+    // Don't use `f` in `y`, but refer to `x` so that the closure substs contain a reference to
+    // `x` that will differ for each instantiation despite polymorphisation of the varying
+    // argument.
+    let y = || x(());
+
+    // Consider `f` used in `foo`.
+    f();
+    // Use `y` so that it is visited in monomorphisation collection.
+    y();
+}
+
+fn entry_a() {
+    foo(|| ());
+}
+
+fn entry_b() {
+    foo(|| ());
+}
+
+fn main() {
+    entry_a();
+    entry_b();
+}

src/test/ui/polymorphization/closure_in_upvar/fnmut.rs

@@ -0,0 +1,34 @@
+// build-pass
+// compile-flags:-Zpolymorphize=on -Zsymbol-mangling-version=v0
+
+fn foo(f: impl Fn()) {
+    // Mutate an upvar from `x` so that it implements `FnMut`.
+    let mut outer = 3;
+    let mut x = |_: ()| {
+        outer = 4;
+        ()
+    };
+
+    // Don't use `f` in `y`, but refer to `x` so that the closure substs contain a reference to
+    // `x` that will differ for each instantiation despite polymorphisation of the varying
+    // argument.
+    let mut y = || x(());
+
+    // Consider `f` used in `foo`.
+    f();
+    // Use `y` so that it is visited in monomorphisation collection.
+    y();
+}
+
+fn entry_a() {
+    foo(|| ());
+}
+
+fn entry_b() {
+    foo(|| ());
+}
+
+fn main() {
+    entry_a();
+    entry_b();
+}

src/test/ui/polymorphization/closure_in_upvar/fnonce.rs

@@ -0,0 +1,34 @@
+// build-pass
+// compile-flags:-Zpolymorphize=on -Zsymbol-mangling-version=v0
+
+fn foo(f: impl Fn()) {
+    // Move a non-copy type into `x` so that it implements `FnOnce`.
+    let outer = Vec::<u32>::new();
+    let x = move |_: ()| {
+        let inner = outer;
+        ()
+    };
+
+    // Don't use `f` in `y`, but refer to `x` so that the closure substs contain a reference to
+    // `x` that will differ for each instantiation despite polymorphisation of the varying
+    // argument.
+    let y = || x(());
+
+    // Consider `f` used in `foo`.
+    f();
+    // Use `y` so that it is visited in monomorphisation collection.
+    y();
+}
+
+fn entry_a() {
+    foo(|| ());
+}
+
+fn entry_b() {
+    foo(|| ());
+}
+
+fn main() {
+    entry_a();
+    entry_b();
+}

src/test/ui/polymorphization/closure_in_upvar/other.rs

@@ -0,0 +1,38 @@
+// build-pass
+// compile-flags:-Zpolymorphize=on -Zsymbol-mangling-version=v0
+
+fn y_uses_f(f: impl Fn()) {
+    let x = |_: ()| ();
+
+    let y = || {
+        f();
+        x(());
+    };
+
+    f();
+    y();
+}
+
+fn x_uses_f(f: impl Fn()) {
+    let x = |_: ()| { f(); };
+
+    let y = || x(());
+
+    f();
+    y();
+}
+
+fn entry_a() {
+    x_uses_f(|| ());
+    y_uses_f(|| ());
+}
+
+fn entry_b() {
+    x_uses_f(|| ());
+    y_uses_f(|| ());
+}
+
+fn main() {
+    entry_a();
+    entry_b();
+}