SMT: Various small improvements for match performance

More ANF expressions can be considered pure

Mutability of let and var expressions is correctly preserved in the
ANF representation. The ANF pretty printer now also shows mutability.

Attributes are not copied incorrectly onto trivial ANF guards. This is
mostly a presentation issue as they cluttered the output.

The --dprofile option now prints nested timing measurements in a more
useful way

src/lib/anf.ml

@@ -458,18 +458,19 @@ let rec pp_aexp (AE_aux (aexp, annot)) =
   | AE_short_circuit (SC_or, aval, aexp) -> pp_aval aval ^^ string " || " ^^ pp_aexp aexp
   | AE_short_circuit (SC_and, aval, aexp) -> pp_aval aval ^^ string " && " ^^ pp_aexp aexp
   | AE_let (mut, id, id_typ, binding, body, typ) ->
+      let keyword = match mut with Mutable -> string "var" | Immutable -> string "let" in
       group
         begin
           match binding with
           | AE_aux (AE_let _, _) ->
-              (pp_annot typ (separate space [string "let"; pp_annot id_typ (pp_id id); string "="])
+              (pp_annot typ (separate space [keyword; pp_annot id_typ (pp_id id); string "="])
               ^^ hardline
               ^^ nest 2 (pp_aexp binding)
               )
               ^^ hardline ^^ string "in" ^^ space ^^ pp_aexp body
           | _ ->
               pp_annot typ
-                (separate space [string "let"; pp_annot id_typ (pp_id id); string "="; pp_aexp binding; string "in"])
+                (separate space [keyword; pp_annot id_typ (pp_id id); string "="; pp_aexp binding; string "in"])
               ^^ hardline ^^ pp_aexp body
         end
   | AE_if (cond, then_aexp, else_aexp, typ) ->
@@ -752,10 +753,14 @@ let rec anf (E_aux (e_aux, (l, tannot)) as exp) =
       mk_aexp (AE_val (AV_ref (id, lvar)))
   | E_match (match_exp, pexps) ->
       let match_aval, match_wrap = to_aval (anf match_exp) in
-      let anf_pexp (Pat_aux (pat_aux, _)) =
+      let anf_pexp (Pat_aux (pat_aux, (l, _))) =
         match pat_aux with
         | Pat_when (pat, guard, body) -> (anf_pat pat, anf guard, anf body)
-        | Pat_exp (pat, body) -> (anf_pat pat, mk_aexp (AE_val (AV_lit (mk_lit L_true, bool_typ))), anf body)
+        | Pat_exp (pat, body) ->
+            ( anf_pat pat,
+              AE_aux (AE_val (AV_lit (mk_lit L_true, bool_typ)), { loc = l; env = env_of body; uannot = empty_uannot }),
+              anf body
+            )
       in
       match_wrap (mk_aexp (AE_match (match_aval, List.map anf_pexp pexps, typ_of exp)))
   | E_try (match_exp, pexps) ->
@@ -766,13 +771,14 @@ let rec anf (E_aux (e_aux, (l, tannot)) as exp) =
         | Pat_exp (pat, body) -> (anf_pat pat, mk_aexp (AE_val (AV_lit (mk_lit L_true, bool_typ))), anf body)
       in
       mk_aexp (AE_try (match_aexp, List.map anf_pexp pexps, typ_of exp))
-  | E_var (LE_aux (LE_id id, _), binding, body)
-  | E_var (LE_aux (LE_typ (_, id), _), binding, body)
-  | E_let (LB_aux (LB_val (P_aux (P_id id, _), binding), _), body)
-  | E_let (LB_aux (LB_val (P_aux (P_typ (_, P_aux (P_id id, _)), _), binding), _), body) ->
+  | ( E_var (LE_aux (LE_id id, _), binding, body)
+    | E_var (LE_aux (LE_typ (_, id), _), binding, body)
+    | E_let (LB_aux (LB_val (P_aux (P_id id, _), binding), _), body)
+    | E_let (LB_aux (LB_val (P_aux (P_typ (_, P_aux (P_id id, _)), _), binding), _), body) ) as binder ->
+      let mut = match binder with E_var _ -> Mutable | E_let _ -> Immutable | _ -> assert false in
       let env = env_of body in
       let lvar = Env.lookup_id id env in
-      mk_aexp (AE_let (Mutable, id, lvar_typ ~loc:l lvar, anf binding, anf body, typ_of exp))
+      mk_aexp (AE_let (mut, id, lvar_typ ~loc:l lvar, anf binding, anf body, typ_of exp))
   | E_var (lexp, _, _) ->
       Reporting.unreachable l __POS__
         ("Encountered complex l-expression " ^ string_of_lexp lexp ^ " when converting to ANF") [@coverage off]

src/lib/jib_compile.ml

@@ -185,6 +185,9 @@ let rec is_pure_aexp ctx (AE_aux (aexp, { uannot; _ })) =
       match aexp with
       | AE_app (f, _, _) -> Effects.function_is_pure f ctx.effect_info
       | AE_let (Immutable, _, _, aexp1, aexp2, _) -> is_pure_aexp ctx aexp1 && is_pure_aexp ctx aexp2
+      | AE_match (_, arms, _) ->
+          List.for_all (fun (_, guard, aexp) -> is_pure_aexp ctx guard && is_pure_aexp ctx aexp) arms
+      | AE_short_circuit (_, _, aexp) -> is_pure_aexp ctx aexp
       | AE_val _ -> true
       | _ -> false
     )

