Auto merge of #101673 - crlf0710:generator_clone, r=oli-obk

Allow generators to impl Clone/Copy

Revives #95137. It's a pity that the original pr didn't land because the implementation is almost complete! All credits goes to `@canndrew,` and i just resolved the merge conflicts and updated the feature gate version number.

r? `@oli-obk`

Author: bors

compiler/rustc_feature/src/active.rs

@@ -394,6 +394,8 @@ declare_features! (
     (active, ffi_returns_twice, "1.34.0", Some(58314), None),
     /// Allows using `#[repr(align(...))]` on function items
     (active, fn_align, "1.53.0", Some(82232), None),
+    /// Allows generators to be cloned.
+    (active, generator_clone, "CURRENT_RUSTC_VERSION", Some(95360), None),
     /// Allows defining generators.
     (active, generators, "1.21.0", Some(43122), None),
     /// Infer generic args for both consts and types.

compiler/rustc_mir_transform/src/shim.rs

@@ -4,7 +4,7 @@ use rustc_hir::lang_items::LangItem;
 use rustc_middle::mir::*;
 use rustc_middle::ty::query::Providers;
 use rustc_middle::ty::subst::{InternalSubsts, Subst};
-use rustc_middle::ty::{self, EarlyBinder, Ty, TyCtxt};
+use rustc_middle::ty::{self, EarlyBinder, GeneratorSubsts, Ty, TyCtxt};
 use rustc_target::abi::VariantIdx;
 
 use rustc_index::vec::{Idx, IndexVec};
@@ -323,6 +323,9 @@ fn build_clone_shim<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, self_ty: Ty<'tcx>) -
             builder.tuple_like_shim(dest, src, substs.as_closure().upvar_tys())
         }
         ty::Tuple(..) => builder.tuple_like_shim(dest, src, self_ty.tuple_fields()),
+        ty::Generator(gen_def_id, substs, hir::Movability::Movable) => {
+            builder.generator_shim(dest, src, *gen_def_id, substs.as_generator())
+        }
         _ => bug!("clone shim for `{:?}` which is not `Copy` and is not an aggregate", self_ty),
     };
 
@@ -388,7 +391,7 @@ impl<'tcx> CloneShimBuilder<'tcx> {
     /// offset=0 will give you the index of the next BasicBlock,
     /// offset=1 will give the index of the next-to-next block,
     /// offset=-1 will give you the index of the last-created block
-    fn block_index_offset(&mut self, offset: usize) -> BasicBlock {
+    fn block_index_offset(&self, offset: usize) -> BasicBlock {
         BasicBlock::new(self.blocks.len() + offset)
     }
 
@@ -461,49 +464,106 @@ impl<'tcx> CloneShimBuilder<'tcx> {
         );
     }
 
-    fn tuple_like_shim<I>(&mut self, dest: Place<'tcx>, src: Place<'tcx>, tys: I)
+    fn clone_fields<I>(
+        &mut self,
+        dest: Place<'tcx>,
+        src: Place<'tcx>,
+        target: BasicBlock,
+        mut unwind: BasicBlock,
+        tys: I,
+    ) -> BasicBlock
     where
         I: IntoIterator<Item = Ty<'tcx>>,
     {
-        let mut previous_field = None;
+        // For an iterator of length n, create 2*n + 1 blocks.
         for (i, ity) in tys.into_iter().enumerate() {
+            // Each iteration creates two blocks, referred to here as block 2*i and block 2*i + 1.
+            //
+            // Block 2*i attempts to clone the field. If successful it branches to 2*i + 2 (the
+            // next clone block). If unsuccessful it branches to the previous unwind block, which
+            // is initially the `unwind` argument passed to this function.
+            //
+            // Block 2*i + 1 is the unwind block for this iteration. It drops the cloned value
+            // created by block 2*i. We store this block in `unwind` so that the next clone block
+            // will unwind to it if cloning fails.
+
             let field = Field::new(i);
             let src_field = self.tcx.mk_place_field(src, field, ity);
 
             let dest_field = self.tcx.mk_place_field(dest, field, ity);
 
-            // #(2i + 1) is the cleanup block for the previous clone operation
-            let cleanup_block = self.block_index_offset(1);
-            // #(2i + 2) is the next cloning block
-            // (or the Return terminator if this is the last block)
+            let next_unwind = self.block_index_offset(1);
             let next_block = self.block_index_offset(2);
+            self.make_clone_call(dest_field, src_field, ity, next_block, unwind);
+            self.block(
+                vec![],
+                TerminatorKind::Drop { place: dest_field, target: unwind, unwind: None },
+                true,
+            );
+            unwind = next_unwind;
+        }
+        // If all clones succeed then we end up here.
+        self.block(vec![], TerminatorKind::Goto { target }, false);
+        unwind
+    }
 
-            // BB #(2i)
-            // `dest.i = Clone::clone(&src.i);`
-            // Goto #(2i + 2) if ok, #(2i + 1) if unwinding happens.
-            self.make_clone_call(dest_field, src_field, ity, next_block, cleanup_block);
-
-            // BB #(2i + 1) (cleanup)
-            if let Some((previous_field, previous_cleanup)) = previous_field.take() {
-                // Drop previous field and goto previous cleanup block.
-                self.block(
-                    vec![],
-                    TerminatorKind::Drop {
-                        place: previous_field,
-                        target: previous_cleanup,
-                        unwind: None,
-                    },
-                    true,
-                );
-            } else {
-                // Nothing to drop, just resume.
-                self.block(vec![], TerminatorKind::Resume, true);
-            }
+    fn tuple_like_shim<I>(&mut self, dest: Place<'tcx>, src: Place<'tcx>, tys: I)
+    where
+        I: IntoIterator<Item = Ty<'tcx>>,
+    {
+        self.block(vec![], TerminatorKind::Goto { target: self.block_index_offset(3) }, false);
+        let unwind = self.block(vec![], TerminatorKind::Resume, true);
+        let target = self.block(vec![], TerminatorKind::Return, false);
 
-            previous_field = Some((dest_field, cleanup_block));
-        }
+        let _final_cleanup_block = self.clone_fields(dest, src, target, unwind, tys);
+    }
 
