Merge last_cached_unwind and diverge_blocks

Author: matthewjasper

src/librustc_mir/build/scope.rs

@@ -107,15 +107,14 @@ struct BreakableScope<'tcx> {
     region_scope: region::Scope,
     /// Where the body of the loop begins. `None` if block
     continue_block: Option<BasicBlock>,
-    /// Block to branch into when the loop or block terminates (either by being `break`-en out
-    /// from, or by having its condition to become false)
+    /// Block to branch into when the loop or block terminates (either by being
+    /// `break`-en out from, or by having its condition to become false)
     break_block: BasicBlock,
-    /// The destination of the loop/block expression itself (i.e., where to put the result of a
-    /// `break` expression)
+    /// The destination of the loop/block expression itself (i.e., where to put
+    /// the result of a `break` expression)
     break_destination: Place<'tcx>,
 }
 
-
 #[derive(Debug)]
 struct DropData {
     /// span where drop obligation was incurred (typically where place was declared)
@@ -131,7 +130,6 @@ struct DropData {
     cached_block: CachedBlock,
 }
 
-
 /// The target of an expression that breaks out of a scope
 #[derive(Clone, Copy, Debug)]
 pub enum BreakableTarget {
@@ -526,25 +524,27 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         // to left reading the cached results but never created anything.
 
         // Find the last cached block
-        debug!("diverge_cleanup_gen(self.scopes = {:?})", self.scopes);
-        let cached_cleanup = self.scopes.last_cached_unwind(generator_drop);
-        let (mut target, first_uncached) = cached_cleanup
-            .unwrap_or_else(|| (self.resume_block(), 0));
-
-        let scopes = self.scopes.diverge_blocks(first_uncached, generator_drop);
-        for (source_scope, drops, cached_unwind) in scopes {
-            target = build_diverge_scope(
-                &mut self.cfg,
-                source_scope,
-                drops,
-                target,
-                generator_drop,
-                self.is_generator,
-            );
-            *cached_unwind = Some(target);
-        }
-
-        target
+        debug!("diverge_cleanup_gen(self.scopes = {:#?})", self.scopes);
+        let resume_block = self.resume_block();
+        let cfg = &mut self.cfg;
+        let is_generator = self.is_generator;
+
+        self.scopes.for_each_diverge_block(
+            generator_drop,
+            resume_block,
+            |source_scope, drops, cached_unwind, mut target| {
+                target = build_diverge_scope(
+                    cfg,
+                    source_scope,
+                    drops,
+                    target,
+                    generator_drop,
+                    is_generator,
+                );
+                *cached_unwind = Some(target);
+                target
+            }
+        )
     }
 
     /// Utility function for *non*-scope code to build their own drops

src/librustc_mir/build/scope/stack.rs

@@ -140,7 +140,11 @@ impl<'tcx> Scopes<'tcx> {
         self.scopes.len()
     }
 
-    pub(super) fn push_scope(&mut self, region_scope: (region::Scope, SourceInfo), vis_scope: SourceScope) {
+    pub(super) fn push_scope(
+        &mut self,
+        region_scope: (region::Scope, SourceInfo),
+        vis_scope: SourceScope,
+    ) {
         debug!("push_scope({:?})", region_scope);
         self.scopes.push(Scope {
             source_scope: vis_scope,
@@ -217,24 +221,31 @@ impl<'tcx> Scopes<'tcx> {
         scope_count
     }
 
-    pub(super) fn last_cached_unwind(
-        &mut self,
-        generator_drop: bool,
-    ) -> Option<(BasicBlock, usize)> {
-        self.scopes.iter_mut().enumerate().rev().find_map(move |(idx, scope)| {
-            let cached_block = scope.cached_unwind.get(generator_drop)?;
-            Some((cached_block, idx + 1))
-        })
-    }
-
-    pub(super) fn diverge_blocks(
+    /// Call the given action for each scope that doesn't already have a cached
+    /// unwind block. Scopes are iterated going up the scope stack.
+    pub(super) fn for_each_diverge_block<F>(
         &mut self,
-        start_at: usize,
         generator_drop: bool,
-    ) -> impl Iterator<Item=(SourceScope, &mut [DropData], &mut Option<BasicBlock>)> {
-        self.scopes[start_at..].iter_mut().map(move |scope| {
-            (scope.source_scope, &mut *scope.drops, scope.cached_unwind.ref_mut(generator_drop))
-        })
+        resume_block: BasicBlock,
+        mut action: F,
+    ) -> BasicBlock
+    where
+        F: FnMut(SourceScope, &mut [DropData], &mut Option<BasicBlock>, BasicBlock) -> BasicBlock
+    {
+        let (mut target, start_at) = self.scopes.iter().enumerate().rev()
+            .find_map(move |(idx, scope)| {
+                let cached_block = scope.cached_unwind.get(generator_drop)?;
+                Some((cached_block, idx + 1))
+            }).unwrap_or((resume_block, 0));
+        self.scopes[start_at..].iter_mut().for_each(|scope| {
+            target = action(
+                scope.source_scope,
+                &mut *scope.drops,
+                scope.cached_unwind.ref_mut(generator_drop),
+                target,
+            )
+        });
+        target
     }
 
     /// Iterate over the [ScopeInfo] for exiting to the given target.