Improved type inference for let statements and assignments with type annotated patterns and lexps

Added get_enum to type checker interface

4 files changed, 37 insertions, 3 deletions

diff --git a/src/parser.mly b/src/parser.mly
index 9f48067f..8e61a0ac 100644
--- a/src/parser.mly
+++ b/src/parser.mly
@@ -209,6 +209,10 @@ id:
     { idl (DeIid($3)) }
   | Lparen Deinfix Lt Rparen
     { idl (DeIid($3)) }
+  | Lparen Deinfix GtUnderS Rparen
+    { idl (DeIid($3)) }
+  | Lparen Deinfix LtUnderS Rparen
+    { idl (DeIid($3)) }
   | Lparen Deinfix Minus Rparen
     { idl (DeIid("-")) }
   | Lparen Deinfix MinusUnderS Rparen
@@ -243,6 +247,8 @@ id:
     { idl (DeIid($3)) }
   | Lparen Deinfix GtEq Rparen
     { idl (DeIid($3)) }
+  | Lparen Deinfix GtEqUnderS Rparen
+    { idl (DeIid($3)) }
   | Lparen Deinfix GtEqPlus Rparen
     { idl (DeIid($3)) }
   | Lparen Deinfix GtGt Rparen
@@ -257,6 +263,8 @@ id:
     { idl (DeIid($3)) }
   | Lparen Deinfix LtEq Rparen
     { idl (DeIid($3)) }
+  | Lparen Deinfix LtEqUnderS Rparen
+    { idl (DeIid($3)) }
   | Lparen Deinfix LtLt Rparen
     { idl (DeIid($3)) }
   | Lparen Deinfix LtLtLt Rparen
diff --git a/src/pretty_print_sail.ml b/src/pretty_print_sail.ml
index 9c33c841..6826087a 100644
--- a/src/pretty_print_sail.ml
+++ b/src/pretty_print_sail.ml
@@ -488,6 +488,7 @@ let doc_rec (Rec_aux(r,_)) = match r with
 
 let doc_tannot_opt (Typ_annot_opt_aux(t,_)) = match t with
   | Typ_annot_opt_some(tq,typ) -> doc_typquant tq (doc_typ typ)
+  | Typ_annot_opt_none -> empty
 
 let doc_effects_opt (Effect_opt_aux(e,_)) = match e with
   | Effect_opt_pure -> string "pure"
diff --git a/src/type_check_new.ml b/src/type_check_new.ml
index 8232fd6a..5be2cd43 100644
--- a/src/type_check_new.ml
+++ b/src/type_check_new.ml
@@ -346,6 +346,7 @@ module Env : sig
   val set_default_order_inc : t -> t
   val set_default_order_dec : t -> t
   val add_enum : id -> id list -> t -> t
+  val get_enum : id -> t -> id list
   val get_casts : t -> id list
   val allow_casts : t -> bool
   val no_casts : t -> t
@@ -518,6 +519,11 @@ end = struct
         { env with enums = Bindings.add id (IdSet.of_list ids) env.enums }
       end
 
+  let get_enum id env =
+    try IdSet.elements (Bindings.find id env.enums)
+    with
+    | Not_found -> typ_error (id_loc id) ("Enumeration " ^ string_of_id id ^ " does not exist")
+
   let is_record id env = Bindings.mem id env.records
 
   let add_record id typq fields env =
