Auto merge of #61572 - Aaron1011:fix/generator-ref, r=varkor

Fix HIR visit order

Fixes #61442

When rustc::middle::region::ScopeTree computes its yield_in_scope
field, it relies on the HIR visitor order to properly compute which
types must be live across yield points. In order for the computed scopes
to agree with the generated MIR, we must ensure that expressions
evaluated before a yield point are visited before the 'yield'
expression.

However, the visitor order for ExprKind::AssignOp
was incorrect. The left-hand side of a compund assignment expression is
evaluated before the right-hand side, but the right-hand expression was
being visited before the left-hand expression. If the left-hand
expression caused a new type to be introduced (e.g. through a
deref-coercion), the new type would be incorrectly seen as occuring
*after* the yield point, instead of before. This leads to a mismatch
between the computed generator types and the MIR, since the MIR will
correctly see the type as being live across the yield point.

To fix this, we correct the visitor order for ExprKind::AssignOp
to reflect the actual evaulation order.

4 files changed, 146 insertions, 2 deletions

diff --git a/src/librustc/hir/intravisit.rs b/src/librustc/hir/intravisit.rs
index 1777d765cc8..9c05f18762d 100644
--- a/src/librustc/hir/intravisit.rs
+++ b/src/librustc/hir/intravisit.rs
@@ -1056,7 +1056,7 @@ pub fn walk_expr<'v, V: Visitor<'v>>(visitor: &mut V, expression: &'v Expr) {
         }
         ExprKind::AssignOp(_, ref left_expression, ref right_expression) => {
             visitor.visit_expr(right_expression);
-            visitor.visit_expr(left_expression)
+            visitor.visit_expr(left_expression);
         }
         ExprKind::Field(ref subexpression, ident) => {
             visitor.visit_expr(subexpression);
diff --git a/src/librustc/middle/region.rs b/src/librustc/middle/region.rs
index d9ccb9d42f2..114684b1524 100644
--- a/src/librustc/middle/region.rs
+++ b/src/librustc/middle/region.rs
@@ -371,7 +371,12 @@ struct RegionResolutionVisitor<'tcx> {
 
     // The number of expressions and patterns visited in the current body
     expr_and_pat_count: usize,
-
+    // When this is `true`, we record the `Scopes` we encounter
+    // when processing a Yield expression. This allows us to fix
+    // up their indices.
+    pessimistic_yield: bool,
+    // Stores scopes when pessimistic_yield is true.
+    fixup_scopes: Vec<Scope>,
     // Generated scope tree:
     scope_tree: ScopeTree,
 
@@ -947,12 +952,107 @@ fn resolve_expr<'tcx>(visitor: &mut RegionResolutionVisitor<'tcx>, expr: &'tcx h
         }
     }
 
