feat: add J/K loop index words and nested control flow tests

CLAUDE.md

@@ -87,7 +87,7 @@ uv run ruff format gpu_test/
 
 - **Stack Type**: `!forth.stack` - untyped stack, programmer ensures type safety
 - **Operations**: All take stack as input and produce stack as output (except `forth.stack`)
-- **Supported Words**: literals, `DUP DROP SWAP OVER ROT NIP TUCK PICK ROLL`, `+ - * / MOD`, `AND OR XOR NOT LSHIFT RSHIFT`, `= < > <> <= >= 0=`, `@ !`, `CELLS`, `IF ELSE THEN`, `BEGIN UNTIL`, `BEGIN WHILE REPEAT`, `DO LOOP I`, `TID-X/Y/Z BID-X/Y/Z BDIM-X/Y/Z GDIM-X/Y/Z GLOBAL-ID` (GPU indexing).
+- **Supported Words**: literals, `DUP DROP SWAP OVER ROT NIP TUCK PICK ROLL`, `+ - * / MOD`, `AND OR XOR NOT LSHIFT RSHIFT`, `= < > <> <= >= 0=`, `@ !`, `CELLS`, `IF ELSE THEN`, `BEGIN UNTIL`, `BEGIN WHILE REPEAT`, `DO LOOP I J K`, `TID-X/Y/Z BID-X/Y/Z BDIM-X/Y/Z GDIM-X/Y/Z GLOBAL-ID` (GPU indexing).
 - **Kernel Parameters**: Declared with `PARAM <name> <size>`, each becomes a `memref<Nxi64>` function argument with `forth.param_name` attribute. Using a param name in code pushes its byte address onto the stack via `forth.param_ref`
 - **Conversion**: `!forth.stack` → `memref<256xi64>` with explicit stack pointer
 - **GPU**: Functions wrapped in `gpu.module`, `main` gets `gpu.kernel` attribute, configured with bare pointers for NVVM conversion

include/warpforth/Dialect/Forth/ForthOps.td

@@ -529,9 +529,20 @@ def Forth_BeginWhileRepeatOp : Forth_Op<"begin_while_repeat",
   let hasCustomAssemblyFormat = 1;
 }
 
