Adding FromRanges (#7)

* adding FromRanges statement

* adding FromRanges statement

* adding test

* adding missing test

Author: weinbe58

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

@@ -17,6 +17,7 @@
 from .concrete import GridInterpreter as GridInterpreter
 from .stmts import (
     FromPositions as FromPositions,
+    FromRanges as FromRanges,
     Get as Get,
     GetSubGrid as GetSubGrid,
     GetXBounds as GetXBounds,

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

@@ -23,6 +23,40 @@ def from_positions(
             ),
         )
 
+    @impl(stmts.FromRanges)
+    def from_ranges(
+        self,
+        interp: Interpreter,
+        frame: Frame,
+        stmt: stmts.FromRanges,
+    ):
+        x_positions = list(
+            map(
+                float,
+                range(
+                    frame.get(stmt.x_start),
+                    frame.get(stmt.x_stop),
+                    frame.get(stmt.x_step),
+                ),
+            )
+        )
+        y_positions = list(
+            map(
+                float,
+                range(
+                    frame.get(stmt.y_start),
+                    frame.get(stmt.y_stop),
+                    frame.get(stmt.y_step),
+                ),
+            )
+        )
+        return (
+            Grid.from_positions(
+                x_positions=x_positions,
+                y_positions=y_positions,
+            ),
+        )
+
     @impl(stmts.New)
     def new(
         self,

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

@@ -21,6 +21,19 @@ class FromPositions(ir.Statement):
     result: ir.ResultValue = info.result(GridType[NumX, NumY])
 
 
+@statement(dialect=dialect)
+class FromRanges(ir.Statement):
+    name = "from_ranges"
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
+
+    x_start: ir.SSAValue = info.argument(types.Int)
+    x_stop: ir.SSAValue = info.argument(types.Int)
+    x_step: ir.SSAValue = info.argument(types.Int)
+    y_start: ir.SSAValue = info.argument(types.Int)
+    y_stop: ir.SSAValue = info.argument(types.Int)
+    y_step: ir.SSAValue = info.argument(types.Int)
+
+
 @statement(dialect=dialect)
 class New(ir.Statement):
     name = "new"

test/grid/test_concrete.py

@@ -50,6 +50,26 @@ def test_new(self):
         assert len(result) == 1
         assert expected_grid_obj.is_equal(result[0])
 
+    def test_from_ranges(self):
+        x_start = 1
+        x_stop = 5
+        x_step = 1
+        y_start = 2
+        y_stop = 6
+        y_step = 1
+
+        expected_grid_obj = grid.Grid.from_positions(
+            [1.0, 2.0, 3.0, 4.0], [2.0, 3.0, 4.0, 5.0]
+        )
+
+        result = self.run_stmt(
+            grid.FromRanges, x_start, x_stop, x_step, y_start, y_stop, y_step
+        )
+
+        assert isinstance(result, tuple)
+        assert len(result) == 1
+        assert expected_grid_obj.is_equal(result[0])
+
     @pytest.mark.parametrize(
         ("stmt_type", "method_name", "args"),
         [
@@ -62,6 +82,7 @@ def test_new(self):
             (grid.Scale, "scale", (1.0, 2.0)),
             (grid.Repeat, "repeat", (1, 2, 0.5, 1.0)),
             (grid.GetSubGrid, "get_view", (ilist.IList((0,)), ilist.IList((1,)))),
+            (grid.Shape, "shape", ()),
         ],
     )
     def test_template(self, stmt_type, method_name, args):