Find label index properly

1 files changed, 30 insertions, 3 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmPatioCodeGen.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmPatioCodeGen.java
index fc723883..1848f727 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmPatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmPatioCodeGen.java
@@ -267,14 +267,13 @@ public class WasmPatioCodeGen {
             }
 
             case PatioLoopBreak patioLoopBreak -> {
-                // TODO: track label indices properly
-                context = context.addInstruction(new WasmBr(2));
+                context = context.addBrToBlockInstruction();
 
                 yield context;
             }
 
             case PatioLoopContinue patioLoopContinue -> {
-                context = context.addInstruction(new WasmBr(0));
+                context = context.addBrToLoopInstruction();
 
                 yield context;
             }
@@ -616,6 +615,10 @@ public class WasmPatioCodeGen {
             return exitStructuredInstruction();
         }
 
+        public Context addBrToBlockInstruction() {
+            return addBr(WasmBlockInstructionBuilder.class);
+        }
+
         public Context enterLoopInstruction() {
             return new Context(
                 this.module,
@@ -632,6 +635,10 @@ public class WasmPatioCodeGen {
             return exitStructuredInstruction();
         }
 
+        public Context addBrToLoopInstruction() {
+            return addBr(WasmLoopInstructionBuilder.class);
+        }
+
         public Context enterIfInstruction(WasmBlockType blockType) {
             return new Context(
                 this.module,
@@ -674,6 +681,26 @@ public class WasmPatioCodeGen {
 
             return context.addInstruction(instruction);
         }
+
+        private Context addBr(Class<? extends WasmStructuredInstructionBuilder> clazz) {
+            var labelIndex = findLabelIndex(clazz);
+
+            return this.addInstruction(new WasmBr(labelIndex));
+        }
+
+        private int findLabelIndex(Class<? extends WasmStructuredInstructionBuilder> clazz) {
+            if (clazz.isInstance(this.structuredInstructionBuilder)) {
+                return 0;
+            }
+
+            for (var i = 0; i < this.structuredInstructionBuilderStack.size(); i++) {
+                if (clazz.isInstance(this.structuredInstructionBuilderStack.get(i))) {
+                    return i + 1;
+                }
+            }
+
+            throw new IllegalStateException("cannot find label index");
+        }
     }
 
     private interface WasmStructuredInstructionBuilder {