Rename trait and type in lattice to avoid collision

Author: johnzl-777

src/bloqade/squin/analysis/nsites/__init__.py

@@ -3,6 +3,6 @@
 from .lattice import (
     NoSites as NoSites,
     AnySites as AnySites,
-    HasNSites as HasNSites,
+    NumberSites as NumberSites,
 )
 from .analysis import NSitesAnalysis as NSitesAnalysis

src/bloqade/squin/analysis/nsites/analysis.py

@@ -5,9 +5,9 @@
 from kirin.analysis.forward import ForwardFrame
 
 from bloqade.squin.op.types import OpType
-from bloqade.squin.op.traits import NSites, HasNSitesTrait
+from bloqade.squin.op.traits import NSites, HasSites
 
-from .lattice import Sites, NoSites, HasNSites
+from .lattice import Sites, NoSites, NumberSites
 
 
 class NSitesAnalysis(Forward[Sites]):
@@ -24,13 +24,13 @@ def eval_stmt(self, frame: ForwardFrame, stmt: ir.Statement):
         method = self.lookup_registry(frame, stmt)
         if method is not None:
             return method(self, frame, stmt)
-        elif stmt.has_trait(HasNSitesTrait):
-            has_n_sites_trait = stmt.get_trait(HasNSitesTrait)
+        elif stmt.has_trait(HasSites):
+            has_n_sites_trait = stmt.get_trait(HasSites)
             sites = has_n_sites_trait.get_sites(stmt)
-            return (HasNSites(sites=sites),)
+            return (NumberSites(sites=sites),)
         elif stmt.has_trait(NSites):
             sites_trait = stmt.get_trait(NSites)
-            return (HasNSites(sites=sites_trait.data),)
+            return (NumberSites(sites=sites_trait.data),)
         else:
             return (NoSites(),)
 

src/bloqade/squin/analysis/nsites/impls.py

@@ -6,7 +6,7 @@
 
 from .lattice import (
     NoSites,
-    HasNSites,
+    NumberSites,
 )
 from .analysis import NSitesAnalysis
 
@@ -18,9 +18,9 @@ class SquinOp(interp.MethodTable):
     def kron(self, interp: NSitesAnalysis, frame: interp.Frame, stmt: op.stmts.Kron):
         lhs = frame.get(stmt.lhs)
         rhs = frame.get(stmt.rhs)
-        if isinstance(lhs, HasNSites) and isinstance(rhs, HasNSites):
+        if isinstance(lhs, NumberSites) and isinstance(rhs, NumberSites):
             new_n_sites = lhs.sites + rhs.sites
-            return (HasNSites(sites=new_n_sites),)
+            return (NumberSites(sites=new_n_sites),)
         else:
             return (NoSites(),)
 
@@ -29,7 +29,7 @@ def mult(self, interp: NSitesAnalysis, frame: interp.Frame, stmt: op.stmts.Mult)
         lhs = frame.get(stmt.lhs)
         rhs = frame.get(stmt.rhs)
 
-        if isinstance(lhs, HasNSites) and isinstance(rhs, HasNSites):
+        if isinstance(lhs, NumberSites) and isinstance(rhs, NumberSites):
             lhs_sites = lhs.sites
             rhs_sites = rhs.sites
             # I originally considered throwing an exception here
@@ -40,7 +40,7 @@ def mult(self, interp: NSitesAnalysis, frame: interp.Frame, stmt: op.stmts.Mult)
             if lhs_sites != rhs_sites:
                 return (NoSites(),)
             else:
-                return (HasNSites(sites=lhs_sites + rhs_sites),)
+                return (NumberSites(sites=lhs_sites + rhs_sites),)
         else:
             return (NoSites(),)
 
@@ -50,11 +50,11 @@ def control(
     ):
         op_sites = frame.get(stmt.op)
 
-        if isinstance(op_sites, HasNSites):
+        if isinstance(op_sites, NumberSites):
             n_sites = op_sites.sites
             n_controls_attr = stmt.get_attr_or_prop("n_controls")
             n_controls = cast(ir.PyAttr[int], n_controls_attr).data
-            return (HasNSites(sites=n_sites + n_controls),)
+            return (NumberSites(sites=n_sites + n_controls),)
         else:
             return (NoSites(),)
 

