Fix rewriter issues

Fixes some re-writer issues that was preventing RISC-V from building
with new flow-typing constraints. Unfortunately because the flow
typing now understands slightly more about boolean variables, the very
large nested case statements with matches predicates produced by the
string-matching end up causing a huge blowup in the overall
compilation time.

4 files changed, 112 insertions, 44 deletions

diff --git a/src/ast_util.ml b/src/ast_util.ml
index 795a41fe..55f8c61c 100644
--- a/src/ast_util.ml
+++ b/src/ast_util.ml
@@ -1173,6 +1173,67 @@ let equal_effects e1 e2 =
   | Effect_aux (Effect_set base_effs1, _), Effect_aux (Effect_set base_effs2, _) ->
      BESet.compare (BESet.of_list base_effs1) (BESet.of_list base_effs2) = 0
 
+let rec kopts_of_nexp (Nexp_aux (nexp,_)) =
+  match nexp with
+  | Nexp_id _
+  | Nexp_constant _ -> KOptSet.empty
+  | Nexp_var kid -> KOptSet.singleton (mk_kopt K_int kid)
+  | Nexp_times (n1,n2)
+  | Nexp_sum (n1,n2)
+  | Nexp_minus (n1,n2) -> KOptSet.union (kopts_of_nexp n1) (kopts_of_nexp n2)
+  | Nexp_exp n
+  | Nexp_neg n -> kopts_of_nexp n
+  | Nexp_app (_, nexps) -> List.fold_left KOptSet.union KOptSet.empty (List.map kopts_of_nexp nexps)
+
+let kopts_of_order (Ord_aux (ord, _)) =
+  match ord with
+  | Ord_var kid -> KOptSet.singleton (mk_kopt K_order kid)
+  | Ord_inc | Ord_dec -> KOptSet.empty
+
+let rec kopts_of_constraint (NC_aux (nc, _)) =
+  match nc with
+  | NC_equal (nexp1, nexp2)
+  | NC_bounded_ge (nexp1, nexp2)
+  | NC_bounded_le (nexp1, nexp2)
+  | NC_not_equal (nexp1, nexp2) ->
+     KOptSet.union (kopts_of_nexp nexp1) (kopts_of_nexp nexp2)
+  | NC_set (kid, _) -> KOptSet.singleton (mk_kopt K_int kid)
+  | NC_or (nc1, nc2)
+  | NC_and (nc1, nc2) ->
+     KOptSet.union (kopts_of_constraint nc1) (kopts_of_constraint nc2)
+  | NC_app (id, args) ->
+     List.fold_left (fun s t -> KOptSet.union s (kopts_of_typ_arg t)) KOptSet.empty args
+  | NC_var kid -> KOptSet.singleton (mk_kopt K_bool kid)
+  | NC_true | NC_false -> KOptSet.empty
+
+and kopts_of_typ (Typ_aux (t,_)) =
+  match t with
+  | Typ_internal_unknown -> KOptSet.empty
+  | Typ_id _ -> KOptSet.empty
+  | Typ_var kid -> KOptSet.singleton (mk_kopt K_type kid)
+  | Typ_fn (ts, t, _) -> List.fold_left KOptSet.union (kopts_of_typ t) (List.map kopts_of_typ ts)
+  | Typ_bidir (t1, t2) -> KOptSet.union (kopts_of_typ t1) (kopts_of_typ t2)
+  | Typ_tup ts ->
+     List.fold_left (fun s t -> KOptSet.union s (kopts_of_typ t))
+       KOptSet.empty ts
+  | Typ_app (_,tas) ->
+     List.fold_left (fun s ta -> KOptSet.union s (kopts_of_typ_arg ta))
+       KOptSet.empty tas
+  | Typ_exist (kopts, nc, t) ->
+     let s = KOptSet.union (kopts_of_typ t) (kopts_of_constraint nc) in
+     KOptSet.diff s (KOptSet.of_list kopts)
+and kopts_of_typ_arg (A_aux (ta,_)) =
+  match ta with
+  | A_nexp nexp -> kopts_of_nexp nexp
+  | A_typ typ -> kopts_of_typ typ
+  | A_order ord -> kopts_of_order ord
+  | A_bool nc -> kopts_of_constraint nc
+
+let kopts_of_quant_item (QI_aux (qi, _)) = match qi with
+  | QI_id kopt ->
+     KOptSet.singleton kopt
+  | QI_const nc -> kopts_of_constraint nc
+
 let rec tyvars_of_nexp (Nexp_aux (nexp,_)) =
   match nexp with
   | Nexp_id _
