Fix an issue with riscv_platform involving flow typing

- Refactor the flow typing implementation in the type-checker. This
  should fix an issue involving riscv_platform. Specifically it should
  now work better when an if statement contains multiple conditions
  combined with and/or, only some of which imply constraints at the
  type level. This change also simplifies the implementation of flow
  typing, and removes some obscure features that were hardly used -
  specifically, flow typing could modify types, but this was fairly
  obscure and doesn't seem to affect any of our specifications. More
  testing is needed to ensure that this change hasn't inadvertantly
  broken anything, but it does pass all our tests and continue to
  typecheck arm, riscv and cheri.

- Also adds a option for generating faster undefined functions for
  enum and variant types. Previously I tried to optimise away such
  functions in the C backend, because they could be slow and cause
  considerable uneccessary allocation, however this was error prone
  and it turns out a much simpler solution is to simply make the
  functions themselves much faster, at the cost of hard-coding certain
  decisions about what undefined means for these types at compile tile
  (which is fine for fast emulation). This almost doubles the
  performance of the generated C code.

- Add a wrapper for right shift to avoid UB when shifting by 64 or
  more places.

5 files changed, 52 insertions, 118 deletions

diff --git a/src/c_backend.ml b/src/c_backend.ml
index 1a150bb5..6daca6b4 100644
--- a/src/c_backend.ml
+++ b/src/c_backend.ml
@@ -62,9 +62,7 @@ let opt_ddump_flow_graphs = ref false
 (* Optimization flags *)
 let optimize_primops = ref false
 let optimize_hoist_allocations = ref false
-let optimize_struct_undefined = ref false
 let optimize_struct_updates = ref false
-let optimize_enum_undefined = ref false
 
 let c_debug str =
   if !c_verbosity > 0 then prerr_endline (Lazy.force str) else ()
@@ -1005,7 +1003,7 @@ let analyze_primop' ctx id args typ =
 
   | "vector_subrange", [AV_C_fragment (vec, _); AV_C_fragment (f, _); AV_C_fragment (t, _)] ->
      let len = F_op (f, "-", F_op (t, "-", v_one)) in
-     AE_val (AV_C_fragment (F_op (F_op (F_raw "UINT64_MAX", ">>", F_op (v_int 64, "-", len)), "&", F_op (vec, ">>", t)), typ))
+     AE_val (AV_C_fragment (F_op (F_call ("safe_rshift", [F_raw "UINT64_MAX"; F_op (v_int 64, "-", len)]), "&", F_op (vec, ">>", t)), typ))
 
   | "vector_access", [AV_C_fragment (vec, _); AV_C_fragment (n, _)] ->
      AE_val (AV_C_fragment (F_op (v_one, "&", F_op (vec, ">>", n)), typ))
@@ -1014,7 +1012,7 @@ let analyze_primop' ctx id args typ =
      AE_val (AV_C_fragment (F_op (a, "==", b), typ))
 
   | "slice", [AV_C_fragment (vec, _); AV_C_fragment (start, _); AV_C_fragment (len, _)] ->
-     AE_val (AV_C_fragment (F_op (F_op (F_raw "UINT64_MAX", ">>", F_op (v_int 64, "-", len)), "&", F_op (vec, ">>", start)), typ))
+     AE_val (AV_C_fragment (F_op (F_call ("safe_rshift", [F_raw "UINT64_MAX"; F_op (v_int 64, "-", len)]), "&", F_op (vec, ">>", start)), typ))
 
   | "undefined_bit", _ ->
      AE_val (AV_C_fragment (F_lit (V_bit Sail_values.B0), typ))
@@ -3411,18 +3409,6 @@ let codegen_def' ctx = function
           ^^ parens (string (sgen_ctyp ret_ctyp ^ " *" ^ sgen_id gs ^ ", ") ^^ string args)
           ^^ hardline
      in
