Auto merge of #48300 - eddyb:split-aggregates, r=<try>

rustc_mir: add a pass for splitting locals into their fields (aka SROA).

**DO NOT MERGE**: based on #48052.

Author: bors

Diff for: 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 for: 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 for: src/librustc_mir/transform/deaggregator.rs

@@ -21,116 +21,90 @@ impl MirPass for Deaggregator {
                           tcx: TyCtxt<'a, 'tcx, 'tcx>,
                           source: MirSource,
                           mir: &mut Mir<'tcx>) {
-        let node_path = tcx.item_path_str(source.def_id);
-        debug!("running on: {:?}", node_path);
-        // we only run when mir_opt_level > 2
-        if tcx.sess.opts.debugging_opts.mir_opt_level <= 2 {
-            return;
-        }
-
         // Don't run on constant MIR, because trans might not be able to
         // evaluate the modified MIR.
         // FIXME(eddyb) Remove check after miri is merged.
         let id = tcx.hir.as_local_node_id(source.def_id).unwrap();
         match (tcx.hir.body_owner_kind(id), source.promoted) {
-            (hir::BodyOwnerKind::Fn, None) => {},
-            _ => return
-        }
-        // In fact, we might not want to trigger in other cases.
-        // Ex: when we could use SROA.  See issue #35259
+            (_, Some(_)) |
+            (hir::BodyOwnerKind::Const, _) |
+            (hir::BodyOwnerKind::Static(_), _) => return,
 
-        for bb in mir.basic_blocks_mut() {
-            let mut curr: usize = 0;
-            while let Some(idx) = get_aggregate_statement_index(curr, &bb.statements) {
-                // do the replacement
-                debug!("removing statement {:?}", idx);
-                let src_info = bb.statements[idx].source_info;
-                let suffix_stmts = bb.statements.split_off(idx+1);
-                let orig_stmt = bb.statements.pop().unwrap();
-                let (lhs, rhs) = match orig_stmt.kind {
-                    StatementKind::Assign(ref lhs, ref rhs) => (lhs, rhs),
-                    _ => span_bug!(src_info.span, "expected assign, not {:?}", orig_stmt),
-                };
-                let (agg_kind, operands) = match rhs {
-                    &Rvalue::Aggregate(ref agg_kind, ref operands) => (agg_kind, operands),
-                    _ => span_bug!(src_info.span, "expected aggregate, not {:?}", rhs),
-                };
-                let (adt_def, variant, substs) = match **agg_kind {
-                    AggregateKind::Adt(adt_def, variant, substs, None)
-                        => (adt_def, variant, substs),
-                    _ => span_bug!(src_info.span, "expected struct, not {:?}", rhs),
-                };
-                let n = bb.statements.len();
-                bb.statements.reserve(n + operands.len() + suffix_stmts.len());
-                for (i, op) in operands.iter().enumerate() {
-                    let ref variant_def = adt_def.variants[variant];
-                    let ty = variant_def.fields[i].ty(tcx, substs);
-                    let rhs = Rvalue::Use(op.clone());
+            (hir::BodyOwnerKind::Fn, _) => {
+                if tcx.is_const_fn(source.def_id) {
+                    // Don't run on const functions, as, again, trans might not be able to evaluate
+                    // the optimized IR.
+                    return
+                }
+            }
+        }
 
-                    let lhs_cast = if adt_def.is_enum() {
-                        Place::Projection(Box::new(PlaceProjection {
-                            base: lhs.clone(),
-                            elem: ProjectionElem::Downcast(adt_def, variant),
-                        }))
+        let (basic_blocks, local_decls) = mir.basic_blocks_and_local_decls_mut();
+        let local_decls = &*local_decls;
+        for bb in basic_blocks {
+            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 {
-                        lhs.clone()
-                    };
-
-                    let lhs_proj = Place::Projection(Box::new(PlaceProjection {
-                        base: lhs_cast,
-                        elem: ProjectionElem::Field(Field::new(i), ty),
-                    }));
-                    let new_statement = Statement {
-                        source_info: src_info,
-                        kind: StatementKind::Assign(lhs_proj, rhs),
-                    };
-                    debug!("inserting: {:?} @ {:?}", new_statement, idx + i);
-                    bb.statements.push(new_statement);
+                        return None;
+                    }
+                } else {
+                    return None;
                 }
 
-                // if the aggregate was an enum, we need to set the discriminant
-                if adt_def.is_enum() {
-                    let set_discriminant = Statement {
-                        kind: StatementKind::SetDiscriminant {
-                            place: lhs.clone(),
-                            variant_index: variant,
-                        },
-                        source_info: src_info,
-                    };
-                    bb.statements.push(set_discriminant);
+                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!()
                 };
 
-                curr = bb.statements.len();
-                bb.statements.extend(suffix_stmts);
-            }
-        }
-    }
-}
+                let mut set_discriminant = None;
+                let active_field_index = match *kind {
+                    AggregateKind::Adt(adt_def, variant_index, _, active_field_index) => {
+                        if adt_def.is_enum() {
+                            set_discriminant = Some(Statement {
+                                kind: StatementKind::SetDiscriminant {
+                                    place: lhs.clone(),
+                                    variant_index,
+                                },
+                                source_info,
+                            });
+                            lhs = lhs.downcast(adt_def, variant_index);
+                        }
+                        active_field_index
+                    }
+                    _ => None
+                };
 
