Make the Transfer Function a Generic, Interchangeable Type for the Monotone Analyzer (#5794)

This change makes the transfer function an object separate from the monotone analyzer. The transfer function class type is a generic template of the monotone analyzer, and an instance of the transfer function is instantiated and then passed into the monotone analyzer via its constructor. The API of the transfer function is enforced by a static assertion.

This change also introduces LivenessTransferFunction as a transfer function for liveness analysis as an example.

4 files changed, 218 insertions, 136 deletions

diff --git a/src/analysis/cfg.h b/src/analysis/cfg.h
index c14cf70f4..650a8c6fe 100644
--- a/src/analysis/cfg.h
+++ b/src/analysis/cfg.h
@@ -69,6 +69,10 @@ struct CFG {
   iterator end() const { return blocks.cend(); }
   size_t size() const { return blocks.size(); }
 
+  using reverse_iterator = std::vector<BasicBlock>::const_reverse_iterator;
+  reverse_iterator rbegin() const { return blocks.rbegin(); }
+  reverse_iterator rend() const { return blocks.rend(); }
+
   static CFG fromFunction(Function* func);
 
   void print(std::ostream& os, Module* wasm = nullptr) const;
diff --git a/src/analysis/liveness-transfer-function.h b/src/analysis/liveness-transfer-function.h
new file mode 100644
index 000000000..69473743f
--- /dev/null
+++ b/src/analysis/liveness-transfer-function.h
@@ -0,0 +1,91 @@
+#ifndef wasm_analysis_liveness_transfer_function_h
+#define wasm_analysis_liveness_transfer_function_h
+
+#include "lattice.h"
+#include "monotone-analyzer.h"
+
+namespace wasm::analysis {
+
+class LivenessTransferFunction : public Visitor<LivenessTransferFunction> {
+  FinitePowersetLattice::Element* currState = nullptr;
+
+public:
+  // Transfer function implementation. A local becomes live before a get
+  // and becomes dead before a set.
+  void visitLocalSet(LocalSet* curr) {
+    assert(currState);
+    currState->set(curr->index, false);
+  }
+  void visitLocalGet(LocalGet* curr) {
+    assert(currState);
+    currState->set(curr->index, true);
+  }
+
+  // Executes the transfer function on all the expressions of the corresponding
+  // CFG node, starting with the node's input state, and changes the input state
+  // to the final output state of the node in place.
+  void transfer(const BasicBlock* cfgBlock,
+                FinitePowersetLattice::Element& inputState) {
+    // If the block is empty, we propagate the state by inputState =
+    // outputState.
+
+    currState = &inputState;
+
+    // This is a backward analysis, so we start from the end of the CFG
+    // and evaluate every expression until the beginning.
+    for (auto cfgIter = cfgBlock->rbegin(); cfgIter != cfgBlock->rend();
+         ++cfgIter) {
+      // Run transfer function.
+      visit(*cfgIter);
+    }
+    currState = nullptr;
+  }
+
+  // Enqueues the worklist before the worklist algorithm is run. For
+  // liveness analysis, we are doing a backward analysis, so we want
+  // to enqueue the worklist backward so that later CFG blocks are
+  // run before earlier CFG blocks. This improves performance by
+  // reducing the number of state propagations needed, since we are
+  // naturally following the backward flow at the beginning.
+  void enqueueWorklist(CFG& cfg, std::queue<const BasicBlock*>& worklist) {
+    for (auto it = cfg.rbegin(); it != cfg.rend(); ++it) {
+      worklist.push(&(*it));
+    }
+  }
+
+  // Predecessors depend on current basic block for information.
+  BasicBlock::Predecessors getDependents(const BasicBlock* currBlock) {
+    return currBlock->preds();
+  }
+
+  // Prints the intermediate states of each BlockState currBlock by applying
+  // the transfer function on each expression of the CFG block. This data is
+  // not stored in the BlockState itself. Requires the cfgBlock, and a temp
+  // copy of the input state to be passed in, where the temp copy is modified
+  // in place to produce the intermediate states.
+  void print(std::ostream& os,
+             const BasicBlock* cfgBlock,
+             FinitePowersetLattice::Element& inputState) {
+    os << "Intermediate States (reverse order): " << std::endl;
+    currState = &inputState;
+    currState->print(os);
+    os << std::endl;
+    auto cfgIter = cfgBlock->rbegin();
+
+    // Since we don't store the intermediate states in the BlockState, we need
+    // to re-run the transfer function on all the CFG node expressions to
+    // reconstruct the intermediate states here.
+    while (cfgIter != cfgBlock->rend()) {
+      os << ShallowExpression{*cfgIter} << std::endl;
+      LivenessTransferFunction::visit(*cfgIter);
+      currState->print(os);
+      os << std::endl;
+      ++cfgIter;
+    }
+    currState = nullptr;
+  }
+};
+
+} // namespace wasm::analysis
+
+#endif // wasm_analysis_liveness_transfer_function_h
diff --git a/src/analysis/monotone-analyzer-impl.h b/src/analysis/monotone-analyzer-impl.h
index a2e059e29..471dc62f6 100644
--- a/src/analysis/monotone-analyzer-impl.h
+++ b/src/analysis/monotone-analyzer-impl.h
@@ -13,132 +13,77 @@ namespace wasm::analysis {
 template<typename Lattice>
 inline BlockState<Lattice>::BlockState(const BasicBlock* underlyingBlock,
                                        Lattice& lattice)
-  : index(underlyingBlock->getIndex()), cfgBlock(underlyingBlock),
-    beginningState(lattice.getBottom()), endState(lattice.getBottom()),
-    currState(lattice.getBottom()) {}
-
-// In our current limited implementation, we just update state on gets
-// and sets of local indices.
-template<typename Lattice>
-inline void BlockState<Lattice>::visitLocalSet(LocalSet* curr) {
-  currState.set(curr->index, false);
-}
-
-template<typename Lattice>
-inline void BlockState<Lattice>::visitLocalGet(LocalGet* curr) {
-  currState.set(curr->index, true);
-}
-
-template<typename Lattice> inline void BlockState<Lattice>::transfer() {
-  // If the block is empty, we propagate the state by endState = currState, then
-  // currState = beginningState.
-
-  // Compute transfer function for all expressions in the CFG block.
-  auto cfgIter = cfgBlock->rbegin();
-  currState = endState;
-
-  while (cfgIter != cfgBlock->rend()) {
-    // Run transfer function.
-    BlockState::visit(*cfgIter);
-    ++cfgIter;
-  }
-  beginningState = std::move(currState);
-}
+  : cfgBlock(underlyingBlock), inputState(lattice.getBottom()) {}
 
+// Prints out inforamtion about a CFG node's state, but not intermediate states.
 template<typename Lattice>
 inline void BlockState<Lattice>::print(std::ostream& os) {
-  os << "State Block: " << index << std::endl;
-  os << "Beginning State: ";
-  beginningState.print(os);
-  os << std::endl << "End State: ";
-  endState.print(os);
+  os << "CFG Block: " << cfgBlock->getIndex() << std::endl;
+  os << "Input State: ";
+  inputState.print(os);
   os << std::endl << "Predecessors:";
-  for (auto pred : predecessors) {
-    os << " " << pred->cfgBlock->getIndex();
+  for (auto pred : cfgBlock->preds()) {
+    os << " " << pred.getIndex();
   }
   os << std::endl << "Successors:";
-  for (auto succ : successors) {
-    os << " " << succ->cfgBlock->getIndex();
+  for (auto succ : cfgBlock->succs()) {
+    os << " " << succ.getIndex();
   }
-  os << std::endl << "Intermediate States (reverse order): " << std::endl;
-
-  currState = endState;
-  currState.print(os);
   os << std::endl;
-  auto cfgIter = cfgBlock->rbegin();
-
-  // Since we don't store the intermediate states in the BlockState, we need to
-  // re-run the transfer function on all the CFG node expressions to reconstruct
-  // the intermediate states here.
-  while (cfgIter != cfgBlock->rend()) {
-    os << ShallowExpression{*cfgIter} << std::endl;
-    BlockState::visit(*cfgIter);
-    currState.print(os);
-    os << std::endl;
-    ++cfgIter;
-  }
 }
 
-template<typename Lattice>
-inline void MonotoneCFGAnalyzer<Lattice>::fromCFG(CFG* cfg) {
-  // Construct BlockStates for each BasicBlock and map each BasicBlock to each
-  // BlockState's index in stateBlocks.
-  std::unordered_map<const BasicBlock*, size_t> basicBlockToState;
-  size_t index = 0;
-  for (auto it = cfg->begin(); it != cfg->end(); it++) {
-    stateBlocks.emplace_back(&(*it), lattice);
-    basicBlockToState[&(*it)] = index++;
-  }
-
-  // Update predecessors and successors of each BlockState object
-  // according to the BasicBlock's predecessors and successors.
-  for (index = 0; index < stateBlocks.size(); ++index) {
-    BlockState<Lattice>& currBlock = stateBlocks.at(index);
-    BasicBlock::Predecessors preds = currBlock.cfgBlock->preds();
-    BasicBlock::Successors succs = currBlock.cfgBlock->succs();
-    for (auto& pred : preds) {
-      currBlock.predecessors.push_back(&stateBlocks[basicBlockToState[&pred]]);
-    }
+template<typename Lattice, typename TransferFunction>
+inline MonotoneCFGAnalyzer<Lattice, TransferFunction>::MonotoneCFGAnalyzer(
+  Lattice& lattice, TransferFunction& transferFunction, CFG& cfg)
+  : lattice(lattice), transferFunction(transferFunction), cfg(cfg) {
 
-    for (auto& succ : succs) {
-      currBlock.successors.push_back(&stateBlocks[basicBlockToState[&succ]]);
-    }
+  // Construct BlockStates for each BasicBlock.
+  for (auto it = cfg.begin(); it != cfg.end(); it++) {
+    stateBlocks.emplace_back(&(*it), lattice);
   }
 }
 
-template<typename Lattice>
-inline void MonotoneCFGAnalyzer<Lattice>::evaluate() {
-  std::queue<Index> worklist;
+template<typename Lattice, typename TransferFunction>
+inline void MonotoneCFGAnalyzer<Lattice, TransferFunction>::evaluate() {
+  std::queue<const BasicBlock*> worklist;
 
-  for (auto it = stateBlocks.rbegin(); it != stateBlocks.rend(); ++it) {
-    worklist.push(it->index);
-  }
+  // Transfer function enqueues the work in some order which is efficient.
+  transferFunction.enqueueWorklist(cfg, worklist);
 
   while (!worklist.empty()) {
-    BlockState<Lattice>& currBlockState = stateBlocks[worklist.front()];
+    BlockState<Lattice>& currBlockState =
+      stateBlocks[worklist.front()->getIndex()];
     worklist.pop();
 
     // For each expression, applies the transfer function, using the expression,
     // on the state of the expression it depends upon (here the next expression)
     // to arrive at the expression's state. The beginning and end states of the
     // CFG block will be updated.
-    currBlockState.transfer();
-
-    // Propagate state to dependents
-    for (size_t j = 0; j < currBlockState.predecessors.size(); ++j) {
-      if (currBlockState.predecessors[j]->getLastState().makeLeastUpperBound(
-            currBlockState.getFirstState())) {
-        worklist.push(currBlockState.predecessors[j]->index);
+    typename Lattice::Element outputState = currBlockState.inputState;
+    transferFunction.transfer(currBlockState.cfgBlock, outputState);
+
+    // Propagate state to dependents of currBlockState.
+    for (auto& dep : transferFunction.getDependents(currBlockState.cfgBlock)) {
+      // If we need to change the input state of a dependent, we need
+      // to enqueue the dependent to recalculate it.
+      if (stateBlocks[dep.getIndex()].inputState.makeLeastUpperBound(
+            outputState)) {
+        worklist.push(&dep);
       }
     }
   }
 }
 
-template<typename Lattice>
-inline void MonotoneCFGAnalyzer<Lattice>::print(std::ostream& os) {
+// Currently prints both the basic information and intermediate states of each
+// BlockState.
+template<typename Lattice, typename TransferFunction>
+inline void
+MonotoneCFGAnalyzer<Lattice, TransferFunction>::print(std::ostream& os) {
   os << "CFG Analyzer" << std::endl;
   for (auto state : stateBlocks) {
     state.print(os);
+    typename Lattice::Element temp = state.inputState;
+    transferFunction.print(os, state.cfgBlock, temp);
   }
   os << "End" << std::endl;
 }
diff --git a/src/analysis/monotone-analyzer.h b/src/analysis/monotone-analyzer.h
index 7c0dcc630..afdf4aebf 100644
--- a/src/analysis/monotone-analyzer.h
+++ b/src/analysis/monotone-analyzer.h
@@ -11,62 +11,104 @@
 
 namespace wasm::analysis {
 
-template<typename Lattice> struct MonotoneCFGAnalyzer;
-
 // A node which contains all the lattice states for a given CFG node.
-template<typename Lattice>
-struct BlockState : public Visitor<BlockState<Lattice>> {
+template<typename Lattice> struct BlockState {
   static_assert(is_lattice<Lattice>);
-  // All states are set to the bottom lattice element in this constructor.
-  BlockState(const BasicBlock* underlyingBlock, Lattice& lattice);
-
-  typename Lattice::Element& getFirstState() { return beginningState; }
-  typename Lattice::Element& getLastState() { return endState; }
 
-  // Transfer function implementation. Modifies the state for a particular
-  // expression type.
-  void visitLocalSet(LocalSet* curr);
-  void visitLocalGet(LocalGet* curr);
+  // CFG node corresponding to this state block.
+  const BasicBlock* cfgBlock;
+  // State at which the analysis flow starts for a CFG. For instance, the ending
+  // state for backward analysis, or the beginning state for forward analysis.
+  typename Lattice::Element inputState;
 
-  // Executes the transfer function on all the expressions of the corresponding
-  // CFG node.
-  void transfer();
+  // All states are set to the bottom lattice element in this constructor.
+  BlockState(const BasicBlock* underlyingBlock, Lattice& lattice);
 
-  // Prints out all intermediate states in the block.
+  // Prints out BlockState information, but not any intermediate states.
   void print(std::ostream& os);
-
-private:
-  // The index of the block is same as the CFG index.
-  Index index;
-  const BasicBlock* cfgBlock;
-  // State at beginning of CFG node.
-  typename Lattice::Element beginningState;
-  // State at the end of the CFG node.
-  typename Lattice::Element endState;
-  // Holds intermediate state values.
-  typename Lattice::Element currState;
-  std::vector<BlockState<Lattice>*> predecessors;
-  std::vector<BlockState<Lattice>*> successors;
-
-  friend MonotoneCFGAnalyzer<Lattice>;
 };
 
-template<typename Lattice> struct MonotoneCFGAnalyzer {
-  MonotoneCFGAnalyzer(Lattice lattice) : lattice(lattice) {}
+// A transfer function using a lattice <Lattice> is required to have the
+// following methods:
+
+// Lattice::Element transfer(const BasicBlock* cfgBlock, Lattice::Element&
+// inputState);
+
+// This function takes in a pointer to a CFG BasicBlock and the input state
+// associated with it and modifies the input state in-place into the ouptut
+// state for the basic block by applying the analysis transfer function to each
+// expression in the CFG BasicBlock. Starting with the input state, the transfer
+// function is used to change the state to new intermediate states based on each
+// expression until we reach the output state. The outuput state will be
+// propagated to dependents of the CFG BasicBlock by the worklist algorithm in
+// MonotoneCFGAnalyzer.
+
+template<typename TransferFunction, typename Lattice>
+constexpr bool has_transfer =
+  std::is_invocable_r<void,
+                      decltype(&TransferFunction::transfer),
+                      TransferFunction,
+                      const BasicBlock*,
+                      typename Lattice::Element&>::value;
+
+// void enqueueWorklist(CFG&, std::queue<const BasicBlock*>& value);
+
+// Loads CFG BasicBlocks in some order into the worklist. Custom specifying the
+// order for each analysis brings performance savings. For example, when doing a
+// backward analysis, loading the BasicBlocks in reverse order will lead to less
+// state propagations, and therefore better performance. The opposite is true
+// for a forward analysis.
+
+template<typename TransferFunction, typename Lattice>
+constexpr bool has_enqueueWorklist =
+  std::is_invocable_r<void,
+                      decltype(&TransferFunction::enqueueWorklist),
+                      TransferFunction,
+                      CFG&,
+                      std::queue<const BasicBlock*>&>::value;
+
+// <iterable> getDependents(const BasicBlock* currBlock);
+
+// Returns an iterable to the CFG BasicBlocks which depend on currBlock for
+// information (e.g. predecessors in a backward analysis). Used to select which
+// blocks to propagate to after applying the transfer function to a block. At
+// present, we allow this function to return any iterable, so we only assert
+// that the method exists with the following parameters.
+
+template<typename TransferFunction>
+constexpr bool has_getDependents =
+  std::is_invocable<decltype(&TransferFunction::getDependents),
+                    TransferFunction,
+                    const BasicBlock*>::value;
+
+// Combined TransferFunction assertions.
+template<typename TransferFunction, typename Lattice>
+constexpr bool is_TransferFunction = has_transfer<TransferFunction, Lattice>&&
+  has_enqueueWorklist<TransferFunction, Lattice>&&
+    has_getDependents<TransferFunction>;
+
+template<typename Lattice, typename TransferFunction>
+class MonotoneCFGAnalyzer {
+  static_assert(is_lattice<Lattice>);
+  static_assert(is_TransferFunction<TransferFunction, Lattice>);
 
-  // Constructs a graph of BlockState objects which parallels
-  // the CFG graph. Each CFG node corresponds to a BlockState node.
-  void fromCFG(CFG* cfg);
+  Lattice& lattice;
+  TransferFunction& transferFunction;
+  CFG& cfg;
+  std::vector<BlockState<Lattice>> stateBlocks;
+
+public:
+  // Will constuct BlockState objects corresponding to BasicBlocks from the
+  // given CFG.
+  MonotoneCFGAnalyzer(Lattice& lattice,
+                      TransferFunction& transferFunction,
+                      CFG& cfg);
 
   // Runs the worklist algorithm to compute the states for the BlockList graph.
   void evaluate();
 
   // Prints out all BlockStates in this analyzer.
   void print(std::ostream& os);
-
-private:
-  Lattice lattice;
-  std::vector<BlockState<Lattice>> stateBlocks;
 };
 
 } // namespace wasm::analysis