Refactor no-cloning validation: enhance error handling and improve test coverage

Author: zhenrongliew

src/bloqade/analysis/validation/nocloning/analysis.py

@@ -10,24 +10,43 @@
     Address,
     AddressAnalysis,
 )
-from bloqade.analysis.address.lattice import AddressReg, AddressQubit
+from bloqade.analysis.address.lattice import (
+    Unknown,
+    AddressReg,
+    UnknownReg,
+    AddressQubit,
+    PartialIList,
+    PartialTuple,
+    UnknownQubit,
+)
 
-from .lattice import QubitValidation
+from .lattice import May, Top, Must, Bottom, QubitValidation
 
 
 class QubitValidationError(ValidationError):
-    """ValidationError that records which qubit and gate caused the violation."""
+    """ValidationError for definite (Must) violations with concrete qubit addresses."""
 
     qubit_id: int
     gate_name: str
 
     def __init__(self, node: ir.IRNode, qubit_id: int, gate_name: str):
-        # message stored in ValidationError so formatting/hint() will include it
         super().__init__(node, f"Qubit[{qubit_id}] cloned at {gate_name} gate.")
         self.qubit_id = qubit_id
         self.gate_name = gate_name
 
 
+class PotentialQubitValidationError(ValidationError):
+    """ValidationError for potential (May) violations with unknown addresses."""
+
+    gate_name: str
+    condition: str
+
+    def __init__(self, node: ir.IRNode, gate_name: str, condition: str):
+        super().__init__(node, f"Potential cloning at {gate_name} gate{condition}.")
+        self.gate_name = gate_name
+        self.condition = condition
+
+
 class NoCloningValidation(Forward[QubitValidation]):
     """
     Validates the no-cloning theorem by tracking qubit addresses.
@@ -40,9 +59,7 @@ class NoCloningValidation(Forward[QubitValidation]):
     _address_frame: ForwardFrame[Address] = field(init=False)
     _type_frame: ForwardFrame = field(init=False)
     method: ir.Method
-    _validation_errors: list[QubitValidationError] = field(
-        default_factory=list, init=False
-    )
+    _validation_errors: list[ValidationError] = field(default_factory=list, init=False)
 
     def __init__(self, mtd: ir.Method):
         """
@@ -88,68 +105,126 @@ def eval_stmt_fallback(
     ) -> tuple[QubitValidation, ...]:
         """
         Default statement evaluation: check for qubit usage violations.
+        Returns Bottom, May, Must, or Top depending on what we can prove.
         """
 
         if not isinstance(stmt, func.Invoke):
-            return tuple(QubitValidation.bottom() for _ in stmt.results)
+            return tuple(Bottom() for _ in stmt.results)
 
         address_frame = self._address_frame
         if address_frame is None:
-            return tuple(QubitValidation.top() for _ in stmt.results)
+            return tuple(Top() for _ in stmt.results)
 
-        has_qubit_args = any(
-            isinstance(address_frame.get(arg), (AddressQubit, AddressReg))
-            for arg in stmt.args
-        )
+        concrete_addrs: list[int] = []
+        has_unknown = False
+        has_qubit_args = False
+        unknown_arg_names: list[str] = []
 
-        if not has_qubit_args:
-            return tuple(QubitValidation.bottom() for _ in stmt.results)
-
-        used_addrs: list[int] = []
         for arg in stmt.args:
             addr = address_frame.get(arg)
-            qubit_addrs = self.get_qubit_addresses(addr)
-            used_addrs.extend(qubit_addrs)
+            match addr:
+                case AddressQubit(data=qubit_addr):
+                    has_qubit_args = True
+                    concrete_addrs.append(qubit_addr)
+                case AddressReg(data=addrs):
+                    has_qubit_args = True
+                    concrete_addrs.extend(addrs)
+                case UnknownQubit() | UnknownReg() | Unknown():
+                    has_qubit_args = True
+                    has_unknown = True
+                    arg_name = self._get_source_name(arg)
+                    unknown_arg_names.append(arg_name)
+                case _:
+                    pass
+
+        if not has_qubit_args:
+            return tuple(Bottom() for _ in stmt.results)
 
         seen: set[int] = set()
