3 files changed, 72 insertions, 0 deletions

diff --git a/src/analysis/cfg.cpp b/src/analysis/cfg.cpp
index 7d770776d..0d651526b 100644
--- a/src/analysis/cfg.cpp
+++ b/src/analysis/cfg.cpp
@@ -103,4 +103,12 @@ void BasicBlock::print(std::ostream& os, Module* wasm, size_t start) const {
   }
 }
 
+CFGBlockIndexes::CFGBlockIndexes(const CFG& cfg) {
+  for (auto& block : cfg) {
+    for (auto* expr : block) {
+      map[expr] = block.getIndex();
+    }
+  }
+}
+
 } // namespace wasm::analysis
diff --git a/src/analysis/cfg.h b/src/analysis/cfg.h
index a7f67c041..ff9b05849 100644
--- a/src/analysis/cfg.h
+++ b/src/analysis/cfg.h
@@ -58,6 +58,7 @@ private:
   std::vector<Expression*> insts;
   std::vector<BasicBlock*> predecessors;
   std::vector<BasicBlock*> successors;
+
   friend CFG;
 };
 
@@ -78,9 +79,32 @@ struct CFG {
 
 private:
   std::vector<BasicBlock> blocks;
+
   friend BasicBlock;
 };
 
+// A helper class that computes block indexes for a CFG and allows querying of
+// them.
+struct CFGBlockIndexes {
+  CFGBlockIndexes(const CFG& cfg);
+
+  // Gets the index of the basic block in which the instruction resides.
+  Index get(Expression* expr) const {
+    auto iter = map.find(expr);
+    if (iter == map.end()) {
+      // There is no entry for this, which can be the case for control flow
+      // structures, or for unreachable code.
+      return InvalidBlock;
+    }
+    return iter->second;
+  }
+
+  enum { InvalidBlock = Index(-1) };
+
+private:
+  std::unordered_map<Expression*, Index> map;
+};
+
 } // namespace wasm::analysis
 
 #include "cfg-impl.h"
diff --git a/test/gtest/cfg.cpp b/test/gtest/cfg.cpp
index 3ee63b9b0..03e333965 100644
--- a/test/gtest/cfg.cpp
+++ b/test/gtest/cfg.cpp
@@ -293,6 +293,46 @@ TEST_F(CFGTest, FinitePowersetLatticeFunctioning) {
   EXPECT_EQ(ss.str(), "101101");
 }
 
+TEST_F(CFGTest, BlockIndexes) {
+  auto moduleText = R"wasm(
+    (module
+      (func $foo
+        (if
+          (i32.const 1)
+          (block
+            (drop
+              (i32.const 2)
+            )
+            (drop
+              (i32.const 3)
+            )
+          )
+        )
+      )
+    )
+  )wasm";
+
+  Module wasm;
+  parseWast(wasm, moduleText);
+
+  auto* func = wasm.getFunction("foo");
+  CFG cfg = CFG::fromFunction(func);
+  CFGBlockIndexes indexes(cfg);
+
+  // The body of the function is an if. An if is a control flow structure and so
+  // it has no basic block (it can contain multiple ones).
+  auto* iff = func->body->cast<If>();
+  EXPECT_EQ(indexes.get(iff), indexes.InvalidBlock);
+
+  // The constant 1 is in the entry block.
+  EXPECT_EQ(indexes.get(iff->condition), Index(0));
+
+  // The dropped constants 2 and three are in another block, together.
+  auto* block = iff->ifTrue->cast<Block>();
+  EXPECT_EQ(indexes.get(block->list[0]), Index(1));
+  EXPECT_EQ(indexes.get(block->list[1]), Index(1));
+}
+
 TEST_F(CFGTest, LinearReachingDefinitions) {
   auto moduleText = R"wasm(
     (module