Add toplevel commands to fix specific register values and simply spec accordingly

3 files changed, 52 insertions, 10 deletions

diff --git a/src/constant_fold.ml b/src/constant_fold.ml
index 0b76a633..166c0ba6 100644
--- a/src/constant_fold.ml
+++ b/src/constant_fold.ml
@@ -191,7 +191,7 @@ let rec run frame =
 let initial_state ast env =
   Interpreter.initial_state ~registers:false ast env safe_primops
 
-let rw_exp target ok not_ok istate =
+let rw_exp fixed_registers target ok not_ok istate =
   let evaluate e_aux annot =
     let initial_monad = Interpreter.return (E_aux (e_aux, annot)) in
     try
@@ -219,6 +219,21 @@ let rw_exp target ok not_ok istate =
     | E_app (id, args) when fold_to_unit id ->
        ok (); E_aux (E_lit (L_aux (L_unit, fst annot)), annot)
 
+    | E_id id ->
+       begin match Bindings.find_opt id fixed_registers with
+       | Some (_, exp) ->
+          ok (); exp
+       | None ->
+          E_aux (e_aux, annot)
+       end
+
+    (* Short-circuit boolean operators with constants *)
+    | E_app (id, [(E_aux (E_lit (L_aux (L_false, _)), _) as false_exp); _]) when string_of_id id = "and_bool" ->
+       ok (); false_exp
+
+    | E_app (id, [(E_aux (E_lit (L_aux (L_false, _)), _) as true_exp); _]) when string_of_id id = "or_bool" ->
+       ok (); true_exp
+
     | E_app (id, args) when List.for_all is_constant args ->
        let env = env_of_annot annot in
        (* We want to fold all primitive operations, but avoid folding
@@ -252,9 +267,9 @@ let rw_exp target ok not_ok istate =
   in
   fold_exp { id_exp_alg with e_aux = (fun (e_aux, annot) -> rw_funcall e_aux annot)}
 
-let rewrite_exp_once target = rw_exp target (fun _ -> ()) (fun _ -> ())
+let rewrite_exp_once target = rw_exp Bindings.empty target (fun _ -> ()) (fun _ -> ())
 
-let rec rewrite_constant_function_calls' target ast =
+let rec rewrite_constant_function_calls' fixed_registers target ast =
   let rewrite_count = ref 0 in
   let ok () = incr rewrite_count in
   let not_ok () = decr rewrite_count in
@@ -262,16 +277,43 @@ let rec rewrite_constant_function_calls' target ast =
 
   let rw_defs = {
       rewriters_base with
-      rewrite_exp = (fun _ -> rw_exp target ok not_ok istate)
+      rewrite_exp = (fun _ -> rw_exp fixed_registers target ok not_ok istate)
     } in
   let ast = rewrite_defs_base rw_defs ast in
   (* We keep iterating until we have no more re-writes to do *)
   if !rewrite_count > 0
-  then rewrite_constant_function_calls' target ast
+  then rewrite_constant_function_calls' fixed_registers target ast
   else ast
 
-let rewrite_constant_function_calls target ast =
+let rewrite_constant_function_calls fixed_registers target ast =
   if !optimize_constant_fold then
-    rewrite_constant_function_calls' target ast
+    rewrite_constant_function_calls' fixed_registers target ast
   else
     ast
+
+let () =
+  let open Interactive in
+  let open Printf in
+
+  ArgString ("target", fun target -> ArgString ("assignments", fun assignments -> Action (fun () ->
+    let assignments = Str.split (Str.regexp " +") assignments in
+    let assignments =
+      List.map (fun assignment ->
+          match String.split_on_char '=' assignment with
+          | [reg; value] ->
+             let reg = mk_id reg in
+             begin match Env.lookup_id reg !env with
+             | Register (_, _, typ) ->
+                let exp = Initial_check.exp_of_string value in
+                let exp = check_exp !env exp typ in
+                (reg, typ, exp)
+             | _ ->
+                failwith (sprintf "Register %s is not defined in the current environment" (string_of_id reg))
+             end
+          | _ -> failwith (sprintf "Could not parse '%s' as an assignment <register>=<value>" assignment)
+        ) assignments in
+    let assignments = List.fold_left (fun m (reg, typ, exp) -> Bindings.add reg (typ, exp) m) Bindings.empty assignments in
+    ast := rewrite_constant_function_calls' assignments target !ast)))
+  |> register_command
+       ~name:"fix_registers"
+       ~help:"Fix the value of specified registers, specified as a list of <register>=<value>"
diff --git a/src/isail.ml b/src/isail.ml
index 01ffe33f..31919f00 100644
--- a/src/isail.ml
+++ b/src/isail.ml
@@ -110,7 +110,7 @@ let sail_logo =
   List.map banner logo @ [""] @ help @ [""]
 
 let sep = "-----------------------------------------------------" |> Util.blue |> Util.clear
-  
+
 let vs_ids = ref (val_spec_ids !Interactive.ast)
 
 let interactive_state = ref (initial_state ~registers:false !Interactive.ast !Interactive.env !Value.primops)
@@ -142,7 +142,7 @@ let setup_sail_scripting () =
   let open Interactive in
 
   let sail_command_name cmd = "sail_" ^ String.sub cmd 1 (String.length cmd - 1) in
-  
+
   let val_specs =
     List.map (fun (cmd, (_, action)) ->
         let name = sail_command_name cmd in
diff --git a/src/rewrites.ml b/src/rewrites.ml
index ad0ed836..273c9eaf 100644
--- a/src/rewrites.ml
+++ b/src/rewrites.ml
@@ -4850,7 +4850,7 @@ let all_rewrites = [
     ("simple_types", Basic_rewriter rewrite_simple_types);
     ("overload_cast", Basic_rewriter rewrite_overload_cast);
     ("top_sort_defs", Basic_rewriter (fun _ -> top_sort_defs));
-    ("constant_fold", String_rewriter (fun target -> Basic_rewriter (fun _ -> Constant_fold.rewrite_constant_function_calls target)));
+    ("constant_fold", String_rewriter (fun target -> Basic_rewriter (fun _ -> Constant_fold.rewrite_constant_function_calls Bindings.empty target)));
     ("split", String_rewriter (fun str -> Basic_rewriter (rewrite_split_fun_ctor_pats str)));
     ("properties", Basic_rewriter (fun _ -> Property.rewrite));
   ]