Parse type-level params for function type

4 files changed, 69 insertions, 1 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionTypeNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionTypeNode.java
index b4a465b3..c1b7ad3d 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionTypeNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionTypeNode.java
@@ -5,9 +5,11 @@ import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.sources.Sources;
 import org.shedlang.compiler.util.P;
 
+import java.util.List;
 import java.util.Optional;
 
 public record UntypedFunctionTypeNode(
+    Optional<List<UntypedTypeLevelParamNode>> typeLevelParams,
     UntypedFunctionTypeParamsNode params,
     Optional<UntypedTypeLevelExpressionNode> effect,
     UntypedTypeLevelExpressionNode returnType,
@@ -15,6 +17,7 @@ public record UntypedFunctionTypeNode(
 ) implements UntypedTypeLevelExpressionNode {
     public static Builder builder() {
         return new Builder(
+            Optional.empty(),
             P.vector(),
             P.vector(),
             Optional.empty(),
@@ -23,6 +26,7 @@ public record UntypedFunctionTypeNode(
     }
 
     public record Builder(
+        Optional<List<UntypedTypeLevelParamNode>> typeLevelParams,
         PVector<UntypedTypeLevelExpressionNode> positionalParams,
         PVector<UntypedFunctionTypeNamedParamNode> namedParams,
         Optional<UntypedTypeLevelExpressionNode> effect,
@@ -30,6 +34,7 @@ public record UntypedFunctionTypeNode(
     ) {
         public UntypedFunctionTypeNode build() {
             return new UntypedFunctionTypeNode(
+                this.typeLevelParams,
                 new UntypedFunctionTypeParamsNode(this.positionalParams, this.namedParams, Sources.NULL),
                 this.effect,
                 this.returnType,
@@ -37,8 +42,19 @@ public record UntypedFunctionTypeNode(
             );
         }
 
+        public Builder withTypeLevelParams(List<UntypedTypeLevelParamNode> typeLevelParams) {
+            return new Builder(
+                Optional.of(typeLevelParams),
+                this.positionalParams,
+                this.namedParams,
+                this.effect,
+                this.returnType
+            );
+        }
+
         public Builder addPositionalParam(UntypedTypeLevelExpressionNode param) {
             return new Builder(
+                this.typeLevelParams,
                 this.positionalParams.plus(param),
                 this.namedParams,
                 this.effect,
@@ -48,6 +64,7 @@ public record UntypedFunctionTypeNode(
 
         public Builder addNamedParam(String name, UntypedTypeLevelExpressionNode type) {
             return new Builder(
+                this.typeLevelParams,
                 this.positionalParams,
                 this.namedParams.plus(new UntypedFunctionTypeNamedParamNode(name, type, Sources.NULL)),
                 this.effect,
@@ -57,6 +74,7 @@ public record UntypedFunctionTypeNode(
 
         public Builder withEffect(UntypedTypeLevelExpressionNode effect) {
             return new Builder(
+                this.typeLevelParams,
                 this.positionalParams,
                 this.namedParams,
                 Optional.of(effect),
@@ -66,6 +84,7 @@ public record UntypedFunctionTypeNode(
 
         public Builder withReturnType(UntypedTypeLevelExpressionNode returnType) {
             return new Builder(
+                this.typeLevelParams,
                 this.positionalParams,
                 this.namedParams,
                 this.effect,
diff --git a/src/main/java/org/shedlang/compiler/parser/Parser.java b/src/main/java/org/shedlang/compiler/parser/Parser.java
index 003d909c..8a55de2d 100644
--- a/src/main/java/org/shedlang/compiler/parser/Parser.java
+++ b/src/main/java/org/shedlang/compiler/parser/Parser.java
@@ -338,7 +338,13 @@ public class Parser {
         var end = tokens.endPosition();
         var source = new FileFragmentSource(start.to(end));
 
-        return new UntypedFunctionTypeNode(params, Optional.empty(), returnType, source);
+        return new UntypedFunctionTypeNode(
+            Optional.empty(),
+            params,
+            Optional.empty(),
+            returnType,
+            source
+        );
     }
 
     private static final Map<TokenType, OperationParser<OperatorPrecedence, UntypedExpressionNode>> OPERATION_PARSERS = Stream.of(
@@ -579,6 +585,7 @@ public class Parser {
         var source = new FileFragmentSource(start.to(end));
 
         return new UntypedFunctionTypeNode(
+            signature.typeLevelParams(),
             signature.params(),
             signature.effect(),
             signature.returnType(),
@@ -591,6 +598,8 @@ public class Parser {
     ) {
         var start = tokens.startPosition();
 
+        var typeLevelParams = parseTypeLevelParamsOptional(tokens);
+
         tokens.skip(TokenType.SYMBOL_PAREN_OPEN);
         var params = parseParamTypes(tokens);
         tokens.skip(TokenType.SYMBOL_PAREN_CLOSE);
@@ -605,6 +614,7 @@ public class Parser {
         var source = new FileFragmentSource(start.to(end));
 
         return new UntypedFunctionTypeNode(
+            typeLevelParams,
             params,
             effect,
             returnType,
diff --git a/src/test/java/org/shedlang/compiler/ast/untyped/UntypedFunctionTypeNodeMatcher.java b/src/test/java/org/shedlang/compiler/ast/untyped/UntypedFunctionTypeNodeMatcher.java
index eb13ced4..2c7da346 100644
--- a/src/test/java/org/shedlang/compiler/ast/untyped/UntypedFunctionTypeNodeMatcher.java
+++ b/src/test/java/org/shedlang/compiler/ast/untyped/UntypedFunctionTypeNodeMatcher.java
@@ -5,6 +5,7 @@ import org.pcollections.PVector;
 import org.shedlang.compiler.testing.MatcherBuilder;
 import org.zwobble.precisely.Matcher;
 
+import java.util.List;
 import java.util.Optional;
 
 import static org.zwobble.precisely.Matchers.has;
@@ -45,4 +46,10 @@ public class UntypedFunctionTypeNodeMatcher extends MatcherBuilder<
     public UntypedFunctionTypeNodeMatcher withReturnType(Matcher<UntypedTypeLevelExpressionNode> returnType) {
         return with(has("returnType", x -> x.returnType(), returnType));
     }
+
+    public UntypedFunctionTypeNodeMatcher withTypeLevelParams(
+        Matcher<Optional<? extends List<UntypedTypeLevelParamNode>>> typeLevelParams
+    ) {
+        return with(has("typeLevelParams", x -> x.typeLevelParams(), typeLevelParams));
+    }
 }
diff --git a/src/test/java/org/shedlang/compiler/parser/ParserFunctionTypeTests.java b/src/test/java/org/shedlang/compiler/parser/ParserFunctionTypeTests.java
index 209220ad..319b0252 100644
--- a/src/test/java/org/shedlang/compiler/parser/ParserFunctionTypeTests.java
+++ b/src/test/java/org/shedlang/compiler/parser/ParserFunctionTypeTests.java
@@ -164,4 +164,36 @@ public class ParserFunctionTypeTests {
                 .withEffect(isOptionalOf(isUntypedTypeLevelReferenceNode("Io")))
         );
     }
+
+    @Test
+    public void canParseSingleTypeLevelParam() {
+        var source = "fn [A]() -> Unit";
+
+        var node = parseString(source, Parser::parseTopLevelTypeLevelExpression);
+
+        assertThat(
+            node,
+            isUntypedFunctionTypeNode()
+                .withTypeLevelParams(isOptionalOf(isSequence(
+                    isUntypedTypeParamNode("A")
+                )))
+        );
+    }
+
+    @Test
+    public void canParseManyTypeLevelParams() {
+        var source = "fn [A, B, C]() -> Unit";
+
+        var node = parseString(source, Parser::parseTopLevelTypeLevelExpression);
+
+        assertThat(
+            node,
+            isUntypedFunctionTypeNode()
+                .withTypeLevelParams(isOptionalOf(isSequence(
+                    isUntypedTypeParamNode("A"),
+                    isUntypedTypeParamNode("B"),
+                    isUntypedTypeParamNode("C")
+                )))
+        );
+    }
 }