-        self.block(vec![], TerminatorKind::Return, false);
+    fn generator_shim(
+        &mut self,
+        dest: Place<'tcx>,
+        src: Place<'tcx>,
+        gen_def_id: DefId,
+        substs: GeneratorSubsts<'tcx>,
+    ) {
+        self.block(vec![], TerminatorKind::Goto { target: self.block_index_offset(3) }, false);
+        let unwind = self.block(vec![], TerminatorKind::Resume, true);
+        // This will get overwritten with a switch once we know the target blocks
+        let switch = self.block(vec![], TerminatorKind::Unreachable, false);
+        let unwind = self.clone_fields(dest, src, switch, unwind, substs.upvar_tys());
+        let target = self.block(vec![], TerminatorKind::Return, false);
+        let unreachable = self.block(vec![], TerminatorKind::Unreachable, false);
+        let mut cases = Vec::with_capacity(substs.state_tys(gen_def_id, self.tcx).count());
+        for (index, state_tys) in substs.state_tys(gen_def_id, self.tcx).enumerate() {
+            let variant_index = VariantIdx::new(index);
+            let dest = self.tcx.mk_place_downcast_unnamed(dest, variant_index);
+            let src = self.tcx.mk_place_downcast_unnamed(src, variant_index);
+            let clone_block = self.block_index_offset(1);
+            let start_block = self.block(
+                vec![self.make_statement(StatementKind::SetDiscriminant {
+                    place: Box::new(Place::return_place()),
+                    variant_index,
+                })],
+                TerminatorKind::Goto { target: clone_block },
+                false,
+            );
+            cases.push((index as u128, start_block));
+            let _final_cleanup_block = self.clone_fields(dest, src, target, unwind, state_tys);
+        }
+        let discr_ty = substs.discr_ty(self.tcx);
+        let temp = self.make_place(Mutability::Mut, discr_ty);
+        let rvalue = Rvalue::Discriminant(src);
+        let statement = self.make_statement(StatementKind::Assign(Box::new((temp, rvalue))));
+        match &mut self.blocks[switch] {
+            BasicBlockData { statements, terminator: Some(Terminator { kind, .. }), .. } => {
+                statements.push(statement);
+                *kind = TerminatorKind::SwitchInt {
+                    discr: Operand::Move(temp),
+                    switch_ty: discr_ty,
+                    targets: SwitchTargets::new(cases.into_iter(), unreachable),
+                };
+            }
+            BasicBlockData { terminator: None, .. } => unreachable!(),
+        }
     }
 }
 

compiler/rustc_span/src/symbol.rs

@@ -763,6 +763,7 @@ symbols! {
         gen_future,
         gen_kill,
         generator,
+        generator_clone,
         generator_state,
         generators,
         generic_arg_infer,

compiler/rustc_trait_selection/src/traits/select/mod.rs

@@ -1928,8 +1928,7 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
             ty::Dynamic(..)
             | ty::Str
             | ty::Slice(..)
-            | ty::Generator(..)
-            | ty::GeneratorWitness(..)
+            | ty::Generator(_, _, hir::Movability::Static)
             | ty::Foreign(..)
             | ty::Ref(_, _, hir::Mutability::Mut) => None,
 
@@ -1938,6 +1937,43 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
                 Where(obligation.predicate.rebind(tys.iter().collect()))
             }
 
