From efe29f417d9c1420af851b526b81e7bee4bb1f32 Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Thu, 20 Jul 2023 12:10:25 -0700
Subject: Static Analysis: Add an API to get the block index of an expression
 (#5822)

---
 src/analysis/cfg.cpp |  8 ++++++++
 src/analysis/cfg.h   | 24 ++++++++++++++++++++++++
 2 files changed, 32 insertions(+)

(limited to 'src/analysis')

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"
-- 
cgit v1.2.3