Compile try node into stack IR

4 files changed, 106 insertions, 12 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedTryNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedTryNode.java
index 5e87a6b7..c3c5d32d 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedTryNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedTryNode.java
@@ -1,8 +1,13 @@
 package org.shedlang.compiler.ast.typed;
 
+import org.pcollections.PVector;
 import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.sources.Sources;
+import org.shedlang.compiler.types.Effects;
+import org.shedlang.compiler.types.NamespaceName;
 import org.shedlang.compiler.types.Type;
 import org.shedlang.compiler.types.UserDefinedEffect;
+import org.shedlang.compiler.util.P;
 
 import java.util.List;
 
@@ -16,4 +21,51 @@ public record TypedTryNode(
     public Type type() {
         return body.type();
     }
+
+    public static Builder builder() {
+        return new Builder(
+            P.vector(),
+            new TypedTypeLevelReferenceNode<>("Effect", Effects.userDefined(NamespaceName.of("Example"), "Effect"), Sources.NULL),
+            P.vector()
+        );
+    }
+
+    public record Builder(
+        PVector<TypedFunctionStatementNode> body,
+        TypedTypeLevelExpressionNode<UserDefinedEffect> handledEffect,
+        PVector<TypedTryHandlerMethodNode> handlerMethods
+    ) {
+        public TypedTryNode build() {
+            return new TypedTryNode(
+                new TypedBlockNode(body, P.vector(), Sources.NULL),
+                handledEffect,
+                handlerMethods,
+                Sources.NULL
+            );
+        }
+
+        public Builder addBodyStatement(TypedFunctionStatementNode statement) {
+            return new Builder(
+                this.body.plus(statement),
+                this.handledEffect,
+                this.handlerMethods
+            );
+        }
+
+        public Builder withHandledEffect(TypedTypeLevelExpressionNode<UserDefinedEffect> handledEffect) {
+            return new Builder(
+                this.body,
+                handledEffect,
+                this.handlerMethods
+            );
+        }
+
+        public Builder addHandlerMethod(TypedTryHandlerMethodNode method) {
+            return new Builder(
+                this.body,
+                this.handledEffect,
+                this.handlerMethods.plus(method)
+            );
+        }
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/stackir/InstructionSerializer.java b/src/main/java/org/shedlang/compiler/stackir/InstructionSerializer.java
index 5cd8478d..10b39832 100644
--- a/src/main/java/org/shedlang/compiler/stackir/InstructionSerializer.java
+++ b/src/main/java/org/shedlang/compiler/stackir/InstructionSerializer.java
@@ -224,8 +224,12 @@ public class InstructionSerializer {
                 treeListIndented(
                     treeIdentifier("try"),
                     userDefinedEffectToTree(tryInstruction.handledEffect()),
-                    instructionsToTree(tryInstruction.body()),
-                    instructionsToTree(tryInstruction.handlerMethods())
+                    treeListIndented(
+                        tryInstruction.handlerMethodNames().stream()
+                            .map(name -> treeString(name))
+                            .toList()
+                    ),
+                    instructionsToTree(tryInstruction.body())
                 );
 
             case TupleCreateInstruction tupleCreateInstruction ->
diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index da20c770..4ca4de8e 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -155,7 +155,7 @@ public class StackCompiler {
                 compileStringLiteral(typedStringLiteralNode, context);
 
             case TypedTryNode typedTryNode ->
-                throw new UnsupportedOperationException("TODO");
+                compileTry(typedTryNode, context);
 
             case TypedTupleNode typedTupleNode ->
                 compileTuple(typedTupleNode, context);
@@ -204,7 +204,7 @@ public class StackCompiler {
 
     private static Context compileFunction(
         String name,
-        TypedFunctionSignatureNode signature,
+        TypedParamsNode params,
         TypedFunctionBodyNode body,
         Source source,
         Context context
@@ -213,14 +213,14 @@ public class StackCompiler {
 
         var c = context;
 
-        var positionalParams = signature.params().positional().stream()
+        var positionalParams = params.positional().stream()
             .map(param -> new DefineFunctionInstruction.Param(
                 c.shedVariableToLocal(param.name()),
                 param.source()
             ))
             .toList();
 
-        var namedParams = signature.params().named().stream()
+        var namedParams = params.named().stream()
             .map(param -> new DefineFunctionInstruction.Param(
                 c.shedVariableToLocal(param.name()),
                 param.source()
@@ -273,7 +273,7 @@ public class StackCompiler {
     ) {
         context = compileFunction(
             node.name(),
-            node.signature(),
+            node.signature().params(),
             node.body(),
             node.source(),
             context
@@ -292,7 +292,7 @@ public class StackCompiler {
         return compileFunction(
             // TODO: include source location?
             "anonymous",
-            node.signature(),
+            node.signature().params(),
             node.body(),
             node.source(),
             context
@@ -609,6 +609,41 @@ public class StackCompiler {
         };
     }
 
+    private static Context compileTry(TypedTryNode node, Context context) {
+        for (var handlerMethod : node.handlerMethods()) {
+            context = compileTryHandlerMethod(handlerMethod, context);
+        }
+
+        context = context.enterBlock(node.body().boundVariableNames());
+        context = compileBlockWithYield(node.body(), context);
+        var bodyInstructions = context.instructions;
+        context = context.exitBlock();
+
+        var instruction = new TryInstruction(
+            node.handledEffect().value(),
+            node.handlerMethods().stream()
+                .map(methodNode -> methodNode.name())
+                .toList(),
+            bodyInstructions,
+            node.source()
+        );
+
+        return context.addInstruction(instruction);
+    }
+
+    private static Context compileTryHandlerMethod(
+        TypedTryHandlerMethodNode node,
+        Context context
+    ) {
+        return compileFunction(
+            node.name(),
+            node.params(),
+            node.body(),
+            node.source(),
+            context
+        );
+    }
+
     private static Context compileTuple(TypedTupleNode node, Context context) {
         for (var elementNode : node.elements()) {
             context = compileExpression(elementNode, context);
diff --git a/src/main/java/org/shedlang/compiler/stackir/instructions/TryInstruction.java b/src/main/java/org/shedlang/compiler/stackir/instructions/TryInstruction.java
index 78a5276d..aa994645 100644
--- a/src/main/java/org/shedlang/compiler/stackir/instructions/TryInstruction.java
+++ b/src/main/java/org/shedlang/compiler/stackir/instructions/TryInstruction.java
@@ -7,15 +7,18 @@ import org.shedlang.compiler.types.UserDefinedEffect;
 import java.util.List;
 
 /**
- * Push an effect handler for the given effect with the given methods onto the
- * effect handler stack, and execute the body. Once the body has exited, pop the
- * effect handler from the effect handler stack.
+ * For each handler method name, pop a value off the stack. The first value
+ * (that is, the first pushed / last popped) corresponds to the first handler
+ * method name, and so on. Push an effect handler for the given effect with the
+ * popped handler methods onto the effect handler stack, and execute the body.
+ * Once the body has exited, pop the effect handler from the effect handler
+ * stack.
  * <p>
  * If the body runs to completion, it should leave a single value on the stack.
  */
 public record TryInstruction(
     UserDefinedEffect handledEffect,
-    List<DefineFunctionInstruction> handlerMethods,
+    List<String> handlerMethodNames,
     List<Instruction> body,
     Source source
 ) implements Instruction {