From 572a6cf9f5c361ba111fe052976627ccd18b58f8 Mon Sep 17 00:00:00 2001
From: David Plankensteiner <da.plankensteiner@gmail.com>
Date: Wed, 6 Aug 2025 10:49:03 +0200
Subject: [PATCH] Add PauliString impl to NsitesAnalysis

---
 src/bloqade/squin/analysis/nsites/impls.py  |  7 +++++++
 test/squin/analysis/test_nsites_analysis.py | 16 ++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/src/bloqade/squin/analysis/nsites/impls.py b/src/bloqade/squin/analysis/nsites/impls.py
index 302c37e7..aaae6ef1 100644
--- a/src/bloqade/squin/analysis/nsites/impls.py
+++ b/src/bloqade/squin/analysis/nsites/impls.py
@@ -81,6 +81,13 @@ def scale(self, interp: NSitesAnalysis, frame: interp.Frame, stmt: op.stmts.Scal
         op_sites = frame.get(stmt.op)
         return (op_sites,)
 
+    @interp.impl(op.stmts.PauliString)
+    def pauli_string(
+        self, interp: NSitesAnalysis, frame: interp.Frame, stmt: op.stmts.PauliString
+    ):
+        s = stmt.string
+        return (NumberSites(sites=len(s)),)
+
 
 @scf.dialect.register(key="op.nsites")
 class ScfSquinOp(ScfTypeInfer):
diff --git a/test/squin/analysis/test_nsites_analysis.py b/test/squin/analysis/test_nsites_analysis.py
index 46d7d9ee..88836419 100644
--- a/test/squin/analysis/test_nsites_analysis.py
+++ b/test/squin/analysis/test_nsites_analysis.py
@@ -248,3 +248,19 @@ def test():
     assert has_n_sites[0].sites == 1
     for n_site in has_n_sites[1:]:
         assert n_site.sites == 2
+
+
+def test_pauli_string():
+    @squin.kernel(fold=False)
+    def main():
+        squin.op.pauli_string(string="XYZ")
+
+    nsites_frame, _ = nsites.NSitesAnalysis(main.dialects).run_analysis(main)
+
+    has_n_sites = []
+    for nsites_type in nsites_frame.entries.values():
+        if isinstance(nsites_type, nsites.NumberSites):
+            has_n_sites.append(nsites_type)
+
+    assert len(has_n_sites) == 1
+    assert has_n_sites[0].sites == 3