-     let instrs =
-       if !optimize_struct_undefined && is_ct_struct ret_ctyp && Str.string_match (Str.regexp_string "undefined_") (string_of_id id) 0 then
-         []
-       else
-         instrs
-     in
-     let instrs =
-       if !optimize_enum_undefined && is_ct_enum ret_ctyp && Str.string_match (Str.regexp_string "undefined_") (string_of_id id) 0 then
-         []
-       else
-         instrs
-     in
      function_header
      (* ^^ string (Printf.sprintf "{ fprintf(stderr, \"%s \"); " (string_of_id id)) *)
      ^^ string "{"
diff --git a/src/initial_check.ml b/src/initial_check.ml
index 1ba1b9e6..d27281fe 100644
--- a/src/initial_check.ml
+++ b/src/initial_check.ml
@@ -55,6 +55,7 @@ module Big_int = Nat_big_num
 
 (* See mli file for details on what these flags do *)
 let opt_undefined_gen = ref false
+let opt_fast_undefined = ref false
 let opt_magic_hash = ref false
 let opt_enum_casts = ref false
 
@@ -1059,8 +1060,11 @@ let generate_undefineds vs_ids (Defs defs) =
        [mk_val_spec (VS_val_spec (typschm, prepend_id "undefined_" id, (fun _ -> None), false));
         mk_fundef [mk_funcl (prepend_id "undefined_" id)
                             (mk_pat (P_lit (mk_lit L_unit)))
-                            (mk_exp (E_app (mk_id "internal_pick",
-                                            [mk_exp (E_list (List.map (fun id -> mk_exp (E_id id)) ids))])))]]
+                            (if !opt_fast_undefined && List.length ids > 0 then
+                               mk_exp (E_id (List.hd ids))
+                             else
+                               mk_exp (E_app (mk_id "internal_pick",
+                                              [mk_exp (E_list (List.map (fun id -> mk_exp (E_id id)) ids))])))]]
     | TD_record (id, _, typq, fields, _) when not (IdSet.mem (prepend_id "undefined_" id) vs_ids) ->
        let pat = mk_pat (P_tup (quant_items typq |> List.map quant_item_param |> List.concat |> List.map (fun id -> mk_pat (P_id id)))) in
        [mk_val_spec (VS_val_spec (undefined_typschm id typq, prepend_id "undefined_" id, (fun _ -> None), false));
@@ -1072,7 +1076,10 @@ let generate_undefineds vs_ids (Defs defs) =
        [mk_val_spec (VS_val_spec (undefined_typschm id typq, prepend_id "undefined_" id, (fun _ -> None), false));
         mk_fundef [mk_funcl (prepend_id "undefined_" id)
                             pat
-                            (mk_exp (E_app (mk_id "internal_pick",
+                            (if !opt_fast_undefined && List.length tus > 0 then
+                               undefined_tu (List.hd tus)
+                             else
+                               mk_exp (E_app (mk_id "internal_pick",
                                             [mk_exp (E_list (List.map undefined_tu tus))])))]]
     | _ -> []
   in
diff --git a/src/initial_check.mli b/src/initial_check.mli
index 757959f7..e6b29216 100644
--- a/src/initial_check.mli
+++ b/src/initial_check.mli
@@ -51,9 +51,20 @@
 open Ast
 open Ast_util
 
-(* Generate undefined_T functions for every type T *)
+(* Generate undefined_T functions for every type T. False by
+   default. *)
 val opt_undefined_gen : bool ref
 
+(* Generate faster undefined_T functions. Rather than generating
+   functions that allow for the undefined values of enums and variants
+   to be picked at runtime using a RNG or similar, this creates
+   undefined_T functions for those types that simply return a specific
+   member of the type chosen at compile time, which is much
+   faster. These functions don't have the right effects, so the
+   -no_effects flag may be needed if this is true. False by
+   default. *)
+val opt_fast_undefined : bool ref
+
 (* Allow # in identifiers when set, like the GHC option of the same name *)
 val opt_magic_hash : bool ref
 
diff --git a/src/sail.ml b/src/sail.ml
index c1ff267d..fd4bdea5 100644
--- a/src/sail.ml
+++ b/src/sail.ml
@@ -107,10 +107,10 @@ let options = Arg.align ([
     Arg.String (fun elf -> opt_process_elf := Some elf),
     " process an elf file so that it can be executed by compiled C code");
   ( "-O",
-    Arg.Tuple [Arg.Set C_backend.optimize_primops;
+    Arg.Tuple [(* Arg.Set C_backend.optimize_primops; *)
                Arg.Set C_backend.optimize_hoist_allocations;
-               (* Arg.Set C_backend.optimize_enum_undefined; *)
-               (* Arg.Set C_backend.optimize_struct_undefined; *)
+               Arg.Set Initial_check.opt_fast_undefined;
+               Arg.Set Type_check.opt_no_effects;
                Arg.Set C_backend.optimize_struct_updates ],
     " turn on optimizations for C compilation");
   ( "-lem_ast",
diff --git a/src/type_check.ml b/src/type_check.ml
index e32a9532..523b0343 100644
--- a/src/type_check.ml
+++ b/src/type_check.ml
@@ -610,7 +610,7 @@ end = struct
     let counter = ref 0 in
     let simplify_nexp (Nexp_aux (nexp_aux, l) as nexp) =
       match nexp_aux with
-      | Nexp_var _ | Nexp_constant _ -> nexp
+      | Nexp_constant _ -> nexp
       | _ -> (incr counter; nexp)
     in
     let typ = map_nexps simplify_nexp typ in
@@ -1968,18 +1968,25 @@ exception Not_a_constraint;;
 
 let rec assert_nexp env exp = destruct_atom_nexp env (typ_of exp)
 
-let rec assert_constraint env (E_aux (exp_aux, _) as exp) =
+let rec combine_constraint b f x y = match b, x, y with
+  | true,  Some x, Some y -> Some (f x y)
+  | true,  Some x, None   -> Some x
+  | true,  None,   Some y -> Some y
+  | false, Some x, Some y -> Some (f x y)
+  | _, _, _ -> None
+
+let rec assert_constraint env b (E_aux (exp_aux, _) as exp) =
   match exp_aux with
   | E_constraint nc ->
      Some nc
   | E_lit (L_aux (L_true, _)) -> Some nc_true
   | E_lit (L_aux (L_false, _)) -> Some nc_false
   | E_let (_,e) ->
-     assert_constraint env e (* TODO: beware of fresh type vars *)
+     assert_constraint env b e (* TODO: beware of fresh type vars *)
   | E_app (op, [x; y]) when string_of_id op = "or_bool" ->
-     option_binop nc_or (assert_constraint env x) (assert_constraint env y)
+     combine_constraint (not b) nc_or (assert_constraint env b x) (assert_constraint env b y)
   | E_app (op, [x; y]) when string_of_id op = "and_bool" ->
-     option_binop nc_and (assert_constraint env x) (assert_constraint env y)
+     combine_constraint b nc_and (assert_constraint env b x) (assert_constraint env b y)
   | E_app (op, [x; y]) when string_of_id op = "gteq_atom" ->
      option_binop nc_gteq (assert_nexp env x) (assert_nexp env y)
   | E_app (op, [x; y]) when string_of_id op = "lteq_atom" ->
@@ -1995,83 +2002,10 @@ let rec assert_constraint env (E_aux (exp_aux, _) as exp) =
   | _ ->
      None
 
-type flow_constraint =
-  | Flow_lteq of Big_int.num * nexp
-  | Flow_gteq of Big_int.num * nexp
-
-let restrict_range_upper c1 nexp1 (Typ_aux (typ_aux, l) as typ) =
-  match typ_aux with
-  | Typ_app (f, [Typ_arg_aux (Typ_arg_nexp nexp, _); Typ_arg_aux (Typ_arg_nexp nexp2, _)])
-     when string_of_id f = "range" ->
-     begin
-       match big_int_of_nexp nexp2 with
-       | Some c2 ->
-          let upper = if (Big_int.less c1 c2) then nexp1 else nexp2 in
-          range_typ nexp upper
-       | _ -> typ
-     end
-  | _ -> typ
-
-let restrict_range_lower c1 nexp1 (Typ_aux (typ_aux, l) as typ) =
-  match typ_aux with
-  | Typ_app (f, [Typ_arg_aux (Typ_arg_nexp nexp2, _); Typ_arg_aux (Typ_arg_nexp nexp, _)])
-     when string_of_id f = "range" ->
-     begin
-       match big_int_of_nexp nexp2 with
-       | Some c2 ->
-          let lower = if (Big_int.greater c1 c2) then nexp1 else nexp2 in
-          range_typ lower nexp
-       | _ -> typ
-     end
-  | _ -> typ
-
-let apply_flow_constraint = function
-  | Flow_lteq (c, nexp) ->
-     (restrict_range_upper c nexp,
-      restrict_range_lower (Big_int.succ c) (nexp_simp (nsum nexp (nint 1))))
-  | Flow_gteq (c, nexp) ->
-     (restrict_range_lower c nexp,
-      restrict_range_upper (Big_int.pred c) (nexp_simp (nminus nexp (nint 1))))
-
-let rec infer_flow env (E_aux (exp_aux, (l, _)) as exp) =
-  match exp_aux with
-  | E_app (f, [E_aux (E_id v, _); y]) when string_of_id f = "lt_range_atom" ->
-     let kid = Env.fresh_kid env in
-     begin
-       match destruct_atom (typ_of y) with
-       | Some (c, nexp) ->
-          [(v, Flow_lteq (Big_int.pred c, nexp_simp (nminus nexp (nint 1))))], []
-       | _ -> [], []
-     end
-  | E_app (f, [E_aux (E_id v, _); y]) when string_of_id f = "lteq_range_atom" ->
-     let kid = Env.fresh_kid env in
-     begin
-       match destruct_atom (typ_of y) with
-       | Some (c, nexp) -> [(v, Flow_lteq (c, nexp))], []
-       | _ -> [], []
-     end
-  | E_app (f, [E_aux (E_id v, _); y]) when string_of_id f = "gt_range_atom" ->
-     let kid = Env.fresh_kid env in
-     begin
-       match destruct_atom (typ_of y) with
-       | Some (c, nexp) ->
-          [(v, Flow_gteq (Big_int.succ c, nexp_simp (nsum nexp (nint 1))))], []
-       | _ -> [], []
-     end
-  | E_app (f, [E_aux (E_id v, _); y]) when string_of_id f = "gteq_range_atom" ->
-     let kid = Env.fresh_kid env in
-     begin
-       match destruct_atom (typ_of y) with
-       | Some (c, nexp) -> [(v, Flow_gteq (c, nexp))], []
-       | _ -> [], []
-     end
-  | _ -> [], option_these [assert_constraint env exp]
-
-let rec add_flows b flows env =
-  match flows with
-  | [] -> env
-  | (id, flow) :: flows when b -> add_flows true flows (Env.add_flow id (fst (apply_flow_constraint flow)) env)
-  | (id, flow) :: flows -> add_flows false flows (Env.add_flow id (snd (apply_flow_constraint flow)) env)
+let rec add_opt_constraint constr env =
+  match constr with
+  | None -> env
+  | Some constr -> Env.add_constraint constr env
 
 let rec add_constraints constrs env =
   List.fold_left (fun env constr -> Env.add_constraint constr env) env constrs
@@ -2169,7 +2103,7 @@ let rec check_exp env (E_aux (exp_aux, (l, ())) as exp : unit exp) (Typ_aux (typ
          | ((E_aux (E_assert (constr_exp, assert_msg), _) as exp) :: exps) ->
             let constr_exp = crule check_exp env constr_exp bool_typ in
             let checked_msg = crule check_exp env assert_msg string_typ in
-            let env = match assert_constraint env constr_exp with
+            let env = match assert_constraint env true constr_exp with
               | Some nc ->
                  typ_print (lazy ("Adding constraint " ^ string_of_n_constraint nc ^ " for assert"));
                  Env.add_constraint nc env
@@ -2293,9 +2227,8 @@ let rec check_exp env (E_aux (exp_aux, (l, ())) as exp : unit exp) (Typ_aux (typ
      type_coercion env inferred_exp typ
   | E_if (cond, then_branch, else_branch), _ ->
      let cond' = crule check_exp env cond (mk_typ (Typ_id (mk_id "bool"))) in
-     let flows, constrs = infer_flow env cond' in
-     let then_branch' = crule check_exp (add_constraints constrs (add_flows true flows env)) then_branch typ in
-     let else_branch' = crule check_exp (add_constraints (List.map nc_negate constrs) (add_flows false flows env)) else_branch typ in
+     let then_branch' = crule check_exp (add_opt_constraint (assert_constraint env true cond') env) then_branch typ in
+     let else_branch' = crule check_exp (add_opt_constraint (option_map nc_negate (assert_constraint env false cond')) env) else_branch typ in
      annot_exp (E_if (cond', then_branch', else_branch')) typ
   | E_exit exp, _ ->
      let checked_exp = crule check_exp env exp (mk_typ (Typ_id (mk_id "unit"))) in
@@ -2327,7 +2260,7 @@ let rec check_exp env (E_aux (exp_aux, (l, ())) as exp : unit exp) (Typ_aux (typ
      let env = match bind_exp with
        | E_aux (E_assert (constr_exp, _), _) ->
           begin
-            match assert_constraint env constr_exp with
+            match assert_constraint env true constr_exp with
             | Some nc ->
                typ_print (lazy ("Adding constraint " ^ string_of_n_constraint nc ^ " for assert"));
                Env.add_constraint nc env
@@ -2367,8 +2300,7 @@ and check_case env pat_typ pexp typ =
        | None -> None, env
        | Some guard ->
           let checked_guard = check_exp env guard bool_typ in
-          let flows, constrs = infer_flow env checked_guard in
-          Some checked_guard, add_constraints constrs (add_flows true flows env)
+          Some checked_guard, add_opt_constraint (assert_constraint env true checked_guard) env
      in
      let checked_case = crule check_exp env' case typ in
      construct_pexp (tpat, checked_guard, checked_case, (l, None))
@@ -2954,8 +2886,7 @@ and infer_exp env (E_aux (exp_aux, (l, ())) as exp) =
   | E_app (f, xs) -> infer_funapp l env f xs None
   | E_loop (loop_type, cond, body) ->
      let checked_cond = crule check_exp env cond bool_typ in
-     let flows, constrs = infer_flow env checked_cond in
-     let checked_body = crule check_exp (add_flows true flows env) body unit_typ in
+     let checked_body = crule check_exp (add_opt_constraint (assert_constraint env true checked_cond) env) body unit_typ in
      annot_exp (E_loop (loop_type, checked_cond, checked_body)) unit_typ
   | E_for (v, f, t, step, ord, body) ->
      begin
@@ -2982,9 +2913,8 @@ and infer_exp env (E_aux (exp_aux, (l, ())) as exp) =
      end
   | E_if (cond, then_branch, else_branch) ->
      let cond' = crule check_exp env cond (mk_typ (Typ_id (mk_id "bool"))) in
-     let flows, constrs = infer_flow env cond' in
-     let then_branch' = irule infer_exp (add_constraints constrs (add_flows true flows env)) then_branch in
-     let else_branch' = crule check_exp (add_constraints (List.map nc_negate constrs) (add_flows false flows env)) else_branch (typ_of then_branch') in
+     let then_branch' = irule infer_exp (add_opt_constraint (assert_constraint env true cond') env) then_branch in
+     let else_branch' = crule check_exp (add_opt_constraint (option_map nc_negate (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')
   | 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, ())))
@@ -3018,7 +2948,7 @@ and infer_exp env (E_aux (exp_aux, (l, ())) as exp) =
      let env = match bind_exp with
        | E_aux (E_assert (constr_exp, _), _) ->
           begin
-            match assert_constraint env constr_exp with
+            match assert_constraint env true constr_exp with
             | Some nc ->
                typ_print (lazy ("Adding constraint " ^ string_of_n_constraint nc ^ " for assert"));
                Env.add_constraint nc env