-        violations: list[str] = []
+        must_violations: list[str] = []
+        gate_name = stmt.callee.sym_name.upper()
 
-        for qubit_addr in used_addrs:
+        for qubit_addr in concrete_addrs:
             if qubit_addr in seen:
-                gate_name = stmt.callee.sym_name.upper()
-                violations.append(self.format_violation(qubit_addr, gate_name))
+                violation = self.format_violation(qubit_addr, gate_name)
+                must_violations.append(violation)
                 self._validation_errors.append(
                     QubitValidationError(stmt, qubit_addr, gate_name)
                 )
             seen.add(qubit_addr)
 
-        if not violations:
-            return tuple(QubitValidation(violations=frozenset()) for _ in stmt.results)
+        if must_violations:
+            usage = Must(violations=frozenset(must_violations))
+        elif has_unknown:
+            args_str = " == ".join(unknown_arg_names)
+            if len(unknown_arg_names) > 1:
+                condition = f", when {args_str}"
+            else:
+                condition = f", with unknown index {args_str}"
+
+            self._validation_errors.append(
+                PotentialQubitValidationError(stmt, gate_name, condition)
+            )
+
+            usage = May(
+                violations=frozenset([f"Unknown qubits at {gate_name} Gate{condition}"])
+            )
+        else:
+            usage = Bottom()
 
-        usage = QubitValidation(violations=frozenset(violations))
         return tuple(usage for _ in stmt.results) if stmt.results else (usage,)
 
+    def _get_source_name(self, value: ir.SSAValue) -> str:
+        """Trace back to get the source variable name for a value.
+
+        For getitem operations like q[a], returns 'a'.
+        For direct values, returns the value's name.
+        """
+        from kirin.dialects.py.indexing import GetItem
+
+        if isinstance(value, ir.ResultValue) and isinstance(value.stmt, GetItem):
+            index_arg = value.stmt.args[1]
+            return self._get_source_name(index_arg)
+
+        if isinstance(value, ir.BlockArgument):
+            return value.name or f"arg{value.index}"
+
+        if hasattr(value, "name") and value.name:
+            return value.name
+
+        return str(value)
+
     def run_method(
         self, method: ir.Method, args: tuple[QubitValidation, ...]
     ) -> tuple[ForwardFrame[QubitValidation], QubitValidation]:
         self_mt = self.method_self(method)
         return self.run_callable(method.code, (self_mt,) + args)
 
     def raise_validation_errors(self):
-        """Raise validation error for each no-cloning violation found.
+        """Raise validation errors for both definite and potential violations.
         Points to source file and line with snippet.
         """
         if not self._validation_errors:
             return
 
-        # If multiple errors, print all with snippets first
-        if len(self._validation_errors) > 1:
-            for err in self._validation_errors:
-                err.attach(self.method)
-                # Print error message before snippet
+        # Print all errors with snippets
+        for err in self._validation_errors:
+            err.attach(self.method)
+
+            # Format error message based on type
+            if isinstance(err, QubitValidationError):
                 print(
-                    f"\033[31mValidation Error\033[0m: Cloned qubit [{err.qubit_id}] at {err.gate_name} gate."
+                    f"\n\033[31mError\033[0m: Cloning qubit [{err.qubit_id}] at {err.gate_name} gate"
                 )
-                print(err.hint())
-            print(f"Raised {len(self._validation_errors)} error(s).")
+            elif isinstance(err, PotentialQubitValidationError):
+                print(
+                    f"\n\033[33mWarning\033[0m: Potential cloning at {err.gate_name} gate{err.condition}"
+                )
+            else:
+                print(
+                    f"\n\033[31mError\033[0m: {err.args[0] if err.args else type(err).__name__}"
+                )
+
+            print(err.hint())
+
         raise

src/bloqade/analysis/validation/nocloning/lattice.py

@@ -1,3 +1,4 @@
+from abc import abstractmethod
 from typing import FrozenSet, final
 from dataclasses import field, dataclass
 
@@ -15,9 +16,11 @@ class QubitValidation(
     SimpleMeetMixin["QubitValidation"],
     BoundedLattice["QubitValidation"],
 ):
-    """Tracks cloning violations detected during analysis."""
+    """Base class for qubit cloning validation lattice.
 
