Generate stack IR for imports

6 files changed, 108 insertions, 5 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedImportNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedImportNode.java
index e9932300..77a290ba 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedImportNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedImportNode.java
@@ -7,4 +7,7 @@ public record TypedImportNode(
     NamespaceName namespaceName,
     Source source
 ) implements TypedNode {
+    public String variableName() {
+        return namespaceName.parts().getLast();
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedNamespaceNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedNamespaceNode.java
index 26da9e61..abc325ce 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedNamespaceNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedNamespaceNode.java
@@ -1,7 +1,10 @@
 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.NamespaceName;
+import org.shedlang.compiler.util.P;
 
 import java.util.List;
 
@@ -13,4 +16,52 @@ public record TypedNamespaceNode(
     List<String> fieldNames,
     Source source
 ) implements TypedNode {
+    public static Builder builder() {
+        return new Builder(
+            NamespaceName.of("Default", "Namespace"),
+            P.vector(),
+            P.vector(),
+            P.vector(),
+            P.vector()
+        );
+    }
+
+    public record Builder(
+        NamespaceName namespaceName,
+        PVector<TypedImportNode> imports,
+        PVector<TypedNamespaceStatementNode> body,
+        PVector<String> boundVariableNames,
+        PVector<String> fieldNames
+    ) {
+        public TypedNamespaceNode build() {
+            return new TypedNamespaceNode(
+                namespaceName,
+                imports,
+                body,
+                boundVariableNames,
+                fieldNames,
+                Sources.NULL
+            );
+        }
+
+        public Builder addImport(TypedImportNode import_) {
+            return new Builder(
+                this.namespaceName,
+                this.imports.plus(import_),
+                this.body,
+                this.boundVariableNames,
+                this.fieldNames
+            );
+        }
+
+        public Builder addBoundVariableName(String variableName) {
+            return new Builder(
+                this.namespaceName,
+                this.imports,
+                this.body,
+                this.boundVariableNames.plus(variableName),
+                this.fieldNames
+            );
+        }
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index 5746fa34..7a9fc502 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -246,7 +246,7 @@ public class StackCompiler {
         context = context.addInstruction(new ReturnInstruction(node.source()));
         return context;
     }
-    
+
     public static Context compileFunctionStatement(
         TypedFunctionStatementNode node,
         Context context
@@ -257,7 +257,7 @@ public class StackCompiler {
 
             case TypedBlockReturnNode blockReturn ->
                 compileBlockReturn(blockReturn, context);
-            
+
             case TypedExpressionStatementNode expressionStatement ->
                 compileExpressionStatement(expressionStatement, context);
 
@@ -342,6 +342,20 @@ public class StackCompiler {
             context = context.addShedVariable(boundVariableName);
         }
 
+        for (var importNode : node.imports()) {
+            var localName = context.shedVariableToLocal(importNode.variableName());
+            context = context
+                .addInstruction(new NamespaceImportInstruction(
+                    importNode.namespaceName(),
+                    importNode.source()
+                ))
+                .addInstruction(new LocalStoreInstruction(
+                    localName,
+                    importNode.source()
+                ));
+
+        }
+
         var statementsAfterHoisting = Iterables.sortedBy(node.body(), statement -> switch (statement) {
             case TypedAssignNode assign ->
                 1;
diff --git a/src/test/java/org/shedlang/compiler/ast/typed/Typed.java b/src/test/java/org/shedlang/compiler/ast/typed/Typed.java
index 057e6062..09e7f080 100644
--- a/src/test/java/org/shedlang/compiler/ast/typed/Typed.java
+++ b/src/test/java/org/shedlang/compiler/ast/typed/Typed.java
@@ -90,6 +90,10 @@ public class Typed {
         );
     }
 
+    public static TypedImportNode import_(NamespaceName namespaceName) {
+        return new TypedImportNode(namespaceName, Sources.NULL);
+    }
+
     public static TypedBinaryOperationNode intAdd(
         TypedExpressionNode left,
         TypedExpressionNode right
diff --git a/src/test/java/org/shedlang/compiler/stackir/StackCompilerNamespaceTests.java b/src/test/java/org/shedlang/compiler/stackir/StackCompilerNamespaceTests.java
index 4ec958cb..e6206b9a 100644
--- a/src/test/java/org/shedlang/compiler/stackir/StackCompilerNamespaceTests.java
+++ b/src/test/java/org/shedlang/compiler/stackir/StackCompilerNamespaceTests.java
@@ -1,9 +1,7 @@
 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.ast.typed.TypedStructDefinitionNode;
+import org.shedlang.compiler.ast.typed.*;
 import org.shedlang.compiler.stackir.instructions.*;
 import org.shedlang.compiler.types.NamespaceName;
 import org.shedlang.compiler.types.Types;
@@ -52,6 +50,30 @@ public class StackCompilerNamespaceTests {
     }
 
     @Test
+    public void importsAreLoaded() {
+        var node = TypedNamespaceNode.builder()
+            .addImport(Typed.import_(NamespaceName.of("A")))
+            .addBoundVariableName("A")
+            .addImport(Typed.import_(NamespaceName.of("B")))
+            .addBoundVariableName("B")
+            .build();
+
+        var stackIrNamespace = StackCompiler.compileNamespace(node);
+
+        assertThat(stackIrNamespace.instructions(), isSequence(
+            // TODO: test the builtin loading properly
+            instanceOf(NamespaceImportInstruction.class),
+            instanceOf(FieldLoadInstruction.class),
+            instanceOf(LocalStoreInstruction.class),
+
+            isNamespaceImportInstruction(NamespaceName.of("A")),
+            isLocalStoreInstruction("A__0"),
+            isNamespaceImportInstruction(NamespaceName.of("B")),
+            isLocalStoreInstruction("B__0")
+        ));
+    }
+
+    @Test
     public void structsAreHoisted() {
         var namespaceName = new NamespaceName(List.of("Example"));
         var structType = Types.struct(namespaceName, "None");
diff --git a/src/test/java/org/shedlang/compiler/stackir/instructions/InstructionMatchers.java b/src/test/java/org/shedlang/compiler/stackir/instructions/InstructionMatchers.java
index 81b35cfe..a33f4f6c 100644
--- a/src/test/java/org/shedlang/compiler/stackir/instructions/InstructionMatchers.java
+++ b/src/test/java/org/shedlang/compiler/stackir/instructions/InstructionMatchers.java
@@ -128,6 +128,15 @@ public class InstructionMatchers {
         );
     }
 
+    public static Matcher<Instruction> isNamespaceImportInstruction(
+        NamespaceName namespaceName
+    ) {
+        return instanceOf(
+            NamespaceImportInstruction.class,
+            has("namespaceName", x -> x.namespaceName(), equalTo(namespaceName))
+        );
+    }
+
     public static Matcher<Instruction> isPushBoolInstruction(boolean value) {
         return instanceOf(
             PushBoolInstruction.class,