1 files changed, 31 insertions, 22 deletions

diff --git a/helpers.cc b/helpers.cc
index c02957e..7bb1215 100644
--- a/helpers.cc
+++ b/helpers.cc
@@ -51,39 +51,48 @@ die_for (Address address, std::string const& message)
   std::exit (1);
 }
 
+static CodeRegion*
+find_region (CodeSource& cs, Address address)
+{
+  if (!cs.isCode (address))
+    die_for (address, "no instruction at");
+  std::set <CodeRegion*> regions;
+  if (cs.findRegions (address, regions) != 1)
+    die_for (address, "not exactly one region found spanning");
+  for (auto* const cr : regions)
+    return cr;
+  std::unreachable ();
+}
+
 /// Find next basic block's entry after given address, reparsing if necessary
 static Block*
-next_block (CodeObject& co, CodeRegion* region, Address address)
+next_block (CodeObject& co, CodeRegion* cr, Address address)
 {
-  auto* blk = co.findBlockByEntry (region, address);
+  auto* blk = co.findBlockByEntry (cr, address);
   if (blk != nullptr)
     return blk;
   co.parse (address, true);
-  blk = co.findBlockByEntry (region, address);
-  return (blk != nullptr) ? blk : co.findNextBlock (region, address);
+  blk = co.findBlockByEntry (cr, address);
+  return (blk != nullptr) ? blk : co.findNextBlock (cr, address);
 }
 
 Block*
 find_block (CodeSource& cs, CodeObject& co, Address address)
 {
-  if (!cs.isCode (address))
-    die_for (address, "no instruction at");
-  std::set <CodeRegion*> regions;
-  if (cs.findRegions (address, regions) != 1)
-    die_for (address, "not exactly 1 region found for instruction at");
-  for (auto* region : regions)
+  auto* const cr = find_region (cs, address);
+  std::set <Block*> blocks;
+  if (co.findBlocks (cr, address, blocks) > 0)
     {
-      std::set <Block*> blocks;
-      if (co.findBlocks (region, address, blocks) > 0)
-        for (auto* blk : blocks) // TODO: choose the best block
-          return blk;
-      auto* blk = next_block (co, region, region->low ());
-      while (blk != nullptr && address > blk->last ())
-        blk = next_block (co, region, blk->end ());
-      if (blk == nullptr)
-        die_for (address, "no block found for instruction at");
-      assert (address >= blk->start () && address < blk->end ());
-      return blk;
+      if (blocks.size () > 1)
+        die_for (address, "more than one block found spanning");
+      for (auto* const blk : blocks)
+        return blk;
     }
-  std::unreachable ();
+  auto* blk = next_block (co, cr, cr->low ());
+  while (blk != nullptr && address > blk->last ())
+    blk = next_block (co, cr, blk->end ());
+  if (blk == nullptr)
+    die_for (address, "no block found spanning");
+  assert (address >= blk->start () && address < blk->end ());
+  return blk;
 }