Getting all loop bodies

backend/Duplicateaux.ml

@@ -270,6 +270,39 @@ let get_inner_loops f code is_loop_header =
     ) (PTree.elements loopmap)
   end
 
+let get_loop_bodies code entrypoint =
+  let predecessors = get_predecessors_rtl code in
+  (* Algorithm from Muchnik, Compiler Design & Implementation, Figure 7.21 page 192 *)
+  let natural_loop n m =
+    debug "Natural Loop from %d to %d\n" (P.to_int n) (P.to_int m);
+    let in_body = ref (PTree.map (fun n b -> false) code) in
+    let body = ref [] in
+    let add_to_body n = begin
+      in_body := PTree.set n true !in_body;
+      body := n :: !body
+    end
+    in let rec process_node p =
+      debug "    Processing node %d\n" (P.to_int p);
+      List.iter (fun pred ->
+        debug "        Looking at predecessor of %d: %d\n" (P.to_int p) (P.to_int pred);
+        let is_in_body = get_some @@ PTree.get pred !in_body in
+        if (not @@ is_in_body) then begin
+          debug "        --> adding to body\n";
+          add_to_body pred;
+          process_node pred
+        end
+      ) (get_some @@ PTree.get p predecessors)
+    in begin
+      add_to_body m;
+      add_to_body n;
+      (if (m != n) then process_node m);
+      !body
+    end
+  in let option_natural_loop n = function
+    | None -> None
+    | Some m -> Some (natural_loop n m)
+  in PTree.map option_natural_loop (LICMaux.get_loop_backedges code entrypoint)
+
 (* Returns a PTree of either None or Some b where b determines the node in the loop body, for a cb instruction *)
 let get_loop_info f is_loop_header bfs_order code =
   let loop_info = ref (PTree.map (fun n i -> None) code) in
@@ -298,6 +331,7 @@ let get_directions f code entrypoint = begin
   let loop_info = get_loop_info f is_loop_header bfs_order code in
   let directions = ref (PTree.map (fun n i -> None) code) in (* None <=> no predicted direction *)
   begin
+    debug_flag := true;
     (* ptree_printbool is_loop_header; *)
     (* debug "\n"; *)
     List.iter (fun n ->
@@ -325,7 +359,7 @@ let get_directions f code entrypoint = begin
           end
           )
       | _ -> ()
-    ) bfs_order;
+    ) bfs_order; debug_flag := false;
     !directions
   end
 end
@@ -929,7 +963,7 @@ let loop_rotate f =
 
 let static_predict f =
   debug_flag := true;
-  let _ = LICMaux.get_loop_backedges f.fn_code f.fn_entrypoint in
+  Printf.printf "Loop bodies: %a" print_ptree_oplist (get_loop_bodies f.fn_code f.fn_entrypoint);
   debug_flag := false;
   let entrypoint = f.fn_entrypoint in
   let code = f.fn_code in

common/DebugPrint.ml

@@ -44,6 +44,20 @@ let print_intlist oc l =
     end
   end
 
+let print_ptree_oplist oc pt =
+  if !debug_flag then
+    let elements = PTree.elements pt in
+    begin
+      Printf.fprintf oc "[";
+      List.iter (fun (n, ol) ->
+        match ol with
+        | None -> ()
+        | Some l -> Printf.fprintf oc "%d -> %a,\n" (P.to_int n) print_intlist l
+      ) elements;
+      Printf.fprintf oc "]\n"
+    end
+  else ()
+
 (* Adapted from backend/PrintRTL.ml: print_function *)
 let print_code code = let open PrintRTL in let open Printf in
   if (!debug_flag) then begin