From edc4e23a2ee455f35e2a77007efd2b963a138b07 Mon Sep 17 00:00:00 2001
From: agle <a.michael1@uq.edu.au>
Date: Fri, 27 Mar 2026 17:04:38 +1000
Subject: [PATCH 01/11] list

---
 lib/analysis/irreducible_loops.ml | 13 ++++++++-----
 lib/lang/block.ml                 | 20 ++++++++++++++++++++
 lib/lang/procedure.ml             | 27 +++++++++++++++++++++------
 lib/transforms/dune               |  1 +
 lib/util/types.ml                 |  2 ++
 5 files changed, 52 insertions(+), 11 deletions(-)

diff --git a/lib/analysis/irreducible_loops.ml b/lib/analysis/irreducible_loops.ml
index 49b7ea0d..47b42075 100644
--- a/lib/analysis/irreducible_loops.ml
+++ b/lib/analysis/irreducible_loops.ml
@@ -50,10 +50,10 @@ module Make (G : GSig) = struct
             (** the set of blocks which are internal to the loop. this set forms
                 a strongly-connected component. note that a block may be
                 internal to multiple loops *)
-        entries : G.E.t Iter.t Lazy.t;
+        entries : G.E.t list Lazy.t;
             (** persistent iterator over edges which enter any header in this
                 loop *)
-        backedges : G.E.t Iter.t Lazy.t;
+        backedges : G.E.t list Lazy.t;
             (** persistent iterator of back-edges to any header of this loop *)
       }  (** Designated primary header of a loop. *)
     | LoopParticipant of {
@@ -105,7 +105,7 @@ module Make (G : GSig) = struct
              |> flat_map (fun h ->
                  diff (G.pred p h |> List.to_iter) (VSet.to_iter nodes))
              |> map (fun src -> G.find_edge p src block))
-          |> Iter.persistent)
+          |> Iter.to_list)
 
       (** Accesses the Basil IR state to compute the set of back-edges. That is,
           the set of edges originating from _inside_ the loop and going towards
@@ -117,7 +117,7 @@ module Make (G : GSig) = struct
              |> flat_map (fun h ->
                  inter (G.pred p h |> List.to_iter) (VSet.to_iter nodes))
              |> map (fun src -> G.find_edge p src block))
-          |> Iter.persistent)
+          |> Iter.to_list)
     end
 
     type block_loop_state = {
@@ -417,6 +417,8 @@ module Make (G : GSig) = struct
              Printf.sprintf "%s -> %s" (G.V.show k) (show_block_loop_state v)))
   end
 
+  (** Returns a list representing the loop forest sorted in reverse-topological
+      order. *)
   let solve g entry =
     let open Implementation in
     let st = create g in
@@ -467,7 +469,8 @@ module ProcIntra = struct
   include Make (BlockGraph)
 
   (** Perform irreducible loop analysis and return a list containing the loop
-      information label for each block in the procedure . *)
+      information label for each block in the procedure. Returns all loop tags
+      for blocks in reverse-topological order. *)
   let solve_proc p =
     Procedure.get_entry_block p
     |> Option.flat_map_l (fun entry -> solve p entry)
diff --git a/lib/lang/block.ml b/lib/lang/block.ml
index 774031a6..42fe68a8 100644
--- a/lib/lang/block.ml
+++ b/lib/lang/block.ml
@@ -103,6 +103,26 @@ let map_fold_forwards ~(phi : 'acc -> 'v phi list -> 'acc * 'v phi list)
 let map ~phi f (b : ('v, 'e) t) : ('vv, 'ee) t =
   { stmts = Vector.map f b.stmts; phis = phi b.phis }
 
