fixing compiler after refactoring circular imports

src/bloqade/builder2/base.py

@@ -1,8 +1,7 @@
 from typing import Optional
-from .compile.trait import CompileJSON
 
 
-class Builder(CompileJSON):
+class Builder:
     __match_args__ = ("__parent__",)
 
     def __init__(

src/bloqade/builder2/compile/ir.py

@@ -1,32 +1,36 @@
-from typing import Tuple, Union
-from .stream import BuilderNode, BuilderStream
-from ... import ir
+from typing import Tuple, Union, TYPE_CHECKING
+
 from ..base import Builder
 
 from ..coupling import LevelCoupling, Rydberg, Hyperfine
 from ..field import Field, Detuning, RabiAmplitude, RabiPhase
 from ..spatial import SpatialModulation, Location, Uniform, Var, Scale
 from ..waveform import (
     WaveformPrimitive,
-    # Linear,
-    # Constant,
-    # Poly,
-    # PiecewiseConstant,
-    # PiecewiseLinear,
-    # Fn,
     Slice,
     Record,
 )
 from ..parallelize import Parallelize, ParallelizeFlatten
 
+if TYPE_CHECKING:
+    from bloqade.ir import Waveform as ir_Waveform
+    from bloqade.ir import SpatialModulation as ir_SpatialModulation
+    from bloqade.ir import Field as ir_Field
+    from bloqade.ir import Sequence as ir_Sequence
+    from bloqade.ir import ParallelRegister as ir_ParallelRegister
+    from bloqade.ir import AtomArrangement as ir_AtomArrangement
+    from .stream import BuilderNode
+
 
 class BuilderCompiler:
     def __init__(self, ast: Builder) -> None:
+        from .stream import BuilderStream
+
         self.stream = BuilderStream.create(ast)
 
 
 class SequenceCompiler(BuilderCompiler):
-    def read_address(self) -> Tuple[LevelCoupling, Field, BuilderNode]:
+    def read_address(self) -> Tuple[LevelCoupling, Field, "BuilderNode"]:
         spatial = self.stream.eat([Location, Uniform, Var], [Scale])
         curr = spatial
 
@@ -52,9 +56,9 @@ def read_address(self) -> Tuple[LevelCoupling, Field, BuilderNode]:
         else:  # only spatial is updated
             return (None, None, spatial)
 
-    def read_waveform(self, head: BuilderNode) -> Tuple[ir.Waveform, BuilderNode]:
+    def read_waveform(self, head: "BuilderNode") -> Tuple["ir_Waveform", "BuilderNode"]:
         curr = head
-        waveform: ir.Waveform = head.node.__bloqade_ir__()
+        waveform = head.node.__bloqade_ir__()
         curr = curr.next
         while curr is not None:
             match curr.node:
@@ -70,8 +74,10 @@ def read_waveform(self, head: BuilderNode) -> Tuple[ir.Waveform, BuilderNode]:
         return waveform, curr
 
     def read_spatial_modulation(
-        self, head: BuilderNode
-    ) -> Tuple[ir.SpatialModulation, BuilderNode]:
+        self, head: "BuilderNode"
+    ) -> Tuple["ir_SpatialModulation", "BuilderNode"]:
+        import bloqade.ir as ir
+
         curr = head
         spatial_modulation = None
         scaled_locations = ir.ScaledLocations({})
