add SetDiscriminant StatementKind to enable deaggregation of enums

3 files changed, 73 insertions, 17 deletions

diff --git a/src/librustc_mir/transform/deaggregator.rs b/src/librustc_mir/transform/deaggregator.rs
index fccd4a607fd..cd6f0ed9cba 100644
--- a/src/librustc_mir/transform/deaggregator.rs
+++ b/src/librustc_mir/transform/deaggregator.rs
@@ -39,7 +39,7 @@ impl<'tcx> MirPass<'tcx> for Deaggregator {
 
         let mut curr: usize = 0;
         for bb in mir.basic_blocks_mut() {
-            let idx = match get_aggregate_statement(curr, &bb.statements) {
+            let idx = match get_aggregate_statement_index(curr, &bb.statements) {
                 Some(idx) => idx,
                 None => continue,
             };
@@ -48,7 +48,11 @@ impl<'tcx> MirPass<'tcx> for Deaggregator {
             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 StatementKind::Assign(ref lhs, ref rhs) = orig_stmt.kind;
+            let (lhs, rhs) = match orig_stmt.kind {
+                StatementKind::Assign(ref lhs, ref rhs) => (lhs, rhs),
+                StatementKind::SetDiscriminant{ .. } =>
+                    span_bug!(src_info.span, "expected aggregate, not {:?}", orig_stmt.kind),
+            };
             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),
@@ -64,10 +68,14 @@ impl<'tcx> MirPass<'tcx> for Deaggregator {
                 let ty = variant_def.fields[i].ty(tcx, substs);
                 let rhs = Rvalue::Use(op.clone());
 
-                // since we don't handle enums, we don't need a cast
-                let lhs_cast = lhs.clone();
-
-                // FIXME we cannot deaggregate enums issue: #35186
+                let lhs_cast = if adt_def.variants.len() > 1 {
+                    Lvalue::Projection(Box::new(LvalueProjection {
+                        base: lhs.clone(),
+                        elem: ProjectionElem::Downcast(adt_def, variant),
+                    }))
+                } else {
+                    lhs.clone()
+                };
 
                 let lhs_proj = Lvalue::Projection(Box::new(LvalueProjection {
                     base: lhs_cast,
@@ -80,18 +88,34 @@ impl<'tcx> MirPass<'tcx> for Deaggregator {
                 debug!("inserting: {:?} @ {:?}", new_statement, idx + i);
                 bb.statements.push(new_statement);
             }
+
+            // if the aggregate was an enum, we need to set the discriminant
+            if adt_def.variants.len() > 1 {
+                let set_discriminant = Statement {
+                    kind: StatementKind::SetDiscriminant {
+                        lvalue: lhs.clone(),
+                        variant_index: variant,
+                    },
+                    source_info: src_info,
+                };
+                bb.statements.push(set_discriminant);
+            };
+
             curr = bb.statements.len();
             bb.statements.extend(suffix_stmts);
         }
     }
 }
 
-fn get_aggregate_statement<'a, 'tcx, 'b>(curr: usize,
+fn get_aggregate_statement_index<'a, 'tcx, 'b>(start: usize,
                                          statements: &Vec<Statement<'tcx>>)
                                          -> Option<usize> {
-    for i in curr..statements.len() {
+    for i in start..statements.len() {
         let ref statement = statements[i];
-        let StatementKind::Assign(_, ref rhs) = statement.kind;
+        let rhs = match statement.kind {
+            StatementKind::Assign(_, ref rhs) => rhs,
+            StatementKind::SetDiscriminant{ .. } => continue,
+        };
         let (kind, operands) = match rhs {
             &Rvalue::Aggregate(ref kind, ref operands) => (kind, operands),
             _ => continue,
@@ -100,9 +124,8 @@ fn get_aggregate_statement<'a, 'tcx, 'b>(curr: usize,
             &AggregateKind::Adt(adt_def, variant, _) => (adt_def, variant),
             _ => continue,
         };
-        if operands.len() == 0 || adt_def.variants.len() > 1 {
+        if operands.len() == 0 {
             // don't deaggregate ()
-            // don't deaggregate enums ... for now
             continue;
         }
         debug!("getting variant {:?}", variant);
diff --git a/src/librustc_mir/transform/promote_consts.rs b/src/librustc_mir/transform/promote_consts.rs
index fa3490cbcf3..eb0d8697f15 100644
--- a/src/librustc_mir/transform/promote_consts.rs
+++ b/src/librustc_mir/transform/promote_consts.rs
@@ -219,7 +219,13 @@ impl<'a, 'tcx> Promoter<'a, 'tcx> {
         let (mut rvalue, mut call) = (None, None);
         let source_info = if stmt_idx < no_stmts {
             let statement = &mut self.source[bb].statements[stmt_idx];
-            let StatementKind::Assign(_, ref mut rhs) = statement.kind;
+            let mut rhs = match statement.kind {
+                StatementKind::Assign(_, ref mut rhs) => rhs,
+                StatementKind::SetDiscriminant{ .. } =>
+                    span_bug!(statement.source_info.span,
+                              "cannot promote SetDiscriminant {:?}",
+                              statement),
+            };
             if self.keep_original {
                 rvalue = Some(rhs.clone());
             } else {
@@ -300,10 +306,16 @@ impl<'a, 'tcx> Promoter<'a, 'tcx> {
         });
         let mut rvalue = match candidate {
             Candidate::Ref(Location { block: bb, statement_index: stmt_idx }) => {
-                match self.source[bb].statements[stmt_idx].kind {
+                let ref mut statement = self.source[bb].statements[stmt_idx];
+                match statement.kind {
                     StatementKind::Assign(_, ref mut rvalue) => {
                         mem::replace(rvalue, Rvalue::Use(new_operand))
                     }
+                    StatementKind::SetDiscriminant{ .. } => {
+                        span_bug!(statement.source_info.span,
+                                  "cannot promote SetDiscriminant {:?}",
+                                  statement);
+                    }
                 }
             }
             Candidate::ShuffleIndices(bb) => {
@@ -340,7 +352,11 @@ pub fn promote_candidates<'a, 'tcx>(mir: &mut Mir<'tcx>,
         let (span, ty) = match candidate {
             Candidate::Ref(Location { block: bb, statement_index: stmt_idx }) => {
                 let statement = &mir[bb].statements[stmt_idx];
-                let StatementKind::Assign(ref dest, _) = statement.kind;
+                let dest = match statement.kind {
+                    StatementKind::Assign(ref dest, _) => dest,
+                    StatementKind::SetDiscriminant{ .. } =>
+                        panic!("cannot promote SetDiscriminant"),
+                };
                 if let Lvalue::Temp(index) = *dest {
                     if temps[index] == TempState::PromotedOut {
                         // Already promoted.
diff --git a/src/librustc_mir/transform/type_check.rs b/src/librustc_mir/transform/type_check.rs
index c8556d5e0d8..026d197cb7f 100644
--- a/src/librustc_mir/transform/type_check.rs
+++ b/src/librustc_mir/transform/type_check.rs
@@ -14,7 +14,7 @@
 use rustc::infer::{self, InferCtxt, InferOk};
 use rustc::traits::{self, ProjectionMode};
 use rustc::ty::fold::TypeFoldable;
-use rustc::ty::{self, Ty, TyCtxt};
+use rustc::ty::{self, Ty, TyCtxt, TypeVariants};
 use rustc::mir::repr::*;
 use rustc::mir::tcx::LvalueTy;
 use rustc::mir::transform::{MirPass, MirSource, Pass};
@@ -360,10 +360,27 @@ impl<'a, 'gcx, 'tcx> TypeChecker<'a, 'gcx, 'tcx> {
                         span_mirbug!(self, stmt, "bad assignment ({:?} = {:?}): {:?}",
                                      lv_ty, rv_ty, terr);
                     }
-                }
-
                 // FIXME: rvalue with undeterminable type - e.g. inline
                 // asm.
+                }
+            }
+            StatementKind::SetDiscriminant{ ref lvalue, variant_index } => {
+                let lvalue_type = lvalue.ty(mir, tcx).to_ty(tcx);
+                let adt = match lvalue_type.sty {
+                    TypeVariants::TyEnum(adt, _) => adt,
+                    _ => {
+                        span_bug!(stmt.source_info.span,
+                                  "bad set discriminant ({:?} = {:?}): lhs is not an enum",
+                                  lvalue,
+                                  variant_index);
+                    }
+                };
+                if variant_index >= adt.variants.len() {
+                     span_bug!(stmt.source_info.span,
+                               "bad set discriminant ({:?} = {:?}): value of of range",
+                               lvalue,
+                               variant_index);
+                };
             }
         }
     }