Migrate Staged_compilation to PPrint.document

lukstafi · lukstafi · commit 657f596d215d · 2025-05-22T13:18:03.000+02:00
Signed-off-by: Lukasz Stafiniak &lt;lukstafi@gmail.com&gt;
diff --git a/arrayjit/lib/c_syntax.ml b/arrayjit/lib/c_syntax.ml
@@ -356,11 +356,7 @@ module C_syntax (B : C_syntax_config) = struct
             ~args_docs:[]
         else string "/* " ^^ string message ^^ string " */"
     | Staged_compilation callback ->
-        (* This is tricky. PPrint needs to generate the document synchronously. We might need to
-           change how Staged_compilation works if it needs to produce dynamic documents. For now,
-           assume it produces no output. *)
-        callback ();
-        empty
+        callback ()
     | Set_local ({ scope_id; tn = { prec; _ } }, value) ->
         let local_defs, value_doc = pp_float (Lazy.force prec) value in
         let assignment =
diff --git a/arrayjit/lib/gcc_backend.ml b/arrayjit/lib/gcc_backend.ml
@@ -439,7 +439,7 @@ let compile_main ~name ~log_functions ~env { ctx; nodes; get_ident; merge_node;
         let value = loop_float ~name ~env ~num_typ local_prec llv in
         Block.assign !current_block lhs value
     | Comment c -> log_comment c
-    | Staged_compilation exp -> exp ()
+    | Staged_compilation _exp -> failwith "gccjit_backend: Staged_compilation not implemented"
   and loop_float ~name ~env ~num_typ prec v_code =
     let loop = loop_float ~name ~env ~num_typ prec in
     match v_code with
diff --git a/arrayjit/lib/low_level.ml b/arrayjit/lib/low_level.ml
@@ -31,7 +31,7 @@ let get_scope =
 type t =
   | Noop
   | Comment of string
-  | Staged_compilation of ((unit -> unit)[@equal.ignore] [@compare.ignore])
+  | Staged_compilation of ((unit -> PPrint.document)[@equal.ignore] [@compare.ignore])
   | Seq of t * t
   | For_loop of { index : Indexing.symbol; from_ : int; to_ : int; body : t; trace_it : bool }
   | Zero_out of Tn.t
@@ -917,7 +917,9 @@ let fprint_cstyle ?name ?static_indices () ppf llc =
           (pp_float @@ Lazy.force p.tn.prec)
           p.llv
     | Comment message -> fprintf ppf "/* %s */" message
-    | Staged_compilation _ -> fprintf ppf "STAGED_COMPILATION_CALLBACK()"
+    | Staged_compilation callback ->
+      let doc = callback () in
+      PPrint.ToFormatter.pretty 1.0 100 ppf doc
     | Set_local (id, llv) ->
         fprintf ppf "@[<2>%a :=@ %a;@]" pp_local id (pp_float @@ Lazy.force id.tn.prec) llv
   and pp_float prec ppf value =
@@ -977,7 +979,9 @@ let fprint_hum ?name ?static_indices () ppf llc =
         p.debug <- asprintf "@[<2>%a[@,%a] :=@ %a;@]" pp_ident p.tn pp_indices p.idcs pp_float p.llv;
         fprintf ppf "@[<2>%a[@,%a] :=@ %a;@]" pp_ident p.tn pp_indices p.idcs pp_float p.llv
     | Comment message -> fprintf ppf "/* %s */" message
-    | Staged_compilation _ -> fprintf ppf "STAGED_COMPILATION_CALLBACK()"
+    | Staged_compilation callback ->
+        let doc = callback () in
+        PPrint.ToFormatter.pretty 1.0 100 ppf doc
     | Set_local (id, llv) -> fprintf ppf "@[<2>%a :=@ %a;@]" pp_local id pp_float llv
   and pp_float ppf value =
     match value with
diff --git a/arrayjit/lib/low_level.mli b/arrayjit/lib/low_level.mli
@@ -18,7 +18,7 @@ type scope_id = Scope_id.t = { tn : Tnode.t; scope_id : int }
 type t =
   | Noop
   | Comment of string
-  | Staged_compilation of (unit -> unit)
+  | Staged_compilation of (unit -> PPrint.document)
   | Seq of t * t
   | For_loop of { index : Indexing.symbol; from_ : int; to_ : int; body : t; trace_it : bool }
   | Zero_out of Tnode.t
diff --git a/arrayjit/lib/lowering_and_inlining.md b/arrayjit/lib/lowering_and_inlining.md
@@ -47,7 +47,7 @@ The low-level representation is a C-like mini-language operating on scalars.
 type t =
   | Noop
   | Comment of string
-  | Staged_compilation of (unit -> unit)
+  | Staged_compilation of (unit -> PPrint.document)
   | Seq of t * t
   | For_loop of { index : Indexing.symbol; from_ : int; to_ : int; body : t; trace_it : bool }
   | Zero_out of Tnode.t
@@ -70,7 +70,7 @@ and float_t =
   | Embed_index of Indexing.axis_index
 ```
 
-The odd part is the `Staged_compilation` element. Backends with an imperative compilation procedure, e.g. using the `Stdlib.Format` module, can use `Staged_compilation` to embed some emitted code within on-the-fly generated `Low_level.t` code.
+The odd part is the `Staged_compilation` element. Backends can use `Staged_compilation` to embed some emitted code within on-the-fly generated `Low_level.t` code. Currently this works only for `PPrint.document` based backends like `C_syntax` derivatives, but this covers almost all backends.
 
 TODO: flesh out explanation.
 
