Handle block returns properly

3 files changed, 46 insertions, 1 deletions

diff --git a/src/test/java/org/shedlang/compiler/stackir/StackCompilerBlockReturnTests.java b/src/test/java/org/shedlang/compiler/stackir/StackCompilerBlockReturnTests.java
new file mode 100644
index 00000000..7ff57a10
--- /dev/null
+++ b/src/test/java/org/shedlang/compiler/stackir/StackCompilerBlockReturnTests.java
@@ -0,0 +1,21 @@
+package org.shedlang.compiler.stackir;
+
+import org.junit.jupiter.api.Test;
+import org.shedlang.compiler.ast.typed.Typed;
+
+import static org.shedlang.compiler.stackir.instructions.InstructionMatchers.isPushStringInstruction;
+import static org.zwobble.precisely.AssertThat.assertThat;
+import static org.zwobble.precisely.Matchers.isSequence;
+
+public class StackCompilerBlockReturnTests {
+    @Test
+    public void expressionIsEvaluated() {
+        var node = Typed.blockReturn(Typed.stringLiteral("hello"));
+
+        var instructions = StackCompiler.compileFunctionStatement(node);
+
+        assertThat(instructions, isSequence(
+            isPushStringInstruction("hello")
+        ));
+    }
+}
diff --git a/src/test/java/org/shedlang/compiler/stackir/StackCompilerFunctionDefinitionTests.java b/src/test/java/org/shedlang/compiler/stackir/StackCompilerFunctionDefinitionTests.java
index 28974ea9..051b0fb6 100644
--- a/src/test/java/org/shedlang/compiler/stackir/StackCompilerFunctionDefinitionTests.java
+++ b/src/test/java/org/shedlang/compiler/stackir/StackCompilerFunctionDefinitionTests.java
@@ -46,4 +46,28 @@ public class StackCompilerFunctionDefinitionTests {
             isLocalStoreInstruction("f")
         ));
     }
+
+    @Test
+    public void whenLastStatementIsBlockReturnThenThatValueIsReturned() {
+        var node = TypedFunctionDefinitionNode.builder()
+            .withName("f")
+            .addPositionalParam("x", Typed.typeLevelReference("String", Types.STRING))
+            .addNamedParam("y", Typed.typeLevelReference("Int", Types.INT))
+            .addFreeVariable("z")
+            .addStatement(Typed.blockReturn(Typed.stringLiteral("hello")))
+            .build();
+
+        var instructions = StackCompiler.compileNamespaceStatement(node);
+
+        assertThat(instructions, isSequence(
+            instanceOf(
+                DefineFunctionInstruction.class,
+                has("body", x -> x.body(), isSequence(
+                    isPushStringInstruction("hello"),
+                    isReturnInstruction()
+                ))
+            ),
+            isLocalStoreInstruction("f")
+        ));
+    }
 }
diff --git a/src/test/java/org/shedlang/compiler/stackir/StackCompilerReturnTests.java b/src/test/java/org/shedlang/compiler/stackir/StackCompilerFunctionReturnTests.java
index dcb1af47..0baf9acb 100644
--- a/src/test/java/org/shedlang/compiler/stackir/StackCompilerReturnTests.java
+++ b/src/test/java/org/shedlang/compiler/stackir/StackCompilerFunctionReturnTests.java
@@ -7,7 +7,7 @@ import static org.shedlang.compiler.stackir.instructions.InstructionMatchers.*;
 import static org.zwobble.precisely.AssertThat.assertThat;
 import static org.zwobble.precisely.Matchers.isSequence;
 
-public class StackCompilerReturnTests {
+public class StackCompilerFunctionReturnTests {
     @Test
     public void expressionIsEvaluatedAndTheResultDiscarded() {
         var node = Typed.functionReturn(Typed.stringLiteral("hello"));