Sanity check the Expr visitation count

2 files changed, 20 insertions, 0 deletions

diff --git a/src/librustc/middle/region.rs b/src/librustc/middle/region.rs
index c2d178e2207..9435b28a013 100644
--- a/src/librustc/middle/region.rs
+++ b/src/librustc/middle/region.rs
@@ -252,6 +252,11 @@ pub struct ScopeTree {
     /// stores the `Span` of the last one and the number of expressions
     /// which came before it in a generator body.
     yield_in_scope: FxHashMap<Scope, (Span, usize)>,
+
+    /// The number of visit_expr calls done in the body.
+    /// Used to sanity check visit_expr call count when
+    /// calculating geneartor interiors.
+    body_expr_count: FxHashMap<hir::BodyId, usize>,
 }
 
 #[derive(Debug, Copy, Clone)]
@@ -619,6 +624,13 @@ impl<'tcx> ScopeTree {
     pub fn yield_in_scope(&self, scope: Scope) -> Option<(Span, usize)> {
         self.yield_in_scope.get(&scope).cloned()
     }
+
+    /// Gives the number of expressions visited in a body.
+    /// Used to sanity check visit_expr call count when
+    /// calculating geneartor interiors.
+    pub fn body_expr_count(&self, body_id: hir::BodyId) -> Option<usize> {
+        self.body_expr_count.get(&body_id).map(|r| *r)
+    }
 }
 
 /// Records the lifetime of a local variable as `cx.var_parent`
@@ -1166,6 +1178,10 @@ impl<'a, 'tcx> Visitor<'tcx> for RegionResolutionVisitor<'a, 'tcx> {
             resolve_local(self, None, Some(&body.value));
         }
 
+        if body.is_generator {
+            self.scope_tree.body_expr_count.insert(body_id, self.expr_count);
+        }
+
         // Restore context we had at the start.
         self.expr_count = outer_ec;
         self.cx = outer_cx;
diff --git a/src/librustc_typeck/check/generator_interior.rs b/src/librustc_typeck/check/generator_interior.rs
index c7971666f44..fc6ea0ad5b0 100644
--- a/src/librustc_typeck/check/generator_interior.rs
+++ b/src/librustc_typeck/check/generator_interior.rs
@@ -77,6 +77,10 @@ pub fn resolve_interior<'a, 'gcx, 'tcx>(fcx: &'a FnCtxt<'a, 'gcx, 'tcx>,
     };
     intravisit::walk_body(&mut visitor, body);
 
+    // Check that we visited the same amount of expressions and the RegionResolutionVisitor
+    let region_expr_count = visitor.region_scope_tree.body_expr_count(body_id).unwrap();
+    assert_eq!(region_expr_count, visitor.expr_count);
+
     let mut types: Vec<_> = visitor.types.drain().collect();
 
     // Sort types by insertion order