@@ -95,22 +101,26 @@ def read_spatial_modulation(
         else:
             return spatial_modulation, curr
 
-    def read_field(self, head) -> Tuple[ir.Field, BuilderNode]:
+    def read_field(self, head) -> "ir_Field":
+        from bloqade.ir import Field
+
         sm, curr = self.read_spatial_modulation(head)
-        wf, curr = self.read_waveform(curr)
-        return ir.Field({sm: wf}), curr
+        wf, _ = self.read_waveform(curr)
+        return Field({sm: wf})
+
+    def compile(self) -> "ir_Sequence":
+        from bloqade.ir import Sequence, Pulse, Field
 
-    def compile(self) -> ir.Sequence:
         coupling_builder, field_builder, spatial_head = self.read_address()
 
-        sequence = ir.Sequence({})
+        sequence = Sequence({})
         coupling_name = coupling_builder.__bloqade_ir__()
         field_name = field_builder.__bloqade_ir__()
 
-        pulse = sequence.pulses.get(coupling_name, ir.Pulse({}))
-        field = pulse.fields.get(field_name, ir.Field({}))
+        pulse = sequence.pulses.get(coupling_name, Pulse({}))
+        field = pulse.fields.get(field_name, Field({}))
 
-        new_field, _ = self.read_field(spatial_head)
+        new_field = self.read_field(spatial_head)
         field = field.add(new_field)
 
         while True:
@@ -121,24 +131,28 @@ def compile(self) -> ir.Sequence:
                 sequence.pulses[coupling_name] = pulse
                 # create/access new pulse
                 coupling_name = coupling_builder.__bloqade_ir__()
-                pulse = sequence.pulses.get(coupling_name, ir.Pulse({}))
+                pulse = sequence.pulses.get(coupling_name, Pulse({}))
 
             if field_builder is not None:
                 # update old field
                 pulse.fields[field_name] = field
                 # create/access new field
                 field_name = field_builder.__bloqade_ir__()
-                field = pulse.fields.get(field_name, ir.Field({}))
+                field = pulse.fields.get(field_name, Field({}))
 
             if spatial_head is None:
                 break
 
-            new_field, _ = self.read_field(spatial_head)
+            new_field = self.read_field(spatial_head)
             field = field.add(new_field)
 
+        return sequence
+
 
 class RegisterCompiler(BuilderCompiler):
-    def compile(self) -> Union[ir.AtomArrangement, ir.ParallelRegister]:
+    def compile(self) -> Union["ir_AtomArrangement", "ir_ParallelRegister"]:
+        from bloqade.ir import ParallelRegister
+
         # register is always head of the stream
         register_block = self.stream.read()
 
@@ -148,6 +162,6 @@ def compile(self) -> Union[ir.AtomArrangement, ir.ParallelRegister]:
 
         if parallel_options is not None:
             parallel_options = parallel_options.node
-            return ir.ParallelRegister(register, parallel_options._cluster_spacing)
+            return ParallelRegister(register, parallel_options._cluster_spacing)
 
         return register

src/bloqade/builder2/compile/json.py

@@ -1,24 +1,25 @@
 from ...codegen.common.json import BloqadeIRSerializer
-from ..start import ProgramStart
-from ..sequence_builder import SequenceBuilder
-from ..base import Builder
-from ..spatial import Location, Scale, Var
-from ..waveform import (
-    Linear,
-    Constant,
-    Poly,
-    Fn,
-    Apply,
-    Slice,
-    Record,
-    Sample,
-    PiecewiseLinear,
-    PiecewiseConstant,
-)
 
 
 class BuilderSerializer(BloqadeIRSerializer):
     def default(self, obj):
+        from ..start import ProgramStart
+        from ..sequence_builder import SequenceBuilder
+        from ..base import Builder
+        from ..spatial import Location, Scale, Var
+        from ..waveform import (
+            Linear,
+            Constant,
+            Poly,
+            Fn,
+            Apply,
+            Slice,
+            Record,
+            Sample,
+            PiecewiseLinear,
+            PiecewiseConstant,
+        )
+
         match obj:
             case Constant(value, duration, parent):
                 return {

src/bloqade/builder2/compile/trait.py

@@ -1,13 +1,12 @@
 from ... import ir
-from .ir import SequenceCompiler, RegisterCompiler
-from .json import BuilderSerializer
-from ...codegen.common.static_assign import StaticAssignProgram
 from typing import Union
 import json
 
 
 class CompileJSON:
     def json(self, **json_options) -> str:
+        from .json import BuilderSerializer
+
         return json.dumps(self, cls=BuilderSerializer, **json_options)
 
     def __repr__(self):
@@ -16,11 +15,15 @@ def __repr__(self):
 
 class CompileRegister:
     def compile_register(self) -> Union[ir.AtomArrangement, ir.ParallelRegister]:
+        from .ir import RegisterCompiler
+
         return RegisterCompiler(self).compile()
 
 
 class CompileSequence:
     def compile_sequence(self):
+        from .ir import SequenceCompiler
+
         return SequenceCompiler(self).compile()
 
 
@@ -31,4 +34,6 @@ def compile_program(self) -> ir.Program:
 
 class Compile(CompileProgram):
     def compile_bloqade_ir(self, **mapping):
+        from ...codegen.common.static_assign import StaticAssignProgram
+
         return StaticAssignProgram(mapping).emit(self.compile_program())

src/bloqade/builder2/spatial.py

@@ -1,8 +1,11 @@
-from typing import Union, Optional
+from typing import Union, Optional, TYPE_CHECKING
 from .waveform import WaveformAttachable
-from .base import Builder, DefaultBuilderEncoder
+from .base import Builder
 from ..ir import Scalar
 
+if TYPE_CHECKING:
+    from bloqade.ir.control.field import UniformModulation
+
 ScalarType = Union[float, str, Scalar]
 
 
@@ -11,7 +14,10 @@ class SpatialModulation(WaveformAttachable):
 
 
 class Uniform(SpatialModulation):
-    pass
+    def __bloqade_ir__(self) -> "UniformModulation":
+        from ..ir import Uniform
+
+        return Uniform
 
 
 class Location(SpatialModulation):
@@ -47,18 +53,3 @@ def __init__(self, name: str, parent: Optional[Builder] = None) -> None:
         assert isinstance(name, str)
         super().__init__(parent)
         self._name = name
-
-
-class SpatialSerializer(DefaultBuilderEncoder):
-    def default(self, obj):
-        match obj:
-            case Uniform(parent):
-                return dict(type="Uniform", parent=self.default(parent))
-            case Var(name, parent):
-                return dict(type="Var", name=name, parent=self.default(parent))
-            case Scale(factor, parent):
-                return dict(type="Scale", factor=factor, parent=self.default(parent))
-            case Location(label, parent):
-                return dict(type="Location", label=label, parent=self.default(parent))
-            case _:
-                return super().default(obj)