diff --git a/src/ast_util.mli b/src/ast_util.mli
index 7a44322d..df7f7efb 100644
--- a/src/ast_util.mli
+++ b/src/ast_util.mli
@@ -354,6 +354,13 @@ val effect_set : effect -> BESet.t
 val equal_effects : effect -> effect -> bool
 val union_effects : effect -> effect -> effect
 
+val kopts_of_order : order -> KOptSet.t
+val kopts_of_nexp : nexp -> KOptSet.t
+val kopts_of_typ : typ -> KOptSet.t
+val kopts_of_typ_arg : typ_arg -> KOptSet.t
+val kopts_of_constraint : n_constraint -> KOptSet.t
+val kopts_of_quant_item : quant_item -> KOptSet.t
+
 val tyvars_of_nexp : nexp -> KidSet.t
 val tyvars_of_typ : typ -> KidSet.t
 val tyvars_of_constraint : n_constraint -> KidSet.t
diff --git a/src/rewrites.ml b/src/rewrites.ml
index ec0ebaa7..e87847bd 100644
--- a/src/rewrites.ml
+++ b/src/rewrites.ml
@@ -1325,7 +1325,7 @@ let contains_bitvector_pexp = function
 
 let remove_bitvector_pat (P_aux (_, (l, _)) as pat) =
 
-  let env = try env_of_pat pat with _ -> Env.empty in
+  let env = try env_of_pat pat with _ -> raise (Reporting.err_unreachable l __POS__ "Pattern without annotation found") in
 
   (* first introduce names for bitvector patterns *)
   let name_bitvector_roots =
@@ -2092,23 +2092,23 @@ let rewrite_split_fun_constr_pats fun_name (Defs defs) =
         | _ ->
            function_typ [args_typ] ret_typ eff
       in
-      let quant_new_tyvars qis =
-        let quant_tyvars = List.fold_left KidSet.union KidSet.empty (List.map tyvars_of_quant_item qis) in
-        let typ_tyvars = tyvars_of_typ fun_typ in
-        let new_tyvars = KidSet.diff typ_tyvars quant_tyvars in
-        List.map (mk_qi_id K_int) (KidSet.elements new_tyvars)
+      let quant_new_kopts qis =
+        let quant_kopts = List.fold_left KOptSet.union KOptSet.empty (List.map kopts_of_quant_item qis) in
+        let typ_kopts = kopts_of_typ fun_typ in
+        let new_kopts = KOptSet.diff typ_kopts quant_kopts in
+        List.map mk_qi_kopt (KOptSet.elements new_kopts)
       in
       let typquant = match typquant with
         | TypQ_aux (TypQ_tq qis, l) ->
            let qis =
              List.filter
-               (fun qi -> KidSet.subset (tyvars_of_quant_item qi) (tyvars_of_typ fun_typ))
+               (fun qi -> KOptSet.subset (kopts_of_quant_item qi) (kopts_of_typ fun_typ))
                qis
-             @ quant_new_tyvars qis
+             @ quant_new_kopts qis
            in
            TypQ_aux (TypQ_tq qis, l)
         | _ ->
-           TypQ_aux (TypQ_tq (List.map (mk_qi_id K_int) (KidSet.elements (tyvars_of_typ fun_typ))), l)
+           TypQ_aux (TypQ_tq (List.map mk_qi_kopt (KOptSet.elements (kopts_of_typ fun_typ))), l)
       in
       let val_spec =
         VS_aux (VS_val_spec
@@ -3016,10 +3016,16 @@ let rec binding_typs_of_pat (P_aux (p_aux, p_annot) as pat) =
 let construct_toplevel_string_append_call env f_id bindings binding_typs guard expr =
   (* s# if match f#(s#) { Some (bindings) => guard, _ => false) } => let Some(bindings) = f#(s#) in expr *)
   let s_id = fresh_stringappend_id () in
+  let hack_typ (Typ_aux (aux, _) as typ) =
+    match aux with
+    | Typ_app (Id_aux (Id "atom_bool", _), [_]) -> bool_typ
+    | Typ_app (Id_aux (Id "atom", _), [_]) -> int_typ
+    | _ -> typ
+  in
   let option_typ = app_typ (mk_id "option") [A_aux (A_typ (match binding_typs with
                                                                          | [] -> unit_typ
-                                                                         | [typ] -> typ
-                                                                         | typs -> tuple_typ typs
+                                                                         | [typ] -> hack_typ typ
+                                                                         | typs -> tuple_typ (List.map hack_typ typs)
                                                               ), unk)]
   in
   let bindings = if bindings = [] then
@@ -3048,11 +3054,22 @@ let construct_toplevel_string_append_func env f_id pat =
                  else
                    bindings
   in
+  (* AA: Pulling the types out of a pattern with binding_typs_of_pat
+     is broken here because they might contain type variables that
+     were bound locally to the pattern, so we can't lift them out to
+     the top-level. As a hacky fix we can generalise types where this
+     is likely to happen. *)
+  let hack_typ (Typ_aux (aux, _) as typ) =
+    match aux with
+    | Typ_app (Id_aux (Id "atom_bool", _), [_]) -> bool_typ
+    | Typ_app (Id_aux (Id "atom", _), [_]) -> int_typ
+    | _ -> typ
+  in
   let option_typ = app_typ (mk_id "option") [A_aux (A_typ (match binding_typs with
                                                                        | [] -> unit_typ
-                                                                       | [typ] -> typ
-                                                                       | typs -> tuple_typ typs
-                                                            ), unk)]
+                                                                       | [typ] -> hack_typ typ
+                                                                       | typs -> tuple_typ (List.map hack_typ typs)
+                                                          ), unk)]
   in
   let fun_typ = (mk_typ (Typ_fn ([string_typ], option_typ, no_effect))) in
   let new_val_spec = VS_aux (VS_val_spec (mk_typschm (TypQ_aux (TypQ_no_forall, unk)) fun_typ, f_id, (fun _ -> None), false), unkt) in
