1 files changed, 36 insertions, 3 deletions
diff --git a/src/constant_propagation.ml b/src/constant_propagation.ml
index 608d25e1..89a4ba82 100644
--- a/src/constant_propagation.ml
+++ b/src/constant_propagation.ml
@@ -321,6 +321,25 @@ let const_props defs ref_vars =
     with
     | _ -> exp
   in
+  let constants =
+    let add m = function
+      | DEF_val (LB_aux (LB_val (P_aux ((P_id id | P_typ (_,P_aux (P_id id,_))),_), exp),_))
+           when Constant_fold.is_constant exp ->
+         Bindings.add id exp m
+      | _ -> m
+    in
+    match defs with
+    | Defs defs ->
+       List.fold_left add Bindings.empty defs
+  in
+  let replace_constant (E_aux (e,annot) as exp) =
+    match e with
+    | E_id id ->
+       (match Bindings.find_opt id constants with
+        | Some e -> e
+        | None -> exp)
+    | _ -> exp
+  in
   let rec const_prop_exp substs assigns ((E_aux (e,(l,annot))) as exp) =
     (* Functions to treat lists and tuples of subexpressions as possibly
        non-deterministic: that is, we stop making any assumptions about
@@ -633,6 +652,8 @@ let const_props defs ref_vars =
      e.g., (length(op : bits(32)) : int(32)) becomes 32 even if op is not constant.
   *)
   and const_prop_try_fn env (id, args) (l, annot) =
+    let exp_orig = E_aux (E_app (id, args), (l, annot)) in
+    let args = List.map replace_constant args in
     let exp = E_aux (E_app (id, args), (l, annot)) in
     let rec is_overload_of f =
       Env.get_overloads f env
@@ -651,15 +672,27 @@ let const_props defs ref_vars =
        (match destruct_atom_nexp env (typ_of exp) with
         | Some (Nexp_aux (Nexp_constant i, _)) ->
            E_aux (E_lit (mk_lit (L_num i)), (l, annot))
-        | _ -> exp)
-    | _ -> exp
+        | _ -> exp_orig)
+    | _ -> exp_orig
 
   and can_match_with_env env (E_aux (e,(l,annot)) as exp0) cases (substs,ksubsts) assigns =
     let rec check_exp_pat (E_aux (e,(l,annot)) as exp) (P_aux (p,(l',_)) as pat) =
       match e, p with
       | _, P_wild -> DoesMatch ([],[])
       | _, P_typ (_,p') -> check_exp_pat exp p'
-      | _, P_id id' when pat_id_is_variable env id' -> DoesMatch ([id',exp],[])
+      | _, P_id id' when pat_id_is_variable env id' ->
+         let exp_typ = typ_of exp in
+         let pat_typ = typ_of_pat pat in
+         let goals = KidSet.diff (tyvars_of_typ pat_typ) (tyvars_of_typ exp_typ) in
+         let unifiers =
+           try Type_check.unify l env goals pat_typ exp_typ
+           with _ -> KBindings.empty in
+         let is_nexp (k,a) = match a with
+           | A_aux (A_nexp n,_) -> Some (k,n)
+           | _ -> None
+         in
+         let kbindings = Util.map_filter is_nexp (KBindings.bindings unifiers) in
+         DoesMatch ([id',exp],kbindings)
       | E_tuple es, P_tup ps ->
          let rec check = function
            | DoesNotMatch -> fun _ -> DoesNotMatch