Wrap rewrite pass

Author: johnzl-777

src/bloqade/squin/analysis/rewrite/stim.py

@@ -1,30 +1,77 @@
 from typing import Dict
 from dataclasses import dataclass
 
+from kirin import ir
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+from kirin.print.printer import Printer
+
+from bloqade.squin import op, wire
 from bloqade.analysis.address import Address
 from bloqade.squin.analysis.nsites import Sites
 
-from kirin import ir
-from kirin.rewrite.abc import RewriteResult, RewriteRule
+# Probably best to move these attributes to a
+# separate file? Keep here for now
+# to get things working first
+
+
+@wire.dialect.register
+@dataclass
+class AddressAttribute(ir.Attribute):
+
+    name = "Address"
+    address: Address
+
+    def __hash__(self) -> int:
+        return hash(self.address)
+
+    def print_impl(self, printer: Printer) -> None:
+        # Can return to implementing this later
+        pass
 
 
+@op.dialect.register
 @dataclass
-class SquinToStim(RewriteRule):
-    
-    # Somehow need to plug in Address and Sites
-    # into the SSAValue Hints field, which only accepts
-    # Attribute types
+class SitesAttribute(ir.Attribute):
 
-    ## Could literally just plug in `ir.Attribute` into
-    ## the Address and Site lattices? 
-    ## Couldn't I just create my own attributes instead?
+    name = "Sites"
+    sites: Sites
+
+    def __hash__(self) -> int:
+        return hash(self.sites)
+
+    def print_impl(self, printer: Printer) -> None:
+        # Can return to implementing this later
+        pass
+
+
+@dataclass
+class WrapSquinAnalysis(RewriteRule):
 
     address_analysis: Dict[ir.SSAValue, Address]
-    op_site_analysis: Dict[ir.SSAValue, Sites] 
+    op_site_analysis: Dict[ir.SSAValue, Sites]
+
+    def wrap(self, value: ir.SSAValue) -> bool:
+        address_analysis_result = self.address_analysis[value]
+        op_site_analysis_result = self.op_site_analysis[value]
+
+        if value.hints["address"] and value.hints["sites"]:
+            return False
+        else:
+            value.hints["address"] = AddressAttribute(address_analysis_result)
+            value.hints["sites"] = SitesAttribute(op_site_analysis_result)
+
+        return True
 
-    # need to plug in data into the SSAValue 
-    # for the rewrite from these passes,
-    # then something should look at those hints 
-    # and generate the corresponding stim statements 
+    def rewrite_Block(self, node: ir.Block) -> RewriteResult:
+        has_done_something = False
+        for arg in node.args:
+            if self.wrap(arg):
+                has_done_something = True
+        return RewriteResult(has_done_something=has_done_something)
 
-    pass
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        has_done_something = False
+        for result in node.results:
+            if self.wrap(result):
+                has_done_something = True
+        return RewriteResult(has_done_something=has_done_something)