feat: add index field to repeat blocks to name loop index (#950)

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

Author: mandel

docs/tutorial.md

@@ -794,7 +794,7 @@ This results in the following output:
 {"Bob": 1, "Carol": 2, "David": 3, "Ernest": 4}
 ```
 
-### While Loop
+### Other loops
 
 The following example shows a while loop in PDL:
 
@@ -804,8 +804,14 @@ The following example shows a while loop in PDL:
 
 The `while` field indicates the looping condition and `repeat` contains the body of the loop.
 
-Notice that `for`, `while`, `until`, and `maxIterations` can all be combined in the same `repeat` block. The loop exits
-as soon as one of the exit conditions is satisfied:
+This loop can be rewritten using `max_iterations` to bound the number of iterations and `index` to name the iteration number.
+
+```yaml
+--8<-- "./examples/tutorial/loop_index.pdl"
+```
+
+
+Notice that `for`, `while`, `until`, and `max_iterations` can all be combined in the same `repeat` block. The loop exits as soon as one of the exit conditions is satisfied:
 
 ```yaml
 --8<-- "./examples/tutorial/repeat.pdl"

examples/tutorial/loop_index.pdl

@@ -0,0 +1,4 @@
+index: i
+repeat:
+  text: ${i + 1}
+max_iterations: 3

examples/tutorial/repeat.pdl

@@ -1,13 +1,14 @@
 description: repeat loop with multiple conditions
 defs:
   numbers:
-    data: [1, 2, 3, 4]
+    data: [42, 2, 4012, 27]
   names:
     data: ["Bob", "Carol", "David", "Ernest"]
 for:
   number: ${ numbers }
   name: ${ names }
+index: i
 repeat:
-  "${ name }'s number is ${ number }\n"
+  "${i}: ${ name }'s number is ${ number }\n"
 until: ${ name == "Carol"}
-max_iterations: 1
+max_iterations: 3

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/results/examples/tutorial/loop_index.0.result

@@ -0,0 +1 @@
+123