Allow slicing to work with address analysis (#374)

This allows slicing to work with the address analysis pass.

@liupengy19 ran into this problem with some of his work and @david-pl
helped me confirm the issue was not specific to my setup (:

Note: the way the tests are written for the address analysis are
definitely a bit outdated, I believe once #370 goes through I updated
the tests there so it will just be a matter of merging in the new
updated stuff

Author: johnzl-777

src/bloqade/analysis/address/impls.py

@@ -70,19 +70,35 @@ def new_ilist(
 class PyIndexing(interp.MethodTable):
     @interp.impl(py.GetItem)
     def getitem(self, interp: AddressAnalysis, frame: interp.Frame, stmt: py.GetItem):
-        # Integer index into the thing being indexed
-        idx = interp.get_const_value(int, stmt.index)
+
+        # determine if the index is an int constant
+        # or a slice
+        hint = stmt.index.hints.get("const")
+        if hint is None:
+            return (NotQubit(),)
+
+        if isinstance(hint, const.Value):
+            idx = hint.data
+        elif isinstance(hint, slice):
+            idx = hint
+        else:
+            return (NotQubit(),)
 
         # The object being indexed into
         obj = frame.get(stmt.obj)
         # The `data` attributes holds onto other Address types
         # so we just extract that here
         if isinstance(obj, AddressTuple):
             return (obj.data[idx],)
-        # an AddressReg is guaranteed to just have some sequence
-        # of integers which is directly pluggable to AddressQubit
+        # If idx is an integer index into an AddressReg,
+        # then it's safe to assume a single qubit is being accessed.
+        # On the other hand, if it's a slice, we return
+        # a new AddressReg to preserve the new sequence.
         elif isinstance(obj, AddressReg):
-            return (AddressQubit(obj.data[idx]),)
+            if isinstance(idx, slice):
+                return (AddressReg(data=obj.data[idx]),)
+            if isinstance(idx, int):
+                return (AddressQubit(obj.data[idx]),)
         else:
             return (NotQubit(),)
 

test/analysis/address/test_analysis.py

@@ -4,6 +4,7 @@
 from kirin.dialects import py, func, ilist
 
 from bloqade import qasm2, squin
+from bloqade.squin import qubit
 from bloqade.analysis import address
 
 
@@ -203,6 +204,38 @@ def test_multiple_wire_apply():
     assert len(address_wire_parent_qubit_1) == 2
 
 
+def test_slice():
+
+    @squin.kernel
+    def main():
+        q = qubit.new(4)
+        # get the middle qubits out and apply to them
+        sub_q = q[1:3]
+        qubit.broadcast(squin.op.x(), sub_q)
+        # get a single qubit out, do some stuff
+        single_q = sub_q[0]
+        qubit.apply(squin.op.h(), single_q)
+
+    address_analysis = address.AddressAnalysis(main.dialects)
+    frame, _ = address_analysis.run_analysis(main, no_raise=False)
+
+    address_regs = [
+        address_reg_type
+        for address_reg_type in frame.entries.values()
+        if isinstance(address_reg_type, address.AddressReg)
+    ]
+    address_qubits = [
+        address_qubit_type
+        for address_qubit_type in frame.entries.values()
+        if isinstance(address_qubit_type, address.AddressQubit)
+    ]
+
+    assert address_regs[0] == address.AddressReg(data=range(0, 4))
+    assert address_regs[1] == address.AddressReg(data=range(1, 3))
+
+    assert address_qubits[0] == address.AddressQubit(data=1)
+
+
 def test_for_loop():
     @squin.kernel
     def main():