rustc_mir: optimize the deaggregator's expansion of statements.

4 files changed, 104 insertions, 57 deletions

diff --git a/src/librustc/lib.rs b/src/librustc/lib.rs
index a7a26195059..2700ef28d67 100644
--- a/src/librustc/lib.rs
+++ b/src/librustc/lib.rs
@@ -69,6 +69,7 @@
 #![feature(underscore_lifetimes)]
 #![feature(universal_impl_trait)]
 #![feature(trace_macros)]
+#![feature(trusted_len)]
 #![feature(catch_expr)]
 #![feature(test)]
 
diff --git a/src/librustc/mir/mod.rs b/src/librustc/mir/mod.rs
index b88dea871ce..475946468fa 100644
--- a/src/librustc/mir/mod.rs
+++ b/src/librustc/mir/mod.rs
@@ -34,13 +34,13 @@ use std::ascii;
 use std::borrow::{Cow};
 use std::cell::Ref;
 use std::fmt::{self, Debug, Formatter, Write};
-use std::{iter, u32};
+use std::{iter, mem, u32};
 use std::ops::{Index, IndexMut};
 use std::rc::Rc;
 use std::vec::IntoIter;
 use syntax::ast::{self, Name};
 use syntax::symbol::InternedString;
-use syntax_pos::Span;
+use syntax_pos::{Span, DUMMY_SP};
 
 mod cache;
 pub mod tcx;
