Store type of function on node

2 files changed, 14 insertions, 0 deletions

diff --git a/src/test/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNodeMatcher.java b/src/test/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNodeMatcher.java
index b1c6e59a..c1d047af 100644
--- a/src/test/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNodeMatcher.java
+++ b/src/test/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNodeMatcher.java
@@ -3,6 +3,7 @@ package org.shedlang.compiler.ast.typed;
 import org.pcollections.Empty;
 import org.pcollections.PVector;
 import org.shedlang.compiler.testing.MatcherBuilder;
+import org.shedlang.compiler.types.FunctionType;
 import org.zwobble.precisely.Matcher;
 
 import java.util.Optional;
@@ -52,4 +53,8 @@ public class TypedFunctionDefinitionNodeMatcher extends MatcherBuilder<TypedName
     public TypedFunctionDefinitionNodeMatcher withReturnType(Matcher<TypedTypeLevelExpressionNode<?>> returnType) {
         return with(has("returnType", x -> x.returnType(), returnType));
     }
+
+    public TypedFunctionDefinitionNodeMatcher withType(FunctionType type) {
+        return with(has("type", x -> x.type(), equalTo(type)));
+    }
 }
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
index 3eb177cc..f3cbf0fe 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
@@ -4,9 +4,12 @@ import org.junit.jupiter.api.Test;
 import org.shedlang.compiler.ast.untyped.Untyped;
 import org.shedlang.compiler.ast.untyped.UntypedFunctionDefinitionNode;
 import org.shedlang.compiler.types.Effects;
+import org.shedlang.compiler.types.FunctionType;
+import org.shedlang.compiler.types.NamedParam;
 import org.shedlang.compiler.types.Types;
 import org.shedlang.compiler.util.P;
 
+import java.util.List;
 import java.util.Map;
 
 import static org.shedlang.compiler.ast.typed.TypedNodeMatchers.*;
@@ -33,6 +36,7 @@ public class TypeCheckerFunctionDefinitionTests {
                 .withEffect(isOptionalEmpty())
                 .withReturnType(isTypedTypeLevelReferenceNode("Unit", Types.UNIT))
                 .withBody(isSequence())
+                .withType(new FunctionType(List.of(), List.of(), Types.UNIT))
         );
     }
 
@@ -61,6 +65,11 @@ public class TypeCheckerFunctionDefinitionTests {
                         isSequence(isTypedParamNode("y", isTypedTypeLevelReferenceNode("Int", Types.INT)))
                     )
                 )
+                .withType(new FunctionType(
+                    List.of(Types.STRING),
+                    List.of(new NamedParam("y", Types.INT)),
+                    Types.UNIT
+                ))
         );
     }
     @Test