src/bloqade/squin/analysis/nsites/lattice.py

@@ -40,10 +40,10 @@ def is_subseteq(self, other: Sites) -> bool:
 
 @final
 @dataclass
-class HasNSites(Sites):
+class NumberSites(Sites):
     sites: int
 
     def is_subseteq(self, other: Sites) -> bool:
-        if isinstance(other, HasNSites):
+        if isinstance(other, NumberSites):
             return self.sites == other.sites
         return False

src/bloqade/squin/op/stmts.py

@@ -2,7 +2,7 @@
 from kirin.decl import info, statement
 
 from .types import OpType
-from .traits import NSites, Unitary, MaybeUnitary, HasNSitesTrait
+from .traits import NSites, Unitary, HasSites, MaybeUnitary
 from .complex import Complex
 from ._dialect import dialect
 
@@ -77,7 +77,7 @@ class Rot(CompositeOp):
 
 @statement(dialect=dialect)
 class Identity(CompositeOp):
-    traits = frozenset({ir.Pure(), ir.FromPythonCall(), Unitary(), HasNSitesTrait()})
+    traits = frozenset({ir.Pure(), ir.FromPythonCall(), Unitary(), HasSites()})
     sites: int = info.attribute()
     result: ir.ResultValue = info.result(OpType)
 

src/bloqade/squin/op/traits.py

@@ -10,7 +10,7 @@ class NSites(ir.StmtTrait):
 
 
 @dataclass(frozen=True)
-class HasNSitesTrait(ir.StmtTrait):
+class HasSites(ir.StmtTrait):
     """An operator with a `sites` attribute."""
 
     def get_sites(self, stmt: ir.Statement):

test/squin/analysis/test_nsites_analysis.py

@@ -69,7 +69,7 @@ def test_primitive_ops():
 
     has_n_sites = []
     for nsites_type in nsites_frame.entries.values():
-        if isinstance(nsites_type, nsites.HasNSites):
+        if isinstance(nsites_type, nsites.NumberSites):
             has_n_sites.append(nsites_type)
             assert nsites_type.sites == 1
 
@@ -98,7 +98,7 @@ def test_control():
 
     has_n_sites = []
     for nsites_type in nsites_frame.entries.values():
-        if isinstance(nsites_type, nsites.HasNSites):
+        if isinstance(nsites_type, nsites.NumberSites):
             has_n_sites.append(nsites_type)
 
     assert len(has_n_sites) == 2
@@ -123,7 +123,7 @@ def test_kron():
 
     has_n_sites = []
     for nsites_type in nsites_frame.entries.values():
-        if isinstance(nsites_type, nsites.HasNSites):
+        if isinstance(nsites_type, nsites.NumberSites):
             has_n_sites.append(nsites_type)
 
     assert len(has_n_sites) == 3
@@ -151,7 +151,7 @@ def test_mult_square_same_sites():
 
     has_n_sites = []
     for nsites_type in nsites_frame.entries.values():
-        if isinstance(nsites_type, nsites.HasNSites):
+        if isinstance(nsites_type, nsites.NumberSites):
             has_n_sites.append(nsites_type)
 
     # should be three HasNSites types
@@ -190,7 +190,7 @@ def test_mult_square_different_sites():
     has_n_sites = []
     no_sites = []
     for nsite_type in nsites_types:
-        if isinstance(nsite_type, nsites.HasNSites):
+        if isinstance(nsite_type, nsites.NumberSites):
             has_n_sites.append(nsite_type)
         elif isinstance(nsite_type, nsites.NoSites):
             no_sites.append(nsite_type)
@@ -221,7 +221,7 @@ def test_rot():
 
     has_n_sites = []
     for nsites_type in nsites_frame.entries.values():
-        if isinstance(nsites_type, nsites.HasNSites):
+        if isinstance(nsites_type, nsites.NumberSites):
             has_n_sites.append(nsites_type)
 
     assert len(has_n_sites) == 2
@@ -247,7 +247,7 @@ def test_scale():
 
     has_n_sites = []
     for nsites_type in nsites_frame.entries.values():
-        if isinstance(nsites_type, nsites.HasNSites):
+        if isinstance(nsites_type, nsites.NumberSites):
             has_n_sites.append(nsites_type)
 
     assert len(has_n_sites) == 2