From 5e678b02bb918ef3e4aff7827495cb2df9b6f404 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Vouillon?= <jerome.vouillon@gmail.com>
Date: Wed, 19 Jun 2024 10:52:07 +0200
Subject: [PATCH] Expose some functions from Freevars

---
 compiler/lib/freevars.ml  | 16 ++++++++++++----
 compiler/lib/freevars.mli |  6 ++++++
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/compiler/lib/freevars.ml b/compiler/lib/freevars.ml
index d65a54c64c..17414f779f 100644
--- a/compiler/lib/freevars.ml
+++ b/compiler/lib/freevars.ml
@@ -100,8 +100,8 @@ type st =
   ; mutable revisited : bool
   }
 
-let find_loops p =
-  let in_loop = ref Addr.Map.empty in
+let find_loops p in_loop pc =
+  let in_loop = ref in_loop in
   let index = ref 0 in
   let state = ref Addr.Map.empty in
   let stack = Stack.create () in
@@ -141,9 +141,17 @@ let find_loops p =
       if st.revisited
       then List.iter !l ~f:(fun pc' -> in_loop := Addr.Map.add pc' pc !in_loop))
   in
-  Code.fold_closures p (fun _ _ (pc, _) () -> traverse pc) ();
+  traverse pc;
   !in_loop
 
+let find_loops_in_closure p pc = find_loops p Addr.Map.empty pc
+
+let find_all_loops p =
+  Code.fold_closures
+    p
+    (fun _ _ (pc, _) (in_loop : _ Addr.Map.t) -> find_loops p in_loop pc)
+    Addr.Map.empty
+
 let mark_variables in_loop p =
   let vars = Var.Tbl.make () (-1) in
   let visited = BitSet.create' p.free_pc in
@@ -245,7 +253,7 @@ let f p =
 let f_mutable p =
   Code.invariant p;
   let t = Timer.make () in
-  let in_loop = find_loops p in
+  let in_loop = find_all_loops p in
   let vars = mark_variables in_loop p in
   let free_vars = free_variables vars in_loop p in
   if times () then Format.eprintf "  free vars 1: %a@." Timer.print t;
diff --git a/compiler/lib/freevars.mli b/compiler/lib/freevars.mli
index f30723923f..ef07c7540e 100644
--- a/compiler/lib/freevars.mli
+++ b/compiler/lib/freevars.mli
@@ -23,6 +23,12 @@ val iter_block_free_vars : (Code.Var.t -> unit) -> Code.block -> unit
 
 val iter_block_bound_vars : (Code.Var.t -> unit) -> Code.block -> unit
 
+val iter_instr_free_vars : (Code.Var.t -> unit) -> Code.instr -> unit
+
+val iter_last_free_var : (Code.Var.t -> unit) -> Code.last -> unit
+
+val find_loops_in_closure : Code.program -> Code.Addr.t -> Code.Addr.t Code.Addr.Map.t
+
 val f_mutable : Code.program -> Code.Var.Set.t Code.Addr.Map.t
 
 val f : Code.program -> Code.Var.Set.t Code.Addr.Map.t