Fix memory leak in string_of_bits

Should hopefully fix memory leak in RISC-V.

Also adds an optimization pass that removes copying structs and allows
some structs to simply alias each other and avoid copying their
contents. This requires knowing certain things about the lifetimes of
the structs involved, as can't free the struct if another variable is
referencing it - therefore we conservatively only apply this
optimization for variables that are lifted outside function
definitions, and should therefore never get freed until the model
exits - however this may cause issues outside ARMv8, as there may be
cases where a struct can exist within a variant type (which are not
yet subject to this lifting optimisation), that would break these
assumptions - therefore this optimisation is only enabled with the
-Oexperimental flag.

4 files changed, 100 insertions, 4 deletions

diff --git a/src/bytecode_util.ml b/src/bytecode_util.ml
index aeb1daf1..c7fdc62d 100644
--- a/src/bytecode_util.ml
+++ b/src/bytecode_util.ml
@@ -82,6 +82,9 @@ let iextern ?loc:(l=Parse_ast.Unknown) clexp id cvals =
 let icopy l clexp cval =
   I_aux (I_copy (clexp, cval), (instr_number (), l))
 
+let ialias l clexp cval =
+  I_aux (I_alias (clexp, cval), (instr_number (), l))
+
 let iclear ?loc:(l=Parse_ast.Unknown) ctyp id =
   I_aux (I_clear (ctyp, id), (instr_number (), l))
 
@@ -167,6 +170,7 @@ let rec instr_rename from_id to_id (I_aux (instr, aux)) =
        I_funcall (clexp_rename from_id to_id clexp, extern, id, List.map (cval_rename from_id to_id) args)
 
     | I_copy (clexp, cval) -> I_copy (clexp_rename from_id to_id clexp, cval_rename from_id to_id cval)
+    | I_alias (clexp, cval) -> I_alias (clexp_rename from_id to_id clexp, cval_rename from_id to_id cval)
 
     | I_clear (ctyp, id) when Id.compare id from_id = 0 -> I_clear (ctyp, to_id)
     | I_clear (ctyp, id) -> I_clear (ctyp, id)
@@ -447,6 +451,8 @@ let rec pp_instr ?short:(short=false) (I_aux (instr, aux)) =
                       string (string_of_id f |> Util.green |> Util.clear) ^^ parens (separate_map (string ", ") pp_cval args) ]
   | I_copy (clexp, cval) ->
      separate space [pp_clexp clexp; string "="; pp_cval cval]
+  | I_alias (clexp, cval) ->
+     pp_keyword "alias" ^^ separate space [pp_clexp clexp; string "="; pp_cval cval]
   | I_clear (ctyp, id) ->
      pp_keyword "kill" ^^ pp_id id ^^ string " : " ^^ pp_ctyp ctyp
   | I_return cval ->
@@ -574,6 +580,7 @@ let instr_deps = function
   | I_jump (cval, label) -> cval_deps cval, NS.singleton (G_label label)
   | I_funcall (clexp, _, _, cvals) -> List.fold_left NS.union NS.empty (List.map cval_deps cvals), clexp_deps clexp
   | I_copy (clexp, cval) -> cval_deps cval, clexp_deps clexp
+  | I_alias (clexp, cval) -> cval_deps cval, clexp_deps clexp
   | I_clear (_, id) -> NS.singleton (G_id id), NS.singleton (G_id id)
   | I_throw cval | I_return cval -> cval_deps cval, NS.empty
   | I_block _ | I_try_block _ -> NS.empty, NS.empty
@@ -698,6 +705,7 @@ let rec map_instr_ctyp f (I_aux (instr, aux)) =
     | I_funcall (clexp, extern, id, cvals) ->
        I_funcall (map_clexp_ctyp f clexp, extern, id, List.map (fun (frag, ctyp) -> frag, f ctyp) cvals)
     | I_copy (clexp, (frag, ctyp)) -> I_copy (map_clexp_ctyp f clexp, (frag, f ctyp))
+    | I_alias (clexp, (frag, ctyp)) -> I_alias (map_clexp_ctyp f clexp, (frag, f ctyp))
     | I_clear (ctyp, id) -> I_clear (f ctyp, id)
     | I_return (frag, ctyp) -> I_return (frag, f ctyp)
     | I_block instrs -> I_block (List.map (map_instr_ctyp f) instrs)
@@ -714,7 +722,7 @@ let rec map_instr_ctyp f (I_aux (instr, aux)) =
 let rec map_instr f (I_aux (instr, aux)) =
   let instr = match instr with
     | I_decl _ | I_init _ | I_reset _ | I_reinit _
-      | I_funcall _ | I_copy _ | I_clear _ | I_jump _ | I_throw _ | I_return _
+      | I_funcall _ | I_copy _ | I_alias _ | I_clear _ | I_jump _ | I_throw _ | I_return _
       | I_comment _ | I_label _ | I_goto _ | I_raw _ | I_match_failure | I_undefined _ -> instr
     | I_if (cval, instrs1, instrs2, ctyp) ->
        I_if (cval, List.map (map_instr f) instrs1, List.map (map_instr f) instrs2, ctyp)
