Control listings in print_accessible_headers

lukstafi · lukstafi · commit 75b5fb751ff7 · 2025-07-14T17:56:52.000+02:00
diff --git a/arrayjit/lib/tnode.ml b/arrayjit/lib/tnode.ml
@@ -755,13 +755,13 @@ let get_values tn =
   let padding = Option.map ~f:fst (Lazy.force tn.padding) in
   Nd.(retrieve_flat_values ?padding @@ Option.value_exn ~here:[%here] @@ Lazy.force tn.array)
 
-let print_accessible_headers () =
+let print_accessible_headers ?(pred = fun _ -> true) () =
   Stdio.printf "Tnode: collecting accessible arrays...%!\n";
   Stdlib.Gc.full_major ();
-  Registry.iter (fun arr -> Stdio.print_endline @@ header arr) registry;
+  Registry.iter (fun arr -> if pred arr then Stdio.print_endline @@ header arr) registry;
   Stdio.printf "Tnode: Finished printing headers.%!\n"
 
-let%debug_sexp log_accessible_headers () =
+let%debug_sexp log_accessible_headers ?(pred = fun _ -> true) () =
   Stdlib.Gc.full_major ();
-  Registry.iter (fun _arr -> [%log header _arr]) registry;
+  Registry.iter (fun arr -> if pred arr then [%log header arr]) registry;
   ()