@@ -3135,7 +3152,7 @@ let construct_toplevel_string_append_func env f_id pat =
        let some_pat = annot_pat (P_app (mk_id "Some",
                                         [tup_arg_pat;
                                          annot_pat (P_id len_id) unk env nat_typ]))
-                        unk env opt_typ in
+                                unk env opt_typ in
        let some_pat, some_pat_env, _ = bind_pat env (strip_pat some_pat) opt_typ in
 
        (* need to add the Some(...) env to tup_arg_pats for pat_to_exp below as it calls the typechecker *)
@@ -3520,7 +3537,7 @@ let rewrite_defs_mapping_patterns =
 
        let false_exp = annot_exp (E_lit (L_aux (L_false, unk))) unk env bool_typ in
        let new_other_guards = annot_exp (E_if (new_guard,
-                                               (annot_exp (E_let (new_letbind, fold_typed_guards env guards)) unk env bool_typ),
+                                               (annot_exp (E_let (new_letbind, annot_exp (E_cast (bool_typ, fold_typed_guards env guards)) unk env bool_typ)) unk env bool_typ),
                                                false_exp)) unk env bool_typ in
 
        annot_pat (P_typ (mapping_in_typ, annot_pat (P_id s_id) unk env mapping_in_typ)) unk env mapping_in_typ, [new_guard; new_other_guards], new_let
diff --git a/src/type_check.ml b/src/type_check.ml
index dedb7015..6ddc31a7 100644
--- a/src/type_check.ml
+++ b/src/type_check.ml
@@ -1300,11 +1300,6 @@ let rec nexp_frees ?exs:(exs=KidSet.empty) (Nexp_aux (nexp, l)) =
   | Nexp_exp n -> nexp_frees ~exs:exs n
   | Nexp_neg n -> nexp_frees ~exs:exs n
 
-let order_frees (Ord_aux (ord_aux, l)) =
-  match ord_aux with
-  | Ord_var kid -> KidSet.singleton kid
-  | _ -> KidSet.empty
-
 let rec typ_nexps (Typ_aux (typ_aux, l)) =
   match typ_aux with
   | Typ_internal_unknown -> []
@@ -1323,24 +1318,6 @@ and typ_arg_nexps (A_aux (typ_arg_aux, l)) =
   | A_typ typ -> typ_nexps typ
   | A_order ord -> []
 
