Stack alloc: option to print in JSON format

The command line option is -print-stack-alloc {plain|json}
jasmin-lang · Nov 29, 2023 · 5c77bd6 · 5c77bd6
1 parent ef263df
commit 5c77bd6
Show file tree

Hide file tree

Showing 4 changed files with 160 additions and 12 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -51,6 +51,11 @@
   `SMULTT`, `SMLABB`, `SMLABT`, `SMLATB`, `SMLATT`, `SMULWB`, and `SMULWT`
   ([PR #644](https://github.com/jasmin-lang/jasmin/pull/644)).
 
+- The results of the stack allocation oracle can now be printed in JSON.
+  The command line option `-print-stack-alloc` can have two possible values:
+  `plain` (textual output as before) and `json` (JSON output).
+  ([PR #656](https://github.com/jasmin-lang/jasmin/pull/656)).
+
 ## Bug fixes
 
 - Type-checking rejects wrongly casted primitive operators
@@ -65,7 +70,7 @@
   ([PR #541](https://github.com/jasmin-lang/jasmin/pull/541);
   fixes [#540](https://github.com/jasmin-lang/jasmin/issues/540)).
 
- - More precise detection of speculative safe loads in the SCT checker
+- More precise detection of speculative safe loads in the SCT checker
   ([PR #556](https://github.com/jasmin-lang/jasmin/pull/556)).
 
 - Fix printing to EasyCrypt of ARMv7 instruction `bic`

diff --git a/compiler/src/dune b/compiler/src/dune
@@ -12,7 +12,7 @@
  (public_name jasmin.jasmin)
  (flags -rectypes)
  (modules :standard \ arch main_compiler x86_safety uint63_63 uint63_31)
- (libraries batteries.unthreaded zarith menhirLib))
+ (libraries batteries.unthreaded zarith menhirLib yojson))
 
 (library
  (name jasminc)

diff --git a/compiler/src/glob_options.ml b/compiler/src/glob_options.ml
@@ -31,11 +31,24 @@ let sct_list = ref None
 let lea = ref false
 let set0 = ref false
 let model = ref Normal
-let print_stack_alloc = ref false
 let introduce_array_copy = ref true
 let print_dependencies = ref false 
 let lazy_regalloc = ref false
 
+type print_stack_alloc = | No | Plain | Json
+let print_stack_alloc = ref No
+
+(* Ideally, -print-stack-alloc without argument would select Plain, but this is
+   not possible with Arg *)
+let set_print_stack_alloc s =
+  let psa =
+    match s with
+    | "plain" -> Plain
+    | "json" -> Json
+    | _ -> assert false
+  in
+  print_stack_alloc := psa
+
 type architecture =
   | X86_64
   | ARM_M4
@@ -222,7 +235,7 @@ let options = [
     "-nowarning", Arg.Unit (nowarning), ": do no print warning";
     "-color", Arg.Symbol (["auto"; "always"; "never"], set_color), ": print messages with color";
     "-help-intrinsics", Arg.Set help_intrinsics, "List the set of intrinsic operators";
-    "-print-stack-alloc", Arg.Set print_stack_alloc, ": print the results of the stack allocation OCaml oracle";
+    "-print-stack-alloc", Arg.Symbol (["plain"; "json"], set_print_stack_alloc), ": print the results of the stack allocation OCaml oracle";
     "-lazy-regalloc", Arg.Set lazy_regalloc, "\tAllocate variables to registers in program order";
     "-pall"    , Arg.Unit set_all_print, "print program after each compilation steps";
     "-print-dependencies", Arg.Set print_dependencies, ": print dependencies and exit";

diff --git a/compiler/src/stackAlloc.ml b/compiler/src/stackAlloc.ml
@@ -57,7 +57,7 @@ let pp_return fmt n =
 
 let pp_sao fmt sao =
   let open Stack_alloc in
-  Format.fprintf fmt "alignment = %s; size = %a; ioff = %a; extra size = %a; max size = %a@;max call depth = %a@;params =@;<2 2>@[<v>%a@]@;return = @[<hov>%a@]@;slots =@;<2 2>@[<v>%a@]@;alloc= @;<2 2>@[<v>%a@]@;saved register = @[<hov>%a@]@;saved stack = %a@;return address = %a"
+  Format.fprintf fmt "alignment = %s; size = %a; ioff = %a; extra size = %a; max size = %a@;max call depth = %a@;params =@;<2 2>@[<v>%a@]@;return = @[<hov>%a@]@;slots =@;<2 2>@[<v>%a@]@;alloc = @;<2 2>@[<v>%a@]@;saved register = @[<hov>%a@]@;saved stack = %a@;return address = %a"
     (string_of_ws sao.sao_align)
     Z.pp_print (Conv.z_of_cz sao.sao_size)
     Z.pp_print (Conv.z_of_cz sao.sao_ioff)
@@ -88,6 +88,127 @@ let pp_oracle up fmt saos =
     (pp_list "@;" pp_slot) ao_global_alloc
     (pp_list "@;" pp_stack_alloc) fs
 
+let json_of_param_info pi =
+  let open Stack_alloc in
+  match pi with
+  | None -> `Null
+  | Some pi ->
+    `Assoc [
+      "writable", `Bool pi.pp_writable;
+      "var", `String (Format.asprintf "%a" pp_var_ty (Conv.var_of_cvar pi.pp_ptr));
+      "alignment", `String (string_of_ws pi.pp_align)
+    ]
+
+let json_of_return n =
+  match n with
+  | None -> `Null
+  | Some n -> `Int (Conv.int_of_nat n)
+
+let json_of_slot ((x, ws), ofs) =
+  `Assoc [
+    "offset", `Intlit (Z.to_string (Conv.z_of_cz ofs));
+    "var", `String (Format.asprintf "%a" pp_var_ty (Conv.var_of_cvar x));
+    "alignment", `String (string_of_ws ws);
+  ]
+
+let json_of_alloc (x, pki) =
+  let json_of_pki pki =
+    let open Stack_alloc in
+    match pki with
+    | PIdirect (v, z, sc) ->
+      "direct",
+      `Assoc [
+        "scope", `String (if sc = Sglob then "global" else "stack");
+        "var", `String (Format.asprintf "%a" pp_var (Conv.var_of_cvar v));
+        "zone", `String (Format.asprintf "%a" pp_zone z);
+      ]
+    | PIregptr v ->
+      "reg ptr",
+      `Assoc [
+        "var", `String (Format.asprintf "%a" pp_var (Conv.var_of_cvar v))
+      ]
+    | PIstkptr (v, z, x) ->
+      "stack ptr",
+      `Assoc [
+        "var", `String (Format.asprintf "%a" pp_var_ty (Conv.var_of_cvar v));
+        "zone", `String (Format.asprintf "%a" pp_zone z);
+        "pseudo-reg", `String (Format.asprintf "%a" pp_var_ty (Conv.var_of_cvar x))
+      ]
+  in
+  `Assoc [
+    "var", `String (Format.asprintf "%a" pp_var (Conv.var_of_cvar x));
+    json_of_pki pki
+  ]
+
+let json_of_to_save (x, ofs) =
+  `Assoc [
+    "var", `String (Format.asprintf "%a" pp_var (Conv.var_of_cvar x));
+    "offset", `Intlit (Z.to_string (Conv.z_of_cz ofs));
+  ]
+
+let json_of_saved_stack ss =
+  let open Expr in
+  match ss with
+  | SavedStackNone -> `Null
+  | SavedStackReg x ->
+    `Assoc ["reg", `String (Format.asprintf "%a" pp_var (Conv.var_of_cvar x))]
+  | SavedStackStk z ->
+    `Assoc ["stack", `Intlit (Z.to_string (Conv.z_of_cz z))]
+
+let json_of_return_address ra =
+  let open Expr in
+  match ra with
+  | RAnone -> `Null
+  | RAreg x ->
+    `Assoc ["reg", `String (Format.asprintf "%a" pp_var (Conv.var_of_cvar x))]
+  | RAstack (x, z) ->
+    `Assoc ["stack",
+      `Assoc [
+        "var",
+          begin match x with
+          | None -> `Null
+          | Some x -> `String (Format.asprintf "%a" pp_var (Conv.var_of_cvar x))
+          end;
+        "offset", `Intlit (Z.to_string (Conv.z_of_cz z))
+      ]
+    ]
+
+let json_of_sao fn sao : Yojson.Safe.t =
+  let open Stack_alloc in
+  `Assoc [
+    "function name", `String fn.fn_name;
+    "alignment", `String (string_of_ws sao.sao_align);
+    "size", `Intlit (Z.to_string (Conv.z_of_cz sao.sao_size));
+    "ioff", `Intlit (Z.to_string (Conv.z_of_cz sao.sao_ioff));
+    "extra size", `Intlit (Z.to_string (Conv.z_of_cz sao.sao_extra_size));
+    "max size", `Intlit (Z.to_string (Conv.z_of_cz sao.sao_max_size));
+    "max call depth", `Intlit (Z.to_string (Conv.z_of_cz sao.sao_max_call_depth));
+    "params", `List (List.map json_of_param_info sao.sao_params);
+    "return", `List (List.map json_of_return sao.sao_return);
+    "slots", `List (List.map json_of_slot sao.sao_slots);
+    "alloc", `List (List.map json_of_alloc sao.sao_alloc);
+    "saved register", `List (List.map json_of_to_save sao.sao_to_save);
+    "saved stack", json_of_saved_stack sao.sao_rsp;
+    "return address", json_of_return_address sao.sao_return_address;
+  ]
+
+let json_of_oracle up saos =
+  let Compiler.{ ao_globals; ao_global_alloc; ao_stack_alloc } = saos in
+  let json_of_globals global =
+    `Intlit (Z.to_string (Conv.z_of_word U8 global))
+  in
+  let json_of_stack_alloc f =
+    let fn = f.f_name in
+    json_of_sao fn (ao_stack_alloc fn)
+  in
+  let _, fs = Conv.prog_of_cuprog up in
+  `Assoc [
+    ("global data", `List (List.map json_of_globals ao_globals));
+    ("global slots", `List (List.map json_of_slot ao_global_alloc));
+    ("stack alloc", `List (List.map json_of_stack_alloc fs));
+  ]
+
+
 module StackAlloc (Arch: Arch_full.Arch) = struct
 
 module Regalloc = Regalloc (Arch)
@@ -152,18 +273,23 @@ let memory_analysis pp_err ~debug up =
       Hf.add atbl fn csao;
       csao in
 
-  if debug && !Glob_options.print_stack_alloc then begin
+  if debug then begin
     let saos =
       Compiler.({
         ao_globals      = gao.gao_data;
         ao_global_alloc = cglobs;
         ao_stack_alloc  = get_sao
       })
     in
-    Format.eprintf
-"(* -------------------------------------------------------------------- *)@.";
-    Format.eprintf "(* Intermediate results of the stack allocation oracle *)@.@.";
-    Format.eprintf "%a@.@.@." (pp_oracle up) saos
+    match !Glob_options.print_stack_alloc with
+    | No -> ()
+    | Plain ->
+      Format.eprintf
+        "(* -------------------------------------------------------------------- *)@.";
+      Format.eprintf "(* Intermediate results of the stack allocation oracle *)@.@.";
+      Format.eprintf "%a@.@.@." (pp_oracle up) saos
+    | Json ->
+      Format.eprintf "%a@." (Yojson.Safe.pretty_print ~std:true) (json_of_oracle up saos)
   end;
 
   let sp' =
@@ -310,11 +436,15 @@ let memory_analysis pp_err ~debug up =
     })
   in
 
-  if !Glob_options.print_stack_alloc then begin
+  begin match !Glob_options.print_stack_alloc with
+  | No -> ()
+  | Plain ->
     Format.eprintf
-"(* -------------------------------------------------------------------- *)@.";
+      "(* -------------------------------------------------------------------- *)@.";
     Format.eprintf "(* Final results of the stack allocation oracle *)@.@.";
     Format.eprintf "%a@.@.@." (pp_oracle up) saos
+  | Json ->
+    Format.eprintf "%a@." (Yojson.Safe.pretty_print ~std:true) (json_of_oracle up saos)
   end;
 
   saos