-fn get_aggregate_statement_index<'a, 'tcx, 'b>(start: usize,
-                                         statements: &Vec<Statement<'tcx>>)
-                                         -> Option<usize> {
-    for i in start..statements.len() {
-        let ref statement = statements[i];
-        let rhs = match statement.kind {
-            StatementKind::Assign(_, ref rhs) => rhs,
-            _ => continue,
-        };
-        let (kind, operands) = match rhs {
-            &Rvalue::Aggregate(ref kind, ref operands) => (kind, operands),
-            _ => continue,
-        };
-        let (adt_def, variant) = match **kind {
-            AggregateKind::Adt(adt_def, variant, _, None) => (adt_def, variant),
-            _ => continue,
-        };
-        if operands.len() == 0 {
-            // don't deaggregate ()
-            continue;
+                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 = 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.
+                            min_length: offset + 1,
+                            from_end: false
+                        })
+                    } else {
+                        let ty = op.ty(local_decls, tcx);
+                        let field = Field::new(active_field_index.unwrap_or(i));
+                        lhs.clone().field(field, ty)
+                    };
+                    Statement {
+                        source_info,
+                        kind: StatementKind::Assign(lhs_field, Rvalue::Use(op)),
+                    }
+                }).chain(set_discriminant))
+            });
         }
-        debug!("getting variant {:?}", variant);
-        debug!("for adt_def {:?}", adt_def);
-        return Some(i);
-    };
-    None
+    }
 }

Diff for: src/librustc_mir/transform/mod.rs

@@ -45,6 +45,7 @@ pub mod generator;
 pub mod inline;
 pub mod lower_128bit;
 pub mod uniform_array_move_out;
+pub mod split_local_fields;
 
 pub(crate) fn provide(providers: &mut Providers) {
     self::qualify_consts::provide(providers);
@@ -258,8 +259,14 @@ fn optimized_mir<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, def_id: DefId) -> &'tcx
 
         // Optimizations begin.
         inline::Inline,
+
+        // Lowering generator control-flow and variables
+        // has to happen before we do anything else to them.
+        generator::StateTransform,
+
         instcombine::InstCombine,
         deaggregator::Deaggregator,
+        split_local_fields::SplitLocalFields,
         copy_prop::CopyPropagation,
         remove_noop_landing_pads::RemoveNoopLandingPads,
         simplify::SimplifyCfg::new("final"),

Diff for: src/librustc_mir/transform/simplify.rs

@@ -42,6 +42,7 @@ use rustc_data_structures::indexed_vec::{Idx, IndexVec};
 use rustc::ty::TyCtxt;
 use rustc::mir::*;
 use rustc::mir::visit::{MutVisitor, Visitor, PlaceContext};
+use rustc::session::config::FullDebugInfo;
 use std::borrow::Cow;
 use transform::{MirPass, MirSource};
 
@@ -281,16 +282,24 @@ pub struct SimplifyLocals;
 
 impl MirPass for SimplifyLocals {
     fn run_pass<'a, 'tcx>(&self,
-                          _: TyCtxt<'a, 'tcx, 'tcx>,
+                          tcx: TyCtxt<'a, 'tcx, 'tcx>,
                           _: MirSource,
                           mir: &mut Mir<'tcx>) {
         let mut marker = DeclMarker { locals: BitVector::new(mir.local_decls.len()) };
         marker.visit_mir(mir);
         // Return pointer and arguments are always live
-        marker.locals.insert(0);
-        for idx in mir.args_iter() {
-            marker.locals.insert(idx.index());
+        marker.locals.insert(RETURN_PLACE.index());
+        for arg in mir.args_iter() {
+            marker.locals.insert(arg.index());
         }
+
+        // We may need to keep dead user variables live for debuginfo.
+        if tcx.sess.opts.debuginfo == FullDebugInfo {
+            for local in mir.vars_iter() {
+                marker.locals.insert(local.index());
+            }
+        }
+
         let map = make_local_map(&mut mir.local_decls, marker.locals);
         // Update references to all vars and tmps now
         LocalUpdater { map: map }.visit_mir(mir);