-let rec typ_frees ?exs:(exs=KidSet.empty) (Typ_aux (typ_aux, l)) =
-  match typ_aux with
-  | Typ_internal_unknown -> KidSet.empty
-  | Typ_id v -> KidSet.empty
-  | Typ_var kid when KidSet.mem kid exs -> KidSet.empty
-  | Typ_var kid -> KidSet.singleton kid
-  | Typ_tup typs -> List.fold_left KidSet.union KidSet.empty (List.map (typ_frees ~exs:exs) typs)
-  | Typ_app (f, args) -> List.fold_left KidSet.union KidSet.empty (List.map (typ_arg_frees ~exs:exs) args)
-  | Typ_exist (kopts, nc, typ) -> typ_frees ~exs:(KidSet.of_list (List.map kopt_kid kopts)) typ
-  | Typ_fn (arg_typs, ret_typ, _) -> List.fold_left KidSet.union (typ_frees ~exs:exs ret_typ) (List.map (typ_frees ~exs:exs) arg_typs)
-  | Typ_bidir (typ1, typ2) -> KidSet.union (typ_frees ~exs:exs typ1) (typ_frees ~exs:exs typ2)
-and typ_arg_frees ?exs:(exs=KidSet.empty) (A_aux (typ_arg_aux, l)) =
-  match typ_arg_aux with
-  | A_nexp n -> nexp_frees ~exs:exs n
-  | A_typ typ -> typ_frees ~exs:exs typ
-  | A_order ord -> order_frees ord
-  | A_bool nc -> tyvars_of_constraint nc
-
 let rec nexp_identical (Nexp_aux (nexp1, _)) (Nexp_aux (nexp2, _)) =
   match nexp1, nexp2 with
   | Nexp_id v1, Nexp_id v2 -> Id.compare v1 v2 = 0
@@ -1791,7 +1768,7 @@ let rec subtyp l env typ1 typ2 =
      typ_debug (lazy "Subtype check with unification");
      let typ1 = canonicalize env typ1 in
      let env = add_typ_vars l kopts env in
-     let kids' = KidSet.elements (KidSet.diff (KidSet.of_list (List.map kopt_kid kopts)) (typ_frees typ2)) in
+     let kids' = KidSet.elements (KidSet.diff (KidSet.of_list (List.map kopt_kid kopts)) (tyvars_of_typ typ2)) in
      if not (kids' = []) then typ_error l "Universally quantified constraint generated" else ();
      let unifiers =
        try unify l env (KidSet.diff (tyvars_of_typ typ2) (tyvars_of_typ typ1)) typ2 typ1 with
@@ -3256,8 +3233,14 @@ and infer_exp env (E_aux (exp_aux, (l, ())) as exp) =
         | None -> typ_error l ("Could not infer type of " ^ string_of_exp else_branch)
         end
      | None ->
-        let else_branch' = crule check_exp (add_opt_constraint (option_map nc_not (assert_constraint env false cond')) env) else_branch (typ_of then_branch') in
-        annot_exp (E_if (cond', then_branch', else_branch')) (typ_of then_branch')
+        begin match typ_of then_branch' with
+        | Typ_aux (Typ_app (f, [_]), _) when string_of_id f = "atom_bool" ->
+           let else_branch' = crule check_exp (add_opt_constraint (option_map nc_not (assert_constraint env false cond')) env) else_branch bool_typ in
+           annot_exp (E_if (cond', then_branch', else_branch')) bool_typ
+        | _ ->
+           let else_branch' = crule check_exp (add_opt_constraint (option_map nc_not (assert_constraint env false cond')) env) else_branch (typ_of then_branch') in
+           annot_exp (E_if (cond', then_branch', else_branch')) (typ_of then_branch')
+        end
      end
   | E_vector_access (v, n) -> infer_exp env (E_aux (E_app (mk_id "vector_access", [v; n]), (l, ())))
   | E_vector_update (v, n, exp) -> infer_exp env (E_aux (E_app (mk_id "vector_update", [v; n; exp]), (l, ())))
@@ -3449,7 +3432,7 @@ and infer_funapp' l env f (typq, f_typ) xs expected_ret_typ =
 
   let universals = KBindings.bindings universals |> List.map fst |> KidSet.of_list in
   let typ_ret =
-    if KidSet.is_empty (KidSet.of_list (List.map kopt_kid existentials)) || KidSet.is_empty (KidSet.diff (typ_frees !typ_ret) universals)
+    if KidSet.is_empty (KidSet.of_list (List.map kopt_kid existentials)) || KidSet.is_empty (KidSet.diff (tyvars_of_typ !typ_ret) universals)
     then !typ_ret
     else mk_typ (Typ_exist (existentials, List.fold_left nc_and nc_true ex_constraints, !typ_ret))
   in
@@ -4386,7 +4369,7 @@ let check_type_union env variant typq (Tu_aux (tu, l)) =
   let ret_typ = app_typ variant (List.fold_left fold_union_quant [] (quant_items typq)) in
   match tu with
   | Tu_ty_id (Typ_aux (Typ_fn (arg_typ, ret_typ, _), _) as typ, v) ->
-     let typq = mk_typquant (List.map (mk_qi_id K_type) (KidSet.elements (typ_frees typ))) in
+     let typq = mk_typquant (List.map (mk_qi_id K_type) (KidSet.elements (tyvars_of_typ typ))) in
      env
      |> Env.add_union_id v (typq, typ)
      |> Env.add_val_spec v (typq, typ)