-def Forth_LoopIndexOp : Forth_StackOpBase<"loop_index"> {
-  let summary = "Push loop index onto stack (I word)";
-  let description = [{ Only valid inside a forth.do_loop body. ( -- i ) }];
+def Forth_LoopIndexOp : Forth_Op<"loop_index", [Pure]> {
+  let summary = "Push loop index onto stack (I/J/K words)";
+  let description = [{
+    Pushes the loop index at the given nesting depth onto the stack.
+    depth=0 is I (innermost), depth=1 is J, depth=2 is K.
+    Only valid inside nested forth.do_loop bodies at sufficient depth.
+    ( -- index )
+  }];
+  let arguments = (ins Forth_StackType:$input_stack,
+                       DefaultValuedAttr<I64Attr, "0">:$depth);
+  let results = (outs Forth_StackType:$output_stack);
+  let assemblyFormat = [{
+    $input_stack attr-dict `:` type($input_stack) `->` type($output_stack)
+  }];
 }
 
 //===----------------------------------------------------------------------===//

lib/Conversion/ForthToMemRef/ForthToMemRef.cpp

@@ -1071,13 +1071,23 @@ struct LoopIndexOpConversion : public OpConversionPattern<forth::LoopIndexOp> {
     Value memref = inputStack[0];
     Value stackPtr = inputStack[1];
 
-    // Find enclosing scf.for
-    auto forOp = op->getParentOfType<scf::ForOp>();
-    if (!forOp)
-      return rewriter.notifyMatchFailure(op, "not inside an scf.for");
+    // Get the depth attribute (default 0 = I)
+    int64_t depth = op.getDepth();
+
+    // Walk up through enclosing scf.for ops
+    Operation *current = op.getOperation();
+    scf::ForOp targetFor;
+    for (int64_t i = 0; i <= depth; ++i) {
+      targetFor = current->getParentOfType<scf::ForOp>();
+      if (!targetFor)
+        return rewriter.notifyMatchFailure(
+            op, "not enough enclosing scf.for ops for depth " +
+                    std::to_string(depth));
+      current = targetFor.getOperation();
+    }
 
     // Get induction variable and cast index to i64
-    Value iv = forOp.getInductionVar();
+    Value iv = targetFor.getInductionVar();
     Value ivI64 =
         rewriter.create<arith::IndexCastOp>(loc, rewriter.getI64Type(), iv);
 

lib/Translation/ForthToMLIR/ForthToMLIR.cpp

@@ -356,12 +356,16 @@ Value ForthParser::emitOperation(StringRef word, Value inputStack,
   } else if (word == "0=") {
     return builder.create<forth::ZeroEqOp>(loc, stackType, inputStack)
         .getResult();
-  } else if (word == "I") {
-    if (doLoopDepth == 0) {
-      (void)emitError("'I' used outside of DO/LOOP");
+  } else if (word == "I" || word == "J" || word == "K") {
+    int64_t depth = (word == "I") ? 0 : (word == "J") ? 1 : 2;
+    if (doLoopDepth < depth + 1) {
+      (void)emitError("'" + word.str() + "' requires " +
+                      std::to_string(depth + 1) + " nested DO/LOOP(s)");
       return nullptr;
     }
-    return builder.create<forth::LoopIndexOp>(loc, stackType, inputStack)
+    return builder
+        .create<forth::LoopIndexOp>(loc, stackType, inputStack,
+                                    builder.getI64IntegerAttr(depth))
         .getResult();
   }
 

test/Conversion/ForthToMemRef/nested-control-flow.mlir

@@ -0,0 +1,149 @@
+// RUN: %warpforth-opt --convert-forth-to-memref %s | %FileCheck %s
+
+// === Nested IF → nested scf.if ===
+// CHECK-LABEL: func.func private @test_nested_if
+// CHECK: scf.if %{{.*}} -> (index) {
+// CHECK:   scf.if %{{.*}} -> (index) {
+// CHECK:     scf.yield
+// CHECK:   } else {
+// CHECK:     scf.yield
+// CHECK:   }
+// CHECK:   scf.yield
+// CHECK: } else {
+// CHECK:   scf.yield
+// CHECK: }
+
+module {
+  func.func private @test_nested_if() {
+    %0 = forth.stack !forth.stack
+    %1 = forth.literal %0 1 : !forth.stack -> !forth.stack
+    %2 = forth.if %1 : !forth.stack -> !forth.stack {
+    ^bb0(%arg0: !forth.stack):
+      %3 = forth.drop %arg0 : !forth.stack -> !forth.stack
+      %4 = forth.literal %3 1 : !forth.stack -> !forth.stack
+      %5 = forth.if %4 : !forth.stack -> !forth.stack {
+      ^bb0(%arg1: !forth.stack):
+        %6 = forth.drop %arg1 : !forth.stack -> !forth.stack
+        %7 = forth.literal %6 42 : !forth.stack -> !forth.stack
+        forth.yield %7 : !forth.stack
+      } else {
+      ^bb0(%arg1: !forth.stack):
+        %6 = forth.drop %arg1 : !forth.stack -> !forth.stack
+        forth.yield %6 : !forth.stack
+      }
+      forth.yield %5 : !forth.stack
+    } else {
+    ^bb0(%arg0: !forth.stack):
+      %3 = forth.drop %arg0 : !forth.stack -> !forth.stack
+      forth.yield %3 : !forth.stack
+    }
+    return
+  }
+
+  // === IF inside DO/LOOP → scf.if inside scf.for ===
+  // CHECK-LABEL: func.func private @test_if_inside_do
+  // CHECK: scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%{{.*}} = %{{.*}}) -> (index) {
+  // CHECK:   scf.if %{{.*}} -> (index) {
+  // CHECK:     scf.yield
+  // CHECK:   } else {
+  // CHECK:     scf.yield
+  // CHECK:   }
+  // CHECK:   scf.yield
+  // CHECK: }
+
+  func.func private @test_if_inside_do() {
+    %0 = forth.stack !forth.stack
+    %1 = forth.literal %0 10 : !forth.stack -> !forth.stack
+    %2 = forth.literal %1 0 : !forth.stack -> !forth.stack
+    %3 = forth.do_loop %2 : !forth.stack -> !forth.stack {
+    ^bb0(%arg0: !forth.stack):
+      %4 = forth.loop_index %arg0 : !forth.stack -> !forth.stack
+      %5 = forth.literal %4 5 : !forth.stack -> !forth.stack
+      %6 = forth.gt %5 : !forth.stack -> !forth.stack
+      %7 = forth.if %6 : !forth.stack -> !forth.stack {
+      ^bb0(%arg1: !forth.stack):
+        %8 = forth.drop %arg1 : !forth.stack -> !forth.stack
+        %9 = forth.literal %8 99 : !forth.stack -> !forth.stack
+        forth.yield %9 : !forth.stack
+      } else {
+      ^bb0(%arg1: !forth.stack):
+        %8 = forth.drop %arg1 : !forth.stack -> !forth.stack
+        forth.yield %8 : !forth.stack
+      }
+      forth.yield %7 : !forth.stack
+    }
+    return
+  }
+
+  // === Nested DO/LOOP with J (depth=1) ===
+  // CHECK-LABEL: func.func private @test_nested_do_j
+  // Outer scf.for
+  // CHECK: scf.for %[[OUTER_IV:.*]] = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%{{.*}} = %{{.*}}) -> (index) {
+  //   Inner scf.for
+  // CHECK:   scf.for %[[INNER_IV:.*]] = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%{{.*}} = %{{.*}}) -> (index) {
+  //     J pushes outer IV, I pushes inner IV
+  // CHECK:     arith.index_cast %[[OUTER_IV]]
+  // CHECK:     arith.index_cast %[[INNER_IV]]
+  // CHECK:     scf.yield
+  // CHECK:   }
+  // CHECK:   scf.yield
+  // CHECK: }
+
+  func.func private @test_nested_do_j() {
+    %0 = forth.stack !forth.stack
+    %1 = forth.literal %0 3 : !forth.stack -> !forth.stack
+    %2 = forth.literal %1 0 : !forth.stack -> !forth.stack
+    %3 = forth.do_loop %2 : !forth.stack -> !forth.stack {
+    ^bb0(%arg0: !forth.stack):
+      %4 = forth.literal %arg0 4 : !forth.stack -> !forth.stack
+      %5 = forth.literal %4 0 : !forth.stack -> !forth.stack
+      %6 = forth.do_loop %5 : !forth.stack -> !forth.stack {
+      ^bb0(%arg1: !forth.stack):
+        %7 = forth.loop_index %arg1 {depth = 1 : i64} : !forth.stack -> !forth.stack
+        %8 = forth.loop_index %7 : !forth.stack -> !forth.stack
+        %9 = forth.add %8 : !forth.stack -> !forth.stack
+        forth.yield %9 : !forth.stack
+      }
+      forth.yield %6 : !forth.stack
+    }
+    return
+  }
+
+  // === BEGIN/WHILE/REPEAT inside IF ===
+  // CHECK-LABEL: func.func private @test_while_inside_if
+  // CHECK: scf.if %{{.*}} -> (index) {
+  // CHECK:   scf.while
+  // CHECK:     scf.condition
+  // CHECK:   } do {
+  // CHECK:     scf.yield
+  // CHECK:   }
+  // CHECK:   scf.yield
+  // CHECK: } else {
+  // CHECK:   scf.yield
+  // CHECK: }
+
+  func.func private @test_while_inside_if() {
+    %0 = forth.stack !forth.stack
+    %1 = forth.literal %0 1 : !forth.stack -> !forth.stack
+    %2 = forth.if %1 : !forth.stack -> !forth.stack {
+    ^bb0(%arg0: !forth.stack):
+      %3 = forth.drop %arg0 : !forth.stack -> !forth.stack
+      %4 = forth.begin_while_repeat %3 : !forth.stack -> !forth.stack {
+      ^bb0(%arg1: !forth.stack):
+        %5 = forth.dup %arg1 : !forth.stack -> !forth.stack
+        forth.yield %5 while_cond : !forth.stack
+      } do {
+      ^bb0(%arg1: !forth.stack):
+        %5 = forth.literal %arg1 1 : !forth.stack -> !forth.stack
+        %6 = forth.sub %5 : !forth.stack -> !forth.stack
+        forth.yield %6 : !forth.stack
+      }
+      forth.yield %4 : !forth.stack
+    } else {
+    ^bb0(%arg0: !forth.stack):
+      %3 = forth.drop %arg0 : !forth.stack -> !forth.stack
+      forth.yield %3 : !forth.stack
+    }
+    return
+  }
+}

test/Pipeline/nested-control-flow.forth

@@ -0,0 +1,15 @@
+\ RUN: %warpforth-translate --forth-to-mlir %s | %warpforth-opt --warpforth-pipeline | %FileCheck %s
+\ RUN: %warpforth-translate --forth-to-mlir %s | %warpforth-opt --convert-forth-to-memref --convert-scf-to-cf --convert-forth-to-gpu | %FileCheck %s --check-prefix=MID
+
+\ Verify nested DO/LOOP with I and J through full pipeline produces gpu.binary
+\ CHECK: gpu.binary @warpforth_module
+
+\ Verify intermediate MLIR: gpu.func with nested loop structure
+\ MID: gpu.module @warpforth_module
+\ MID: gpu.func @main(%arg0: memref<4xi64> {forth.param_name = "DATA"}) kernel
+\ MID: cf.br
+\ MID: cf.cond_br
+\ MID: gpu.return
+
+PARAM DATA 4
+3 0 DO 4 0 DO J I + LOOP LOOP DATA 0 CELLS + !

test/Translation/Forth/loop-index-depth-error.forth

@@ -0,0 +1,3 @@
+\ RUN: %not %warpforth-translate --forth-to-mlir %s 2>&1 | %FileCheck %s
+\ CHECK: 'J' requires 2 nested DO/LOOP(s)
+10 0 DO J LOOP

test/Translation/Forth/nested-control-flow.forth

@@ -0,0 +1,83 @@
+\ RUN: %warpforth-translate --forth-to-mlir %s | %FileCheck %s
+
+\ === Nested IF ===
+\ CHECK: %[[S0:.*]] = forth.stack
+\ CHECK: %[[S1:.*]] = forth.literal %[[S0]] 1
+\ CHECK: %[[IF1:.*]] = forth.if %[[S1]]
+\ CHECK:   ^bb0(%[[A1:.*]]: !forth.stack):
+\ CHECK:   %[[D1:.*]] = forth.drop %[[A1]]
+\ CHECK:   %[[L2:.*]] = forth.literal %[[D1]] 2
+\ CHECK:   %[[IF2:.*]] = forth.if %[[L2]]
+\ CHECK:     ^bb0(%[[A2:.*]]: !forth.stack):
+\ CHECK:     forth.drop %[[A2]]
+\ CHECK:     forth.literal %{{.*}} 3
+\ CHECK:     forth.yield
+\ CHECK:   } else {
+\ CHECK:     ^bb0(%[[A3:.*]]: !forth.stack):
+\ CHECK:     forth.drop %[[A3]]
+\ CHECK:     forth.yield
+\ CHECK:   }
+\ CHECK:   forth.yield
+\ CHECK: } else {
+\ CHECK:   ^bb0(%[[A4:.*]]: !forth.stack):
+\ CHECK:   forth.drop %[[A4]]
+\ CHECK:   forth.yield
+\ CHECK: }
+1 IF 2 IF 3 THEN THEN
+
+\ === IF inside DO ===
+\ CHECK: %[[S2:.*]] = forth.literal %[[IF1]] 10
+\ CHECK: %[[S3:.*]] = forth.literal %[[S2]] 0
+\ CHECK: %[[LOOP1:.*]] = forth.do_loop %[[S3]]
+\ CHECK:   ^bb0(%[[BA:.*]]: !forth.stack):
+\ CHECK:   %[[LI:.*]] = forth.loop_index %[[BA]]
+\ CHECK:   %[[L5:.*]] = forth.literal %[[LI]] 5
+\ CHECK:   %[[GT:.*]] = forth.gt %[[L5]]
+\ CHECK:   forth.if %[[GT]]
+\ CHECK:     forth.loop_index
+\ CHECK:   forth.yield
+10 0 DO I 5 > IF I THEN LOOP
+
+\ === Nested DO with J ===
+\ CHECK: forth.do_loop
+\ CHECK:   forth.do_loop
+\ CHECK:     forth.loop_index %{{.*}} {depth = 1 : i64}
+\ CHECK:     forth.loop_index %{{.*}}
+\ CHECK:     forth.add
+3 0 DO 4 0 DO J I + LOOP LOOP
+
+\ === Triple-nested DO with K ===
+\ CHECK: forth.do_loop
+\ CHECK:   forth.do_loop
+\ CHECK:     forth.do_loop
+\ CHECK:       forth.loop_index %{{.*}} {depth = 2 : i64}
+\ CHECK:       forth.loop_index %{{.*}} {depth = 1 : i64}
+\ CHECK:       forth.loop_index %{{.*}}
+\ CHECK:       forth.add
+\ CHECK:       forth.add
+2 0 DO 2 0 DO 2 0 DO K J I + + LOOP LOOP LOOP
+
+\ === BEGIN/WHILE inside IF ===
+\ CHECK: forth.if
+\ CHECK:   forth.begin_while_repeat
+\ CHECK:     forth.dup
+\ CHECK:     forth.yield %{{.*}} while_cond
+\ CHECK:   } do {
+\ CHECK:     forth.literal %{{.*}} 1
+\ CHECK:     forth.sub
+\ CHECK:     forth.yield
+5 IF BEGIN DUP WHILE 1 - REPEAT THEN
+
+\ === IF inside BEGIN/UNTIL ===
+\ CHECK: forth.begin_until
+\ CHECK:   forth.dup
+\ CHECK:   forth.literal %{{.*}} 10
+\ CHECK:   forth.lt
+\ CHECK:   forth.if
+\ CHECK:     forth.literal %{{.*}} 1
+\ CHECK:     forth.add
+\ CHECK:   forth.dup
+\ CHECK:   forth.literal %{{.*}} 20
+\ CHECK:   forth.eq
+\ CHECK:   forth.yield
+BEGIN DUP 10 < IF 1 + THEN DUP 20 = UNTIL