Type check native function declarations

4 files changed, 156 insertions, 7 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedNamespaceStatementNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedNamespaceStatementNode.java
index 79c69548..d38cff5e 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedNamespaceStatementNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedNamespaceStatementNode.java
@@ -1,4 +1,4 @@
 package org.shedlang.compiler.ast.typed;
 
-public sealed interface TypedNamespaceStatementNode extends TypedNode permits TypedAssignNode, TypedEffectDefinitionNode, TypedFunctionDefinitionNode, TypedStructDefinitionNode, TypedSumTypeDefinitionNode {
+public sealed interface TypedNamespaceStatementNode extends TypedNode permits TypedAssignNode, TypedEffectDefinitionNode, TypedFunctionDefinitionNode, TypedNativeFunctionDeclarationNode, TypedStructDefinitionNode, TypedSumTypeDefinitionNode {
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedNativeFunctionDeclarationNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedNativeFunctionDeclarationNode.java
new file mode 100644
index 00000000..e22bc311
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedNativeFunctionDeclarationNode.java
@@ -0,0 +1,88 @@
+package org.shedlang.compiler.ast.typed;
+
+import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.sources.Sources;
+import org.shedlang.compiler.types.*;
+
+import java.util.List;
+import java.util.Optional;
+
+public record TypedNativeFunctionDeclarationNode(
+    String name,
+    TypedFunctionSignatureNode signature,
+    FunctionType type,
+    Source source
+) implements TypedNamespaceStatementNode {
+    public static Builder builder() {
+        return new Builder(
+            "f",
+            TypedFunctionSignatureNode.builder(),
+            new FunctionType(Optional.empty(), List.of(), List.of(), Effects.NONE, Types.UNIT)
+        );
+    }
+
+    public record Builder(
+        String name,
+        TypedFunctionSignatureNode.Builder signature,
+        FunctionType type
+    ) {
+        public TypedNativeFunctionDeclarationNode build() {
+            return new TypedNativeFunctionDeclarationNode(
+                name,
+                signature.build(),
+                type,
+                Sources.NULL
+            );
+        }
+
+        public TypedNativeFunctionDeclarationNode.Builder withName(String name) {
+            return new TypedNativeFunctionDeclarationNode.Builder(
+                name,
+                this.signature,
+                this.type
+            );
+        }
+
+        public TypedNativeFunctionDeclarationNode.Builder addPositionalParam(
+            String name,
+            TypedTypeLevelExpressionNode<Type> type
+        ) {
+            return addPositionalParam(name, Optional.of(type));
+        }
+
+        public TypedNativeFunctionDeclarationNode.Builder addPositionalParam(
+            String name,
+            Optional<TypedTypeLevelExpressionNode<Type>> type
+        ) {
+            return new TypedNativeFunctionDeclarationNode.Builder(
+                this.name,
+                this.signature.addPositionalParam(name, type),
+                this.type
+            );
+        }
+
+        public TypedNativeFunctionDeclarationNode.Builder addNamedParam(String name, TypedTypeLevelExpressionNode<Type> type) {
+            return new TypedNativeFunctionDeclarationNode.Builder(
+                this.name,
+                this.signature.addNamedParam(name, type),
+                this.type
+            );
+        }
+
+        public TypedNativeFunctionDeclarationNode.Builder withEffect(TypedTypeLevelExpressionNode<Effect> effect) {
+            return new TypedNativeFunctionDeclarationNode.Builder(
+                this.name,
+                this.signature.withEffect(effect),
+                this.type
+            );
+        }
+
+        public TypedNativeFunctionDeclarationNode.Builder withReturnType(TypedTypeLevelExpressionNode<Type> returnType) {
+            return new TypedNativeFunctionDeclarationNode.Builder(
+                this.name,
+                this.signature.withReturnType(returnType),
+                this.type
+            );
+        }
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java b/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java
index d6626b37..194d4708 100644
--- a/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java
@@ -551,7 +551,7 @@ public class ShidCodeGen {
 
         var statementsAfterHoisting = Iterables.sortedBy(node.body(), statement -> switch (statement) {
             case TypedAssignNode assign ->
-                1;
+                2;
 
             case TypedEffectDefinitionNode effectDefinition ->
                 0;
@@ -559,6 +559,9 @@ public class ShidCodeGen {
             case TypedFunctionDefinitionNode functionDefinition ->
                 1;
 
+            case TypedNativeFunctionDeclarationNode nativeFunctionDeclaration ->
+                1;
+
             case TypedStructDefinitionNode structDefinition ->
                 0;
 
@@ -605,6 +608,9 @@ public class ShidCodeGen {
             case TypedFunctionDefinitionNode functionDefinition ->
                 compileFunctionDefinition(functionDefinition, context);
 
+            case TypedNativeFunctionDeclarationNode nativeFunctionDeclaration ->
+                throw new UnsupportedOperationException("TODO");
+
             case TypedStructDefinitionNode structDefinition ->
                 compileStructDefinition(structDefinition, context);
 
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 00960bd0..fa32a440 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -848,11 +848,27 @@ public class TypeChecker {
         Optional<Type> thisType,
         TypeCheckerContext context
     ) {
+        return declareFunction(
+            untyped.name(),
+            untyped.signature(),
+            untyped.source(),
+            thisType,
+            context
+        );
+    }
+
+    private static TypeCheckerResult<TypeCheckFunctionDefinitionState> declareFunction(
+        String name,
+        UntypedFunctionSignatureNode signature,
+        Source source,
+        Optional<Type> thisType,
+        TypeCheckerContext context
+    ) {
         var namespaceName = context.namespaceName()
-            .orElseThrow(() -> new CannotDefineFunctionOutsideOfNamespaceError(untyped.source()));
+            .orElseThrow(() -> new CannotDefineFunctionOutsideOfNamespaceError(source));
         var typedSignature = typeCheckFunctionSignature(
-            untyped.signature(),
-            Optional.of(new Scope.Function(namespaceName, untyped.name())),
+            signature,
+            Optional.of(new Scope.Function(namespaceName, name)),
             context
         );
 
@@ -888,7 +904,7 @@ public class TypeChecker {
         var effect = functionSignatureToEffect(typedSignature);
 
         var returnType = functionSignatureToReturnType(typedSignature)
-            .orElseThrow(() -> new FunctionSignatureMissingReturnTypeError(untyped.signature().source()));
+            .orElseThrow(() -> new FunctionSignatureMissingReturnTypeError(signature.source()));
 
         var partialSignature = new PartialFunctionSignature(
             typedSignature.typeLevelParams(),
@@ -1564,7 +1580,7 @@ public class TypeChecker {
                 typeCheckStepsFunctionDefinition(functionDefinition);
 
             case UntypedNativeFunctionDeclarationNode nativeFunctionDeclaration ->
-                throw new UnsupportedOperationException("TODO");
+                typeCheckStepsNativeFunctionDeclaration(nativeFunctionDeclaration);
 
             case UntypedStructDefinitionNode structDefinition ->
                 typeCheckStepsStructDefinition(structDefinition);
@@ -1574,6 +1590,45 @@ public class TypeChecker {
         };
     }
 
+    private static TypeCheckSteps<TypedNativeFunctionDeclarationNode> typeCheckStepsNativeFunctionDeclaration(
+        UntypedNativeFunctionDeclarationNode untyped
+    ) {
+        return new TypeCheckSteps<TypedNativeFunctionDeclarationNode>() {
+            @Override
+            public TypeCheckerContext declareFunctions(TypeCheckerContext context) {
+                var declareResult = declareFunction(
+                    untyped.name(),
+                    untyped.signature(),
+                    untyped.source(),
+                    Optional.empty(),
+                    context
+                );
+                var state = declareResult.value();
+                context = declareResult.context();
+
+                context = context.declareVariableType(
+                    untyped.name(),
+                    state.type(),
+                    untyped.source()
+                );
+
+                return context.addLexicalState(untyped.name(), state);
+            }
+
+            @Override
+            public TypeCheckerResult<TypedNativeFunctionDeclarationNode> finish(TypeCheckerContext context) {
+                var state = context.getLexicalState(TypeCheckFunctionDefinitionState.class, untyped.name());
+
+                return context.toResult(new TypedNativeFunctionDeclarationNode(
+                    untyped.name(),
+                    state.typedSignatureNode,
+                    state.type,
+                    untyped.source()
+                ));
+            }
+        };
+    }
+
     private static TypedParamNode typeCheckParam(
         UntypedParamNode untyped,
         TypeCheckerContext context