From 63fdd46f9af41398fd2191135336e63d118114c5 Mon Sep 17 00:00:00 2001
From: Scott Olson <scott@solson.me>
Date: Mon, 28 Mar 2016 21:08:08 -0600
Subject: [PATCH] Handle custom discriminant values and detect invalid
 discriminants.

---
 src/error.rs       |  3 +++
 src/interpreter.rs | 34 ++++++++++++++++++++++++----------
 test/c_enums.rs    | 21 +++++++++++++++++++++
 3 files changed, 48 insertions(+), 10 deletions(-)
 create mode 100755 test/c_enums.rs

diff --git a/src/error.rs b/src/error.rs
index 08ef5a46ff2..65dd187fa0b 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -5,6 +5,7 @@
 pub enum EvalError {
     DanglingPointerDeref,
     InvalidBool,
+    InvalidDiscriminant,
     PointerOutOfBounds,
     ReadPointerAsBytes,
     ReadBytesAsPointer,
@@ -21,6 +22,8 @@ fn description(&self) -> &str {
                 "dangling pointer was dereferenced",
             EvalError::InvalidBool =>
                 "invalid boolean value read",
+            EvalError::InvalidDiscriminant =>
+                "invalid enum discriminant value read",
             EvalError::PointerOutOfBounds =>
                 "pointer offset outside bounds of allocation",
             EvalError::ReadPointerAsBytes =>
diff --git a/src/interpreter.rs b/src/interpreter.rs
index 969a0fad4d0..f2e7702717d 100644
--- a/src/interpreter.rs
+++ b/src/interpreter.rs
@@ -228,7 +228,7 @@ fn eval_terminator(&mut self, terminator: &mir::Terminator<'tcx>)
                 TerminatorTarget::Block(target_block)
             }
 
-            Switch { ref discr, ref targets, .. } => {
+            Switch { ref discr, ref targets, adt_def } => {
                 let adt_ptr = try!(self.eval_lvalue(discr)).to_ptr();
                 let adt_repr = self.lvalue_repr(discr);
                 let discr_size = match *adt_repr {
@@ -236,7 +236,14 @@ fn eval_terminator(&mut self, terminator: &mir::Terminator<'tcx>)
                     _ => panic!("attmpted to switch on non-aggregate type"),
                 };
                 let discr_val = try!(self.memory.read_uint(adt_ptr, discr_size));
-                TerminatorTarget::Block(targets[discr_val as usize])
+
+                let matching = adt_def.variants.iter()
+                    .position(|v| discr_val == v.disr_val.to_u64_unchecked());
+
+                match matching {
+                    Some(i) => TerminatorTarget::Block(targets[i]),
+                    None => return Err(EvalError::InvalidDiscriminant),
+                }
             }
 
             Call { ref func, ref args, ref destination, .. } => {
@@ -481,13 +488,18 @@ fn call_c_abi(&mut self, def_id: DefId, args: &[mir::Operand<'tcx>], dest: Point
         Ok(TerminatorTarget::Call)
     }
 
-    fn assign_to_aggregate(&mut self, dest: Pointer, dest_repr: &Repr, variant: usize,
-                         operands: &[mir::Operand<'tcx>]) -> EvalResult<()> {
+    fn assign_to_aggregate(
+        &mut self,
+        dest: Pointer,
+        dest_repr: &Repr,
+        variant: usize,
+        discr: Option<u64>,
+        operands: &[mir::Operand<'tcx>],
+    ) -> EvalResult<()> {
         match *dest_repr {
             Repr::Aggregate { discr_size, ref variants, .. } => {
                 if discr_size > 0 {
-                    let discr = variant as u64;
-                    try!(self.memory.write_uint(dest, discr, discr_size));
+                    try!(self.memory.write_uint(dest, discr.unwrap(), discr_size));
                 }
                 let after_discr = dest.offset(discr_size as isize);
                 for (field, operand) in variants[variant].iter().zip(operands) {
@@ -538,10 +550,12 @@ fn eval_assignment(&mut self, lvalue: &mir::Lvalue<'tcx>, rvalue: &mir::Rvalue<'
                 use rustc::mir::repr::AggregateKind::*;
                 match *kind {
                     Tuple | Closure(..) =>
-                        try!(self.assign_to_aggregate(dest, &dest_repr, 0, operands)),
+                        try!(self.assign_to_aggregate(dest, &dest_repr, 0, None, operands)),
 
-                    Adt(_, variant_idx, _) =>
-                        try!(self.assign_to_aggregate(dest, &dest_repr, variant_idx, operands)),
+                    Adt(adt_def, variant, _) => {
+                        let discr = Some(adt_def.variants[variant].disr_val.to_u64_unchecked());
+                        try!(self.assign_to_aggregate(dest, &dest_repr, variant, discr, operands));
+                    }
 
                     Vec => if let Repr::Array { elem_size, length } = *dest_repr {
                         assert_eq!(length, operands.len());
@@ -668,7 +682,7 @@ fn lvalue_repr(&self, lvalue: &mir::Lvalue<'tcx>) -> &'arena Repr {
         use rustc::mir::tcx::LvalueTy;
         match self.mir().lvalue_ty(self.tcx, lvalue) {
             LvalueTy::Ty { ty } => self.ty_to_repr(ty),
-            LvalueTy::Downcast { ref adt_def, substs, variant_index } => {
+            LvalueTy::Downcast { adt_def, substs, variant_index } => {
                 let field_tys = adt_def.variants[variant_index].fields.iter()
                     .map(|f| f.ty(self.tcx, substs));
                 self.repr_arena.alloc(self.make_aggregate_repr(iter::once(field_tys)))
diff --git a/test/c_enums.rs b/test/c_enums.rs
new file mode 100755
index 00000000000..190e82ac01c
--- /dev/null
+++ b/test/c_enums.rs
@@ -0,0 +1,21 @@
+#![feature(custom_attribute)]
+#![allow(dead_code, unused_attributes)]
+
+enum Foo {
+    Bar = 42,
+    Baz,
+    Quux = 100,
+}
+
+#[miri_run]
+fn foo() -> [u8; 3] {
+    [Foo::Bar as u8, Foo::Baz as u8, Foo::Quux as u8]
+}
+
+#[miri_run]
+fn unsafe_match() -> bool {
+    match unsafe { std::mem::transmute::<u8, Foo>(43) } {
+        Foo::Baz => true,
+        _ => false,
+    }
+}
-- 
2.44.0