@@ -756,7 +764,7 @@ let rec map_instrs f (I_aux (instr, aux)) =
     | I_decl _ | I_init _ | I_reset _ | I_reinit _ -> instr
     | I_if (cval, instrs1, instrs2, ctyp) ->
        I_if (cval, f (List.map (map_instrs f) instrs1), f (List.map (map_instrs f) instrs2), ctyp)
-    | I_funcall _ | I_copy _ | I_clear _ | I_jump _ | I_throw _ | I_return _ -> instr
+    | I_funcall _ | I_copy _ | I_alias _ | I_clear _ | I_jump _ | I_throw _ | I_return _ -> instr
     | I_block instrs -> I_block (f (List.map (map_instrs f) instrs))
     | I_try_block instrs -> I_try_block (f (List.map (map_instrs f) instrs))
     | I_comment _ | I_label _ | I_goto _ | I_raw _ | I_match_failure | I_undefined _ -> instr
@@ -774,6 +782,28 @@ let rec instr_ids (I_aux (instr, _)) =
   |> Util.option_these
   |> IdSet.of_list
 
+let rec instr_reads (I_aux (instr, _)) =
+  let reads, _ = instr_deps instr in
+  let get_id = function
+    | G_id id -> Some id
+    | _ -> None
+  in
+  NS.elements reads
+  |> List.map get_id
+  |> Util.option_these
+  |> IdSet.of_list
+
+let rec instr_writes (I_aux (instr, _)) =
+  let _, writes = instr_deps instr in
+  let get_id = function
+    | G_id id -> Some id
+    | _ -> None
+  in
+  NS.elements writes
+  |> List.map get_id
+  |> Util.option_these
+  |> IdSet.of_list
+
 let rec filter_instrs f instrs =
   let filter_instrs' = function
     | I_aux (I_block instrs, aux) -> I_aux (I_block (filter_instrs f instrs), aux)
diff --git a/src/c_backend.ml b/src/c_backend.ml
index 31a989f8..5003e432 100644
--- a/src/c_backend.ml
+++ b/src/c_backend.ml
@@ -72,6 +72,7 @@ let optimize_primops = ref false
 let optimize_hoist_allocations = ref false
 let optimize_struct_updates = ref false
 let optimize_alias = ref false
+let optimize_experimental = ref false
 
 let c_debug str =
   if !c_verbosity > 0 then prerr_endline (Lazy.force str) else ()
@@ -621,7 +622,7 @@ let rec instr_ctyps (I_aux (instr, aux)) =
      ctyp :: cval_ctyp cval :: List.concat (List.map instr_ctyps instrs1 @ List.map instr_ctyps instrs2)
   | I_funcall (clexp, _, _, cvals) ->
      clexp_ctyp clexp :: List.map cval_ctyp cvals
-  | I_copy (clexp, cval) -> [clexp_ctyp clexp; cval_ctyp cval]
+  | I_copy (clexp, cval) | I_alias (clexp, cval) -> [clexp_ctyp clexp; cval_ctyp cval]
   | I_block instrs | I_try_block instrs -> List.concat (List.map instr_ctyps instrs)
   | I_throw cval | I_jump (cval, _) | I_return cval -> [cval_ctyp cval]
   | I_comment _ | I_label _ | I_goto _ | I_raw _ | I_match_failure -> []
@@ -1546,7 +1547,7 @@ let rec map_try_block f (I_aux (instr, aux)) =
     | I_decl _ | I_reset _ | I_init _ | I_reinit _ -> instr
     | I_if (cval, instrs1, instrs2, ctyp) ->
        I_if (cval, List.map (map_try_block f) instrs1, List.map (map_try_block f) instrs2, ctyp)
-    | I_funcall _ | I_copy _ | I_clear _ | I_throw _ | I_return _ -> instr
+    | I_funcall _ | I_copy _ | I_alias _ | I_clear _ | I_throw _ | I_return _ -> instr
     | I_block instrs -> I_block (List.map (map_try_block f) instrs)
     | I_try_block instrs -> I_try_block (f (List.map (map_try_block f) instrs))
     | I_comment _ | I_label _ | I_goto _ | I_raw _ | I_jump _ | I_match_failure | I_undefined _ -> instr
@@ -1800,6 +1801,7 @@ let rec instrs_rename from_id to_id =
   | I_aux (I_funcall (clexp, extern, id, cvals), aux) :: instrs ->
      I_aux (I_funcall (lrename clexp, extern, rename id, List.map crename cvals), aux) :: irename instrs
   | I_aux (I_copy (clexp, cval), aux) :: instrs -> I_aux (I_copy (lrename clexp, crename cval), aux) :: irename instrs
+  | I_aux (I_alias (clexp, cval), aux) :: instrs -> I_aux (I_alias (lrename clexp, crename cval), aux) :: irename instrs
   | I_aux (I_clear (ctyp, id), aux) :: instrs -> I_aux (I_clear (ctyp, rename id), aux) :: irename instrs
   | I_aux (I_return cval, aux) :: instrs -> I_aux (I_return (crename cval), aux) :: irename instrs
   | I_aux (I_block block, aux) :: instrs -> I_aux (I_block (irename block), aux) :: irename instrs