@@ -1533,19 +1539,24 @@ let rec check_exp env (E_aux (exp_aux, (l, ())) as exp : unit exp) (Typ_aux (typ
      begin
        match letbind with
        | LB_val_explicit (typschm, pat, bind) -> assert false
+       | LB_val_implicit (P_aux (P_typ (ptyp, _), _) as pat, bind) ->
+          let checked_bind = crule check_exp env bind ptyp in
+          let tpat, env = bind_pat env pat (typ_of checked_bind) in
+          annot_exp (E_let (LB_aux (LB_val_implicit (tpat, checked_bind), (let_loc, None)), crule check_exp env exp typ)) typ
        | LB_val_implicit (pat, bind) ->
           let inferred_bind = irule infer_exp env bind in
           let tpat, env = bind_pat env pat (typ_of inferred_bind) in
           annot_exp (E_let (LB_aux (LB_val_implicit (tpat, inferred_bind), (let_loc, None)), crule check_exp env exp typ)) typ
      end
-  | E_app_infix (x, op, y), _ when List.length (Env.get_overloads (deinfix op) env) > 0 -> check_exp env (E_aux (E_app (deinfix op, [x; y]), (l, ()))) typ
+  | E_app_infix (x, op, y), _ when List.length (Env.get_overloads (deinfix op) env) > 0 ->
+     check_exp env (E_aux (E_app (deinfix op, [x; y]), (l, ()))) typ
   | E_app (f, xs), _ when List.length (Env.get_overloads f env) > 0 ->
      let rec try_overload = function
        | [] -> typ_error l ("No valid overloading for " ^ string_of_exp exp)
        | (f :: fs) -> begin
            typ_print ("Overload: " ^ string_of_id f ^ "(" ^ string_of_list ", " string_of_exp xs ^ ")");
            try crule check_exp env (E_aux (E_app (f, xs), (l, ()))) typ with
-           | Type_error (_, m2) -> try_overload fs
+           | Type_error (_, m) -> typ_print ("Error : " ^ m); try_overload fs
          end
      in
      try_overload (Env.get_overloads f env)
@@ -1789,6 +1800,10 @@ and bind_assignment env (LEXP_aux (lexp_aux, _) as lexp) (E_aux (_, (l, ())) as
      end
   | LEXP_memory (f, xs) ->
      check_exp env (E_aux (E_app (f, xs @ [exp]), (l, ()))) unit_typ, env
+  | LEXP_cast (typ_annot, v) ->
+     let checked_exp = crule check_exp env exp typ_annot in
+     let tlexp, env' = bind_lexp env lexp (typ_of checked_exp) in
+     annot_assign tlexp checked_exp, env'
   | _ ->
      let inferred_exp = irule infer_exp env exp in
      let tlexp, env' = bind_lexp env lexp (typ_of inferred_exp) in
@@ -1963,7 +1978,7 @@ and infer_exp env (E_aux (exp_aux, (l, ())) as exp) =
        | (f :: fs) -> begin
            typ_print ("Overload: " ^ string_of_id f ^ "(" ^ string_of_list ", " string_of_exp xs ^ ")");
            try irule infer_exp env (E_aux (E_app (f, xs), (l, ()))) with
-           | Type_error (_, m2) -> try_overload fs
+           | Type_error (_, m) -> typ_print ("Error: " ^ m); try_overload fs
          end
      in
      try_overload (Env.get_overloads f env)
@@ -2043,6 +2058,12 @@ and infer_funapp' l env f (typq, f_typ) xs ret_ctx_typ =
          else typ_error l ("Quantifiers " ^ string_of_list ", " string_of_quant_item quants
                            ^ " not resolved during application of " ^ string_of_id f)
        end
+    | (utyps, (typ :: typs)), (uargs, ((n, arg) :: args)) when KidSet.is_empty (typ_frees typ) ->
+       begin
+         let carg = crule check_exp env arg typ in
+         let (iargs, ret_typ') = instantiate quants (utyps, typs) ret_typ (uargs, args) in
+         ((n, carg) :: iargs, ret_typ')
+       end
     | (utyps, (typ :: typs)), (uargs, ((n, arg) :: args)) ->
        begin
          typ_debug ("INSTANTIATE: " ^ string_of_exp arg ^ " with " ^ string_of_typ typ ^ " NF " ^ string_of_tnf (normalize_typ env typ));
diff --git a/src/type_check_new.mli b/src/type_check_new.mli
index f55ccf3a..cdacf523 100644
--- a/src/type_check_new.mli
+++ b/src/type_check_new.mli
@@ -64,6 +64,10 @@ module Env : sig
 
   val get_regtyp : id -> t -> int * int * (index_range * id) list
 
+  (* Return all the identifiers in an enumeration. Throws a type error
+     if the enumeration doesn't exist. *)
+  val get_enum : id -> t -> id list
+
   (* Returns true if id is a register type, false otherwise *)
   val is_regtyp : id -> t -> bool