Added code to populate loop latch blocks in a loop

Author: VedantParanjape

include/blocks/loops.h

@@ -4,6 +4,7 @@
 #include "blocks/basic_blocks.h"
 #include "blocks/dominance.h"
 #include "blocks/stmt.h"
+#include <unordered_set>
 
 using namespace block;
 class loop {
@@ -23,9 +24,10 @@ class loop {
     } loop_bounds;
 
     basic_block::cfg_block blocks;
+    std::unordered_set<int> blocks_id_map;
     std::shared_ptr<loop> parent_loop;
     std::shared_ptr<basic_block> header_block;
-    std::shared_ptr<basic_block> backedge_block;
+    basic_block::cfg_block loop_latch_blocks;
     std::vector<std::shared_ptr<loop>> subloops;
 };
 

src/blocks/loops.cpp

@@ -84,6 +84,7 @@ void loop_info::analyze() {
             }
             new_loop->subloops.reserve(num_subloops);
             new_loop->blocks.reserve(num_blocks);
+            new_loop->blocks_id_map.reserve(num_blocks);
         }
     }
 
@@ -108,7 +109,21 @@ void loop_info::analyze() {
 
         while (subloop) {
             subloop->blocks.push_back(dta.cfg_[bb_id]);
+            subloop->blocks_id_map.insert(dta.cfg_[bb_id]->id);
             subloop = subloop->parent_loop;
         }
     }
+
+    // Populate the loop latches
+    for (auto loop: loops) {
+        if (!loop->header_block)
+            continue;
+
+        std::shared_ptr<basic_block> header = loop->header_block;
+        for (auto children: header->predecessor) {
+            if (loop->blocks_id_map.count(children->id)) {
+                loop->loop_latch_blocks.push_back(children);
+            }
+        }
+    }
 }

src/builder/builder_context.cpp

@@ -371,11 +371,17 @@ block::stmt::Ptr builder_context::extract_ast_from_function_impl(void) {
 		std::cerr << "++++++ loop " << loop_num++ << " ++++++ \n";
 
 		std::cerr << "loop headers: " << loop->header_block->id << "\n";
+		
 		std::cerr << "blocks: ";
 		for (auto bb: loop->blocks) std::cerr << bb->id << " ";
 		std::cerr << "\n";
-		// std::cerr << "backedge: " << loop->backedge_block->id << "\n";
+
+		std::cerr << "loop latches: ";
+		for (auto bb: loop->loop_latch_blocks) std::cerr << bb->id << " ";
+		std::cerr << "\n";
+		
 		std::cerr << "parent loop: (loop header: " << (loop->parent_loop ? (int)loop->parent_loop->header_block->id : -1) << ")\n";
+		
 		std::cerr << "subloops: ";
 		for (auto subl: loop->subloops) std::cerr << "(loop header: " << subl->header_block->id << ") ";
 		std::cerr << "\n";