+(** Modify stmt list by creating a mutable copy of the underlying vector *)
+let fmap_stmts_copy (f : (('a, 'a, 'b) Stmt.t, 'c) Vector.t -> unit) b =
+  let v = Vector.copy b.stmts in
+  f v;
+  { b with stmts = Vector.freeze v }
+
+(** prepend statements to block statement list (copies underlying vector) *)
+let prepend_stmts (b : ('v, 'e) t) (nstmts : ('v, 'v, 'e) Stmt.t list) :
+    ('v, 'e) t =
+  let stmts = Vector.create () in
+  Vector.append_list stmts nstmts;
+  Vector.append stmts b.stmts;
+  let stmts = Vector.freeze stmts in
+  { b with stmts }
+
+(** append statements to block statement list (copies underlying vector) *)
+let append_stmts (b : ('v, 'e) t) (stmt : ('v, 'v, 'e) Stmt.t list) :
+    ('vv, 'ee) t =
+  fmap_stmts_copy (fun v -> Vector.append_list v stmt) b
+
 let flat_map ~phi f (b : ('v, 'e) t) : ('vv, 'ee) t =
   {
     stmts =
diff --git a/lib/lang/procedure.ml b/lib/lang/procedure.ml
index 588fb54a..540de80f 100644
--- a/lib/lang/procedure.ml
+++ b/lib/lang/procedure.ml
@@ -338,15 +338,21 @@ let get_entry_block p =
         Some (List.hd id))
   with Not_found -> None
 
+(** Get the block for an id
+
+    @raises {! Not_found} when the block does not exist.
+*)
+let find_block p id =
+  let open Edge in
+  let open G in
+  let g = graph p |> function Some e -> e | _ -> raise Not_found in
+  let _, e, _ = G.find_edge g (Begin id) (End id) in
+  match e with Block b -> b | Jump -> raise Not_found
+
 let get_block p id =
   let open Edge in
   let open G in
-  try
-    graph p
-    |> Option.flat_map (fun g ->
-        let _, e, _ = G.find_edge g (Begin id) (End id) in
-        match e with Block b -> Some b | Jump -> None)
-  with Not_found -> None
+  try Some (find_block p id) with Not_found -> None
 
 let decl_block_exn p name ?(phis = [])
     ~(stmts : ('var, 'var, 'expr) Stmt.t list) ?(successors = []) () =
@@ -365,6 +371,15 @@ let update_block p id (block : (Var.t, BasilExpr.t) Block.t) =
       let g = G.add_edge_e g (Begin id, Block block, End id) in
       g)
 
+let replace_block_succs p id succs =
+  let open Edge in
+  let open G in
+  p
+  |> map_graph (fun g ->
+      let g = G.succ_e g (End id) |> List.fold_left G.remove_edge_e g in
+      let succs = List.map (fun s -> Vert.(Begin id, Jump, End s)) succs in
+      List.fold_left G.add_edge_e g succs)
+
 let replace_edge p id (block : (Var.t, BasilExpr.t) Block.t) =
   update_block p id block
 
diff --git a/lib/transforms/dune b/lib/transforms/dune
index 920b25c5..4f9107a2 100644
--- a/lib/transforms/dune
+++ b/lib/transforms/dune
@@ -3,6 +3,7 @@
  (name transforms)
  (flags -w -27)
  (modules
+  irreducible_loop
   ssa
   cf_tx
   livevars
diff --git a/lib/util/types.ml b/lib/util/types.ml
index 2e311705..25617ab2 100644
--- a/lib/util/types.ml
+++ b/lib/util/types.ml
@@ -85,6 +85,8 @@ let mk_adt name (variants : (string * field list) list) =
   Sort
     (name, variants |> List.map (fun (variant, fields) -> { variant; fields }))
 
+let bv_min_width_for_nat n = Bitvector (Z.of_int n |> Z.numbits)
+
 let get_field field_name record : record_field =
   match record with
   | Record fields -> (

From b65986cf01593f49318805e6d96f895ad2da0560 Mon Sep 17 00:00:00 2001
From: agle <a.michael1@uq.edu.au>
Date: Fri, 27 Mar 2026 17:07:49 +1000
Subject: [PATCH 02/11] irred transform

---
 lib/transforms/irreducible_loop.ml | 125 +++++++++++++++++++++++++++++
 1 file changed, 125 insertions(+)
 create mode 100644 lib/transforms/irreducible_loop.ml

diff --git a/lib/transforms/irreducible_loop.ml b/lib/transforms/irreducible_loop.ml
new file mode 100644
index 00000000..6f59ad2f
--- /dev/null
+++ b/lib/transforms/irreducible_loop.ml
@@ -0,0 +1,125 @@
+open Lang
+open Common
+
+type state = {
+  new_header : ID.t;
+  from_variable : Var.t;
+  entry_indexes : int IDMap.t;
+  preceding_indices : int list IDMap.t;
+}
+
+open Analysis.Irreducible_loops.ProcIntra
+
+let transform_loop p l =
+  let open Option in
+  let header, next_h, headers, nodes, entries, backedges =
+    match l with
+    | {
+     block;
+     loop = PrimaryHeader { primary_header; headers; nodes; entries; backedges };
+    } ->
+        ( block,
+          primary_header,
+          headers,
+          nodes,
+          Lazy.force entries,
+          Lazy.force backedges )
+    | _ -> raise (Invalid_argument "called on non-primary header")
+  in
+  let dest = BlockGraph.E.dst in
+  let src = BlockGraph.E.src in
+  let backedges_to_primary_header =
+    List.filter (BlockGraph.E.dst %> ID.equal header) entries
+  in
+  let entry_indexes =
+    entries @ backedges
+    |> List.mapi (fun i v -> (BlockGraph.E.src v, i))
+    |> IDMap.of_list
+  in
+
+  (* included entry blocks should be a superset of (externalEntries ++ backEdgesToFirstHeader).
+     in particular, it additionally includes internal edges to alternative headers.*)
+  assert (
+    Iter.append (List.to_iter entries)
+      (List.to_iter backedges_to_primary_header)
+    |> Iter.map BlockGraph.E.src
+    |> Iter.subset ~eq:ID.equal (IDMap.keys entry_indexes));
+
+  let preceding_indices =
+    entries @ backedges
+    |> List.group_by ~hash:(dest %> ID.hash) ~eq:(fun a b ->
+        ID.equal (dest a) (dest b))
+    |> List.map (function
+      | h :: tl ->
+          ( dest h,
+            h :: tl |> List.map src
+            |> List.sort_uniq ~cmp:(fun a b ->
+                Int.compare
+                  (IDMap.find a entry_indexes)
+                  (IDMap.find a entry_indexes)) )
+      | _ -> failwith "emtpy")
+  in
+  let ctrl_sz = Types.bv_min_width_for_nat (IDSet.cardinal headers) in
+  let loop_crtl_v =
+    Procedure.fresh_var ~pure:true
+      ~name:(ID.to_string header ^ "loop_from")
+      p ctrl_sz
+  in
+  let open Lang.Expr in
+  let bvali idx =
+    let size = match ctrl_sz with Bitvector i -> i | _ -> assert false in
+    BasilExpr.bvconst (Bitvec.of_int ~size idx)
+  in
+  let bval m =
+    let idx = IDMap.find m entry_indexes in
+    bvali idx
+  in
+  (* create new primary header which jumps to all existing headers *)
+  let p, n_header =
+    Procedure.fresh_block
+      ~name:(ID.to_string header ^ "header_loop_N")
+      ~successors:(VSet.to_list headers) p ~stmts:[] ()
+  in
+  ( ( p |> fun p ->
+      (* add guards to old headers to restrict their predecessors to original predecessor set *)
+      List.fold_left
+        (fun p (h, indices) ->
+          let preds =
+            List.map
+              (fun i ->
+                BasilExpr.binexp ~op:`EQ (BasilExpr.rvar loop_crtl_v) (bval i))
+              indices
+          in
+          let e = BasilExpr.applyintrin ~op:`OR preds in
+          let b = Procedure.find_block p h in
+          let b =
+            Block.prepend_stmts b [ Instr_Assume { body = e; branch = false } ]
+          in
+          Procedure.update_block p h b)
+        p preceding_indices )
+  |> fun p ->
+    (* set control var index in predecessors *)
+    VMap.fold
+      (fun bid idx p ->
+        let b = Procedure.find_block p bid in
+        let b =
+          Block.append_stmts b [ Instr_Assign [ (loop_crtl_v, bvali idx) ] ]
+        in
+        Procedure.update_block p bid b)
+      entry_indexes p )
+  |> fun p ->
+  (* make all headers jump to the new primary header *)
+  List.fold_left
+    (fun p (src, dest) -> Procedure.replace_block_succs p src [ n_header ])
+    p
+    (entries @ backedges_to_primary_header)
+
+let transform (p : Program.proc) =
+  (* NOTE: we get the result sorted in reverse-topological order *)
+  let loops =
+    solve_proc p
+    |> List.filter (function
+      | { loop = PrimaryHeader _; _ } -> true
+      | _ -> false)
+  in
+  List.fold_left transform_loop p loops

From 7d90075fd0804bf9efe6df460786cb805396e26d Mon Sep 17 00:00:00 2001
From: agle <a.michael1@uq.edu.au>
Date: Fri, 27 Mar 2026 17:52:58 +1000
Subject: [PATCH 03/11] wip

---
 lib/analysis/irreducible_loops.ml       | 60 +++++++++++++------------
 lib/transforms/irreducible_loop.ml      | 13 +++---
 test/analysis/test_irreducible_loops.ml | 56 ++++++++++++++++++++++-
 3 files changed, 93 insertions(+), 36 deletions(-)

diff --git a/lib/analysis/irreducible_loops.ml b/lib/analysis/irreducible_loops.ml
index 47b42075..60352eca 100644
--- a/lib/analysis/irreducible_loops.ml
+++ b/lib/analysis/irreducible_loops.ml
@@ -38,22 +38,26 @@ module Make (G : GSig) = struct
   module VSet = Set.Make (G.V)
   module VMap = Map.Make (G.V)
 
+  let vset_pp fmt m =
+    Format.pp_print_string fmt
+      (VSet.to_string ~sep:", " ~start:"{" ~stop:"}" G.V.show m)
+
   (** Loop context information of a given block. *)
   type loop_info =
     | PrimaryHeader of {
         primary_header : G.V.t option;
             (** next-outermost header to this loop, if we are inside a loop *)
-        headers : VSet.t;
+        headers : VSet.t; [@printer vset_pp]
             (** the set of blocks which can be used to enter the loop. a header
                 is defined as dominating all non-header nodes of the loop. *)
-        nodes : VSet.t;
+        nodes : VSet.t; [@printer vset_pp]
             (** the set of blocks which are internal to the loop. this set forms
                 a strongly-connected component. note that a block may be
                 internal to multiple loops *)
-        entries : G.E.t list Lazy.t;
+        entries : G.E.t list; [@equal fun a b -> Equal.poly a b]
             (** persistent iterator over edges which enter any header in this
                 loop *)
-        backedges : G.E.t list Lazy.t;
+        backedges : G.E.t list; [@equal fun a b -> Equal.poly a b]
             (** persistent iterator of back-edges to any header of this loop *)
       }  (** Designated primary header of a loop. *)
     | LoopParticipant of {
@@ -64,8 +68,18 @@ module Make (G : GSig) = struct
         (** Member of a loop that is not a primary header (may be a non-primary
             header) *)
     | NonLoop
+  [@@deriving show { with_path = false }, eq]
 
   type block_info = { block : G.V.t; dfs_pos : int; loop : loop_info }
+  [@@deriving eq, show { with_path = false }]
+
+  let classify_block b =
+    match b with
+    | { loop = NonLoop } -> `NonLoop
+    | { loop = LoopParticipant _ } -> `LoopNode
+    | { loop = PrimaryHeader { headers } } when VSet.cardinal headers = 1 ->
+        `ReducibleHeader
+    | { loop = PrimaryHeader { headers } } -> `IrreducibleHeader
 
   module Implementation = struct
     (** Internal implementation of irreducible loop forest algorithm.
@@ -99,25 +113,23 @@ module Make (G : GSig) = struct
           originating from outside the loop and going towards *any* header of
           the loop. *)
       let compute_entries p block headers nodes =
-        lazy
-          (Iter.(
-             VSet.to_iter headers
-             |> flat_map (fun h ->
-                 diff (G.pred p h |> List.to_iter) (VSet.to_iter nodes))
-             |> map (fun src -> G.find_edge p src block))
-          |> Iter.to_list)
+        Iter.(
+          VSet.to_iter headers
+          |> flat_map (fun h ->
+              diff (G.pred p h |> List.to_iter) (VSet.to_iter nodes))
+          |> map (fun src -> G.find_edge p src block))
+        |> Iter.to_list
 
       (** Accesses the Basil IR state to compute the set of back-edges. That is,
           the set of edges originating from _inside_ the loop and going towards
           *any* header of the loop. *)
       let compute_backedges p block headers nodes =
-        lazy
-          (Iter.(
-             VSet.to_iter headers
-             |> flat_map (fun h ->
-                 inter (G.pred p h |> List.to_iter) (VSet.to_iter nodes))
-             |> map (fun src -> G.find_edge p src block))
-          |> Iter.to_list)
+        Iter.(
+          VSet.to_iter headers
+          |> flat_map (fun h ->
+              inter (G.pred p h |> List.to_iter) (VSet.to_iter nodes))
+          |> map (fun src -> G.find_edge p src block))
+        |> Iter.to_list
     end
 
     type block_loop_state = {
@@ -133,11 +145,7 @@ module Make (G : GSig) = struct
               maintained even after visiting is finished. *)
       mutable is_traversed : bool;
           (** whether visiting this block has _started_. *)
-      mutable headers : VSet.t;
-          [@printer
-            fun fmt m ->
-              Format.pp_print_string fmt
-                (VSet.to_string ~sep:", " ~start:"{" ~stop:"}" G.V.show m)]
+      mutable headers : VSet.t; [@printer vset_pp]
           (** headers of the loop headed by this block. if this block heads a
               loop, this always contains `b` as the primary header. for
               irreducible loops, it also contains secondary headers. *)
@@ -284,10 +292,8 @@ module Make (G : GSig) = struct
     (** Sets [h] as the loop header for the block `b` and all containing loops.
     *)
     let tag_lhead s (b : block_loop_state) (h : block_loop_state option) =
-      print_endline "tag lhead";
       match h with
       | Some h when not @@ G.V.equal b.block h.block -> begin
-          print_endline (show_block_loop_state h);
           let rec loop ~(cur1 : block_loop_state) ~(cur2 : block_loop_state) =
             match cur1.iloop_header with
             | Some ih when G.V.equal ih cur2.block -> ()
@@ -368,10 +374,6 @@ module Make (G : GSig) = struct
         (* iterate the remaining blocks *)
         iter_tails
           (fun it b ->
-            (*print_endline
-            (Printf.sprintf "%s %d"
-               (ID.to_string (b : block_loop_state).block)
-               (List.length it));*)
             match b with
             | { is_traversed = false } ->
                 raise
diff --git a/lib/transforms/irreducible_loop.ml b/lib/transforms/irreducible_loop.ml
index 6f59ad2f..486e63d9 100644
--- a/lib/transforms/irreducible_loop.ml
+++ b/lib/transforms/irreducible_loop.ml
@@ -17,14 +17,18 @@ let transform_loop p l =
     | {
      block;
      loop = PrimaryHeader { primary_header; headers; nodes; entries; backedges };
-    } ->
+    }
+      when VSet.cardinal headers > 1 ->
         ( block,
           primary_header,
           headers,
           nodes,
           Lazy.force entries,
           Lazy.force backedges )
-    | _ -> raise (Invalid_argument "called on non-primary header")
+    | _ ->
+        raise
+          (Invalid_argument
+             "called on non-primary header / non-irreducible loop ")
   in
   let dest = BlockGraph.E.dst in
   let src = BlockGraph.E.src in
@@ -118,8 +122,7 @@ let transform (p : Program.proc) =
   (* NOTE: we get the result sorted in reverse-topological order *)
   let loops =
     solve_proc p
-    |> List.filter (function
-      | { loop = PrimaryHeader _; _ } -> true
-      | _ -> false)
+    |> List.filter (fun b ->
+        match classify_block b with `IrreducibleHeader -> true | _ -> false)
   in
   List.fold_left transform_loop p loops
diff --git a/test/analysis/test_irreducible_loops.ml b/test/analysis/test_irreducible_loops.ml
index f496a15b..e386b91a 100644
--- a/test/analysis/test_irreducible_loops.ml
+++ b/test/analysis/test_irreducible_loops.ml
@@ -12,6 +12,8 @@ open struct
   }
   [@@deriving eq, show]
 
+  let block_info = Alcotest.testable pp_block_info equal_block_info
+
   let id_map equal str =
     Alcotest.testable
       (fun f p ->
@@ -63,6 +65,16 @@ open struct
     let checked = { iloop_headers = headers; headers = members } in
     check_test_comparison expect checked
 
+  let run_transform prog =
+    let p = (Loader.Loadir.ast_of_string prog).prog in
+    let p =
+      IDMap.find (Option.get_exn_or "no entry proc" p.entry_proc) p.procs
+    in
+    let before = solve_proc p in
+    let p' = Transforms.Irreducible_loop.transform p in
+    let after = solve_proc p' in
+    (before, after)
+
   let check_loop_result name prog ~header_ptrs ~all_loop_headers =
     let p = (Loader.Loadir.ast_of_string prog).prog in
     let p =
@@ -186,7 +198,6 @@ proc @main () -> ()
 prog entry @main;
 proc @main () -> ()
 [
- 
   block %S [ goto(%loop); ]; 
   block %loop [ goto(%loop2); ]; 
   block %loop2 [ goto(%loop3); ]; 
@@ -211,7 +222,6 @@ proc @main () -> ()
 prog entry @main;
 proc @main () -> ()
 [
- 
   block %S [ goto(%loop); ]; 
   block %loop [ goto(%loop2); ]; 
   block %loop2 [ goto(%loop3, %loop2); ]; 
@@ -226,6 +236,47 @@ proc @main () -> ()
       [ ("%loop", [ "%loop" ]); ("%loop2", [ "%loop2" ]) ]
     in
     check_loop_result name p ~header_ptrs ~all_loop_headers
+
+  let loops_irreducible p =
+    List.filter
+      (fun b ->
+        match classify_block b with `IrreducibleHeader -> true | _ -> false)
+      p
+
+  let irrloops_fixed after =
+    Alcotest.(check (list block_info))
+      "all irreducible loops fixed" [] (loops_irreducible after)
+
+  let sub_cycles_transform =
+    let p =
+      {|
+prog entry @main;                           
+                                            
+proc @main () -> ()                         
+  { .name = "main"; .returnBlock = "exit" } 
+[                                           
+  block %S [                                
+    goto(%h1, %h2);                         
+  ];                                        
+  block %h1 [                               
+    goto(%h2);                              
+  ];                                        
+  block %h2 [                               
+    goto(%h1, %h3);                         
+  ];                                        
+  block %h3 [                               
+    goto(%h2, %exit);                       
+  ];                                        
+  block %exit [                             
+    return ();                              
+  ]                                         
+];                                          
+                                            
+    |}
+    in
+    let _, after = run_transform p in
+    Alcotest.test_case "sub cycles transform" `Quick (fun () ->
+        irrloops_fixed after)
 end
 
 let tests =
@@ -238,5 +289,6 @@ let tests =
         one_long_loop;
         nested_loop;
         nested_self_loop;
+        sub_cycles_transform;
       ] );
   ]

From 44c47dd927a7b0805756293c307c5af1a24fa493 Mon Sep 17 00:00:00 2001
From: agle <a.michael1@uq.edu.au>
Date: Fri, 27 Mar 2026 20:42:47 +1000
Subject: [PATCH 04/11] fix

---
 lib/analysis/irreducible_loops.ml  | 16 ++++++++++------
 lib/transforms/irreducible_loop.ml |  4 ++--
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/lib/analysis/irreducible_loops.ml b/lib/analysis/irreducible_loops.ml
index 60352eca..2bc30001 100644
--- a/lib/analysis/irreducible_loops.ml
+++ b/lib/analysis/irreducible_loops.ml
@@ -116,9 +116,11 @@ module Make (G : GSig) = struct
         Iter.(
           VSet.to_iter headers
           |> flat_map (fun h ->
-              diff (G.pred p h |> List.to_iter) (VSet.to_iter nodes))
-          |> map (fun src -> G.find_edge p src block))
-        |> Iter.to_list
+              G.pred p h |> List.to_iter 
+              |> map (fun p -> (p, h))
+              |> filter (fun (s,_) -> not @@ VSet.mem s nodes))
+          |> map (fun (a,b) -> G.find_edge p a b)
+          |> to_list)
 
       (** Accesses the Basil IR state to compute the set of back-edges. That is,
           the set of edges originating from _inside_ the loop and going towards
@@ -127,9 +129,11 @@ module Make (G : GSig) = struct
         Iter.(
           VSet.to_iter headers
           |> flat_map (fun h ->
-              inter (G.pred p h |> List.to_iter) (VSet.to_iter nodes))
-          |> map (fun src -> G.find_edge p src block))
-        |> Iter.to_list
+              G.pred p h |> List.to_iter 
+              |> map (fun p -> (p, h))
+              |> filter (fun (s,_) -> VSet.mem s nodes))
+          |> map (fun (a,b) -> G.find_edge p a b)
+        |> to_list)
     end
 
     type block_loop_state = {
diff --git a/lib/transforms/irreducible_loop.ml b/lib/transforms/irreducible_loop.ml
index 486e63d9..9a09b829 100644
--- a/lib/transforms/irreducible_loop.ml
+++ b/lib/transforms/irreducible_loop.ml
@@ -23,8 +23,8 @@ let transform_loop p l =
           primary_header,
           headers,
           nodes,
-          Lazy.force entries,
-          Lazy.force backedges )
+          entries,
+          backedges )
     | _ ->
         raise
           (Invalid_argument

From 89ab629068169e2f35071e904d63e572e3ec109e Mon Sep 17 00:00:00 2001
From: agle <a.michael1@uq.edu.au>
Date: Mon, 30 Mar 2026 12:23:59 +1000
Subject: [PATCH 05/11] check in irred loop transform tests

---
 lib/analysis/irreducible_loops.ml       |  14 +-
 test/analysis/test_irreducible_loops.ml | 394 +++++++++++++++++++++++-
 2 files changed, 388 insertions(+), 20 deletions(-)

diff --git a/lib/analysis/irreducible_loops.ml b/lib/analysis/irreducible_loops.ml
index 2bc30001..d1c6e5c3 100644
--- a/lib/analysis/irreducible_loops.ml
+++ b/lib/analysis/irreducible_loops.ml
@@ -116,10 +116,10 @@ module Make (G : GSig) = struct
         Iter.(
           VSet.to_iter headers
           |> flat_map (fun h ->
-              G.pred p h |> List.to_iter 
+              G.pred p h |> List.to_iter
               |> map (fun p -> (p, h))
-              |> filter (fun (s,_) -> not @@ VSet.mem s nodes))
-          |> map (fun (a,b) -> G.find_edge p a b)
+              |> filter (fun (s, _) -> not @@ VSet.mem s nodes))
+          |> map (fun (a, b) -> G.find_edge p a b)
           |> to_list)
 
       (** Accesses the Basil IR state to compute the set of back-edges. That is,
@@ -129,11 +129,11 @@ module Make (G : GSig) = struct
         Iter.(
           VSet.to_iter headers
           |> flat_map (fun h ->
-              G.pred p h |> List.to_iter 
+              G.pred p h |> List.to_iter
               |> map (fun p -> (p, h))
-              |> filter (fun (s,_) -> VSet.mem s nodes))
-          |> map (fun (a,b) -> G.find_edge p a b)
-        |> to_list)
+              |> filter (fun (s, _) -> VSet.mem s nodes))
+          |> map (fun (a, b) -> G.find_edge p a b)
+          |> to_list)
     end
 
     type block_loop_state = {
diff --git a/test/analysis/test_irreducible_loops.ml b/test/analysis/test_irreducible_loops.ml
index e386b91a..87f7c12c 100644
--- a/test/analysis/test_irreducible_loops.ml
+++ b/test/analysis/test_irreducible_loops.ml
@@ -1,8 +1,12 @@
 open Bincaml_util.Common
 open Analysis.Irreducible_loops.ProcIntra
 
+(** Tests for irreducible loop forest analysis and transform. Here "paper"
+    refers to T. Wei et al. {:http://dx.doi.org/10.1007/978-3-540-74061-2_11}.
+*)
+
 open struct
-  (** Put all the implementation in a hidden struct so not exported and we get
+  (* Put all the implementation in a hidden struct so not exported and we get
       unused function warnings if we define a test and dont add it to the suite
   *)
 
@@ -31,8 +35,7 @@ open struct
            (StringSet.to_string ~stop:"}" ~start:"{" Fun.id)))
       "loop header->participant sets equal" a.headers b.headers
 
-  let assert_loop_detector p iloop_headers headers =
-    let loops = solve_proc p in
+  let assert_loop_detector loops iloop_headers headers =
     let headers =
       List.map (Pair.map Fun.id StringSet.of_list) headers |> StringMap.of_list
     in
@@ -80,7 +83,11 @@ open struct
     let p =
       IDMap.find (Option.get_exn_or "no entry proc" p.entry_proc) p.procs
     in
-    let c = fun () -> assert_loop_detector p header_ptrs all_loop_headers in
+    let c =
+     fun () ->
+      let loops = solve_proc p in
+      assert_loop_detector loops header_ptrs all_loop_headers
+    in
     Alcotest.test_case name `Quick c
 
   let paper_fig2 =
@@ -237,18 +244,47 @@ proc @main () -> ()
     in
     check_loop_result name p ~header_ptrs ~all_loop_headers
 
+  let loops_reducible p =
+    List.filter
+      (fun b ->
+        match classify_block b with
+        | `IrreducibleHeader -> false
+        | `ReducibleHeader -> true
+        | _ -> false)
+      p
+
   let loops_irreducible p =
     List.filter
       (fun b ->
         match classify_block b with `IrreducibleHeader -> true | _ -> false)
       p
 
-  let irrloops_fixed after =
-    Alcotest.(check (list block_info))
-      "all irreducible loops fixed" [] (loops_irreducible after)
+  let check_transform_fixed name ~num_irr_loops ~num_red_loops ?header_ptrs
+      ?all_headers p =
+    let checks () =
+      let before, after = run_transform p in
+
+      let check_x =
+        let open Option in
+        let* hdrs = header_ptrs in
+        let* headers = all_headers in
+        Some (fun () -> assert_loop_detector before hdrs headers)
+      in
+      Option.iter (fun x -> x ()) check_x;
+      Alcotest.(check int)
+        "number of irreducible loops present" num_irr_loops
+        (List.length @@ loops_irreducible before);
+      Alcotest.(check int)
+        "number of reducible loops present" num_red_loops
+        (List.length @@ loops_reducible before);
+      Alcotest.(check (list block_info))
+        "all irreducible loops fixed" [] (loops_irreducible after)
+    in
+    Alcotest.test_case name `Quick checks
 
   let sub_cycles_transform =
-    let p =
+    check_transform_fixed "subcycles applying transform" ~num_irr_loops:1
+      ~num_red_loops:1
       {|
 prog entry @main;                           
                                             
@@ -271,14 +307,337 @@ proc @main () -> ()
     return ();                              
   ]                                         
 ];                                          
-                                            
     |}
-    in
-    let _, after = run_transform p in
-    Alcotest.test_case "sub cycles transform" `Quick (fun () ->
-        irrloops_fixed after)
+
+  let crossover =
+    check_transform_fixed "crossover" ~num_irr_loops:2 ~num_red_loops:0
+      {|
+prog entry @main;
+proc @main () -> ()
+  { .name = "main"; .returnBlock = "exit" }
+[
+  block %S [
+    goto(%h1, %h2);
+  ];
+  block %h1 [
+    goto(%x);
+  ];
+  block %x [
+    goto(%h2, %h1);
+  ];
+  block %h2 [
+    goto(%y);
+  ];
+  block %y [
+    goto(%x, %exit);
+  ];
+  block %exit [
+    return ();
+  ]
+];
+    |}
+
+  let paper_fig4a =
+    check_transform_fixed "paper fig4a" ~num_irr_loops:0 ~num_red_loops:0
+      {|
+prog entry @main;
+
+proc @main () -> ()
+  { .name = "main"; .returnBlock = "exit" }
+[
+  block %S [
+    goto(%1);
+  ];
+  block %1 [
+    goto(%2);
+  ];
+  block %2 [
+    goto(%b0);
+  ];
+  block %b0 [
+    goto(%b);
+  ];
+  block %b [
+    goto(%exit);
+  ];
+  block %exit [
+    return ();
+  ]
+];
+
+    |}
+
+  let paper_fig4b =
+    check_transform_fixed "paper fig4b" ~num_irr_loops:0 ~num_red_loops:1
+      ~header_ptrs:[ ("%b0", "%b"); ("%x", "%b") ]
+      ~all_headers:[ ("%b", [ "%b" ]) ]
+      {|
+prog entry @main;
+
+proc @main () -> ()
+  { .name = "main"; .returnBlock = "exit" }
+[
+  block %S [
+    goto(%1);
+  ];
+  block %1 [
+    goto(%b);
+  ];
+  block %b [
+    goto(%x);
+  ];
+  block %x [
+    goto(%b0);
+  ];
+  block %b0 [
+    goto(%exit, %b);
+  ];
+  block %exit [
+    return ();
+  ]
+];
+    |}
+
+  let paper_fig4c =
+    (* TODO: sus there are no loops before *)
+    check_transform_fixed "paper fig4c" ~num_irr_loops:0 ~num_red_loops:0
+      {|
+prog entry @main;
+
+proc @main () -> ()
+  { .name = "main"; .returnBlock = "exit" }
+[
+  block %S [
+    goto(%1);
+  ];
+  block %1 [
+    goto(%h);
+  ];
+  block %h [
+    goto(%x, %b0);
+  ];
+  block %x [
+    goto(%b);
+  ];
+  block %b0 [
+    goto(%b);
+  ];
+  block %b [
+    goto(%z);
+  ];
+  block %z [
+    goto(%exit);
+  ];
+  block %exit [
+    return ();
+  ]
+];
+  |}
+
+  let paper_fig4d =
+    check_transform_fixed "paper fig4d" ~num_irr_loops:0 ~num_red_loops:1
+      ~header_ptrs:
+        [
+          ("%b", "%h"); ("%b0", "%h"); ("%x", "%h"); ("%y", "%h"); ("%z", "%h");
+        ]
+      ~all_headers:[ ("%h", [ "%h" ]) ]
+      {|
+prog entry @main;
+
+proc @main () -> ()
+  { .name = "main"; .returnBlock = "exit" }
+[
+  block %S [
+    goto(%1);
+  ];
+  block %1 [
+    goto(%h);
+  ];
+  block %h [
+    goto(%x);
+  ];
+  block %x [
+    goto(%b0, %y);
+  ];
+  block %y [
+    goto(%b);
+  ];
+  block %b0 [
+    goto(%b);
+  ];
+  block %b [
+    goto(%z);
+  ];
+  block %z [
+    goto(%exit, %h);
+  ];
+  block %exit [
+    return ();
+  ]
+];
+
+    |}
+
+  let paper_fig6a =
+    (* TODO: wrong; does scala impl identifies 2 irreducible loops *)
+    check_transform_fixed "paper fig6a" ~num_irr_loops:2 ~num_red_loops:1
+      ~header_ptrs:[] ~all_headers:[]
+      {|
+prog entry @main;
+
+proc @main () -> ()
+  { .name = "main"; .returnBlock = "exit" }
+[
+  block %S [
+    goto(%h3);
+  ];
+  block %h3 [
+    goto(%x);
+  ];
+  block %x [
+    goto(%h2, %y);
+  ];
+  block %y [
+    goto(%b0);
+  ];
+  block %b0 [
+    goto(%b);
+  ];
+  block %h2 [
+    goto(%h1);
+  ];
+  block %h1 [
+    goto(%b);
+  ];
+  block %b [
+    goto(%z);
+  ];
+  block %z [
+    goto(%h1, %a);
+  ];
+  block %a [
+    goto(%h2, %back);
+  ];
+  block %back [
+    goto(%h3, %exit);
+  ];
+  block %exit [
+    return ();
+  ]
+];
+
+    |}
+
+  let paper_fig6b =
+    (* FIXME: shouldn't we identify 4 reducible loops? *)
+    check_transform_fixed "paper fig6b" ~num_irr_loops:0 ~num_red_loops:1
+      {|
+prog entry @main;
+
+proc @main () -> ()
+  { .name = "main"; .returnBlock = "exit" }
+[
+  block %S [
+    goto(%h4);
+  ];
+  block %h4 [
+    goto(%h3);
+  ];
+  block %h3 [
+    goto(%h2);
+  ];
+  block %h2 [
+    goto(%h1);
+  ];
+  block %h1 [
+    goto(%x);
+  ];
+  block %x [
+    goto(%y, %h4);
+  ];
+  block %y [
+    goto(%z, %h3);
+  ];
+  block %z [
+    goto(%back, %h2);
+  ];
+  block %back [
+    goto(%h1, %exit);
+  ];
+  block %exit [
+    return ();
+  ]
+];
+|}
+
+  let paper_fig4e =
+    check_transform_fixed "paper fig4e" ~num_irr_loops:1 ~num_red_loops:1
+      ~header_ptrs:[] ~all_headers:[]
+      {|
+prog entry @main;
+
+proc @main () -> ()
+  { .name = "main"; .returnBlock = "exit" }
+[
+  block %S [
+    goto(%h1);
+  ];
+  block %h1 [
+    goto(%y, %z);
+  ];
+  block %y [
+    goto(%h);
+  ];
+  block %h [
+    goto(%b);
+  ];
+  block %z [
+    goto(%b0);
+  ];
+  block %b0 [
+    goto(%b);
+  ];
+  block %b [
+    goto(%a);
+  ];
+  block %a [
+    goto(%h, %h1, %exit);
+  ];
+  block %exit [
+    return ();
+  ]
+];
+    |}
 end
 
+let triforce =
+  check_transform_fixed "triforce" ~num_irr_loops:2 ~num_red_loops:1
+    {|
+prog entry @main;
+
+proc @main () -> ()
+  { .name = "main"; .returnBlock = "exit" }
+[
+  block %S [
+    goto(%h1, %h2, %h3);
+  ];
+  block %h1 [
+    goto(%h1, %h2, %h3);
+  ];
+  block %h2 [
+    goto(%h1, %h2, %h3);
+  ];
+  block %h3 [
+    goto(%h1, %h2, %h3, %exit);
+  ];
+  block %exit [
+    return ();
+  ]
+];
+
+
+  |}
+
 let tests =
   [
     ( "loop identification",
@@ -290,5 +649,14 @@ let tests =
         nested_loop;
         nested_self_loop;
         sub_cycles_transform;
+        crossover;
+        paper_fig4a;
+        paper_fig4b;
+        paper_fig4c;
+        paper_fig4d;
+        paper_fig6a;
+        paper_fig6b;
+        paper_fig4e;
+        triforce;
       ] );
   ]

From 3929877e33b8f73ae388f5daed4b5755bcf0a658 Mon Sep 17 00:00:00 2001
From: agle <a.michael1@uq.edu.au>
Date: Mon, 30 Mar 2026 16:23:30 +1000
Subject: [PATCH 06/11] irred transform and more fixes

---
 lib/analysis/irreducible_loops.ml       |  97 +++++++++++-----------
 lib/lang/procedure.ml                   |  20 ++++-
 lib/passes.ml                           |   8 ++
 lib/transforms/irreducible_loop.ml      |  37 +++------
 test/analysis/dune                      |   2 +
 test/analysis/test_irreducible_loops.ml | 104 +++++++++++++++---------
 6 files changed, 157 insertions(+), 111 deletions(-)

diff --git a/lib/analysis/irreducible_loops.ml b/lib/analysis/irreducible_loops.ml
index d1c6e5c3..7e9eed47 100644
--- a/lib/analysis/irreducible_loops.ml
+++ b/lib/analysis/irreducible_loops.ml
@@ -54,11 +54,6 @@ module Make (G : GSig) = struct
             (** the set of blocks which are internal to the loop. this set forms
                 a strongly-connected component. note that a block may be
                 internal to multiple loops *)
-        entries : G.E.t list; [@equal fun a b -> Equal.poly a b]
-            (** persistent iterator over edges which enter any header in this
-                loop *)
-        backedges : G.E.t list; [@equal fun a b -> Equal.poly a b]
-            (** persistent iterator of back-edges to any header of this loop *)
       }  (** Designated primary header of a loop. *)
     | LoopParticipant of {
         primary_header : G.V.t;
@@ -81,6 +76,33 @@ module Make (G : GSig) = struct
         `ReducibleHeader
     | { loop = PrimaryHeader { headers } } -> `IrreducibleHeader
 
+  (** Accesses the Basil IR state to compute the set of entry edges originating
+      from outside the loop and going towards *any* header of the loop. *)
+  let compute_entries p = function
+    | { block; loop = PrimaryHeader { headers; nodes } } ->
+        Iter.(
+          VSet.to_iter headers
+          |> flat_map (fun h ->
+              G.pred p h |> List.to_iter
+              |> filter (fun a -> not @@ VSet.mem a nodes)
+              |> map (fun a -> G.find_edge p a h))
+          |> to_list)
+    | _ -> []
+
+  (** Accesses the Basil IR state to compute the set of back-edges. That is, the
+      set of edges originating from _inside_ the loop and going towards *any*
+      header of the loop. *)
+  let compute_backedges p = function
+    | { block; loop = PrimaryHeader { headers; nodes } } ->
+        Iter.(
+          VSet.to_iter headers
+          |> flat_map (fun h ->
+              G.pred p h |> List.to_iter
+              |> filter (fun a -> VSet.mem a nodes)
+              |> map (fun a -> G.find_edge p a h))
+          |> to_list)
+    | _ -> []
+
   module Implementation = struct
     (** Internal implementation of irreducible loop forest algorithm.
 
@@ -108,34 +130,6 @@ module Make (G : GSig) = struct
         - Loop: {:https://llvm.org/docs/LoopTerminology.html}
         - Cycle: {:https://llvm.org/docs/CycleTerminology.html} *)
 
-    open struct
-      (** Accesses the Basil IR state to compute the set of entry edges
-          originating from outside the loop and going towards *any* header of
-          the loop. *)
-      let compute_entries p block headers nodes =
-        Iter.(
-          VSet.to_iter headers
-          |> flat_map (fun h ->
-              G.pred p h |> List.to_iter
-              |> map (fun p -> (p, h))
-              |> filter (fun (s, _) -> not @@ VSet.mem s nodes))
-          |> map (fun (a, b) -> G.find_edge p a b)
-          |> to_list)
-
-      (** Accesses the Basil IR state to compute the set of back-edges. That is,
-          the set of edges originating from _inside_ the loop and going towards
-          *any* header of the loop. *)
-      let compute_backedges p block headers nodes =
-        Iter.(
-          VSet.to_iter headers
-          |> flat_map (fun h ->
-              G.pred p h |> List.to_iter
-              |> map (fun p -> (p, h))
-              |> filter (fun (s, _) -> VSet.mem s nodes))
-          |> map (fun (a, b) -> G.find_edge p a b)
-          |> to_list)
-    end
-
     type block_loop_state = {
       block : G.V.t;  (** the block which this loop information concerns. *)
       mutable iloop_header : G.V.t option;
@@ -167,13 +161,7 @@ module Make (G : GSig) = struct
       let loop =
         if is_loop_header e then
           PrimaryHeader
-            {
-              primary_header = e.iloop_header;
-              headers = e.headers;
-              nodes;
-              entries = compute_entries p e.block e.headers nodes;
-              backedges = compute_backedges p e.block e.headers nodes;
-            }
+            { primary_header = e.iloop_header; headers = e.headers; nodes }
         else if is_loop_participant e then
           LoopParticipant
             {
@@ -183,6 +171,7 @@ module Make (G : GSig) = struct
       in
       { block = e.block; dfs_pos = e.dfsp_pos_max; loop }
 
+    (** return loop forest in {b reverse} topological order.*)
     let compute_block_loop_info p (block_states : block_loop_state VMap.t) :
         block_info list =
       let open Iter in
@@ -197,12 +186,14 @@ module Make (G : GSig) = struct
           all_blocks
         |> Iter.persistent
       in
+      (* 1 *)
       let (forest : VSet.t VMap.t) =
         header_blocks
         |> map (fun b -> (b.block, VSet.singleton b.block))
         |> VMap.of_iter
       in
 
+      (* 2 *)
       (* NOTE: iterates the forest in *bottom-up* topological order. this
        ensures that node-sets of sub-cycles are fully populated before
        processing their parent cycle. this avoids us having to compute
@@ -221,6 +212,7 @@ module Make (G : GSig) = struct
                      forest))
              forest
       in
+      (* 3 *)
       let forest =
         header_blocks
         |> fold
@@ -234,6 +226,7 @@ module Make (G : GSig) = struct
                      forest)
              forest
       in
+      (* 4 *)
       header_blocks
       |> iter (fun b ->
           let nodes = VMap.find b.block forest in
@@ -249,6 +242,7 @@ module Make (G : GSig) = struct
             ^ " bad headers: "
             ^ VSet.to_string G.V.show bad_headers);
 
+      (* 5 *)
       let new_loops =
         all_blocks
         |> map (fun b ->
@@ -373,14 +367,14 @@ module Make (G : GSig) = struct
             end
           | Return nh, [] ->
               (* this is the outermost call : we are done*)
-              raise (Recurse (Return nh, []))
+              raise_notrace (Recurse (Return nh, []))
         in
         (* iterate the remaining blocks *)
         iter_tails
           (fun it b ->
             match b with
             | { is_traversed = false } ->
-                raise
+                raise_notrace
                   (Recurse
                      ( Call { block = b; dfsp_pos = dfsp_pos + 1 },
                        ({ block = b0; dfsp_pos }, it) :: continuations ))
@@ -410,7 +404,7 @@ module Make (G : GSig) = struct
           it;
         b0.dfsp_pos <- 0;
         let result = Option.map st.l b0.iloop_header in
-        raise (Recurse (Return result, continuations))
+        raise_notrace (Recurse (Return result, continuations))
       in
       try run () with
       | Recurse (Return c, []) -> c
@@ -420,16 +414,23 @@ module Make (G : GSig) = struct
       print_endline
       @@ (VMap.to_iter st.loops
          |> Iter.to_string (fun (k, v) ->
-             Printf.sprintf "%s -> %s" (G.V.show k) (show_block_loop_state v)))
+             Printf.sprintf "%s -> %s\n" (G.V.show k) (show_block_loop_state v))
+         )
+
+    let dbg_show_r r =
+      print_endline
+      @@ (List.to_iter r
+         |> Iter.to_string (fun v -> Printf.sprintf "%s\n" (show_block_info v))
+         )
   end
 
-  (** Returns a list representing the loop forest sorted in reverse-topological
-      order. *)
+  (** Returns a list representing the loop forest in topological order. *)
   let solve g entry =
     let open Implementation in
     let st = create g in
     ignore @@ trav_loops st (Call { block = st.l entry; dfsp_pos = 1 }) [];
-    compute_block_loop_info g st.loops
+    let r = List.rev @@ compute_block_loop_info g st.loops in
+    r
 end
 
 module ProcIntra = struct
@@ -465,7 +466,7 @@ module ProcIntra = struct
       let open Option in
       Procedure.blocks_succ p src |> Iter.find_pred (fst %> ID.equal dest)
       |> function
-      | Some (e, b) -> (e, dest)
+      | Some _ -> (src, dest)
       | _ -> raise Not_found
 
     let succ p v = Procedure.blocks_succ p v |> Iter.map fst |> Iter.to_list
diff --git a/lib/lang/procedure.ml b/lib/lang/procedure.ml
index 540de80f..f0e29244 100644
--- a/lib/lang/procedure.ml
+++ b/lib/lang/procedure.ml
@@ -371,14 +371,30 @@ let update_block p id (block : (Var.t, BasilExpr.t) Block.t) =
       let g = G.add_edge_e g (Begin id, Block block, End id) in
       g)
 
+let modify_succs p id ~remove ~add =
+  let open Edge in
+  let open G in
+  p
+  |> map_graph (fun g ->
+      let g =
+        G.succ_e g (End id)
+        |> List.filter
+             ( G.E.dst %> function
+               | Begin e -> List.exists (ID.equal e) remove
+               | _ -> false )
+        |> List.fold_left G.remove_edge_e g
+      in
+      let new_succs = List.map (fun s -> Vert.(End id, Jump, Begin s)) add in
+      List.fold_left G.add_edge_e g new_succs)
+
 let replace_block_succs p id succs =
   let open Edge in
   let open G in
   p
   |> map_graph (fun g ->
       let g = G.succ_e g (End id) |> List.fold_left G.remove_edge_e g in
-      let succs = List.map (fun s -> Vert.(Begin id, Jump, End s)) succs in
-      List.fold_left G.add_edge_e g succs)
+      let new_succs = List.map (fun s -> Vert.(End id, Jump, Begin s)) succs in
+      List.fold_left G.add_edge_e g new_succs)
 
 let replace_edge p id (block : (Var.t, BasilExpr.t) Block.t) =
   update_block p id block
diff --git a/lib/passes.ml b/lib/passes.ml
index 19aa5efc..820ce3ca 100644
--- a/lib/passes.ml
+++ b/lib/passes.ml
@@ -146,6 +146,13 @@ module PassManager = struct
       doc = "Remove blocks unreachable from entry";
     }
 
+  let irreducible_loop =
+    {
+      name = "irreducible-loops";
+      apply = Proc Transforms.Irreducible_loop.transform;
+      doc = "Remove blocks unreachable from entry";
+    }
+
   let full_ssa =
     {
       name = "ssa";
@@ -189,6 +196,7 @@ module PassManager = struct
 
   let passes =
     [
+      irreducible_loop;
       cleanup_cfg;
       dfg_bool;
       dfg_ival_wint_product;
diff --git a/lib/transforms/irreducible_loop.ml b/lib/transforms/irreducible_loop.ml
index 9a09b829..b62d0834 100644
--- a/lib/transforms/irreducible_loop.ml
+++ b/lib/transforms/irreducible_loop.ml
@@ -14,17 +14,11 @@ let transform_loop p l =
   let open Option in
   let header, next_h, headers, nodes, entries, backedges =
     match l with
-    | {
-     block;
-     loop = PrimaryHeader { primary_header; headers; nodes; entries; backedges };
-    }
+    | { block; loop = PrimaryHeader { primary_header; headers; nodes } }
       when VSet.cardinal headers > 1 ->
-        ( block,
-          primary_header,
-          headers,
-          nodes,
-          entries,
-          backedges )
+        let entries = compute_entries p l in
+        let backedges = compute_backedges p l in
+        (block, primary_header, headers, nodes, entries, backedges)
     | _ ->
         raise
           (Invalid_argument
@@ -33,7 +27,7 @@ let transform_loop p l =
   let dest = BlockGraph.E.dst in
   let src = BlockGraph.E.src in
   let backedges_to_primary_header =
-    List.filter (BlockGraph.E.dst %> ID.equal header) entries
+    List.filter (BlockGraph.E.dst %> ID.equal header) backedges
   in
   let entry_indexes =
     entries @ backedges
@@ -43,12 +37,12 @@ let transform_loop p l =
 
   (* included entry blocks should be a superset of (externalEntries ++ backEdgesToFirstHeader).
      in particular, it additionally includes internal edges to alternative headers.*)
-  assert (
+  (*assert (
     Iter.append (List.to_iter entries)
       (List.to_iter backedges_to_primary_header)
     |> Iter.map BlockGraph.E.src
     |> Iter.subset ~eq:ID.equal (IDMap.keys entry_indexes));
-
+    *)
   let preceding_indices =
     entries @ backedges
     |> List.group_by ~hash:(dest %> ID.hash) ~eq:(fun a b ->
@@ -57,10 +51,8 @@ let transform_loop p l =
       | h :: tl ->
           ( dest h,
             h :: tl |> List.map src
-            |> List.sort_uniq ~cmp:(fun a b ->
-                Int.compare
-                  (IDMap.find a entry_indexes)
-                  (IDMap.find a entry_indexes)) )
+            |> List.map (fun a -> IDMap.find a entry_indexes)
+            |> List.sort_uniq ~cmp:Int.compare )
       | _ -> failwith "emtpy")
   in
   let ctrl_sz = Types.bv_min_width_for_nat (IDSet.cardinal headers) in
@@ -74,15 +66,11 @@ let transform_loop p l =
     let size = match ctrl_sz with Bitvector i -> i | _ -> assert false in
     BasilExpr.bvconst (Bitvec.of_int ~size idx)
   in
-  let bval m =
-    let idx = IDMap.find m entry_indexes in
-    bvali idx
-  in
   (* create new primary header which jumps to all existing headers *)
   let p, n_header =
     Procedure.fresh_block
       ~name:(ID.to_string header ^ "header_loop_N")
-      ~successors:(VSet.to_list headers) p ~stmts:[] ()
+      ~successors:(List.map dest entries) p ~stmts:[] ()
   in
   ( ( p |> fun p ->
       (* add guards to old headers to restrict their predecessors to original predecessor set *)
@@ -91,7 +79,7 @@ let transform_loop p l =
           let preds =
             List.map
               (fun i ->
-                BasilExpr.binexp ~op:`EQ (BasilExpr.rvar loop_crtl_v) (bval i))
+                BasilExpr.binexp ~op:`EQ (BasilExpr.rvar loop_crtl_v) (bvali i))
               indices
           in
           let e = BasilExpr.applyintrin ~op:`OR preds in
@@ -114,7 +102,8 @@ let transform_loop p l =
   |> fun p ->
   (* make all headers jump to the new primary header *)
   List.fold_left
-    (fun p (src, dest) -> Procedure.replace_block_succs p src [ n_header ])
+    (fun p (src, dest) ->
+      Procedure.modify_succs p src ~remove:[ dest ] ~add:[ n_header ])
     p
     (entries @ backedges_to_primary_header)
 
diff --git a/test/analysis/dune b/test/analysis/dune
index 437b0576..00410e7e 100644
--- a/test/analysis/dune
+++ b/test/analysis/dune
@@ -20,6 +20,8 @@
   test_irreducible_loops
   test_lattice_collections)
  (name analysis_tests)
+ (preprocess
+  (pps ppx_here))
  (libraries
   bincaml.loader
   bincaml.ast
diff --git a/test/analysis/test_irreducible_loops.ml b/test/analysis/test_irreducible_loops.ml
index 87f7c12c..0cb5cdf5 100644
--- a/test/analysis/test_irreducible_loops.ml
+++ b/test/analysis/test_irreducible_loops.ml
@@ -26,16 +26,7 @@ open struct
           |> Iter.to_string ~sep:", " (fun (k, v) -> k ^ "->" ^ str v)))
       (StringMap.equal equal)
 
-  let check_test_comparison a b =
-    Alcotest.(check @@ id_map String.equal Fun.id)
-      "loop participant->header ptrs equal" a.iloop_headers b.iloop_headers;
-    Alcotest.(
-      check
-        (id_map StringSet.equal
-           (StringSet.to_string ~stop:"}" ~start:"{" Fun.id)))
-      "loop header->participant sets equal" a.headers b.headers
-
-  let assert_loop_detector loops iloop_headers headers =
+  let assert_loop_detector here loops iloop_headers headers =
     let headers =
       List.map (Pair.map Fun.id StringSet.of_list) headers |> StringMap.of_list
     in
@@ -66,7 +57,15 @@ open struct
       |> StringMap.of_list
     in
     let checked = { iloop_headers = headers; headers = members } in
-    check_test_comparison expect checked
+    (let open Alcotest in
+     check ~here @@ id_map String.equal Fun.id)
+      "loop participant->header ptrs equal" expect.iloop_headers
+      checked.iloop_headers;
+    (let open Alcotest in
+     check ~here
+       (id_map StringSet.equal
+          (StringSet.to_string ~stop:"]" ~start:"[" ~sep:";" Fun.id)))
+      "loop header->participant sets equal" expect.headers checked.headers
 
   let run_transform prog =
     let p = (Loader.Loadir.ast_of_string prog).prog in
@@ -86,7 +85,7 @@ open struct
     let c =
      fun () ->
       let loops = solve_proc p in
-      assert_loop_detector loops header_ptrs all_loop_headers
+      assert_loop_detector [%here] loops header_ptrs all_loop_headers
     in
     Alcotest.test_case name `Quick c
 
@@ -250,7 +249,8 @@ proc @main () -> ()
         match classify_block b with
         | `IrreducibleHeader -> false
         | `ReducibleHeader -> true
-        | _ -> false)
+        | `LoopNode -> false
+        | `NonLoop -> false)
       p
 
   let loops_irreducible p =
@@ -259,7 +259,7 @@ proc @main () -> ()
         match classify_block b with `IrreducibleHeader -> true | _ -> false)
       p
 
-  let check_transform_fixed name ~num_irr_loops ~num_red_loops ?header_ptrs
+  let check_transform_fixed here name ~num_irr_loops ~num_red_loops ?header_ptrs
       ?all_headers p =
     let checks () =
       let before, after = run_transform p in
@@ -268,23 +268,35 @@ proc @main () -> ()
         let open Option in
         let* hdrs = header_ptrs in
         let* headers = all_headers in
-        Some (fun () -> assert_loop_detector before hdrs headers)
+        Some (fun () -> assert_loop_detector here before hdrs headers)
       in
-      Option.iter (fun x -> x ()) check_x;
-      Alcotest.(check int)
+      print_endline @@ "before transform: ";
+      Implementation.dbg_show_r before;
+      print_endline @@ "after transform: ";
+      Implementation.dbg_show_r after;
+      print_endline @@ "irreducible: "
+      ^ List.to_string show_block_info (loops_irreducible before);
+      print_endline @@ "reducible: "
+      ^ List.to_string show_block_info (loops_reducible before);
+      Alcotest.(check ~here int)
         "number of irreducible loops present" num_irr_loops
         (List.length @@ loops_irreducible before);
-      Alcotest.(check int)
+      Alcotest.(check ~here int)
         "number of reducible loops present" num_red_loops
         (List.length @@ loops_reducible before);
-      Alcotest.(check (list block_info))
-        "all irreducible loops fixed" [] (loops_irreducible after)
+      Option.iter (fun x -> x ()) check_x;
+      Alcotest.(check ~here (list block_info))
+        "all irreducible loops fixed" [] (loops_irreducible after);
+      Alcotest.(check ~here bool)
+        "have at least one loop left"
+        (num_irr_loops + num_red_loops > 0)
+        (List.length @@ loops_reducible after > 0)
     in
     Alcotest.test_case name `Quick checks
 
   let sub_cycles_transform =
-    check_transform_fixed "subcycles applying transform" ~num_irr_loops:1
-      ~num_red_loops:1
+    check_transform_fixed [%here] "subcycles applying transform"
+      ~num_irr_loops:1 ~num_red_loops:1
       {|
 prog entry @main;                           
                                             
@@ -310,7 +322,7 @@ proc @main () -> ()
     |}
 
   let crossover =
-    check_transform_fixed "crossover" ~num_irr_loops:2 ~num_red_loops:0
+    check_transform_fixed [%here] "crossover" ~num_irr_loops:2 ~num_red_loops:0
       {|
 prog entry @main;
 proc @main () -> ()
@@ -338,7 +350,8 @@ proc @main () -> ()
     |}
 
   let paper_fig4a =
-    check_transform_fixed "paper fig4a" ~num_irr_loops:0 ~num_red_loops:0
+    check_transform_fixed [%here] "paper fig4a" ~num_irr_loops:0
+      ~num_red_loops:0
       {|
 prog entry @main;
 
@@ -368,7 +381,8 @@ proc @main () -> ()
     |}
 
   let paper_fig4b =
-    check_transform_fixed "paper fig4b" ~num_irr_loops:0 ~num_red_loops:1
+    check_transform_fixed [%here] "paper fig4b" ~num_irr_loops:0
+      ~num_red_loops:1
       ~header_ptrs:[ ("%b0", "%b"); ("%x", "%b") ]
       ~all_headers:[ ("%b", [ "%b" ]) ]
       {|
@@ -399,8 +413,8 @@ proc @main () -> ()
     |}
 
   let paper_fig4c =
-    (* TODO: sus there are no loops before *)
-    check_transform_fixed "paper fig4c" ~num_irr_loops:0 ~num_red_loops:0
+    check_transform_fixed [%here] "paper fig4c" ~num_irr_loops:0
+      ~num_red_loops:0
       {|
 prog entry @main;
 
@@ -435,7 +449,8 @@ proc @main () -> ()
   |}
 
   let paper_fig4d =
-    check_transform_fixed "paper fig4d" ~num_irr_loops:0 ~num_red_loops:1
+    check_transform_fixed [%here] "paper fig4d" ~num_irr_loops:0
+      ~num_red_loops:1
       ~header_ptrs:
         [
           ("%b", "%h"); ("%b0", "%h"); ("%x", "%h"); ("%y", "%h"); ("%z", "%h");
@@ -479,9 +494,14 @@ proc @main () -> ()
     |}
 
   let paper_fig6a =
-    (* TODO: wrong; does scala impl identifies 2 irreducible loops *)
-    check_transform_fixed "paper fig6a" ~num_irr_loops:2 ~num_red_loops:1
-      ~header_ptrs:[] ~all_headers:[]
+    (* FIXME: scala impl identifies 2 irreducible loops 
+
+  + BlockLoopInfo(%h2,Some(%h3),4,Set(%h2, %b),HashSet(%b, %h1, %h2, %z, %a)) 
+  + BlockLoopInfo(%h1,Some(%h2),5,Set(%h1, %b),Set(%h1, %z, %b)) 
+
+       *)
+    check_transform_fixed [%here] "paper fig6a" ~num_irr_loops:2
+      ~num_red_loops:1
       {|
 prog entry @main;
 
@@ -529,8 +549,8 @@ proc @main () -> ()
     |}
 
   let paper_fig6b =
-    (* FIXME: shouldn't we identify 4 reducible loops? *)
-    check_transform_fixed "paper fig6b" ~num_irr_loops:0 ~num_red_loops:1
+    check_transform_fixed [%here] "paper fig6b" ~num_irr_loops:0
+      ~num_red_loops:4
       {|
 prog entry @main;
 
@@ -571,8 +591,18 @@ proc @main () -> ()
 |}
 
   let paper_fig4e =
-    check_transform_fixed "paper fig4e" ~num_irr_loops:1 ~num_red_loops:1
-      ~header_ptrs:[] ~all_headers:[]
+    check_transform_fixed [%here] "paper fig4e" ~num_irr_loops:1
+      ~num_red_loops:1
+      ~header_ptrs:
+        [
+          ("%a", "%h");
+          ("%b", "%h");
+          ("%b0", "%h1");
+          ("%h", "%h1");
+          ("%y", "%h1");
+          ("%z", "%h1");
+        ]
+      ~all_headers:[ ("%h", [ "%b"; "%h" ]); ("%h1", [ "%h1" ]) ]
       {|
 prog entry @main;
 
@@ -611,7 +641,7 @@ proc @main () -> ()
 end
 
 let triforce =
-  check_transform_fixed "triforce" ~num_irr_loops:2 ~num_red_loops:1
+  check_transform_fixed [%here] "triforce" ~num_irr_loops:2 ~num_red_loops:1
     {|
 prog entry @main;
 
@@ -654,8 +684,8 @@ let tests =
         paper_fig4b;
         paper_fig4c;
         paper_fig4d;
-        paper_fig6a;
         paper_fig6b;
+        paper_fig6a;
         paper_fig4e;
         triforce;
       ] );

From de94aa0334faedb9ad77d95bd5dd7678ee78de15 Mon Sep 17 00:00:00 2001
From: agle <a.michael1@uq.edu.au>
Date: Mon, 30 Mar 2026 17:19:06 +1000
Subject: [PATCH 07/11] add assert

---
 lib/analysis/irreducible_loops.ml       |  5 -----
 lib/transforms/irreducible_loop.ml      | 13 +++++++------
 test/analysis/test_irreducible_loops.ml | 10 +++++++---
 3 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/lib/analysis/irreducible_loops.ml b/lib/analysis/irreducible_loops.ml
index 7e9eed47..18d8e823 100644
--- a/lib/analysis/irreducible_loops.ml
+++ b/lib/analysis/irreducible_loops.ml
@@ -186,14 +186,12 @@ module Make (G : GSig) = struct
           all_blocks
         |> Iter.persistent
       in
-      (* 1 *)
       let (forest : VSet.t VMap.t) =
         header_blocks
         |> map (fun b -> (b.block, VSet.singleton b.block))
         |> VMap.of_iter
       in
 
-      (* 2 *)
       (* NOTE: iterates the forest in *bottom-up* topological order. this
        ensures that node-sets of sub-cycles are fully populated before
        processing their parent cycle. this avoids us having to compute
@@ -212,7 +210,6 @@ module Make (G : GSig) = struct
                      forest))
              forest
       in
-      (* 3 *)
       let forest =
         header_blocks
         |> fold
@@ -226,7 +223,6 @@ module Make (G : GSig) = struct
                      forest)
              forest
       in
-      (* 4 *)
       header_blocks
       |> iter (fun b ->
           let nodes = VMap.find b.block forest in
@@ -242,7 +238,6 @@ module Make (G : GSig) = struct
             ^ " bad headers: "
             ^ VSet.to_string G.V.show bad_headers);
 
-      (* 5 *)
       let new_loops =
         all_blocks
         |> map (fun b ->
diff --git a/lib/transforms/irreducible_loop.ml b/lib/transforms/irreducible_loop.ml
index b62d0834..85cd30cc 100644
--- a/lib/transforms/irreducible_loop.ml
+++ b/lib/transforms/irreducible_loop.ml
@@ -37,12 +37,13 @@ let transform_loop p l =
 
   (* included entry blocks should be a superset of (externalEntries ++ backEdgesToFirstHeader).
      in particular, it additionally includes internal edges to alternative headers.*)
-  (*assert (
-    Iter.append (List.to_iter entries)
-      (List.to_iter backedges_to_primary_header)
-    |> Iter.map BlockGraph.E.src
-    |> Iter.subset ~eq:ID.equal (IDMap.keys entry_indexes));
-    *)
+  assert (
+    let e =
+      Iter.append (List.to_iter entries)
+        (List.to_iter backedges_to_primary_header)
+      |> Iter.map BlockGraph.E.src
+    in
+    Iter.subset e (IDMap.keys entry_indexes));
   let preceding_indices =
     entries @ backedges
     |> List.group_by ~hash:(dest %> ID.hash) ~eq:(fun a b ->
diff --git a/test/analysis/test_irreducible_loops.ml b/test/analysis/test_irreducible_loops.ml
index 0cb5cdf5..bbd10b58 100644
--- a/test/analysis/test_irreducible_loops.ml
+++ b/test/analysis/test_irreducible_loops.ml
@@ -274,10 +274,14 @@ proc @main () -> ()
       Implementation.dbg_show_r before;
       print_endline @@ "after transform: ";
       Implementation.dbg_show_r after;
-      print_endline @@ "irreducible: "
+      print_endline @@ "irreducible before: "
       ^ List.to_string show_block_info (loops_irreducible before);
-      print_endline @@ "reducible: "
+      print_endline @@ "reducible before: "
       ^ List.to_string show_block_info (loops_reducible before);
+      print_endline @@ "irreducible after: "
+      ^ List.to_string show_block_info (loops_irreducible after);
+      print_endline @@ "reducible after: "
+      ^ List.to_string show_block_info (loops_reducible after);
       Alcotest.(check ~here int)
         "number of irreducible loops present" num_irr_loops
         (List.length @@ loops_irreducible before);
@@ -500,7 +504,7 @@ proc @main () -> ()
   + BlockLoopInfo(%h1,Some(%h2),5,Set(%h1, %b),Set(%h1, %z, %b)) 
 
        *)
-    check_transform_fixed [%here] "paper fig6a" ~num_irr_loops:2
+    check_transform_fixed [%here] "paper fig6a" ~num_irr_loops:1
       ~num_red_loops:1
       {|
 prog entry @main;

From 8eac246a38bab72ae9085d7c291d13526fb30938 Mon Sep 17 00:00:00 2001
From: agle <a.michael1@uq.edu.au>
Date: Mon, 30 Mar 2026 17:28:57 +1000
Subject: [PATCH 08/11] doc

---
 lib/lang/procedure.ml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/lang/procedure.ml b/lib/lang/procedure.ml
index f0e29244..2892282d 100644
--- a/lib/lang/procedure.ml
+++ b/lib/lang/procedure.ml
@@ -340,7 +340,7 @@ let get_entry_block p =
 
 (** Get the block for an id
 
-    @raises {! Not_found} when the block does not exist.
+   @raises {! Not_found } when the block does not exist.
 *)
 let find_block p id =
   let open Edge in

From ed27f48f816f38e5b76af55adc14d09fc5ba4b10 Mon Sep 17 00:00:00 2001
From: agle <a.michael1@uq.edu.au>
Date: Wed, 1 Apr 2026 16:39:27 +1000
Subject: [PATCH 09/11] idk

---
 lib/lang/procedure.ml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/lang/procedure.ml b/lib/lang/procedure.ml
index 2892282d..8cd4bf9d 100644
--- a/lib/lang/procedure.ml
+++ b/lib/lang/procedure.ml
@@ -340,7 +340,7 @@ let get_entry_block p =
 
 (** Get the block for an id
 
-   @raises {! Not_found } when the block does not exist.
+   @raise {! Not_found } when the block does not exist.
 *)
 let find_block p id =
   let open Edge in

From cc743e41b355fae5012cf978d54fcf51b590caae Mon Sep 17 00:00:00 2001
From: agle <a.michael1@uq.edu.au>
Date: Wed, 1 Apr 2026 16:45:39 +1000
Subject: [PATCH 10/11] giveup

---
 lib/lang/procedure.ml | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/lib/lang/procedure.ml b/lib/lang/procedure.ml
index 8cd4bf9d..bb6627ef 100644
--- a/lib/lang/procedure.ml
+++ b/lib/lang/procedure.ml
@@ -340,8 +340,7 @@ let get_entry_block p =
 
 (** Get the block for an id
 
-   @raise {! Not_found } when the block does not exist.
-*)
+    raise Not_found when the block does not exist. *)
 let find_block p id =
   let open Edge in
   let open G in

From f927533aa4e363f9e99aaf0ccc32448aa7a06495 Mon Sep 17 00:00:00 2001
From: agle <a.michael1@uq.edu.au>
Date: Wed, 1 Apr 2026 16:52:18 +1000
Subject: [PATCH 11/11] tingy

---
 lib/lang/procedure.ml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/lang/procedure.ml b/lib/lang/procedure.ml
index bb6627ef..68566673 100644
--- a/lib/lang/procedure.ml
+++ b/lib/lang/procedure.ml
@@ -340,7 +340,7 @@ let get_entry_block p =
 
 (** Get the block for an id
 
-    raise Not_found when the block does not exist. *)
+    @raise Not_found when the block does not exist. *)
 let find_block p id =
   let open Edge in
   let open G in