-    violations: FrozenSet[str] = field(default_factory=frozenset)
+    Lattice ordering:
+    Bottom ⊑ May{...} ⊑ Must{...} ⊑ Top
+    """
 
     @classmethod
     def bottom(cls) -> "QubitValidation":
@@ -26,52 +29,95 @@ def bottom(cls) -> "QubitValidation":
 
     @classmethod
     def top(cls) -> "QubitValidation":
-        """Unknown state - assume potential violations"""
+        """Unknown state"""
         return Top()
 
+    @abstractmethod
     def is_subseteq(self, other: "QubitValidation") -> bool:
-        """Check if this state is a subset of another.
-
-        Lattice ordering:
-        Bottom ⊑ {{'Qubit[1] at CX Gate'}} ⊑ {{'Qubit[0] at CX Gate'},{'Qubit[1] at CX Gate'}} ⊑ Top
-        """
-        if isinstance(other, Top):
-            return True
-        if isinstance(self, Bottom):
-            return True
-        if isinstance(other, Bottom):
-            return False
-
-        return self.violations.issubset(other.violations)
-
-    def __repr__(self) -> str:
-        """Custom repr to show violations clearly."""
-        if not self.violations:
-            return "QubitValidation()"
-        return f"QubitValidation(violations={self.violations})"
+        """Check if this state is a subset of another."""
+        ...
 
 
 @final
 class Bottom(QubitValidation, metaclass=SingletonMeta):
-    """Bottom element representing no violations."""
+    """Bottom element: no violations detected (safe)."""
 
     def is_subseteq(self, other: QubitValidation) -> bool:
         """Bottom is subset of everything."""
         return True
 
     def __repr__(self) -> str:
-        """Cleaner printing."""
-        return "⊥ (Bottom)"
+        return "⊥ (No Errors)"
 
 
 @final
 class Top(QubitValidation, metaclass=SingletonMeta):
-    """Top element representing unknown state with potential violations."""
+    """Top element: unknown state (worst case - assume violations possible)."""
 
     def is_subseteq(self, other: QubitValidation) -> bool:
         """Top is only subset of Top."""
         return isinstance(other, Top)
 
     def __repr__(self) -> str:
-        """Cleaner printing."""
-        return "⊤ (Top)"
+        return "⊤ (Unknown)"
+
+
+@final
+@dataclass
+class May(QubitValidation):
+    """Potential violations that may occur depending on runtime values.
+
+    Used when we have unknown addresses (UnknownQubit, UnknownReg, Unknown).
+    """
+
+    violations: FrozenSet[str] = field(default_factory=frozenset)
+
+    def is_subseteq(self, other: QubitValidation) -> bool:
+        """May ⊑ May' if violations ⊆ violations'
+        May ⊑ Must (any may is less precise than must)
+        May ⊑ Top
+        """
+        match other:
+            case Bottom():
+                return False
+            case May(violations=other_violations):
+                return self.violations.issubset(other_violations)
+            case Must():
+                return True  # May is less precise than Must
+            case Top():
+                return True
+        return False
+
+    def __repr__(self) -> str:
+        if not self.violations:
+            return "MayError(∅)"
+        return f"MayError({self.violations})"
+
+
+@final
+@dataclass
+class Must(QubitValidation):
+    """Definite violations with concrete qubit addresses.
+
+    These are violations we can prove will definitely occur.
+    """
+
+    violations: FrozenSet[str] = field(default_factory=frozenset)
+
+    def is_subseteq(self, other: QubitValidation) -> bool:
+        """Must ⊑ Must' if violations ⊆ violations'
+        Must ⊑ Top
+        """
+        match other:
+            case Bottom() | May():
+                return False
+            case Must(violations=other_violations):
+                return self.violations.issubset(other_violations)
+            case Top():
+                return True
+        return False
+
+    def __repr__(self) -> str:
+        if not self.violations:
+            return "MustError(∅)"
+        return f"MustError({self.violations})"