+    let prev_pessimistic = visitor.pessimistic_yield;
+
+    // Ordinarily, we can rely on the visit order of HIR intravisit
+    // to correspond to the actual execution order of statements.
+    // However, there's a weird corner case with compund assignment
+    // operators (e.g. `a += b`). The evaluation order depends on whether
+    // or not the operator is overloaded (e.g. whether or not a trait
+    // like AddAssign is implemented).
+
+    // For primitive types (which, despite having a trait impl, don't actually
+    // end up calling it), the evluation order is right-to-left. For example,
+    // the following code snippet:
+    //
+    //    let y = &mut 0;
+    //    *{println!("LHS!"); y} += {println!("RHS!"); 1};
+    //
+    // will print:
+    //
+    // RHS!
+    // LHS!
+    //
+    // However, if the operator is used on a non-primitive type,
+    // the evaluation order will be left-to-right, since the operator
+    // actually get desugared to a method call. For example, this
+    // nearly identical code snippet:
+    //
+    //     let y = &mut String::new();
+    //    *{println!("LHS String"); y} += {println!("RHS String"); "hi"};
+    //
+    // will print:
+    // LHS String
+    // RHS String
+    //
+    // To determine the actual execution order, we need to perform
+    // trait resolution. Unfortunately, we need to be able to compute
+    // yield_in_scope before type checking is even done, as it gets
+    // used by AST borrowcheck.
+    //
+    // Fortunately, we don't need to know the actual execution order.
+    // It suffices to know the 'worst case' order with respect to yields.
+    // Specifically, we need to know the highest 'expr_and_pat_count'
+    // that we could assign to the yield expression. To do this,
+    // we pick the greater of the two values from the left-hand
+    // and right-hand expressions. This makes us overly conservative
+    // about what types could possibly live across yield points,
+    // but we will never fail to detect that a type does actually
+    // live across a yield point. The latter part is critical -
+    // we're already overly conservative about what types will live
+    // across yield points, as the generated MIR will determine
+    // when things are actually live. However, for typecheck to work
+    // properly, we can't miss any types.
+
+
     match expr.node {
         // Manually recurse over closures, because they are the only
         // case of nested bodies that share the parent environment.
         hir::ExprKind::Closure(.., body, _, _) => {
             let body = visitor.tcx.hir().body(body);
             visitor.visit_body(body);
+        },
+        hir::ExprKind::AssignOp(_, ref left_expr, ref right_expr) => {
+            debug!("resolve_expr - enabling pessimistic_yield, was previously {}",
+                   prev_pessimistic);
+
+            let start_point = visitor.fixup_scopes.len();
+            visitor.pessimistic_yield = true;
+
+            // If the actual execution order turns out to be right-to-left,
+            // then we're fine. However, if the actual execution order is left-to-right,
+            // then we'll assign too low a count to any `yield` expressions
+            // we encounter in 'right_expression' - they should really occur after all of the
+            // expressions in 'left_expression'.
+            visitor.visit_expr(&right_expr);
+            visitor.pessimistic_yield = prev_pessimistic;
+
+            debug!("resolve_expr - restoring pessimistic_yield to {}", prev_pessimistic);
+            visitor.visit_expr(&left_expr);
+            debug!("resolve_expr - fixing up counts to {}", visitor.expr_and_pat_count);
+
+            // Remove and process any scopes pushed by the visitor
+            let target_scopes = visitor.fixup_scopes.drain(start_point..);
+
+            for scope in target_scopes {
+                let mut yield_data = visitor.scope_tree.yield_in_scope.get_mut(&scope).unwrap();
+                let count = yield_data.expr_and_pat_count;
+                let span = yield_data.span;
+
+                // expr_and_pat_count never decreases. Since we recorded counts in yield_in_scope
+                // before walking the left-hand side, it should be impossible for the recorded
+                // count to be greater than the left-hand side count.
+                if count > visitor.expr_and_pat_count {
+                    bug!("Encountered greater count {} at span {:?} - expected no greater than {}",
+                         count, span, visitor.expr_and_pat_count);
+                }
+                let new_count = visitor.expr_and_pat_count;
+                debug!("resolve_expr - increasing count for scope {:?} from {} to {} at span {:?}",
+                       scope, count, new_count, span);
+
+                yield_data.expr_and_pat_count = new_count;
+            }
+
         }
 
         _ => intravisit::walk_expr(visitor, expr)
@@ -972,6 +1072,10 @@ fn resolve_expr<'tcx>(visitor: &mut RegionResolutionVisitor<'tcx>, expr: &'tcx h
                 source: *source,
             };
             visitor.scope_tree.yield_in_scope.insert(scope, data);
+            if visitor.pessimistic_yield {
+                debug!("resolve_expr in pessimistic_yield - marking scope {:?} for fixup", scope);
+                visitor.fixup_scopes.push(scope);
+            }
 
             // Keep traversing up while we can.
             match visitor.scope_tree.parent_map.get(&scope) {
@@ -1360,6 +1464,8 @@ fn region_scope_tree<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> &'tcx ScopeTree
                 var_parent: None,
             },
             terminating_scopes: Default::default(),
+            pessimistic_yield: false,
+            fixup_scopes: vec![],
         };
 
         let body = tcx.hir().body(body_id);
diff --git a/src/librustc_typeck/check/generator_interior.rs b/src/librustc_typeck/check/generator_interior.rs
index 72f42c85ead..0bd078dace4 100644
--- a/src/librustc_typeck/check/generator_interior.rs
+++ b/src/librustc_typeck/check/generator_interior.rs
@@ -28,6 +28,10 @@ impl<'a, 'tcx> InteriorVisitor<'a, 'tcx> {
               source_span: Span) {
         use syntax_pos::DUMMY_SP;
 
+        debug!("generator_interior: attempting to record type {:?} {:?} {:?} {:?}",
+               ty, scope, expr, source_span);
+
+
         let live_across_yield = scope.map(|s| {
             self.region_scope_tree.yield_in_scope(s).and_then(|yield_data| {
                 // If we are recording an expression that is the last yield
diff --git a/src/test/run-pass/generator/addassign-yield.rs b/src/test/run-pass/generator/addassign-yield.rs
new file mode 100644
index 00000000000..6a417936384
--- /dev/null
+++ b/src/test/run-pass/generator/addassign-yield.rs
@@ -0,0 +1,34 @@
+// Regression test for broken MIR error (#61442)
+// Due to the two possible evaluation orders for
+// a '+=' expression (depending on whether or not the 'AddAssign' trait
+// is being used), we were failing to account for all types that might
+// possibly be live across a yield point.
+
+#![feature(generators)]
+
+fn foo() {
+    let _x = static || {
+        let mut s = String::new();
+        s += { yield; "" };
+    };
+
+    let _y = static || {
+        let x = &mut 0;
+        *{ yield; x } += match String::new() { _ => 0 };
+    };
+
+    // Please don't ever actually write something like this
+    let _z = static || {
+        let x = &mut 0;
+        *{
+            let inner = &mut 1;
+            *{ yield (); inner } += match String::new() { _ => 1};
+            yield;
+            x
+        } += match String::new() { _ => 2 };
+    };
+}
+
+fn main() {
+    foo()
+}