1 files changed, 20 insertions, 12 deletions

diff --git a/backend/LICMaux.ml b/backend/LICMaux.ml
index 1f6b8817..b88dbc2d 100644
--- a/backend/LICMaux.ml
+++ b/backend/LICMaux.ml
@@ -41,24 +41,25 @@ let rtl_successors = function
   *
   * If we come accross an edge to a Processed node, it's a loop!
   *)
-let get_loop_headers code entrypoint = begin
-  debug "get_loop_headers\n";
+let get_loop_backedges code entrypoint = begin
+  debug "get_loop_backedges\n";
   let visited = ref (PTree.map (fun n i -> Unvisited) code)
-  and is_loop_header = ref (PTree.map (fun n i -> false) code)
-  in let rec dfs_visit code = function
+  and loop_backedge = ref (PTree.map (fun n i -> None) code)
+  in let rec dfs_visit code origin = function
   | [] -> ()
   | node :: ln ->
       debug "ENTERING node %d, REM are %a\n" (P.to_int node) print_intlist ln;
       match (get_some @@ PTree.get node !visited) with
       | Visited -> begin
           debug "\tNode %d is already Visited, skipping\n" (P.to_int node);
-          dfs_visit code ln
+          dfs_visit code origin ln
         end
       | Processed -> begin
           debug "Node %d is a loop header\n" (P.to_int node);
-          is_loop_header := PTree.set node true !is_loop_header;
+          debug "The backedge is from %d\n" (P.to_int @@ get_some origin);
+          loop_backedge := PTree.set node origin !loop_backedge;
           visited := PTree.set node Visited !visited;
-          dfs_visit code ln
+          dfs_visit code origin ln
         end
       | Unvisited -> begin
           visited := PTree.set node Processed !visited;
@@ -67,19 +68,26 @@ let get_loop_headers code entrypoint = begin
           | None -> failwith "No such node"
           | Some i -> let next_visits = rtl_successors i in begin
               debug "About to visit: %a\n" print_intlist next_visits;
-              dfs_visit code next_visits
+              dfs_visit code (Some node) next_visits
             end);
           debug "Node %d is Visited!\n" (P.to_int node);
           visited := PTree.set node Visited !visited;
-          dfs_visit code ln
+          dfs_visit code origin ln
         end
   in begin
-    dfs_visit code [entrypoint];
-    debug "LOOP HEADERS: %a\n" print_ptree_bool !is_loop_header;
-    !is_loop_header
+    dfs_visit code None [entrypoint];
+    debug "LOOP BACKEDGES: %a\n" print_ptree_opint !loop_backedge;
+    !loop_backedge
   end
 end
 
+let get_loop_headers code entrypoint =
+  let backedges = get_loop_backedges code entrypoint in
+  PTree.map (fun _ ob ->
+    match ob with
+    | None -> false
+    | Some _ -> true
+  ) backedges
 
 module Dominator =
   struct