Tidy up generation of WASM structured instructions

1 files changed, 145 insertions, 47 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 bbf8298d..fc723883 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmPatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmPatioCodeGen.java
@@ -202,23 +202,19 @@ public class WasmPatioCodeGen {
             case PatioIf patioIf -> {
                 context = generateOperand(patioIf.condition(), context);
 
-                var trueBodyResult = generateBlock(patioIf.trueBody(), context);
-                var trueBody = trueBodyResult.value;
-                context = trueBodyResult.context;
-
-                var falseBodyResult = generateBlock(patioIf.falseBody(), context);
-                var falseBody = falseBodyResult.value;
-                context = falseBodyResult.context;
-
-                context = context.addInstruction(new WasmIf(
-                    new WasmBlockType(
-                        patioIf.targets().stream()
-                            .map(patioTarget -> patioOperandToWasmValueType(patioTarget))
-                            .toList()
-                    ),
-                    trueBody,
-                    falseBody
-                ));
+                var wasmBlockType = new WasmBlockType(
+                    patioIf.targets().stream()
+                        .map(patioTarget -> patioOperandToWasmValueType(patioTarget))
+                        .toList()
+                );
+
+                context = context.enterIfInstruction(wasmBlockType);
+                context = generateBlockBody(patioIf.trueBody(), context);
+
+                context = context.enterElse();
+                context = generateBlockBody(patioIf.falseBody(), context);
+
+                context = context.exitIfInstruction();
 
                 for (var patioTarget : Lists.reversed(patioIf.targets())) {
                     context = popIntoPatioLocal(patioTarget, context);
@@ -261,20 +257,11 @@ public class WasmPatioCodeGen {
                 );
 
             case PatioLoop patioLoop -> {
-                context = context.enterBlock();
+                context = context.enterBlockInstruction();
+                context = context.enterLoopInstruction();
                 context = generateInstructions(patioLoop.body(), context);
-                var bodyInstructions = context.instructions;
-                context = context.exitBlock();
-
-                context = context.addInstruction(new WasmBlock(
-                    new WasmBlockType(List.of()),
-                    List.of(
-                        new WasmLoop(
-                            new WasmBlockType(List.of()),
-                            bodyInstructions
-                        )
-                    )
-                ));
+                context = context.exitLoopInstruction();
+                context = context.exitBlockInstruction();
 
                 yield context;
             }
@@ -342,21 +329,17 @@ public class WasmPatioCodeGen {
         };
     }
 
-    private static Context.Result<List<WasmInstruction>> generateBlock(
+    private static Context generateBlockBody(
         PatioBlock block,
         Context context
     ) {
-        context = context.enterBlock();
         context = generateInstructions(block.body(), context);
 
         for (var patioResult : block.results()) {
             context = generateOperand(patioResult, context);
         }
 
-        var wasmBody = context.instructions;
-        context = context.exitBlock();
-
-        return context.toResult(wasmBody);
+        return context;
     }
 
     private static Context generateBinaryOperation(
@@ -562,7 +545,9 @@ public class WasmPatioCodeGen {
                 module,
                 P.vector(),
                 P.stack(),
-                patioLocalNameToWasmLocalIndex
+                patioLocalNameToWasmLocalIndex,
+                new NullStructuredInstructionBuilder(),
+                P.stack()
             );
         }
 
@@ -570,17 +555,23 @@ public class WasmPatioCodeGen {
         private final PVector<WasmInstruction> instructions;
         private final PStack<PVector<WasmInstruction>> instructionsStack;
         private final PMap<String, WasmLocalIndex> patioLocalNameToWasmLocalIndex;
+        private final WasmStructuredInstructionBuilder structuredInstructionBuilder;
+        private final PStack<WasmStructuredInstructionBuilder> structuredInstructionBuilderStack;
 
         private Context(
             WasmModule module,
             PVector<WasmInstruction> instructions,
             PStack<PVector<WasmInstruction>> instructionsStack,
-            PMap<String, WasmLocalIndex> patioLocalNameToWasmLocalIndex
+            PMap<String, WasmLocalIndex> patioLocalNameToWasmLocalIndex,
+            WasmStructuredInstructionBuilder structuredInstructionBuilder,
+            PStack<WasmStructuredInstructionBuilder> structuredInstructionBuilderStack
         ) {
             this.module = module;
             this.instructions = instructions;
             this.instructionsStack = instructionsStack;
             this.patioLocalNameToWasmLocalIndex = patioLocalNameToWasmLocalIndex;
+            this.structuredInstructionBuilder = structuredInstructionBuilder;
+            this.structuredInstructionBuilderStack = structuredInstructionBuilderStack;
         }
 
         private Context addType(WasmFuncType type) {
@@ -588,7 +579,9 @@ public class WasmPatioCodeGen {
                 this.module.addType(type),
                 this.instructions,
                 this.instructionsStack,
-                this.patioLocalNameToWasmLocalIndex
+                this.patioLocalNameToWasmLocalIndex,
+                this.structuredInstructionBuilder,
+                this.structuredInstructionBuilderStack
             );
         }
 
@@ -597,7 +590,9 @@ public class WasmPatioCodeGen {
                 this.module,
                 this.instructions.plus(instruction),
                 this.instructionsStack,
-                this.patioLocalNameToWasmLocalIndex
+                this.patioLocalNameToWasmLocalIndex,
+                this.structuredInstructionBuilder,
+                this.structuredInstructionBuilderStack
             );
         }
 
@@ -605,29 +600,132 @@ public class WasmPatioCodeGen {
             return this.patioLocalNameToWasmLocalIndex.get(patioLocalRef.name());
         }
 
-        public Context enterBlock() {
+        public Context enterBlockInstruction() {
             return new Context(
                 this.module,
                 P.vector(),
                 this.instructionsStack.plus(this.instructions),
-                this.patioLocalNameToWasmLocalIndex
+                this.patioLocalNameToWasmLocalIndex,
+                new WasmBlockInstructionBuilder(),
+                this.structuredInstructionBuilderStack.plus(this.structuredInstructionBuilder)
             );
         }
 
-        public Context exitBlock() {
+        public Context exitBlockInstruction() {
+            // TODO: verify matching enter/exit?
+            return exitStructuredInstruction();
+        }
+
+        public Context enterLoopInstruction() {
             return new Context(
                 this.module,
+                P.vector(),
+                this.instructionsStack.plus(this.instructions),
+                this.patioLocalNameToWasmLocalIndex,
+                new WasmLoopInstructionBuilder(),
+                this.structuredInstructionBuilderStack.plus(this.structuredInstructionBuilder)
+            );
+        }
+
+        public Context exitLoopInstruction() {
+            // TODO: verify matching enter/exit?
+            return exitStructuredInstruction();
+        }
+
+        public Context enterIfInstruction(WasmBlockType blockType) {
+            return new Context(
+                this.module,
+                P.vector(),
+                this.instructionsStack.plus(this.instructions),
+                this.patioLocalNameToWasmLocalIndex,
+                new WasmIfInstructionBuilder(blockType),
+                this.structuredInstructionBuilderStack.plus(this.structuredInstructionBuilder)
+            );
+        }
+
+        public Context enterElse() {
+            return new Context(
+                this.module,
+                P.vector(),
+                this.instructionsStack,
+                this.patioLocalNameToWasmLocalIndex,
+                this.structuredInstructionBuilder.enterElse(this.instructions),
+                this.structuredInstructionBuilderStack
+            );
+        }
+
+        public Context exitIfInstruction() {
+            // TODO: verify matching enter/exit?
+            return exitStructuredInstruction();
+        }
+
+        public Context exitStructuredInstruction() {
+            var context = new Context(
+                this.module,
                 this.instructionsStack.get(0),
                 this.instructionsStack.minus(0),
-                this.patioLocalNameToWasmLocalIndex
+                this.patioLocalNameToWasmLocalIndex,
+                this.structuredInstructionBuilderStack.get(0),
+                this.structuredInstructionBuilderStack.minus(0)
             );
+
+            var instruction = this.structuredInstructionBuilder
+                .exit(this.instructions);
+
+            return context.addInstruction(instruction);
+        }
+    }
+
+    private interface WasmStructuredInstructionBuilder {
+        WasmInstruction exit(PVector<WasmInstruction> instructions);
+
+        default WasmStructuredInstructionBuilder enterElse(PVector<WasmInstruction> instructions) {
+            throw new UnsupportedOperationException();
         }
+    }
+
+    private static class NullStructuredInstructionBuilder implements WasmStructuredInstructionBuilder {
+        @Override
+        public WasmInstruction exit(PVector<WasmInstruction> instructions) {
+            throw new UnsupportedOperationException();
+        }
+    }
+
+    private static class WasmBlockInstructionBuilder implements WasmStructuredInstructionBuilder {
+        @Override
+        public WasmInstruction exit(PVector<WasmInstruction> instructions) {
+            return new WasmBlock(new WasmBlockType(List.of()), instructions);
+        }
+    }
 
-        private <T> Result<T> toResult(T value) {
-            return new Result<>(value, this);
+    private static class WasmLoopInstructionBuilder implements WasmStructuredInstructionBuilder {
+        @Override
+        public WasmInstruction exit(PVector<WasmInstruction> instructions) {
+            return new WasmLoop(new WasmBlockType(List.of()), instructions);
         }
+    }
+
+    private record WasmIfInstructionBuilder(
+        WasmBlockType blockType
+    ) implements WasmStructuredInstructionBuilder {
+        @Override
+        public WasmInstruction exit(PVector<WasmInstruction> instructions) {
+            throw new IllegalStateException("if instruction without else not supported");
+        }
+
+        @Override
+        public WasmStructuredInstructionBuilder enterElse(PVector<WasmInstruction> instructions) {
+            return new WasmIfElseInstructionBuilder(blockType, instructions);
+        }
+    }
 
-        private record Result<T>(T value, Context context) {
+    private record WasmIfElseInstructionBuilder(
+        WasmBlockType blockType,
+        PVector<WasmInstruction> ifTrue
+    ) implements WasmStructuredInstructionBuilder {
+        @Override
+        public WasmInstruction exit(PVector<WasmInstruction> instructions) {
+            return new WasmIf(blockType, ifTrue, instructions);
         }
     }
 }