Fix instruction extractor

5 files changed, 19 insertions, 2 deletions

diff --git a/src/lem_interp/instruction_extractor.lem b/src/lem_interp/instruction_extractor.lem
index f78e2d29..3507a5ee 100644
--- a/src/lem_interp/instruction_extractor.lem
+++ b/src/lem_interp/instruction_extractor.lem
@@ -30,6 +30,8 @@ let rec extract_ityp t tag = match (t,tag) with
   | (T_app "range" _,_) -> IRange Nothing
   | (T_app i (T_args []),Tag_enum max) ->
     IEnum i (natFromInteger max)
+  | (T_id i,Tag_enum max) ->
+    IEnum i (natFromInteger max)
   | _ -> IOther
 end 
 
diff --git a/src/lem_interp/interp_inter_imp.lem b/src/lem_interp/interp_inter_imp.lem
index aa802641..66bfd6b7 100644
--- a/src/lem_interp/interp_inter_imp.lem
+++ b/src/lem_interp/interp_inter_imp.lem
@@ -396,6 +396,7 @@ let instruction_to_istate (top_level:context) (((name, parms, _) as instr):instr
 	match typ with
 	  | Bit -> match bits with | [b] -> b | _ -> Assert_extra.failwith "Expected a bitvector of length 1" end
 	  | Range _ -> Interp_lib.to_num Interp_lib.Unsigned vec
+	  | Enum _ _ -> Interp_lib.to_num Interp_lib.Unsigned vec
 	  | _ -> vec
        end 
     end in
diff --git a/src/type_check.ml b/src/type_check.ml
index a4466acb..105748a3 100644
--- a/src/type_check.ml
+++ b/src/type_check.ml
@@ -248,8 +248,11 @@ let rec check_pattern envs emp_tag expect_t (P_aux(p,(l,annot))) : ((tannot pat)
     | P_id id -> 
       let i = id_to_string id in
       let default_bounds = extract_bounds d_env i expect_t in
-      let default = let tannot = Base(([],expect_t),emp_tag,cs,pure_e,default_bounds) in
-		    (P_aux(p,(l,tannot)),Envmap.from_list [(i,tannot)],cs,default_bounds,expect_t) in
+      let default = let id_annot = Base(([],expect_t),emp_tag,cs,pure_e,default_bounds) in
+		    let pat_annot = match is_enum_typ d_env expect_t with
+		      | None -> id_annot
+		      | Some n -> Base(([],expect_t), Enum n, cs,pure_e,default_bounds) in
+		    (P_aux(p,(l,pat_annot)),Envmap.from_list [(i,id_annot)],cs,default_bounds,expect_t) in
       (match Envmap.apply t_env i with
 	| Some(Base((params,t),Constructor,cs,ef,bounds)) ->
 	  let t,cs,ef,_ = subst params false t cs ef in
diff --git a/src/type_internal.ml b/src/type_internal.ml
index a760c9d4..fb5e3489 100644
--- a/src/type_internal.ml
+++ b/src/type_internal.ml
@@ -1975,6 +1975,15 @@ let rec get_abbrev d_env t =
 	| _ -> t,[])
     | _ -> t,[]
 
+let is_enum_typ d_env t = 
+  let t,_ = get_abbrev d_env t in
+  let t_actual = match t.t with | Tabbrev(_,ta) -> ta | _ -> t in
+  match t.t with
+    | Tid i -> (match Envmap.apply d_env.enum_env i with
+	| Some(ns) -> Some(List.length ns)
+	| _ -> None)
+    | _ -> None
+
 let eq_error l msg = raise (Reporting_basic.err_typ l msg)
 
 let compare_effect (BE_aux(e1,_)) (BE_aux(e2,_)) =
diff --git a/src/type_internal.mli b/src/type_internal.mli
index fac69bb1..9e8b27c8 100644
--- a/src/type_internal.mli
+++ b/src/type_internal.mli
@@ -192,6 +192,8 @@ val type_param_consistent : Parse_ast.l -> t_arg emap -> t_arg emap -> nexp_rang
 
 val get_abbrev : def_envs -> t -> (t * nexp_range list)
 
+val is_enum_typ : def_envs -> t -> int option
+
 val extract_bounds : def_envs -> string -> t -> bounds_env
 val merge_bounds : bounds_env -> bounds_env -> bounds_env
 val find_var_from_nexp : nexp -> bounds_env -> (string option * string) option