+            ty::Generator(_, substs, hir::Movability::Movable) => {
+                if self.tcx().features().generator_clone {
+                    let resolved_upvars =
+                        self.infcx.shallow_resolve(substs.as_generator().tupled_upvars_ty());
+                    let resolved_witness =
+                        self.infcx.shallow_resolve(substs.as_generator().witness());
+                    if resolved_upvars.is_ty_var() || resolved_witness.is_ty_var() {
+                        // Not yet resolved.
+                        Ambiguous
+                    } else {
+                        let all = substs
+                            .as_generator()
+                            .upvar_tys()
+                            .chain(iter::once(substs.as_generator().witness()))
+                            .collect::<Vec<_>>();
+                        Where(obligation.predicate.rebind(all))
+                    }
+                } else {
+                    None
+                }
+            }
+
+            ty::GeneratorWitness(binder) => {
+                let witness_tys = binder.skip_binder();
+                for witness_ty in witness_tys.iter() {
+                    let resolved = self.infcx.shallow_resolve(witness_ty);
+                    if resolved.is_ty_var() {
+                        return Ambiguous;
+                    }
+                }
+                // (*) binder moved here
+                let all_vars = self.tcx().mk_bound_variable_kinds(
+                    obligation.predicate.bound_vars().iter().chain(binder.bound_vars().iter()),
+                );
+                Where(ty::Binder::bind_with_vars(witness_tys.to_vec(), all_vars))
+            }
+
             ty::Closure(_, substs) => {
                 // (*) binder moved here
                 let ty = self.infcx.shallow_resolve(substs.as_closure().tupled_upvars_ty());

compiler/rustc_ty_utils/src/instance.rs

@@ -263,7 +263,10 @@ fn resolve_associated_item<'tcx>(
                     let is_copy = self_ty.is_copy_modulo_regions(tcx.at(DUMMY_SP), param_env);
                     match self_ty.kind() {
                         _ if is_copy => (),
-                        ty::Closure(..) | ty::Tuple(..) => {}
+                        ty::Generator(..)
+                        | ty::GeneratorWitness(..)
+                        | ty::Closure(..)
+                        | ty::Tuple(..) => {}
                         _ => return Ok(None),
                     };
 

src/test/ui/generator/clone-impl-async.rs

@@ -0,0 +1,71 @@
+// edition:2021
+// gate-test-generator_clone
+// Verifies that feature(generator_clone) doesn't allow async blocks to be cloned/copied.
+
+#![feature(generators, generator_clone)]
+
+use std::future::ready;
+
+struct NonClone;
+
+fn main() {
+    let inner_non_clone = async {
+        let non_clone = NonClone;
+        let () = ready(()).await;
+        drop(non_clone);
+    };
+    check_copy(&inner_non_clone);
+    //~^ ERROR the trait bound `impl Future<Output = ()>: Copy` is not satisfied
+    check_clone(&inner_non_clone);
+    //~^ ERROR the trait bound `impl Future<Output = ()>: Clone` is not satisfied
+
+    let non_clone = NonClone;
+    let outer_non_clone = async move {
+        drop(non_clone);
+    };
+    check_copy(&outer_non_clone);
+    //~^ ERROR the trait bound `impl Future<Output = ()>: Copy` is not satisfied
+    check_clone(&outer_non_clone);
+    //~^ ERROR the trait bound `impl Future<Output = ()>: Clone` is not satisfied
+
+    let maybe_copy_clone = async move {};
+    check_copy(&maybe_copy_clone);
+    //~^ ERROR the trait bound `impl Future<Output = ()>: Copy` is not satisfied
+    check_clone(&maybe_copy_clone);
+    //~^ ERROR the trait bound `impl Future<Output = ()>: Clone` is not satisfied
+
+    let inner_non_clone_fn = the_inner_non_clone_fn();
+    check_copy(&inner_non_clone_fn);
+    //~^ ERROR the trait bound `impl Future<Output = ()>: Copy` is not satisfied
+    check_clone(&inner_non_clone_fn);
+    //~^ ERROR the trait bound `impl Future<Output = ()>: Clone` is not satisfied
+
+    let outer_non_clone_fn = the_outer_non_clone_fn(NonClone);
+    check_copy(&outer_non_clone_fn);
+    //~^ ERROR the trait bound `impl Future<Output = ()>: Copy` is not satisfied
+    check_clone(&outer_non_clone_fn);
+    //~^ ERROR the trait bound `impl Future<Output = ()>: Clone` is not satisfied
+
+    let maybe_copy_clone_fn = the_maybe_copy_clone_fn();
+    check_copy(&maybe_copy_clone_fn);
+    //~^ ERROR the trait bound `impl Future<Output = ()>: Copy` is not satisfied
+    check_clone(&maybe_copy_clone_fn);
+    //~^ ERROR the trait bound `impl Future<Output = ()>: Clone` is not satisfied
+}
+
+async fn the_inner_non_clone_fn() {
+    let non_clone = NonClone;
+    let () = ready(()).await;
+    drop(non_clone);
+}
+
+async fn the_outer_non_clone_fn(non_clone: NonClone) {
+    let () = ready(()).await;
+    drop(non_clone);
+}
+
+async fn the_maybe_copy_clone_fn() {
+}
+
+fn check_copy<T: Copy>(_x: &T) {}
+fn check_clone<T: Clone>(_x: &T) {}