Hoist structs

2 files changed, 57 insertions, 5 deletions

diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index dc171e00..cca640b4 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -5,6 +5,7 @@ import org.shedlang.compiler.ShedBuiltins;
 import org.shedlang.compiler.ast.typed.*;
 import org.shedlang.compiler.stackir.instructions.*;
 import org.shedlang.compiler.types.StructType;
+import org.shedlang.compiler.util.Iterables;
 import org.shedlang.compiler.util.P;
 
 import java.util.List;
@@ -337,7 +338,17 @@ public class StackCompiler {
             context = context.addShedVariable(boundVariableName);
         }
 
-        for (var statement : node.body()) {
+        var statementsAfterHoisting = Iterables.sortedBy(node.body(), statement -> switch (statement) {
+            case TypedAssignNode assign ->
+                1;
+
+            case TypedFunctionDefinitionNode functionDefinition ->
+                1;
+
+            case TypedStructDefinitionNode structDefinition ->
+                0;
+        });
+        for (var statement : statementsAfterHoisting) {
             context = compileNamespaceStatement(statement, context);
         }
 
diff --git a/src/test/java/org/shedlang/compiler/stackir/StackCompilerNamespaceTests.java b/src/test/java/org/shedlang/compiler/stackir/StackCompilerNamespaceTests.java
index b4524e90..3557ac20 100644
--- a/src/test/java/org/shedlang/compiler/stackir/StackCompilerNamespaceTests.java
+++ b/src/test/java/org/shedlang/compiler/stackir/StackCompilerNamespaceTests.java
@@ -3,10 +3,8 @@ package org.shedlang.compiler.stackir;
 import org.junit.jupiter.api.Test;
 import org.shedlang.compiler.ast.typed.Typed;
 import org.shedlang.compiler.ast.typed.TypedFunctionDefinitionNode;
-import org.shedlang.compiler.stackir.instructions.DefineFunctionInstruction;
-import org.shedlang.compiler.stackir.instructions.FieldLoadInstruction;
-import org.shedlang.compiler.stackir.instructions.LocalStoreInstruction;
-import org.shedlang.compiler.stackir.instructions.NamespaceImportInstruction;
+import org.shedlang.compiler.ast.typed.TypedStructDefinitionNode;
+import org.shedlang.compiler.stackir.instructions.*;
 import org.shedlang.compiler.types.NamespaceName;
 import org.shedlang.compiler.types.Types;
 
@@ -52,4 +50,47 @@ public class StackCompilerNamespaceTests {
             isNamespaceFieldStore(namespaceName, "f")
         ));
     }
+
+    @Test
+    public void structsAreHoisted() {
+        var namespaceName = new NamespaceName(List.of("Example"));
+        var structType = Types.struct(namespaceName, "None");
+        var node = Typed.namespace(
+            namespaceName,
+            List.of(
+                TypedFunctionDefinitionNode.builder()
+                    .withName("f")
+                    .build(),
+                TypedStructDefinitionNode.builder()
+                    .withType(structType)
+                    .build()
+            ),
+            List.of("f", "None"),
+            List.of("f", "None")
+        );
+
+        var stackIrNamespace = StackCompiler.compileNamespace(node);
+
+        assertThat(stackIrNamespace.instructions(), isSequence(
+            // TODO: test the builtin loading properly
+            instanceOf(NamespaceImportInstruction.class),
+            instanceOf(FieldLoadInstruction.class),
+            instanceOf(LocalStoreInstruction.class),
+
+            instanceOf(
+                DefineStructInstruction.class,
+                has("type", x -> x.type(), equalTo(structType))
+            ),
+            isLocalStoreInstruction("None__0"),
+            instanceOf(
+                DefineFunctionInstruction.class,
+                has("name", x -> x.name(), equalTo("f"))
+            ),
+            isLocalStoreInstruction("f__0"),
+            isLocalLoadInstruction("f__0"),
+            isNamespaceFieldStore(namespaceName, "f"),
+            isLocalLoadInstruction("None__0"),
+            isNamespaceFieldStore(namespaceName, "None")
+        ));
+    }
 }