3 files changed, 72 insertions, 21 deletions
diff --git a/src/analysis/lattice.h b/src/analysis/lattice.h
index 4fb794f1b..6b468c2fc 100644
--- a/src/analysis/lattice.h
+++ b/src/analysis/lattice.h
@@ -2,6 +2,7 @@
 #define wasm_analysis_lattice_h
 
 #include <iostream>
+#include <unordered_map>
 #include <vector>
 
 #include "wasm.h"
@@ -39,18 +40,17 @@ constexpr bool is_lattice = has_getBottom<Lattice>&& has_compare<Lattice>&&
   has_makeLeastUpperBound<typename Lattice::Element>&& has_isTop<
     typename Lattice::Element>&& has_isBottom<typename Lattice::Element>;
 
-// Represents a powerset lattice which is constructed from a finite set which
-// can be represented by a bitvector. Set elements are represented by
-// FinitePowersetLattice::Element, which represents members present in each
-// element by bits in the bitvector.
-class FinitePowersetLattice {
-
+// Represents a powerset lattice constructed from a finite set of consecutive
+// integers from 0 to n which can be represented by a bitvector. Set elements
+// are represented by FiniteIntPowersetLattice::Element, which represents
+// members present in each element by bits in the bitvector.
+class FiniteIntPowersetLattice {
   // The size of the set that the powerset lattice was created from. This is
   // equivalent to the size of the Top lattice element.
   size_t setSize;
 
 public:
-  FinitePowersetLattice(size_t setSize) : setSize(setSize) {}
+  FiniteIntPowersetLattice(size_t setSize) : setSize(setSize) {}
 
   // This represents an element of a powerset lattice. The element is itself a
   // set which has set members. The bitvector tracks which possible members of
@@ -90,7 +90,7 @@ public:
     // Prints out the bits in the bitvector for a lattice element.
     void print(std::ostream& os);
 
-    friend FinitePowersetLattice;
+    friend FiniteIntPowersetLattice;
   };
 
   // Compares two lattice elements and returns a result indicating the
@@ -101,7 +101,58 @@ public:
   Element getBottom();
 };
 
-static_assert(is_lattice<FinitePowersetLattice>);
+// A layer of abstraction over FiniteIntPowersetLattice which maps
+// set members of some type T to indices in the bitvector. Allows
+// the finite powerset lattice to be generalized to arbitrary types.
+template<typename T> class FinitePowersetLattice {
+  FiniteIntPowersetLattice intLattice;
+
+  // Maps a bitvector index to some element member of type T.
+  // Used to produce initial ordering of element members.
+  std::vector<T> members;
+
+  // Maps an element member of type T to a bitvector index.
+  std::unordered_map<T, size_t> memberIndices;
+
+public:
+  using Element = FiniteIntPowersetLattice::Element;
+
+  // Takes in an ordered list of all elements of the set to create
+  // the powerset lattice from (i.e. the powerset lattice top element). This
+  // is used for mapping the elements to bitvector indices.
+  FinitePowersetLattice(std::vector<T>&& setMembers)
+    : intLattice(setMembers.size()), members(std::move(setMembers)) {
+    for (size_t i = 0; i < members.size(); ++i) {
+      memberIndices[members[i]] = i;
+    }
+  }
+
+  T indexToMember(size_t index) { return members[index]; }
+
+  size_t memberToIndex(T member) { return memberIndices[member]; }
+
+  // Adds member to a powerset lattice element.
+  void add(Element* element, T member) {
+    element->set(memberIndices[member], true);
+  }
+
+  // Removes member from a powerset lattice element.
+  void remove(Element* element, T member) {
+    element->set(memberIndices[member], false);
+  }
+
+  // Checks if member is included in the element set.
+  bool exists(Element* element, T member) {
+    return element->get(memberIndices[member]);
+  }
+
+  // We use implementations from FiniteIntPowersetLattice here.
+  static LatticeComparison compare(const Element& left, const Element& right) {
+    return FiniteIntPowersetLattice::compare(left, right);
+  }
+
+  Element getBottom() { return intLattice.getBottom(); }
+};
 
 } // namespace wasm::analysis
 
