Temp: Caching

Author: matthewjasper

src/librustc_mir_build/build/scope.rs

@@ -105,8 +105,6 @@ pub struct Scopes<'tcx> {
 
     /// Drops that need to be done on paths to the `GeneratorDrop` terminator.
     generator_drops: DropTree,
-    // TODO: implement caching
-    // cached_unwind_drop: DropIdx,
 }
 
 #[derive(Debug)]
@@ -127,6 +125,14 @@ struct Scope {
     drops: Vec<DropData>,
 
     moved_locals: Vec<Local>,
+
+    /// The drop index that will drop everything in and below this scope on an
+    /// unwind path.
+    cached_unwind_block: Option<DropIdx>,
+
+    /// The drop index that will drop everything in and below this scope on a
+    /// generator drop path.
+    cached_generator_drop_block: Option<DropIdx>,
 }
 
 #[derive(Clone, Copy, Debug)]
@@ -213,6 +219,11 @@ impl Scope {
             DropKind::Storage => false,
         })
     }
+
+    fn invalidate_cache(&mut self) {
+        self.cached_unwind_block = None;
+        self.cached_generator_drop_block = None;
+    }
 }
 
 /// A trait that determined how [DropTree::lower_to_mir] creates its blocks and
@@ -229,11 +240,7 @@ impl DropTree {
             DropData { source_info: fake_source_info, local: Local::MAX, kind: DropKind::Storage };
         let drop_idx = DropIdx::MAX;
         let drops = IndexVec::from_elem_n((fake_data, drop_idx), 1);
-        Self {
-            drops,
-            entry_points: Vec::new(),
-            previous_drops: FxHashMap::default(),
-        }
+        Self { drops, entry_points: Vec::new(), previous_drops: FxHashMap::default() }
     }
 
     fn add_drop(&mut self, drop: DropData, next: DropIdx) -> DropIdx {
@@ -383,6 +390,8 @@ impl<'tcx> Scopes<'tcx> {
             region_scope_span: region_scope.1.span,
             drops: vec![],
             moved_locals: vec![],
+            cached_unwind_block: None,
+            cached_generator_drop_block: None,
         });
     }
 
@@ -399,10 +408,6 @@ impl<'tcx> Scopes<'tcx> {
             .unwrap_or_else(|| span_bug!(span, "region_scope {:?} does not enclose", region_scope))
     }
 
-    fn iter_mut(&mut self) -> impl DoubleEndedIterator<Item = &mut Scope> + '_ {
-        self.scopes.iter_mut().rev()
-    }
-
     /// Returns the topmost active scope, which is known to be alive until
     /// the next scope expression.
     fn topmost(&self) -> region::Scope {
@@ -609,10 +614,11 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         } else {
             self.scopes.breakable_scopes[break_index].continue_drops.as_mut().unwrap()
         };
-
-        let mut drop_idx = DropIdx::from_u32(destination.is_none() as u32);
-        for drop in scope_drops {
-            drop_idx = drops.add_drop(*drop, drop_idx);
+        let mut drop_idx = ROOT_NODE;
+        for scope in &self.scopes.scopes[scope_index + 1..] {
+            for drop in &scope.drops {
+                drop_idx = drops.add_drop(*drop, drop_idx);
+            }
         }
         drops.add_entry(block, drop_idx);
         // `build_drop_tree` doesn't have access to our source_info, so we
@@ -668,19 +674,6 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         ))
     }
 
