Implement rearrange_impl_horizontal_vertical

.gitignore

@@ -166,3 +166,4 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+scratch/

.pre-commit-config.yaml

@@ -2,24 +2,24 @@
 # See https://pre-commit.com/hooks.html for more hooks
 repos:
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v5.0.0
+  rev: v6.0.0
   hooks:
     - id: check-yaml
       args: ['--unsafe']
     - id: end-of-file-fixer
     - id: trailing-whitespace
 - repo: https://github.com/pycqa/isort
-  rev: 6.0.1
+  rev: 6.1.0
   hooks:
     - id: isort
       name: isort (python)
 - repo: https://github.com/psf/black
-  rev: 25.1.0
+  rev: 25.9.0
   hooks:
     - id: black
 - repo: https://github.com/charliermarsh/ruff-pre-commit
   # Ruff version.
-  rev: "v0.12.1"
+  rev: "v0.13.3"
   hooks:
     - id: ruff
 - repo: local

src/bloqade/shuttle/stdlib/layouts/two_col_zone.py

@@ -2,6 +2,7 @@
 from typing import TypeVar
 
 from bloqade.geometry.dialects import grid
+from kirin import ir
 from kirin.dialects import ilist
 
 from bloqade.shuttle import action, schedule, spec
@@ -118,18 +119,96 @@ def parking_y_end(index: int):
     action.turn_off(action.ALL, action.ALL)
 
 
+@tweezer
+def rearrange_impl_horizontal_vertical(
+    src_x: ilist.IList[int, NumX],
+    src_y: ilist.IList[int, NumY],
+    dst_x: ilist.IList[int, NumX],
+    dst_y: ilist.IList[int, NumY],
+):
+    assert len(src_x) == len(
+        dst_x
+    ), "Source and destination x indices must have the same length."
+    assert len(src_y) == len(
+        dst_y
+    ), "Source and destination y indices must have the same length."
+
+    assert_sorted(src_x)
+    assert_sorted(src_y)
+    assert_sorted(dst_x)
+    assert_sorted(dst_y)
+
+    zone = spec.get_static_trap(zone_id="traps")
+
+    start = grid.sub_grid(zone, src_x, src_y)
+    end = grid.sub_grid(zone, dst_x, dst_y)
+
+    x_positions = grid.get_xpos(zone)
+
+    # the direction for x displacement is decided based on the left (index % 2 == 0) or right site (index % 2 == 1)
+    def parking_x(index: int):
+        return x_positions[index] + 3.0 * (2 * (index % 2) - 1)
+
+    def parking_y_end(index: int):
+        start_y = grid.get_ypos(start)[index]
+        end_y = grid.get_ypos(end)[index]
+        if start_y < end_y:
+            end_y = end_y - 3.0
+        else:
+            end_y = end_y + 3.0
+
+        return end_y
+
+    num_y = len(src_y)
+
+    src_horizontal_parking = grid.from_positions(
+        ilist.map(parking_x, src_x), grid.get_ypos(start)
+    )
+    parking_y = ilist.map(parking_y_end, ilist.range(num_y))
+    mid_pos_after_vertical_move = grid.from_positions(
+        grid.get_xpos(src_horizontal_parking), parking_y
+    )
+    mid_pos_after_horizontal_move = grid.from_positions(grid.get_xpos(end), parking_y)
+
+    action.set_loc(start)
+    action.turn_on(action.ALL, action.ALL)
+    action.move(src_horizontal_parking)
+    action.move(mid_pos_after_vertical_move)
+    action.move(mid_pos_after_horizontal_move)
+    action.move(end)
+    action.turn_off(action.ALL, action.ALL)
+
+
 @move
 def rearrange(
     src_x: ilist.IList[int, NumX],
     src_y: ilist.IList[int, NumY],
     dst_x: ilist.IList[int, NumX],
     dst_y: ilist.IList[int, NumY],
+    tweezer_kernal: ir.Method = rearrange_impl,
 ):
     if len(src_x) < 1 or len(dst_x) < 1:
         return
 
     x_tones = ilist.range(len(src_x))
     y_tones = ilist.range(len(src_y))
 
-    device_fn = schedule.device_fn(rearrange_impl, x_tones, y_tones)
+    device_fn = schedule.device_fn(tweezer_kernal, x_tones, y_tones)
     device_fn(src_x, src_y, dst_x, dst_y)
+
+
+@move
+def get_device_fn(
+    src_x: ilist.IList[int, NumX],
+    src_y: ilist.IList[int, NumY],
+    dst_x: ilist.IList[int, NumX],
+    dst_y: ilist.IList[int, NumY],
+    tweezer_kernal: ir.Method = rearrange_impl,
+):
+    if len(src_x) < 1 or len(dst_x) < 1:
+        return
+
+    x_tones = ilist.range(len(src_x))
+    y_tones = ilist.range(len(src_y))
+
+    return schedule.device_fn(tweezer_kernal, x_tones, y_tones)

test/stdlib/two_col_zone/test_move.py

@@ -0,0 +1,64 @@
+from kirin.dialects import ilist
+
+from bloqade.shuttle import move
+from bloqade.shuttle.arch import ArchSpecInterpreter
+from bloqade.shuttle.codegen import TraceInterpreter, taskgen
+from bloqade.shuttle.stdlib.layouts.two_col_zone import (
+    get_device_fn,
+    get_spec,
+    rearrange_impl_horizontal_vertical,
+)
+
+
+def test_aom_move():
+    spec_value = get_spec(10, 2)
+    zone = spec_value.layout.static_traps["traps"]
+    qubit_init_x = ilist.IList([0, 2, 4, 6])
+    qubit_init_y = ilist.IList([0])
+
+    qubit_final_x = ilist.IList([1, 3, 5, 7])
+    qubit_final_y = ilist.IList([1])
+
+    device_fn = ArchSpecInterpreter(
+        dialects=move, arch_spec=(arch_spec := spec_value)
+    ).run(
+        get_device_fn,
+        (
+            qubit_init_x,
+            qubit_init_y,
+            qubit_final_x,
+            qubit_final_y,
+            rearrange_impl_horizontal_vertical,
+        ),
+    )
+    actions = TraceInterpreter(arch_spec=arch_spec).run_trace(
+        device_fn.move_fn,
+        (qubit_init_x, qubit_init_y, qubit_final_x, qubit_final_y),
+        kwargs={},
+    )
+    set_waypoint_init, turn_on, movement, turn_off, set_waypoint_end = actions
+
+    start_pos = zone.get_view(qubit_init_x, qubit_init_y)
+    end_pos = zone.get_view(qubit_final_x, qubit_final_y)
+
+    first_pos = start_pos.shift(-3.0, 0.0)
+    second_pos = first_pos.shift(0.0, 7.0)
+    third_pos = first_pos.shift(5.0, 7.0)
+
+    expected_movement = taskgen.WayPointsAction(
+        [start_pos, first_pos, second_pos, third_pos, end_pos]
+    )
+
+    assert set_waypoint_init == taskgen.WayPointsAction([start_pos])
+    assert set_waypoint_end == taskgen.WayPointsAction([end_pos])
+    assert (
+        isinstance(turn_on, taskgen.TurnOnXYSliceAction)
+        and turn_on.x_tone_indices == slice(None)
+        and turn_on.y_tone_indices == slice(None)
+    )
+    assert (
+        isinstance(turn_off, taskgen.TurnOffXYSliceAction)
+        and turn_off.x_tone_indices == slice(None)
+        and turn_off.y_tone_indices == slice(None)
+    )
+    assert movement == expected_movement