Static Analysis: Add an API to get the block index of an expression (#5822)

2 files changed, 32 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"