Only specialize Jib definitions containing relevant types

If we are specializing a polymorphic type, previously we would specialize
all functions even if they didn't contain references to this type. While
this would not change the definition it essentially caused it to be copied
which caused the GC to have to do a bunch of extra work.

src/bin/dune

@@ -13,7 +13,7 @@
  (public_name sail)
  (package sail)
  (link_flags -linkall)
- (libraries libsail linenoise dynlink))
+ (libraries libsail linenoise dynlink memtrace))
 
 ; For legacy reasons install all the Sail files in lib as part of this package
 

src/bin/sail.ml

@@ -390,6 +390,7 @@ let parse_config_file file =
     None
 
 let main () =
+  Memtrace.trace_if_requested ();
   begin
     match Sys.getenv_opt "SAIL_NO_PLUGINS" with
     | Some _ -> ()

src/lib/jib_compile.ml

@@ -1691,6 +1691,10 @@ module Make (C : CONFIG) = struct
     if IdSet.is_empty (IdSet.diff polymorphic_functions unreachable_polymorphic_functions) then (cdefs, ctx)
     else specialize_functions ~specialized_calls ctx cdefs
 
+  let is_struct id = function CT_struct (id', _) -> Id.compare id id' = 0 | _ -> false
+
+  let is_variant id = function CT_variant (id', _) -> Id.compare id id' = 0 | _ -> false
+
   let map_structs_and_variants f = function
     | ( CT_lint | CT_fint _ | CT_constant _ | CT_lbits | CT_fbits _ | CT_sbits _ | CT_bit | CT_unit | CT_bool | CT_real
       | CT_string | CT_poly _ | CT_enum _ | CT_float _ | CT_rounding_mode ) as ctyp ->
@@ -1845,8 +1849,8 @@ module Make (C : CONFIG) = struct
           |> List.concat
         in
 
-        let prior = List.map (cdef_map_ctyp (fix_variants ctx var_id)) prior in
-        let cdefs = List.map (cdef_map_ctyp (fix_variants ctx var_id)) cdefs in
+        let prior = Util.map_if (cdef_ctyps_has (is_variant var_id)) (cdef_map_ctyp (fix_variants ctx var_id)) prior in
+        let cdefs = Util.map_if (cdef_ctyps_has (is_variant var_id)) (cdef_map_ctyp (fix_variants ctx var_id)) cdefs in
 
         let ctx =
           {
@@ -1900,8 +1904,12 @@ module Make (C : CONFIG) = struct
           |> List.concat
         in
 
-        let prior = List.map (cdef_map_ctyp (fix_variants ctx struct_id)) prior in
-        let cdefs = List.map (cdef_map_ctyp (fix_variants ctx struct_id)) cdefs in
+        let prior =
+          Util.map_if (cdef_ctyps_has (is_struct struct_id)) (cdef_map_ctyp (fix_variants ctx struct_id)) prior
+        in
+        let cdefs =
+          Util.map_if (cdef_ctyps_has (is_struct struct_id)) (cdef_map_ctyp (fix_variants ctx struct_id)) cdefs
+        in
         let ctx =
           {
             ctx with

src/lib/jib_util.ml

@@ -458,6 +458,18 @@ let rec map_ctyp f = function
   | CT_struct (id, fields) -> f (CT_struct (id, List.map (fun (id, ctyp) -> (id, map_ctyp f ctyp)) fields))
   | CT_variant (id, ctors) -> f (CT_variant (id, List.map (fun (id, ctyp) -> (id, map_ctyp f ctyp)) ctors))
 
+let rec ctyp_has pred ctyp =
+  pred ctyp
+  ||
+  match ctyp with
+  | CT_lint | CT_fint _ | CT_constant _ | CT_lbits | CT_fbits _ | CT_sbits _ | CT_float _ | CT_rounding_mode | CT_bit
+  | CT_unit | CT_bool | CT_real | CT_string | CT_poly _ | CT_enum _ ->
+      false
+  | CT_tup ctyps -> List.exists pred ctyps
+  | CT_ref ctyp | CT_vector ctyp | CT_fvector (_, ctyp) | CT_list ctyp -> pred ctyp
+  | CT_struct (id, fields) -> List.exists (fun (_, ctyp) -> pred ctyp) fields
+  | CT_variant (id, ctors) -> List.exists (fun (_, ctyp) -> pred ctyp) ctors
+
 let rec ctyp_equal ctyp1 ctyp2 =
   match (ctyp1, ctyp2) with
   | CT_lint, CT_lint -> true
@@ -1101,6 +1113,21 @@ let rec instr_ctyps (I_aux (instr, aux)) =
 
 and instrs_ctyps instrs = List.fold_left CTSet.union CTSet.empty (List.map instr_ctyps instrs)
 
+let rec instr_ctyps_exist pred (I_aux (instr, aux)) =
+  match instr with
+  | I_decl (ctyp, _) | I_reset (ctyp, _) | I_clear (ctyp, _) | I_undefined ctyp -> pred ctyp
+  | I_init (ctyp, _, cval) | I_reinit (ctyp, _, cval) -> pred ctyp || pred (cval_ctyp cval)
+  | I_if (cval, instrs1, instrs2, ctyp) ->
+      pred (cval_ctyp cval) || instrs_ctyps_exist pred instrs1 || instrs_ctyps_exist pred instrs2 || pred ctyp
+  | I_funcall (clexp, _, (_, ctyps), cvals) ->
+      pred (clexp_ctyp clexp) || List.exists pred ctyps || Util.map_exists pred cval_ctyp cvals
+  | I_copy (clexp, cval) -> pred (clexp_ctyp clexp) || pred (cval_ctyp cval)
+  | I_block instrs | I_try_block instrs -> instrs_ctyps_exist pred instrs
+  | I_throw cval | I_jump (cval, _) | I_return cval -> pred (cval_ctyp cval)
+  | I_comment _ | I_label _ | I_goto _ | I_raw _ | I_exit _ | I_end _ -> false
+
+and instrs_ctyps_exist pred instrs = List.exists (instr_ctyps_exist pred) instrs
+
 let ctype_def_ctyps = function
   | CTD_enum _ -> []
   | CTD_struct (_, fields) -> List.map snd fields
@@ -1116,6 +1143,17 @@ let cdef_ctyps = function
       |> CTSet.union (instrs_ctyps instrs)
   | CDEF_pragma (_, _) -> CTSet.empty
 
+let rec cdef_ctyps_exist pred = function
+  | CDEF_register (_, ctyp, instrs) -> pred ctyp || instrs_ctyps_exist pred instrs
+  | CDEF_val (_, _, ctyps, ctyp) -> List.exists pred ctyps || pred ctyp
+  | CDEF_fundef (_, _, _, instrs) | CDEF_startup (_, instrs) | CDEF_finish (_, instrs) -> instrs_ctyps_exist pred instrs
+  | CDEF_type tdef -> List.exists pred (ctype_def_ctyps tdef)
+  | CDEF_let (_, bindings, instrs) ->
+      List.exists (fun (_, ctyp) -> pred ctyp) bindings || instrs_ctyps_exist pred instrs
+  | CDEF_pragma (_, _) -> false
+
+let cdef_ctyps_has pred cdef = cdef_ctyps_exist (ctyp_has pred) cdef
+
 let rec c_ast_registers = function
   | CDEF_register (id, ctyp, instrs) :: ast -> (id, ctyp, instrs) :: c_ast_registers ast
   | _ :: ast -> c_ast_registers ast

src/lib/jib_util.mli

@@ -186,6 +186,8 @@ val cval_ctyp : cval -> ctyp
 val clexp_ctyp : clexp -> ctyp
 val cdef_ctyps : cdef -> CTSet.t
 
+val cdef_ctyps_has : (ctyp -> bool) -> cdef -> bool
+
 (** {1 Functions for mapping over and extracting information from instructions, values, and definitions} *)
 
 val instr_ids : instr -> NameSet.t

src/lib/util.ml

@@ -421,6 +421,15 @@ let fold_left_index_last f init xs =
   let rec go n acc = function [] -> acc | [x] -> f n true acc x | x :: xs -> go (n + 1) (f n false acc x) xs in
   go 0 init xs
 
+let map_if pred f xs =
+  let rec go acc = function
+    | x :: xs -> begin match pred x with true -> go (f x :: acc) xs | false -> go (x :: acc) xs end
+    | [] -> List.rev acc
+  in
+  go [] xs
+
+let rec map_exists pred f = function x :: xs -> if pred (f x) then true else map_exists pred f xs | [] -> false
+
 let rec take n xs = match (n, xs) with 0, _ -> [] | _, [] -> [] | n, x :: xs -> x :: take (n - 1) xs
 
 let rec drop n xs = match (n, xs) with 0, xs -> xs | _, [] -> [] | n, _ :: xs -> drop (n - 1) xs

src/lib/util.mli

@@ -168,6 +168,10 @@ val map_split : ('a -> ('b, 'c) result) -> 'a list -> 'b list * 'c list
     the [Some] function is removed from the list. *)
 val map_all : ('a -> 'b option) -> 'a list -> 'b list option
 
+val map_if : ('a -> bool) -> ('a -> 'a) -> 'a list -> 'a list
+
+val map_exists : ('b -> bool) -> ('a -> 'b) -> 'a list -> bool
+
 (** [list_to_front i l] resorts the list [l] by bringing the element at index [i]
     to the front. 
     @throws Failure if [i] is not smaller than the length of [l]*)