[generator] Special cases for match guard when analyzing interior types in generators

compiler/rustc_middle/src/middle/region.rs

@@ -283,23 +283,27 @@ pub struct ScopeTree {
     /// To see that this method works, consider:
     ///
     /// Let `D` be our binding/temporary and `U` be our other HIR node, with
-    /// `HIR-postorder(U) < HIR-postorder(D)` (in our example, U would be
-    /// the yield and D would be one of the calls). Let's show that
-    /// `D` is storage-dead at `U`.
+    /// `HIR-postorder(U) < HIR-postorder(D)`. Suppose, as in our example,
+    /// U is the yield and D is one of the calls.
+    /// Let's show that `D` is storage-dead at `U`.
     ///
     /// Remember that storage-live/storage-dead refers to the state of
     /// the *storage*, and does not consider moves/drop flags.
     ///
     /// Then:
+    ///
     ///     1. From the ordering guarantee of HIR visitors (see
     ///     `rustc_hir::intravisit`), `D` does not dominate `U`.
+    ///
     ///     2. Therefore, `D` is *potentially* storage-dead at `U` (because
     ///     we might visit `U` without ever getting to `D`).
+    ///
     ///     3. However, we guarantee that at each HIR point, each
     ///     binding/temporary is always either always storage-live
     ///     or always storage-dead. This is what is being guaranteed
     ///     by `terminating_scopes` including all blocks where the
     ///     count of executions is not guaranteed.
+    ///
     ///     4. By `2.` and `3.`, `D` is *statically* storage-dead at `U`,
     ///     QED.
     ///

compiler/rustc_mir_build/src/build/matches/simplify.rs

@@ -28,8 +28,9 @@ use std::mem;
 impl<'a, 'tcx> Builder<'a, 'tcx> {
     /// Simplify a candidate so that all match pairs require a test.
     ///
-    /// This method will also split a candidate where the only match-pair is an
-    /// or-pattern into multiple candidates. This is so that
+    /// This method will also split a candidate, in which the only
+    /// match-pair is an or-pattern, into multiple candidates.
+    /// This is so that
     ///
     /// match x {
     ///     0 | 1 => { ... },

compiler/rustc_typeck/src/check/generator_interior.rs

@@ -8,11 +8,13 @@ use rustc_data_structures::fx::{FxHashSet, FxIndexSet};
 use rustc_hir as hir;
 use rustc_hir::def::{CtorKind, DefKind, Res};
 use rustc_hir::def_id::DefId;
+use rustc_hir::hir_id::HirIdSet;
 use rustc_hir::intravisit::{self, NestedVisitorMap, Visitor};
-use rustc_hir::{Expr, ExprKind, Pat, PatKind};
+use rustc_hir::{Arm, Expr, ExprKind, Guard, HirId, Pat, PatKind};
 use rustc_middle::middle::region::{self, YieldData};
 use rustc_middle::ty::{self, Ty};
 use rustc_span::Span;
+use smallvec::SmallVec;
 
 struct InteriorVisitor<'a, 'tcx> {
     fcx: &'a FnCtxt<'a, 'tcx>,
@@ -21,6 +23,13 @@ struct InteriorVisitor<'a, 'tcx> {
     expr_count: usize,
     kind: hir::GeneratorKind,
     prev_unresolved_span: Option<Span>,
+    /// Match arm guards have temporary borrows from the pattern bindings.
+    /// In case there is a yield point in a guard with a reference to such bindings,
+    /// such borrows can span across this yield point.
+    /// As such, we need to track these borrows and record them despite of the fact
+    /// that they may succeed the said yield point in the post-order.
+    guard_bindings: SmallVec<[SmallVec<[HirId; 4]>; 1]>,
+    guard_bindings_set: HirIdSet,
 }
 
 impl<'a, 'tcx> InteriorVisitor<'a, 'tcx> {
@@ -30,6 +39,7 @@ impl<'a, 'tcx> InteriorVisitor<'a, 'tcx> {
         scope: Option<region::Scope>,
         expr: Option<&'tcx Expr<'tcx>>,
         source_span: Span,
+        guard_borrowing_from_pattern: bool,
     ) {
         use rustc_span::DUMMY_SP;
 
@@ -53,7 +63,12 @@ impl<'a, 'tcx> InteriorVisitor<'a, 'tcx> {
                         yield_data.expr_and_pat_count, self.expr_count, source_span
                     );
 
-                    if yield_data.expr_and_pat_count >= self.expr_count {
+                    // If it is a borrowing happening in the guard,
+                    // it needs to be recorded regardless because they
+                    // do live across this yield point.
+                    if guard_borrowing_from_pattern
+                        || yield_data.expr_and_pat_count >= self.expr_count
+                    {
                         Some(yield_data)
                     } else {
                         None
@@ -134,6 +149,8 @@ pub fn resolve_interior<'a, 'tcx>(
         expr_count: 0,
         kind,
         prev_unresolved_span: None,
+        guard_bindings: <_>::default(),
+        guard_bindings_set: <_>::default(),
     };
     intravisit::walk_body(&mut visitor, body);
 
@@ -210,6 +227,38 @@ impl<'a, 'tcx> Visitor<'tcx> for InteriorVisitor<'a, 'tcx> {
         NestedVisitorMap::None
     }
 
+    fn visit_arm(&mut self, arm: &'tcx Arm<'tcx>) {
+        let Arm { guard, pat, body, .. } = arm;
+        self.visit_pat(pat);
+        if let Some(ref g) = guard {
+            self.guard_bindings.push(<_>::default());
+            ArmPatCollector {
+                guard_bindings_set: &mut self.guard_bindings_set,
+                guard_bindings: self
+                    .guard_bindings
+                    .last_mut()
+                    .expect("should have pushed at least one earlier"),
+            }
+            .visit_pat(pat);
+
+            match g {
+                Guard::If(ref e) => {
+                    self.visit_expr(e);
+                }
+            }
+
+            let mut scope_var_ids =
+                self.guard_bindings.pop().expect("should have pushed at least one earlier");
+            for var_id in scope_var_ids.drain(..) {
+                assert!(
+                    self.guard_bindings_set.remove(&var_id),
+                    "variable should be placed in scope earlier"
+                );
+            }
+        }
+        self.visit_expr(body);
+    }
+
     fn visit_pat(&mut self, pat: &'tcx Pat<'tcx>) {
         intravisit::walk_pat(self, pat);
 
@@ -218,11 +267,12 @@ impl<'a, 'tcx> Visitor<'tcx> for InteriorVisitor<'a, 'tcx> {
         if let PatKind::Binding(..) = pat.kind {
             let scope = self.region_scope_tree.var_scope(pat.hir_id.local_id);
             let ty = self.fcx.typeck_results.borrow().pat_ty(pat);
-            self.record(ty, Some(scope), None, pat.span);
+            self.record(ty, Some(scope), None, pat.span, false);
         }
     }
 
     fn visit_expr(&mut self, expr: &'tcx Expr<'tcx>) {
+        let mut guard_borrowing_from_pattern = false;
         match &expr.kind {
             ExprKind::Call(callee, args) => match &callee.kind {
                 ExprKind::Path(qpath) => {
@@ -249,6 +299,16 @@ impl<'a, 'tcx> Visitor<'tcx> for InteriorVisitor<'a, 'tcx> {
                 }
                 _ => intravisit::walk_expr(self, expr),
             },
+            ExprKind::Path(qpath) => {
+                intravisit::walk_expr(self, expr);
+                let res = self.fcx.typeck_results.borrow().qpath_res(qpath, expr.hir_id);
+                match res {
+                    Res::Local(id) if self.guard_bindings_set.contains(&id) => {
+                        guard_borrowing_from_pattern = true;
+                    }
+                    _ => {}
+                }
+            }
             _ => intravisit::walk_expr(self, expr),
         }
 
@@ -259,18 +319,18 @@ impl<'a, 'tcx> Visitor<'tcx> for InteriorVisitor<'a, 'tcx> {
         // If there are adjustments, then record the final type --
         // this is the actual value that is being produced.
         if let Some(adjusted_ty) = self.fcx.typeck_results.borrow().expr_ty_adjusted_opt(expr) {
-            self.record(adjusted_ty, scope, Some(expr), expr.span);
+            self.record(adjusted_ty, scope, Some(expr), expr.span, guard_borrowing_from_pattern);
         }
 
         // Also record the unadjusted type (which is the only type if
         // there are no adjustments). The reason for this is that the
         // unadjusted value is sometimes a "temporary" that would wind
         // up in a MIR temporary.
         //
-        // As an example, consider an expression like `vec![].push()`.
+        // As an example, consider an expression like `vec![].push(x)`.
         // Here, the `vec![]` would wind up MIR stored into a
         // temporary variable `t` which we can borrow to invoke
-        // `<Vec<_>>::push(&mut t)`.
+        // `<Vec<_>>::push(&mut t, x)`.
         //
         // Note that an expression can have many adjustments, and we
         // are just ignoring those intermediate types. This is because
@@ -287,9 +347,30 @@ impl<'a, 'tcx> Visitor<'tcx> for InteriorVisitor<'a, 'tcx> {
         // The type table might not have information for this expression
         // if it is in a malformed scope. (#66387)
         if let Some(ty) = self.fcx.typeck_results.borrow().expr_ty_opt(expr) {
-            self.record(ty, scope, Some(expr), expr.span);
+            self.record(ty, scope, Some(expr), expr.span, guard_borrowing_from_pattern);
         } else {
             self.fcx.tcx.sess.delay_span_bug(expr.span, "no type for node");
         }
     }
 }
+
+struct ArmPatCollector<'a> {
+    guard_bindings_set: &'a mut HirIdSet,
+    guard_bindings: &'a mut SmallVec<[HirId; 4]>,
+}
+
+impl<'a, 'tcx> Visitor<'tcx> for ArmPatCollector<'a> {
+    type Map = intravisit::ErasedMap<'tcx>;
+
+    fn nested_visit_map(&mut self) -> NestedVisitorMap<Self::Map> {
+        NestedVisitorMap::None
+    }
+
+    fn visit_pat(&mut self, pat: &'tcx Pat<'tcx>) {
+        intravisit::walk_pat(self, pat);
+        if let PatKind::Binding(_, id, ..) = pat.kind {
+            self.guard_bindings.push(id);
+            self.guard_bindings_set.insert(id);
+        }
+    }
+}

src/test/ui/generator/yielding-in-match-guards.rs

@@ -0,0 +1,24 @@
+// check-pass
+// edition:2018
+
+// This test is derived from
+// https://github.com/rust-lang/rust/issues/72651#issuecomment-668720468
+
+// This test demonstrates that, in `async fn g()`,
+// indeed a temporary borrow `y` from `x` is live
+// while `f().await` is being evaluated.
+// Thus, `&'_ u8` should be included in type signature
+// of the underlying generator.
+
+async fn f() -> u8 { 1 }
+
+pub async fn g(x: u8) {
+    match x {
+        y if f().await == y => (),
+        _ => (),
+    }
+}
+
+fn main() {
+    let _ = g(10);
+}