Add analysis to compute mappings between instructions and basic blocks.

This analysis builds a map from instructions to the basic block that
contains them.  It is accessed via get_instr_block().  Once built, it is kept
up-to-date by the IRContext, as long as instructions are removed via
KillInst.

I have not yet marked passes that preserve this analysis. I will do it
in a separate change.

Other changes:

- Add documentation about analysis values requirement to be powers of 2.
- Force a re-build of the def-use manager in tests.
- Fix AllPreserveFirstOnlyAfterPassWithChange to use the
  DummyPassPreservesFirst pass.
- Fix sentinel value for IRContext::Analysis enum.

- Fix logic for checking if the instr<->block mapping is valid in KillInst.

6 files changed, 72 insertions, 38 deletions

diff --git a/source/opt/aggressive_dead_code_elim_pass.cpp b/source/opt/aggressive_dead_code_elim_pass.cpp
index fb11f13a..7be3f9ae 100644
--- a/source/opt/aggressive_dead_code_elim_pass.cpp
+++ b/source/opt/aggressive_dead_code_elim_pass.cpp
@@ -371,7 +371,6 @@ bool AggressiveDCEPass::AggressiveDCE(ir::Function* func) {
 
 void AggressiveDCEPass::Initialize(ir::IRContext* c) {
   InitializeProcessing(c);
-  InitializeCFGCleanup(c);
 
   // Clear collections
   worklist_ = std::queue<ir::Instruction*>{};
diff --git a/source/opt/cfg_cleanup_pass.cpp b/source/opt/cfg_cleanup_pass.cpp
index 946c85ca..a74b6d35 100644
--- a/source/opt/cfg_cleanup_pass.cpp
+++ b/source/opt/cfg_cleanup_pass.cpp
@@ -29,7 +29,6 @@ namespace opt {
 
 void CFGCleanupPass::Initialize(ir::IRContext* c) {
   InitializeProcessing(c);
-  InitializeCFGCleanup(c);
 }
 
 Pass::Status CFGCleanupPass::Process(ir::IRContext* c) {
diff --git a/source/opt/ir_context.cpp b/source/opt/ir_context.cpp
index a4824026..d7c92800 100644
--- a/source/opt/ir_context.cpp
+++ b/source/opt/ir_context.cpp
@@ -22,14 +22,24 @@ void IRContext::BuildInvalidAnalyses(IRContext::Analysis set) {
   if (set & kAnalysisDefUse) {
     BuildDefUseManager();
   }
+  if (set & kAnalysisInstrToBlockMapping) {
+    BuildInstrToBlockMapping();
+  }
 }
 
 void IRContext::InvalidateAnalysesExceptFor(
     IRContext::Analysis preserved_analyses) {
   uint32_t analyses_to_invalidate = valid_analyses_ & (~preserved_analyses);
+  InvalidateAnalyses(static_cast<IRContext::Analysis>(analyses_to_invalidate));
+}
+
+void IRContext::InvalidateAnalyses(IRContext::Analysis analyses_to_invalidate) {
   if (analyses_to_invalidate & kAnalysisDefUse) {
     def_use_mgr_.reset(nullptr);
   }
+  if (analyses_to_invalidate & kAnalysisInstrToBlockMapping) {
+    instr_to_block_.clear();
+  }
   valid_analyses_ = Analysis(valid_analyses_ & ~analyses_to_invalidate);
 }
 
@@ -41,6 +51,10 @@ void IRContext::KillInst(ir::Instruction* inst) {
   if (AreAnalysesValid(kAnalysisDefUse)) {
     get_def_use_mgr()->ClearInst(inst);
   }
+  if (AreAnalysesValid(kAnalysisInstrToBlockMapping)) {
+    instr_to_block_.erase(inst);
+  }
+
   inst->ToNop();
 }
 
@@ -118,5 +132,6 @@ void IRContext::AnalyzeUses(Instruction* inst) {
     get_def_use_mgr()->AnalyzeInstUse(inst);
   }
 }
+
 }  // namespace ir
 }  // namespace spvtools
diff --git a/source/opt/ir_context.h b/source/opt/ir_context.h
index 907ff964..464749ec 100644
--- a/source/opt/ir_context.h
+++ b/source/opt/ir_context.h
@@ -25,11 +25,24 @@ namespace ir {
 
 class IRContext {
  public:
+  // Available analyses.
+  //
+  // When adding a new analysis:
+  //
+  // 1. Enum values should be powers of 2. These are cast into uint32_t
+  //    bitmasks, so we can have at most 31 analyses represented.
+  //
+  // 2. Make sure it gets invalidated or preserved by IRContext methods that add
+  //    or remove IR elements (e.g., KillDef, KillInst, ReplaceAllUsesWith).
+  //
+  // 3. Add handling code in BuildInvalidAnalyses and
+  //    InvalidateAnalysesExceptFor.
   enum Analysis {
-    kAnalysisNone = 0x0,
-    kAnalysisBegin = 0x1,
+    kAnalysisNone = 0 << 0,
+    kAnalysisBegin = 1 << 0,
     kAnalysisDefUse = kAnalysisBegin,
-    kAnalysisEnd = 0x2
+    kAnalysisInstrToBlockMapping = 1 << 1,
+    kAnalysisEnd = 1 << 2
   };
 
   friend inline Analysis operator|(Analysis lhs, Analysis rhs);
@@ -147,10 +160,14 @@ class IRContext {
     return def_use_mgr_.get();
   }
 
-  // Builds the def-use manager from scratch, even if it was already valid.
-  void BuildDefUseManager() {
-    def_use_mgr_.reset(new opt::analysis::DefUseManager(module()));
-    valid_analyses_ = valid_analyses_ | kAnalysisDefUse;
+  // Returns the basic block for instruction |instr|. Re-builds the instruction
+  // block map, if needed.
+  ir::BasicBlock* get_instr_block(ir::Instruction* instr) {
+    if (!AreAnalysesValid(kAnalysisInstrToBlockMapping)) {
+      BuildInstrToBlockMapping();
+    }
+    auto entry = instr_to_block_.find(instr);
+    return (entry != instr_to_block_.end()) ? entry->second : nullptr;
   }
 
   // Sets the message consumer to the given |consumer|. |consumer| which will be
@@ -168,6 +185,9 @@ class IRContext {
   // Invalidates all of the analyses except for those in |preserved_analyses|.
   void InvalidateAnalysesExceptFor(Analysis preserved_analyses);
 
+  // Invalidates the analyses marked in |analyses_to_invalidate|.
+  void InvalidateAnalyses(Analysis analyses_to_invalidate);
+
   // Turns the instruction defining the given |id| into a Nop. Returns true on
   // success, false if the given |id| is not defined at all. This method also
   // erases both the uses of |id| and the information of this |id|-generating
@@ -202,10 +222,36 @@ class IRContext {
   void AnalyzeUses(Instruction* inst);
 
  private:
+  // Builds the def-use manager from scratch, even if it was already valid.
+  void BuildDefUseManager() {
+    def_use_mgr_.reset(new opt::analysis::DefUseManager(module()));
+    valid_analyses_ = valid_analyses_ | kAnalysisDefUse;
+  }
+
+  // Builds the instruction-block map for the whole module.
+  void BuildInstrToBlockMapping() {
+    instr_to_block_.clear();
+    for (auto& fn : *module_) {
+      for (auto& block : fn) {
+        block.ForEachInst([this, &block](ir::Instruction* inst) {
+          instr_to_block_[inst] = &block;
+        });
+      }
+    }
+    valid_analyses_ = valid_analyses_ | kAnalysisInstrToBlockMapping;
+  }
+
   std::unique_ptr<Module> module_;
   spvtools::MessageConsumer consumer_;
   std::unique_ptr<opt::analysis::DefUseManager> def_use_mgr_;
 
+  // A map from instructions the the basic block they belong to. This mapping is
+  // built on-demand when get_instr_block() is called.
+  //
+  // NOTE: Do not traverse this map. Ever. Use the function and basic block
+  // iterators to traverse instructions.
+  std::unordered_map<ir::Instruction*, ir::BasicBlock*> instr_to_block_;
+
   // A bitset indicating which analyes are currently valid.
   Analysis valid_analyses_;
 };
diff --git a/source/opt/mem_pass.cpp b/source/opt/mem_pass.cpp
index 812caf03..e0179979 100644
--- a/source/opt/mem_pass.cpp
+++ b/source/opt/mem_pass.cpp
@@ -731,14 +731,15 @@ void MemPass::RemovePhiOperands(
 
     // In all other cases, the operand must be kept but may need to be changed.
     uint32_t arg_id = phi->GetSingleWordOperand(i);
-    ir::BasicBlock* def_block = def_block_[arg_id];
+    ir::Instruction *arg_def_instr = get_def_use_mgr()->GetDef(arg_id);
+    ir::BasicBlock* def_block = context()->get_instr_block(arg_def_instr);
     if (def_block &&
-        reachable_blocks.find(def_block_[arg_id]) == reachable_blocks.end()) {
+        reachable_blocks.find(def_block) == reachable_blocks.end()) {
       // If the current |phi| argument was defined in an unreachable block, it
       // means that this |phi| argument is no longer defined. Replace it with
       // |undef_id|.
       if (!undef_id) {
-        type_id = get_def_use_mgr()->GetDef(arg_id)->type_id();
+        type_id = arg_def_instr->type_id();
         undef_id = Type2Undef(type_id);
       }
       keep_operands.push_back(
@@ -850,23 +851,5 @@ bool MemPass::CFGCleanup(ir::Function* func) {
   return modified;
 }
 
-void MemPass::InitializeCFGCleanup(ir::IRContext* c) {
-  // Build a map between SSA names to the block they are defined in.
-  //
-  // TODO(dnovillo): This is expensive and unnecessary if ir::Instruction
-  // instances could figure out what basic block they belong to. Remove this
-  // once this is possible.
-  for (auto& fn : *c->module()) {
-    for (auto& block : fn) {
-      block.ForEachInst([this, &block](ir::Instruction* inst) {
-        uint32_t result_id = inst->result_id();
-        if (result_id > 0) {
-          def_block_[result_id] = &block;
-        }
-      });
-    }
-  }
-}
-
 }  // namespace opt
 }  // namespace spvtools
diff --git a/source/opt/mem_pass.h b/source/opt/mem_pass.h
index 29beeccd..2d6e2c75 100644
--- a/source/opt/mem_pass.h
+++ b/source/opt/mem_pass.h
@@ -109,9 +109,6 @@ class MemPass : public Pass {
   // |loadInst|.
   void ReplaceAndDeleteLoad(ir::Instruction* loadInst, uint32_t replId);
 
-  // Initialize CFG Cleanup variables
-  void InitializeCFGCleanup(ir::IRContext* context);
-
   // Call all the cleanup helper functions on |func|.
   bool CFGCleanup(ir::Function* func);
 
@@ -231,11 +228,6 @@ class MemPass : public Pass {
   // The Ids of OpPhi instructions that are in a loop header and which require
   // patching of the value for the loop back-edge.
   std::unordered_set<uint32_t> phis_to_patch_;
-
-  // Map from an instruction result ID to the block that holds it.
-  // TODO(dnovillo): This would be unnecessary if ir::Instruction instances
-  // knew what basic block they belong to.
-  std::unordered_map<uint32_t, ir::BasicBlock*> def_block_;
 };
 
 }  // namespace opt