New Address Analysis. (#563)

~In this PR I implement a Joint analysis with constant prop and Address
analysis to try to get the address analysis to work properly going into
the call stack.~

I do not use a Cartesian product for the lattic because I need to
implement partial lambda and partial IList and tuple to get the analysis
to work. In order to support constant folding, however, I opted into
adding an extra lattice element to wrap the constant prop results.

---------

Co-authored-by: David Plankensteiner <[email protected]>

Author: weinbe58

src/bloqade/analysis/address/__init__.py

@@ -1,11 +1,15 @@
 from . import impls as impls
 from .lattice import (
+    Bottom as Bottom,
     Address as Address,
-    NotQubit as NotQubit,
+    Unknown as Unknown,
     AddressReg as AddressReg,
-    AnyAddress as AnyAddress,
-    AddressWire as AddressWire,
+    UnknownReg as UnknownReg,
+    ConstResult as ConstResult,
     AddressQubit as AddressQubit,
-    AddressTuple as AddressTuple,
+    PartialIList as PartialIList,
+    PartialTuple as PartialTuple,
+    UnknownQubit as UnknownQubit,
+    PartialLambda as PartialLambda,
 )
 from .analysis import AddressAnalysis as AddressAnalysis

src/bloqade/analysis/address/analysis.py

@@ -1,13 +1,13 @@
-from typing import TypeVar
+from typing import Any, Type, TypeVar
 from dataclasses import field
 
-from kirin import ir, interp
+from kirin import ir, types, interp
 from kirin.analysis import Forward, const
+from kirin.dialects.ilist import IList
 from kirin.analysis.forward import ForwardFrame
+from kirin.analysis.const.lattice import PartialLambda
 
-from bloqade.types import QubitType
-
-from .lattice import Address
+from .lattice import Address, AddressReg, ConstResult, PartialIList, PartialTuple
 
 
 class AddressAnalysis(Forward[Address]):
@@ -16,12 +16,15 @@ class AddressAnalysis(Forward[Address]):
     """
 
     keys = ["qubit.address"]
+    _const_prop: const.Propagate
     lattice = Address
     next_address: int = field(init=False)
 
     def initialize(self):
         super().initialize()
         self.next_address: int = 0
+        self._const_prop = const.Propagate(self.dialects)
+        self._const_prop.initialize()
         return self
 
     @property
@@ -31,30 +34,117 @@ def qubit_count(self) -> int:
 
     T = TypeVar("T")
 
-    def get_const_value(self, typ: type[T], value: ir.SSAValue) -> T:
-        if isinstance(hint := value.hints.get("const"), const.Value):
-            data = hint.data
-            if isinstance(data, typ):
-                return hint.data
-            raise interp.InterpreterError(
-                f"Expected constant value <type = {typ}>, got {data}"
-            )
-        raise interp.InterpreterError(
-            f"Expected constant value <type = {typ}>, got {value}"
-        )
-
-    def eval_stmt_fallback(
-        self, frame: ForwardFrame[Address], stmt: ir.Statement
-    ) -> tuple[Address, ...] | interp.SpecialValue[Address]:
-        return tuple(
-            (
-                self.lattice.top()
-                if result.type.is_subseteq(QubitType)
-                else self.lattice.bottom()
-            )
-            for result in stmt.results
-        )
+    def to_address(self, result: const.Result):
+        return ConstResult(result)
+
+    def try_eval_const_prop(
+        self,
+        frame: ForwardFrame[Address],
+        stmt: ir.Statement,
+        args: tuple[ConstResult, ...],
+    ) -> interp.StatementResult[Address]:
+        _frame = self._const_prop.initialize_frame(frame.code)
+        _frame.set_values(stmt.args, tuple(x.result for x in args))
+        result = self._const_prop.eval_stmt(_frame, stmt)
+
+        match result:
+            case interp.ReturnValue(constant_ret):
+                return interp.ReturnValue(self.to_address(constant_ret))
+            case interp.YieldValue(constant_values):
+                return interp.YieldValue(tuple(map(self.to_address, constant_values)))
+            case interp.Successor(block, block_args):
+                return interp.Successor(block, *map(self.to_address, block_args))
+            case tuple():
+                return tuple(map(self.to_address, result))
+            case _:
+                return result
+
+    def unpack_iterable(self, iterable: Address):
+        """Extract the values of a container lattice element.
+
+        Args:
+            iterable: The lattice element representing a container.
+
+        Returns:
+            A tuple of the container type and the contained values.
+
+        """
+
+        def from_constant(constant: const.Result) -> Address:
+            return ConstResult(constant)
+
+        def from_literal(literal: Any) -> Address:
+            return ConstResult(const.Value(literal))
+
+        match iterable:
+            case PartialIList(data):
+                return PartialIList, data
+            case PartialTuple(data):
+                return PartialTuple, data
+            case AddressReg():
+                return PartialIList, iterable.qubits
+            case ConstResult(const.Value(IList() as data)):
+                return PartialIList, tuple(map(from_literal, data))
+            case ConstResult(const.Value(tuple() as data)):
+                return PartialTuple, tuple(map(from_literal, data))
+            case ConstResult(const.PartialTuple(data)):
+                return PartialTuple, tuple(map(from_constant, data))
+            case _:
+                return None, ()
+
+    def run_lattice(
+        self,
+        callee: Address,
+        inputs: tuple[Address, ...],
+        kwargs: tuple[str, ...],
+    ) -> Address:
+        """Run a callable lattice element with the given inputs and keyword arguments.
+
+        Args:
+            callee (Address): The lattice element representing the callable.
+            inputs (tuple[Address, ...]): The input lattice elements.
+            kwargs (tuple[str, ...]): The keyword argument names.
+
+        Returns:
+            Address: The resulting lattice element after invoking the callable.
+
+        """
+
+        match callee:
+            case PartialLambda(code=code, argnames=argnames):
+                _, ret = self.run_callable(
+                    code, (callee,) + self.permute_values(argnames, inputs, kwargs)
+                )
+                return ret
+            case ConstResult(const.Value(ir.Method() as method)):
+                _, ret = self.run_method(
+                    method,
+                    self.permute_values(method.arg_names, inputs, kwargs),
+                )
+                return ret
+            case _:
+                return Address.top()
+
+    def get_const_value(self, addr: Address, typ: Type[T]) -> T | None:
+        if not isinstance(addr, ConstResult):
+            return None
+
+        if not isinstance(result := addr.result, const.Value):
+            return None
+
+        if not isinstance(value := result.data, typ):
+            return None
+
+        return value
+
+    def eval_stmt_fallback(self, frame: ForwardFrame[Address], stmt: ir.Statement):
+        args = frame.get_values(stmt.args)
+        if types.is_tuple_of(args, ConstResult):
+            return self.try_eval_const_prop(frame, stmt, args)
+
+        return tuple(Address.from_type(result.type) for result in stmt.results)
 
     def run_method(self, method: ir.Method, args: tuple[Address, ...]):
         # NOTE: we do not support dynamic calls here, thus no need to propagate method object
-        return self.run_callable(method.code, (self.lattice.bottom(),) + args)
+        self_mt = ConstResult(const.Value(method))
+        return self.run_callable(method.code, (self_mt,) + args)