-    /// Sets up a path that performs all required cleanup for dropping a generator.
-    ///
-    /// This path terminates in GeneratorDrop. Returns the start of the path.
-    /// None indicates there’s no cleanup to do at this point.
-    crate fn generator_drop_cleanup(&mut self, yield_block: BasicBlock) {
-        let drops = self.scopes.scopes.iter().flat_map(|scope| &scope.drops);
-        let mut next_drop = ROOT_NODE;
-        for drop in drops {
-            next_drop = self.scopes.generator_drops.add_drop(*drop, next_drop);
-        }
-        self.scopes.generator_drops.add_entry(yield_block, next_drop);
-    }
-
     /// Creates a new source scope, nested in the current one.
     crate fn new_source_scope(
         &mut self,
@@ -777,8 +770,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         local: Local,
         drop_kind: DropKind,
     ) {
-        // TODO: add back in caching.
-        let _needs_drop = match drop_kind {
+        let needs_drop = match drop_kind {
             DropKind::Value => {
                 if !self.hir.needs_drop(self.local_decls[local].ty) {
                     return;
@@ -798,23 +790,29 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             }
         };
 
-        let scope = self
-            .scopes
-            .iter_mut()
-            .find(|scope| scope.region_scope == region_scope)
-            .unwrap_or_else(|| {
-                span_bug!(span, "region scope {:?} not in scope to drop {:?}", region_scope, local);
-            });
-
-        let region_scope_span = region_scope.span(self.hir.tcx(), &self.hir.region_scope_tree);
-        // Attribute scope exit drops to scope's closing brace.
-        let scope_end = self.hir.tcx().sess.source_map().end_point(region_scope_span);
-
-        scope.drops.push(DropData {
-            source_info: SourceInfo { span: scope_end, scope: scope.source_scope },
-            local,
-            kind: drop_kind,
-        });
+        let invalidate_caches = needs_drop || self.is_generator;
+        for scope in self.scopes.scopes.iter_mut().rev() {
+            if invalidate_caches {
+                scope.invalidate_cache();
+            }
+
+            if scope.region_scope == region_scope {
+                let region_scope_span =
+                    region_scope.span(self.hir.tcx(), &self.hir.region_scope_tree);
+                // Attribute scope exit drops to scope's closing brace.
+                let scope_end = self.hir.tcx().sess.source_map().end_point(region_scope_span);
+
+                scope.drops.push(DropData {
+                    source_info: SourceInfo { span: scope_end, scope: scope.source_scope },
+                    local,
+                    kind: drop_kind,
+                });
+
+                return;
+            }
+        }
+
+        span_bug!(span, "region scope {:?} not in scope to drop {:?}", region_scope, local);
     }
 
     /// Indicates that the "local operand" stored in `local` is
@@ -861,9 +859,10 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             }
 
             Some(local_scope) => self
+                .scopes
                 .scopes
                 .iter_mut()
-                .find(|scope| scope.region_scope == local_scope)
+                .rfind(|scope| scope.region_scope == local_scope)
                 .unwrap_or_else(|| bug!("scope {:?} not found in scope list!", local_scope)),
         };
 
@@ -913,6 +912,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                     // Manually drop the condition on both branches.
                     let top_scope = self.scopes.scopes.last_mut().unwrap();
                     let top_drop_data = top_scope.drops.pop().unwrap();
+                    if self.is_generator {
+                        top_scope.invalidate_cache();
+                    }
 
                     match top_drop_data.kind {
                         DropKind::Value { .. } => {
@@ -943,17 +945,25 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
 
     fn diverge_cleanup(&mut self) -> DropIdx {
         let is_generator = self.is_generator;
-        let drops = self
+        let (uncached_scope, mut cached_drop) = self
             .scopes
             .scopes
             .iter()
-            .flat_map(|scope| &scope.drops)
-            .filter(|drop| is_generator || drop.kind == DropKind::Value);
-        let mut next_drop = ROOT_NODE;
-        for drop in drops {
-            next_drop = self.scopes.unwind_drops.add_drop(*drop, next_drop);
+            .enumerate()
+            .rev()
+            .find_map(|(scope_idx, scope)| {
+                scope.cached_unwind_block.map(|cached_block| (scope_idx + 1, cached_block))
+            })
+            .unwrap_or((0, ROOT_NODE));
+        for scope in &mut self.scopes.scopes[uncached_scope..] {
+            for drop in &scope.drops {
+                if is_generator || drop.kind == DropKind::Value {
+                    cached_drop = self.scopes.unwind_drops.add_drop(*drop, cached_drop);
+                }
+            }
+            scope.cached_unwind_block = Some(cached_drop);
         }
-        next_drop
+        cached_drop
     }
 
     /// Prepares to create a path that performs all required cleanup for
@@ -966,6 +976,30 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         self.scopes.unwind_drops.add_entry(start, next_drop);
     }
 
+    /// Sets up a path that performs all required cleanup for dropping a generator.
+    ///
+    /// This path terminates in GeneratorDrop. Returns the start of the path.
+    /// None indicates there’s no cleanup to do at this point.
+    crate fn generator_drop_cleanup(&mut self, yield_block: BasicBlock) {
+        let (uncached_scope, mut cached_drop) = self
+            .scopes
+            .scopes
+            .iter()
+            .enumerate()
+            .rev()
+            .find_map(|(scope_idx, scope)| {
+                scope.cached_generator_drop_block.map(|cached_block| (scope_idx + 1, cached_block))
+            })
+            .unwrap_or((0, ROOT_NODE));
+        for scope in &mut self.scopes.scopes[uncached_scope..] {
+            for drop in &scope.drops {
+                cached_drop = self.scopes.generator_drops.add_drop(*drop, cached_drop);
+            }
+            scope.cached_generator_drop_block = Some(cached_drop);
+        }
+        self.scopes.generator_drops.add_entry(yield_block, cached_drop);
+    }
+
     /// Utility function for *non*-scope code to build their own drops
     crate fn build_drop_and_replace(
         &mut self,
@@ -1022,6 +1056,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         assert_eq!(top_scope.region_scope, region_scope);
 
         top_scope.drops.clear();
+        top_scope.invalidate_cache();
     }
 }