src/lib/mappings.ml

@@ -215,7 +215,7 @@ let match_completeness c (E_aux (aux, (l, uannot))) =
   in
   match aux with
   | E_match _ -> E_aux (aux, (l, uannot))
-  | _ -> Reporting.unreachable l __POS__ "Non-match in match_complete"
+  | _ -> Reporting.unreachable l __POS__ "Non-match in match_completeness"
 
 let match_complete = match_completeness "complete"
 let match_incomplete = match_completeness "incomplete"

src/lib/profile.ml

@@ -86,12 +86,18 @@ let start () =
   Sys.time ()
 
 let finish msg t =
-  if !opt_profile then begin
+  let open Printf in
+  if !opt_profile then (
+    let depth = 2 * (List.length !profile_stack - 1) in
     match !profile_stack with
     | p :: ps ->
-        prerr_endline (Printf.sprintf "%s %s: %fs" Util.("Profiled" |> magenta |> clear) msg (Sys.time () -. t));
-        prerr_endline (Printf.sprintf "  SMT calls: %d, SMT time: %fs" p.smt_calls p.smt_time);
+        let indent =
+          if depth > 0 then String.init depth (fun i -> if i land 1 = 0 then '|' else ' ') |> Util.magenta |> Util.clear
+          else ""
+        in
+        (* Note ksprintf prerr_endline flushes unlike eprintf so the profiling output occurs immediately *)
+        ksprintf prerr_endline "%s%s %s: %fs" indent Util.("Profiled" |> magenta |> clear) msg (Sys.time () -. t);
+        ksprintf prerr_endline "%s  SMT calls: %d, SMT time: %fs" indent p.smt_calls p.smt_time;
         profile_stack := ps
     | [] -> ()
-  end
-  else ()
+  )

src/sail_smt_backend/jib_smt.ml

@@ -739,7 +739,10 @@ module Make (Config : CONFIG) = struct
     let stack = Stack.create () in
 
     let open Jib_ssa in
+    let t = Profile.start () in
     let start, _, cfg = ssa ?debug_prefix:(Option.map (fun _ -> name) debug_attr) instrs in
+    Profile.finish (Printf.sprintf "SSA conversion (%s)" name) t;
+
     let visit_order =
       try topsort cfg
       with Not_a_DAG n ->
@@ -819,6 +822,7 @@ module Make (Config : CONFIG) = struct
     | CDEF_val (function_id, _, arg_ctyps, ret_ctyp) when Bindings.mem function_id props -> begin
         match find_function [] function_id all_cdefs with
         | intervening_lets, Some (None, args, instrs, function_def_annot) ->
+            let function_id_string = string_of_id function_id in
             let debug_attr = get_def_attribute "jib_debug" function_def_annot in
             let prop_type, prop_args, pragma_l, vs = Bindings.find function_id props in
 
@@ -844,18 +848,20 @@ module Make (Config : CONFIG) = struct
             in
 
             if Option.is_some debug_attr then (
-              prerr_endline Util.("Pre-SMT IR for " ^ string_of_id function_id ^ ":" |> yellow |> bold |> clear);
+              prerr_endline Util.("Pre-SMT IR for " ^ function_id_string ^ ":" |> yellow |> bold |> clear);
               List.iter (fun instr -> prerr_endline (string_of_instr instr)) instrs
             );
 
+            let t = Profile.start () in
             let (stack, state), _ =
-              Smt_gen.run (smt_instr_list debug_attr (string_of_id function_id) ctx all_cdefs instrs) pragma_l
+              Smt_gen.run (smt_instr_list debug_attr function_id_string ctx all_cdefs instrs) pragma_l
             in
+            Profile.finish (Printf.sprintf "SMT conversion (%s)" function_id_string) t;
 
             let query = smt_query state pragma.query in
             push_smt_defs stack [Assert (Fn ("not", [query]))];
 
-            let fname = name_file (string_of_id function_id) in
+            let fname = name_file function_id_string in
             let out_chan = open_out fname in
             if prop_type = "counterexample" then output_string out_chan "(set-option :produce-models true)\n";
 
@@ -872,7 +878,7 @@ module Make (Config : CONFIG) = struct
 
             let t = Profile.start () in
             let queue = Queue_optimizer.optimize stack in
-            Profile.finish "SMT optimization" t;
+            Profile.finish (Printf.sprintf "SMT optimization (%s)" function_id_string) t;
 
             (* let queue = Queue.of_seq (List.to_seq (List.rev (List.of_seq (Stack.to_seq stack)))) in *)
             Queue.iter

src/sail_smt_backend/sail_plugin_smt.ml

@@ -166,7 +166,9 @@ let smt_target out_file { ast; effect_info; env; _ } =
   let module Counterexample = Smt_exp.Counterexample (struct
     let max_unknown_integer_width = !opt_smt_unknown_integer_width
   end) in
+  let t = Profile.start () in
   let generated_smt = SMTGen.generate_smt ~properties ~name_file ~smt_includes:!opt_smt_includes ctx cdefs in
+  Profile.finish "Generating SMT" t;
   if !opt_smt_auto then
     List.iter
       (fun ({ file_name; function_id; args; arg_ctyps; arg_smt_names } : SMTGen.generated_smt_info) ->