Adding statement that generates all positions of the grid (#16)

* Adding  statement that generates all positions of the grid

* updating pre-commit

* run linter

Author: weinbe58

.pre-commit-config.yaml

@@ -14,7 +14,7 @@ repos:
     - id: isort
       name: isort (python)
 - repo: https://github.com/psf/black
-  rev: 24.3.0
+  rev: 25.1.0
   hooks:
     - id: black
 - repo: https://github.com/charliermarsh/ruff-pre-commit

src/bloqade/geometry/dialects/grid/__init__.py

@@ -5,6 +5,7 @@
     get_xpos as get_xpos,
     get_ypos as get_ypos,
     new as new,
+    positions as positions,
     repeat as repeat,
     scale as scale,
     shape as shape,
@@ -25,6 +26,7 @@
     GetYBounds as GetYBounds,
     GetYPos as GetYPos,
     New as New,
+    Positions as Positions,
     Repeat as Repeat,
     Scale as Scale,
     Shape as Shape,

src/bloqade/geometry/dialects/grid/_interface.py

@@ -12,6 +12,7 @@
     GetYBounds,
     GetYPos,
     New,
+    Positions,
     Repeat,
     Scale,
     Shape,
@@ -164,6 +165,22 @@ def y_bounds(grid: Grid[typing.Any, typing.Any]) -> tuple[float, float]:
     ...
 
 
+@_wraps(Positions)
+def positions(
+    grid: Grid[typing.Any, typing.Any],
+) -> ilist.IList[tuple[float, float], typing.Any]:
+    """Get the positions of a grid as a list of (x, y) tuples.
+
+    Args:
+        grid (Grid): a grid object
+
+    Returns:
+        ilist.IList[tuple[float, float], typing.Any]: a list of (x, y) tuples representing the positions of the grid points
+
+    """
+    ...
+
+
 @_wraps(Repeat)
 def repeat(
     grid: Grid, x_times: int, y_times: int, x_spacing: float, y_spacing: float

src/bloqade/geometry/dialects/grid/concrete.py

@@ -181,3 +181,13 @@ def repeat(
         y_gap = frame.get_casted(stmt.y_gap, float)
 
         return (grid.repeat(x_times, y_times, x_gap, y_gap),)
+
+    @impl(stmts.Positions)
+    def positions(
+        self,
+        interp: Interpreter,
+        frame: Frame,
+        stmt: stmts.Positions,
+    ):
+        grid = frame.get_casted(stmt.zone, Grid)
+        return (grid.positions,)

src/bloqade/geometry/dialects/grid/stmts.py

@@ -52,6 +52,16 @@ class New(ir.Statement):
     )
 
 
+@statement(dialect=dialect)
+class Positions(ir.Statement):
+    name = "positions"
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
+    zone: ir.SSAValue = info.argument(type=GridType[types.Any, types.Any])
+    result: ir.ResultValue = info.result(
+        ilist.IListType[types.Tuple[types.Float, types.Float], types.Any]
+    )
+
+
 # Maybe do this with hints?
 @statement(dialect=dialect)
 class Shape(ir.Statement):

src/bloqade/geometry/dialects/grid/types.py

@@ -1,6 +1,6 @@
 import dataclasses
 from functools import cached_property
-from itertools import chain
+from itertools import chain, product
 from typing import Any, Generic, Sequence, TypeVar
 
 from kirin import ir, types
@@ -118,6 +118,10 @@ def y_positions(self) -> tuple[float, ...]:
             )
         )
 
+    @cached_property
+    def positions(self) -> ilist.IList[tuple[float, float], Any]:
+        return ilist.IList(tuple(product(self.x_positions, self.y_positions)))
+
     def get(self, idx: tuple[int, int]) -> tuple[float, float]:
         return (self.x_positions[idx[0]], self.y_positions[idx[1]])