diff --git a/src/analysis/liveness-transfer-function.h b/src/analysis/liveness-transfer-function.h
index 69473743f..c24e67eab 100644
--- a/src/analysis/liveness-transfer-function.h
+++ b/src/analysis/liveness-transfer-function.h
@@ -7,7 +7,7 @@
 namespace wasm::analysis {
 
 class LivenessTransferFunction : public Visitor<LivenessTransferFunction> {
-  FinitePowersetLattice::Element* currState = nullptr;
+  FiniteIntPowersetLattice::Element* currState = nullptr;
 
 public:
   // Transfer function implementation. A local becomes live before a get
@@ -25,7 +25,7 @@ public:
   // 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) {
+                FiniteIntPowersetLattice::Element& inputState) {
     // If the block is empty, we propagate the state by inputState =
     // outputState.
 
@@ -65,7 +65,7 @@ public:
   // in place to produce the intermediate states.
   void print(std::ostream& os,
              const BasicBlock* cfgBlock,
-             FinitePowersetLattice::Element& inputState) {
+             FiniteIntPowersetLattice::Element& inputState) {
     os << "Intermediate States (reverse order): " << std::endl;
     currState = &inputState;
     currState->print(os);
diff --git a/src/analysis/powerset-lattice-impl.h b/src/analysis/powerset-lattice-impl.h
index 24931618a..e4e6dd4cf 100644
--- a/src/analysis/powerset-lattice-impl.h
+++ b/src/analysis/powerset-lattice-impl.h
@@ -5,9 +5,9 @@
 
 namespace wasm::analysis {
 
-inline LatticeComparison
-FinitePowersetLattice::compare(const FinitePowersetLattice::Element& left,
-                               const FinitePowersetLattice::Element& right) {
+inline LatticeComparison FiniteIntPowersetLattice::compare(
+  const FiniteIntPowersetLattice::Element& left,
+  const FiniteIntPowersetLattice::Element& right) {
   // Both must be from the powerset lattice of the same set.
   assert(left.bitvector.size() == right.bitvector.size());
 
@@ -41,14 +41,14 @@ FinitePowersetLattice::compare(const FinitePowersetLattice::Element& left,
   return NO_RELATION;
 }
 
-inline FinitePowersetLattice::Element FinitePowersetLattice::getBottom() {
-  FinitePowersetLattice::Element result(setSize);
+inline FiniteIntPowersetLattice::Element FiniteIntPowersetLattice::getBottom() {
+  FiniteIntPowersetLattice::Element result(setSize);
   return result;
 }
 
 // We count the number of element members present in the element by counting the
 // trues in the bitvector.
-inline size_t FinitePowersetLattice::Element::count() {
+inline size_t FiniteIntPowersetLattice::Element::count() {
   size_t count = 0;
   for (auto it : bitvector) {
     count += it;
@@ -59,8 +59,8 @@ inline size_t FinitePowersetLattice::Element::count() {
 // Least upper bound is implemented as a logical OR between the bitvectors on
 // both sides. We return true if a bit is flipped in-place on the left so the
 // worklist algorithm will know if when to enqueue more work.
-inline bool FinitePowersetLattice::Element::makeLeastUpperBound(
-  const FinitePowersetLattice::Element& other) {
+inline bool FiniteIntPowersetLattice::Element::makeLeastUpperBound(
+  const FiniteIntPowersetLattice::Element& other) {
   // Both must be from powerset lattice of the same set.
   assert(other.bitvector.size() == bitvector.size());
 
@@ -75,7 +75,7 @@ inline bool FinitePowersetLattice::Element::makeLeastUpperBound(
   return modified;
 }
 
-inline void FinitePowersetLattice::Element::print(std::ostream& os) {
+inline void FiniteIntPowersetLattice::Element::print(std::ostream& os) {
   // Element member 0 is on the left, element member N is on the right.
   for (auto it : bitvector) {
     os << it;