add predecessor and union types to lattice

Author: neelay893

src/kirin/analysis/const/__init__.py

@@ -13,6 +13,8 @@
     Value as Value,
     Bottom as Bottom,
     Result as Result,
+    Union as Union,
+    Predecessor as Predecessor,
     Unknown as Unknown,
     PartialConst as PartialConst,
     PartialTuple as PartialTuple,

src/kirin/analysis/const/lattice.py

@@ -264,3 +264,69 @@ def is_structurally_equal(
                 for x, y in zip(self.captured, other.captured)
             )
         )
+
+@final
+@dataclass
+class Predecessor(Result):
+    """Predecessor block in CFG."""
+
+    block: ir.Block
+    value: Result
+
+    def __hash__(self) -> int:
+        return id(self)
+
+    def is_subseteq(self, other: Result) -> bool:
+        if isinstance(other, Predecessor):
+            return self.value.is_subseteq(other.value)
+        else:
+            return self.value.is_subseteq(other)
+    
+    def join(self, other: Result) -> Result:
+        if isinstance(other, Predecessor):
+            if self.is_subseteq(other):
+                return other.value
+            elif other.is_subseteq(self):
+                return self.value
+            else:
+                return Union(predecessors=frozenset({self, other}))
+        elif isinstance(other, Union):
+            return other.join(self)
+        else:
+            return self.value.join(other)
+                
+    def meet(self, other: Result) -> Result:
+        if isinstance(other, Predecessor):
+            if self.is_subseteq(other):
+                return self.value
+            elif other.is_subseteq(self):
+                return other.value
+            else:
+                return self.bottom()
+        elif isinstance(other, Union):
+            return other.meet(self)
+        else:
+            return self.value.meet(other)
+
+@final
+@dataclass
+class Union(Result):
+
+    predecessors: frozenset[Predecessor]
+
+    def join(self, other: Result) -> Result:
+        if isinstance(other, Union):
+            union_preds = self.predecessors.union(other.predecessors)
+            return Union(predecessors=union_preds)
+        elif isinstance(other, Predecessor):
+            union_preds = self.predecessors.union({other})
+            return Union(predecessors=union_preds)
+    
+    def meet(self, other: Result) -> Result:
+        if isinstance(other, Union):
+            common_preds = self.predecessors.intersection(other.predecessors)
+            return Union(predecessors=common_preds)
+        elif isinstance(other, Predecessor):
+            common_preds = self.predecessors.intersection({other})
+            return Union(predecessors=common_preds)
+

src/kirin/dialects/cf/constprop.py

@@ -10,7 +10,8 @@ class ConstPropMethodTable(MethodTable):
     @impl(Branch)
     def branch(self, interp: const.Propagate, frame: const.Frame, stmt: Branch):
         interp.state.current_frame.worklist.append(
-            Successor(stmt.successor, *frame.get_values(stmt.arguments))
+            Successor(stmt.successor, *(const.Predecessor(stmt.parent_block, arg) 
+                                        for arg in frame.get_values(stmt.arguments)))
         )
         return ()
 
@@ -25,10 +26,12 @@ def conditional_branch(
         cond = frame.get(stmt.cond)
         if isinstance(cond, const.Value):
             else_successor = Successor(
-                stmt.else_successor, *frame.get_values(stmt.else_arguments)
+                stmt.else_successor, *(const.Predecessor(stmt.parent_block, arg) 
+                                       for arg in frame.get_values(stmt.else_arguments))
             )
             then_successor = Successor(
-                stmt.then_successor, *frame.get_values(stmt.then_arguments)
+                stmt.then_successor, *(const.Predecessor(stmt.parent_block, arg) 
+                                       for arg in frame.get_values(stmt.then_arguments))
             )
             if cond.data:
                 frame.worklist.append(then_successor)
@@ -37,13 +40,15 @@ def conditional_branch(
         else:
             frame.entries[stmt.cond] = const.Value(True)
             then_successor = Successor(
-                stmt.then_successor, *frame.get_values(stmt.then_arguments)
+                stmt.then_successor, *(const.Predecessor(stmt.parent_block, arg) 
+                                       for arg in frame.get_values(stmt.then_arguments))
             )
             frame.worklist.append(then_successor)
 
             frame.entries[stmt.cond] = const.Value(False)
             else_successor = Successor(
-                stmt.else_successor, *frame.get_values(stmt.else_arguments)
+                stmt.else_successor, *(const.Predecessor(stmt.parent_block, arg) 
+                                       for arg in frame.get_values(stmt.else_arguments))
             )
             frame.worklist.append(else_successor)