@@ -984,11 +984,62 @@ impl<'tcx> BasicBlockData<'tcx> {
     pub fn retain_statements<F>(&mut self, mut f: F) where F: FnMut(&mut Statement) -> bool {
         for s in &mut self.statements {
             if !f(s) {
-                s.kind = StatementKind::Nop;
+                s.make_nop();
             }
         }
     }
 
+    pub fn expand_statements<F, I>(&mut self, mut f: F)
+        where F: FnMut(&mut Statement<'tcx>) -> Option<I>,
+              I: iter::TrustedLen<Item = Statement<'tcx>>
+    {
+        // Gather all the iterators we'll need to splice in, and their positions.
+        let mut splices: Vec<(usize, I)> = vec![];
+        let mut extra_stmts = 0;
+        for (i, s) in self.statements.iter_mut().enumerate() {
+            if let Some(mut new_stmts) = f(s) {
+                if let Some(first) = new_stmts.next() {
+                    // We can already store the first new statement.
+                    *s = first;
+
+                    // Save the other statements for optimized splicing.
+                    let remaining = new_stmts.size_hint().0;
+                    if remaining > 0 {
+                        splices.push((i + 1 + extra_stmts, new_stmts));
+                        extra_stmts += remaining;
+                    }
+                } else {
+                    s.make_nop();
+                }
+            }
+        }
+
+        // Splice in the new statements, from the end of the block.
+        // FIXME(eddyb) This could be more efficient with a "gap buffer"
+        // where a range of elements ("gap") is left uninitialized, with
+        // splicing adding new elements to the end of that gap and moving
+        // existing elements from before the gap to the end of the gap.
+        // For now, this is safe code, emulating a gap but initializing it.
+        let mut gap = self.statements.len()..self.statements.len()+extra_stmts;
+        self.statements.resize(gap.end, Statement {
+            source_info: SourceInfo {
+                span: DUMMY_SP,
+                scope: ARGUMENT_VISIBILITY_SCOPE
+            },
+            kind: StatementKind::Nop
+        });
+        for (splice_start, new_stmts) in splices.into_iter().rev() {
+            let splice_end = splice_start + new_stmts.size_hint().0;
+            while gap.end > splice_end {
+                gap.start -= 1;
+                gap.end -= 1;
+                self.statements.swap(gap.start, gap.end);
+            }
+            self.statements.splice(splice_start..splice_end, new_stmts);
+            gap.end = splice_start;
+        }
+    }
+
     pub fn visitable(&self, index: usize) -> &dyn MirVisitable<'tcx> {
         if index < self.statements.len() {
             &self.statements[index]
@@ -1157,6 +1208,14 @@ impl<'tcx> Statement<'tcx> {
     pub fn make_nop(&mut self) {
         self.kind = StatementKind::Nop
     }
+
+    /// Changes a statement to a nop and returns the original statement.
+    pub fn replace_nop(&mut self) -> Self {
+        Statement {
+            source_info: self.source_info,
+            kind: mem::replace(&mut self.kind, StatementKind::Nop)
+        }
+    }
 }
 
 #[derive(Clone, Debug, RustcEncodable, RustcDecodable)]
diff --git a/src/librustc_mir/transform/deaggregator.rs b/src/librustc_mir/transform/deaggregator.rs
index 8ffcc4025f1..503354ebc4f 100644
--- a/src/librustc_mir/transform/deaggregator.rs
+++ b/src/librustc_mir/transform/deaggregator.rs
@@ -39,32 +39,27 @@ impl MirPass for Deaggregator {
             }
         }
 
-        let can_deaggregate = |statement: &Statement| {
-            if let StatementKind::Assign(_, ref rhs) = statement.kind {
-                if let Rvalue::Aggregate(ref kind, _) = *rhs {
-                    // FIXME(#48193) Deaggregate arrays when it's cheaper to do so.
-                    if let AggregateKind::Array(_) = **kind {
-                        return false;
-                    }
-                    return true;
-                }
-            }
-
-            false
-        };
-
         let (basic_blocks, local_decls) = mir.basic_blocks_and_local_decls_mut();
+        let local_decls = &*local_decls;
         for bb in basic_blocks {
-            let mut start = 0;
-            while let Some(i) = bb.statements[start..].iter().position(&can_deaggregate) {
-                let i = start + i;
+            bb.expand_statements(|stmt| {
+                // FIXME(eddyb) don't match twice on `stmt.kind` (post-NLL).
+                if let StatementKind::Assign(_, ref rhs) = stmt.kind {
+                    if let Rvalue::Aggregate(ref kind, _) = *rhs {
+                        // FIXME(#48193) Deaggregate arrays when it's cheaper to do so.
+                        if let AggregateKind::Array(_) = **kind {
+                            return None;
+                        }
+                    } else {
+                        return None;
+                    }
+                } else {
+                    return None;
+                }
 
-                // FIXME(eddyb) this is probably more expensive than it should be.
-                // Ideally we'd move the block's statements all at once.
-                let suffix_stmts = bb.statements.split_off(i + 1);
-                let orig_stmt = bb.statements.pop().unwrap();
-                let source_info = orig_stmt.source_info;
-                let (mut lhs, kind, operands) = match orig_stmt.kind {
+                let stmt = stmt.replace_nop();
+                let source_info = stmt.source_info;
+                let (mut lhs, kind, operands) = match stmt.kind {
                     StatementKind::Assign(lhs, Rvalue::Aggregate(kind, operands))
                         => (lhs, kind, operands),
                     _ => bug!()
@@ -88,17 +83,11 @@ impl MirPass for Deaggregator {
                     _ => None
                 };
 
-                let new_total_count = bb.statements.len() +
-                    operands.len() +
-                    (set_discriminant.is_some() as usize) +
-                    suffix_stmts.len();
-                bb.statements.reserve(new_total_count);
-
-                for (j, op) in operands.into_iter().enumerate() {
+                Some(operands.into_iter().enumerate().map(move |(i, op)| {
                     let lhs_field = if let AggregateKind::Array(_) = *kind {
                         // FIXME(eddyb) `offset` should be u64.
-                        let offset = j as u32;
-                        assert_eq!(offset as usize, j);
+                        let offset = i as u32;
+                        assert_eq!(offset as usize, i);
                         lhs.clone().elem(ProjectionElem::ConstantIndex {
                             offset,
                             // FIXME(eddyb) `min_length` doesn't appear to be used.
@@ -107,21 +96,15 @@ impl MirPass for Deaggregator {
                         })
                     } else {
                         let ty = op.ty(local_decls, tcx);
-                        let field = Field::new(active_field_index.unwrap_or(j));
+                        let field = Field::new(active_field_index.unwrap_or(i));
                         lhs.clone().field(field, ty)
                     };
-                    bb.statements.push(Statement {
+                    Statement {
                         source_info,
                         kind: StatementKind::Assign(lhs_field, Rvalue::Use(op)),
-                    });
-                }
-
-                // If the aggregate was an enum, we need to set the discriminant.
-                bb.statements.extend(set_discriminant);
-
-                start = bb.statements.len();
-                bb.statements.extend(suffix_stmts);
-            }
+                    }
+                }).chain(set_discriminant))
+            });
         }
     }
 }
diff --git a/src/test/mir-opt/copy_propagation_arg.rs b/src/test/mir-opt/copy_propagation_arg.rs
index 9e4c06b7d7b..feadec6bbf7 100644
--- a/src/test/mir-opt/copy_propagation_arg.rs
+++ b/src/test/mir-opt/copy_propagation_arg.rs
@@ -78,6 +78,7 @@ fn main() {
 // bb1: {
 //     StorageDead(_3);
 //     _1 = const 5u8;
+//     ...
 //     return;
 // }
 // END rustc.bar.CopyPropagation.before.mir
@@ -91,6 +92,7 @@ fn main() {
 //     ...
 //     _1 = const 5u8;
 //     ...
+//     return;
 // }
 // END rustc.bar.CopyPropagation.after.mir
 // START rustc.baz.CopyPropagation.before.mir
@@ -99,6 +101,7 @@ fn main() {
 //     _2 = _1;
 //     _1 = move _2;
 //     StorageDead(_2);
+//     ...
 //     return;
 // }
 // END rustc.baz.CopyPropagation.before.mir
@@ -108,21 +111,22 @@ fn main() {
 //     _2 = _1;
 //     _1 = move _2;
 //     ...
+//     return;
 // }
 // END rustc.baz.CopyPropagation.after.mir
 // START rustc.arg_src.CopyPropagation.before.mir
 // bb0: {
-//       ...
-//       _3 = _1;
-//       _2 = move _3;
-//       ...
-//       _1 = const 123i32;
-//       ...
-//       _4 = _2;
-//       _0 = move _4;
-//       ...
-//       return;
-//   }
+//      ...
+//      _3 = _1;
+//      _2 = move _3;
+//      ...
+//      _1 = const 123i32;
+//      ...
+//      _4 = _2;
+//      _0 = move _4;
+//      ...
+//      return;
+//  }
 // END rustc.arg_src.CopyPropagation.before.mir
 // START rustc.arg_src.CopyPropagation.after.mir
 // bb0: {