instance: polymorphize upvar closures/generators

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();
+}