Auto merge of #75213 - dingxiangfei2009:yield-point-in-match-guard, r=tmandry

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

Fix #72651

This proposes one of ways to fix the mentioned issue. One cause of #72651 is that the interior type analysis misses out types of match pattern locals. Those locals are manifested as temporary borrows in the scopes of match arm guards. If uses of these locals appear after yield points, the borrows from them were not considered live across the yield points. However, this is not the case since the borrowing always happens at the very beginning of the match guard.

This calls for special treatment to analysis of types appearing in the match guard. Those borrows are recorded as the HIR tree is walked by `InteriorVisitor` and their uses are recorded whenever a yield point is crossed.

Author: bors

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