feat: add index field to repeat blocks to name loop index

Signed-off-by: Louis Mandel <[email protected]>

Author: mandel

pdl-live-react/src/pdl_ast.d.ts

@@ -1483,6 +1483,11 @@ export type Kind10 = "repeat"
 export type For = {
   [k: string]: LocalizedExpression | unknown[] | string
 } | null
+/**
+ * Name of the variable containing the loop iteration.
+ *
+ */
+export type Index = string | null
 /**
  * Condition to stay at the beginning of the loop.
  *
@@ -3834,6 +3839,16 @@ export interface Defs9 {
  * repeat:
  *     "${ name }'s number is ${ number }\n"
  * ```
+ *
+ * Bounded loop:
+ * ```PDL
+ * index: i
+ * max_iterations: 5
+ * repeat:
+ *     ${ i }
+ * join:
+ *   as: array
+ * ```
  */
 export interface RepeatBlock {
   description?: Description10
@@ -3873,6 +3888,7 @@ export interface RepeatBlock {
   context?: IndependentEnum5
   kind?: Kind10
   for?: For
+  index?: Index
   while?: While
   repeat: Repeat
   until?: Until

pdl-live-react/src/pdl_code_cleanup.ts

@@ -123,6 +123,9 @@ function clean_repeat_block(block: RepeatBlock) {
   if (block.for === null) {
     delete block.for
   }
+  if (block.index === null) {
+    delete block.index
+  }
   if (block.while === true) {
     delete block.while
   }

src/pdl/pdl-schema.json

@@ -10642,7 +10642,7 @@
     },
     "RepeatBlock": {
       "additionalProperties": false,
-      "description": "Repeat the execution of a block.\n\nFor loop example:\n```PDL\nfor:\n    number: [1, 2, 3, 4]\n    name: [\"Bob\", \"Carol\", \"David\", \"Ernest\"]\nrepeat:\n    \"${ name }'s number is ${ number }\\n\"\n```",
+      "description": "Repeat the execution of a block.\n\nFor loop example:\n```PDL\nfor:\n    number: [1, 2, 3, 4]\n    name: [\"Bob\", \"Carol\", \"David\", \"Ernest\"]\nrepeat:\n    \"${ name }'s number is ${ number }\\n\"\n```\n\nBounded loop:\n```PDL\nindex: i\nmax_iterations: 5\nrepeat:\n    ${ i }\njoin:\n  as: array\n```",
       "properties": {
         "description": {
           "anyOf": [
@@ -11042,6 +11042,19 @@
           "description": "Arrays to iterate over.\n    ",
           "title": "For"
         },
+        "index": {
+          "anyOf": [
+            {
+              "type": "string"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "default": null,
+          "description": "Name of the variable containing the loop iteration.\n    ",
+          "title": "Index"
+        },
         "while": {
           "anyOf": [
             {

src/pdl/pdl_ast.py

@@ -872,12 +872,25 @@ class RepeatBlock(StructuredBlock):
     repeat:
         "${ name }'s number is ${ number }\\n"
     ```
+
+    Bounded loop:
+    ```PDL
+    index: i
+    max_iterations: 5
+    repeat:
+        ${ i }
+    join:
+      as: array
+    ```
     """
 
     kind: Literal[BlockKind.REPEAT] = BlockKind.REPEAT
     for_: Optional[dict[str, ExpressionType[list]]] = Field(default=None, alias="for")
     """Arrays to iterate over.
     """
+    index: Optional[str] = None
+    """Name of the variable containing the loop iteration.
+    """
     while_: ExpressionType[bool] = Field(default=True, alias="while")
     """Condition to stay at the beginning of the loop.
     """

src/pdl/pdl_dumper.py

@@ -231,6 +231,8 @@ def block_to_dict(  # noqa: C901
         case RepeatBlock():
             if block.for_ is not None:
                 d["for"] = expr_to_dict(block.for_, json_compatible)
+            if block.index is not None:
+                d["index"] = block.index
             if block.while_ is not None:
                 d["while"] = expr_to_dict(block.while_, json_compatible)
             d["repeat"] = block_to_dict(block.repeat, json_compatible)

src/pdl/pdl_interpreter.py

@@ -807,6 +807,8 @@ def process_block_body(
                 first = True
                 saved_background: PdlLazy[list[dict[str, Any]]] = PdlList([])
                 while True:
+                    if block.index is not None:
+                        scope = scope | {block.index: iidx}
                     if max_iterations is not None and iidx >= max_iterations:
                         break
                     if lengths is not None and iidx >= lengths[0]:
@@ -859,8 +861,8 @@ def process_block_body(
                     results.append(iteration_result)
                     iter_trace.append(body_trace)
                     iteration_state = iteration_state.with_pop()
-                    iidx = iidx + 1
                     stop, _ = process_condition_of(block, "until", scope, loc)
+                    iidx = iidx + 1
                     if stop:
                         break
             except PDLRuntimeError as exc:

tests/test_repeat.py

@@ -174,6 +174,20 @@ def test_for_max_iterations2():
     assert result == "012"
 
 
+def _for_max_iterations_index_prog(n):
+    return f"""
+    index: i
+    repeat: {'${i}'}
+    max_iterations: {'${'}{n}{'}'}
+    """
+
+
+def test_for_max_iterations_index1():
+    prog = _for_max_iterations_index_prog(10)
+    result = exec_str(prog)
+    assert result == "0123456789"
+
+
 def _for_until_prog(n):
     return f"""
     for: