instance: polymorphize upvar closures/generators

This commit modifies how instances are polymorphized so that closures
and generators have any closures or generators captured within their
upvars also polymorphized - this avoids symbol clashes with the new
symbol mangling scheme.

Signed-off-by: David Wood <[email protected]>

Author: davidtwco

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,92 @@ 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);
+
+    if unused.is_empty() {
+        // Exit early if every parameter was used.
+        return substs;
+    }
+
+    // If this is a closure or generator then we need to handle the case where another closure
+    // from the function is captured as an upvar and hasn't been polymorphized. In this case,
+    // the unpolymorphized upvar closure would result in a polymorphized closure producing
+    // multiple mono items (and eventually symbol clashes).
+    let upvars_ty = if tcx.is_closure(def_id) {
+        Some(substs.as_closure().tupled_upvars_ty())
+    } else if tcx.type_of(def_id).is_generator() {
+        Some(substs.as_generator().tupled_upvars_ty())
+    } else {
+        None
+    };
+    let has_upvars = upvars_ty.map(|ty| ty.tuple_fields().count() > 0).unwrap_or(false);
+    debug!("polymorphize: upvars_ty={:?} has_upvars={:?}", upvars_ty, has_upvars);
+
+    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);
+                    self.tcx.mk_closure(def_id, polymorphized_substs)
+                }
+                ty::Generator(def_id, substs, movability) => {
+                    let polymorphized_substs = polymorphize(self.tcx, def_id, substs);
+                    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 {
+            // Upvar case: If parameter is a type parameter..
+            ty::GenericParamDefKind::Type { .. } if
+                // ..and has upvars..
+                has_upvars &&
+                // ..and this param has the same type as the tupled upvars..
+                upvars_ty == Some(substs[param.index as usize].expect_ty()) => {
+                    // ..then double-check that polymorphization marked it used..
+                    debug_assert!(!is_unused);
+                    // ..and polymorphize any closures/generators captured as upvars.
+                    let upvars_ty = upvars_ty.unwrap();
+                    let polymorphized_upvars_ty = upvars_ty.fold_with(
+                        &mut PolymorphizationFolder { tcx });
+                    debug!("polymorphize: polymorphized_upvars_ty={:?}", polymorphized_upvars_ty);
+                    ty::GenericArg::from(polymorphized_upvars_ty)
+                },
+
+            // Simple case: 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.
+            _ => substs[param.index as usize],
+        }
+    })
+}
+
 fn needs_fn_once_adapter_shim(
     actual_closure_kind: ty::ClosureKind,
     trait_closure_kind: ty::ClosureKind,

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