@@ -2259,6 +2261,62 @@ let combine_variables ctx =
      [CDEF_fundef (function_id, heap_return, args, opt body)]
   | cdef -> [cdef]
 
+(** hoist_alias looks for patterns like
+
+    recreate x; y = x; // no furthner mentions of x
+
+    Provided x has a certain type, then we can make y an alias to x
+    (denoted in the IR as 'alias y = x'). This only works if y also has
+    a lifespan that also spans the entire function body. It's possible
+    we may need to do a more thorough lifetime evaluation to get this
+    to be 100% correct - so it's behind the -Oexperimental flag
+    for now. Some benchmarking shows that this kind of optimization
+    is very valuable however! *)
+let hoist_alias ctx =
+  (* Must return true for a subset of the types hoist_ctyp would return true for. *)
+  let is_struct = function
+    | CT_struct _ -> true
+    | _ -> false
+  in
+  let pattern heap_return id ctyp instrs =
+    let rec scan instrs =
+      match instrs with
+      (* The only thing that has a longer lifetime than id is the
+         function return, so we want to make sure we avoid that
+         case. *)
+      | (I_aux (I_copy (clexp, (F_id id', ctyp')), aux) as instr) :: instrs
+           when not (IdSet.mem heap_return (instr_writes instr)) && Id.compare id id' = 0
+                && ctyp_equal (clexp_ctyp clexp) ctyp && ctyp_equal ctyp ctyp' ->
+         if List.exists (IdSet.mem id) (List.map instr_ids instrs) then
+           instr :: scan instrs
+         else
+           I_aux (I_alias (clexp, (F_id id', ctyp')), aux) :: instrs
+
+      | instr :: instrs -> instr :: scan instrs
+      | [] -> []
+    in
+    scan instrs
+  in
+  let optimize heap_return =
+    let rec opt = function
+      | (I_aux (I_reset (ctyp, id), _) as instr) :: instrs when not (is_stack_ctyp ctyp) && is_struct ctyp ->
+         instr :: opt (pattern heap_return id ctyp instrs)
+
+      | I_aux (I_block block, aux) :: instrs -> I_aux (I_block (opt block), aux) :: opt instrs
+      | I_aux (I_try_block block, aux) :: instrs -> I_aux (I_try_block (opt block), aux) :: opt instrs
+      | I_aux (I_if (cval, then_instrs, else_instrs, ctyp), aux) :: instrs ->
+         I_aux (I_if (cval, opt then_instrs, opt else_instrs, ctyp), aux) :: opt instrs
+
+      | instr :: instrs ->
+         instr :: opt instrs
+      | [] -> []
+    in
+    opt
+  in
+  function
+  | CDEF_fundef (function_id, Some heap_return, args, body) ->
+     [CDEF_fundef (function_id, Some heap_return, args, optimize heap_return body)]
+  | cdef -> [cdef]
 
 let concatMap f xs = List.concat (List.map f xs)
 
@@ -2269,6 +2327,7 @@ let optimize ctx cdefs =
   |> (if !optimize_alias then concatMap (remove_alias ctx) else nothing)
   |> (if !optimize_alias then concatMap (combine_variables ctx) else nothing)
   |> (if !optimize_hoist_allocations then concatMap (hoist_allocations ctx) else nothing)
+  |> (if !optimize_hoist_allocations && !optimize_experimental then concatMap (hoist_alias ctx) else nothing)
 
 (**************************************************************************)
 (* 6. Code generation                                                     *)
@@ -2395,6 +2454,9 @@ let rec codegen_instr fid ctx (I_aux (instr, (_, l))) =
 
   | I_copy (clexp, cval) -> codegen_conversion l clexp cval
 
+  | I_alias (clexp, cval) ->
+     ksprintf string "  %s = %s;" (sgen_clexp_pure clexp) (sgen_cval cval)
+
   | I_jump (cval, label) ->
      ksprintf string "  if (%s) goto %s;" (sgen_cval cval) label
 
diff --git a/src/c_backend.mli b/src/c_backend.mli
index b39f58e2..24f6e03b 100644
--- a/src/c_backend.mli
+++ b/src/c_backend.mli
@@ -65,6 +65,7 @@ val optimize_primops : bool ref
 val optimize_hoist_allocations : bool ref
 val optimize_struct_updates : bool ref
 val optimize_alias : bool ref
+val optimize_experimental : bool ref
 
 (** The compilation context. *)
 type ctx
diff --git a/src/sail.ml b/src/sail.ml
index 44766926..417d3ef4 100644
--- a/src/sail.ml
+++ b/src/sail.ml
@@ -134,6 +134,9 @@ let options = Arg.align ([
   ( "-Oconstant_fold",
     Arg.Set Constant_fold.optimize_constant_fold,
     " Apply constant folding optimizations");
+  ( "-Oexperimental",
+    Arg.Set C_backend.optimize_experimental,
+    " turn on additional, experimental optimisations");
   ( "-static",
     Arg.Set C_backend.opt_static,
     " Make generated C functions static");