Pass hint when type checking expression

1 files changed, 61 insertions, 38 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 333afb2b..ff416edc 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -72,42 +72,52 @@ public class TypeChecker {
             );
         }
 
-        var positionalResult = context.mapWithContext(args.positional(), TypeChecker::typeCheckExpression);
-        var typedPositionalArgs = positionalResult.value();
-        context = positionalResult.context();
-
-        var namedResult = context.mapWithContext(args.named(), (arg, currentContext) -> {
-            return typeCheckExpression(arg.expression(), currentContext)
-                .map(typedExpression -> new TypedNamedArgNode(
-                    arg.name(),
-                    typedExpression,
-                    arg.source()
-                ));
-        });
-        var typedNamedArgs = namedResult.value();
-        context = namedResult.context();
+        var typedPositionalArgs = new ArrayList<TypedExpressionNode>();
+        for (var argIndex = 0; argIndex < args.positional().size(); argIndex++) {
+            var argNode = args.positional().get(argIndex);
+            var paramType = receiverType.positionalParams().get(argIndex);
+
+            var typedPositionalArgResult = typeCheckExpression(
+                argNode,
+                Optional.of(paramType),
+                context
+            );
+            var typedPositionalArg = typedPositionalArgResult.value();
+            context = typedPositionalArgResult.context();
 
-        for (var i = 0; i < typedPositionalArgs.size(); i++) {
-            var typedArg = typedPositionalArgs.get(i);
-            var paramType = receiverType.positionalParams().get(i);
+            checkExpressionType(typedPositionalArg, paramType, context);
 
-            checkExpressionType(typedArg, paramType, context);
+            typedPositionalArgs.add(typedPositionalArg);
         }
 
+        var typedNamedArgs = new ArrayList<TypedNamedArgNode>();
         var seenNamedArgNames = new HashSet<String>();
-        for (var typedArg : typedNamedArgs) {
+        for (var untypedNamedArgNode : args.named()) {
             var param = receiverType.namedParams().stream()
-                .filter(namedParam -> namedParam.name().equals(typedArg.name()))
+                .filter(namedParam -> namedParam.name().equals(untypedNamedArgNode.name()))
                 .findFirst()
-                .orElseThrow(() -> new UnexpectedNamedArgError(typedArg.name(), typedArg.source()));
+                .orElseThrow(() -> new UnexpectedNamedArgError(untypedNamedArgNode.name(), untypedNamedArgNode.source()));
+
+            var typedExpressionResult = typeCheckExpression(
+                untypedNamedArgNode.expression(),
+                Optional.of(param.type()),
+                context
+            );
+            var typedExpression = typedExpressionResult.value();
+            context = typedExpressionResult.context();
 
-            checkExpressionType(typedArg.expression(), param.type(), context);
+            checkExpressionType(typedExpression, param.type(), context);
 
-            if (seenNamedArgNames.contains(typedArg.name())) {
-                throw new CannotPassNamedArgTwiceError(typedArg.name(), typedArg.source());
+            if (seenNamedArgNames.contains(untypedNamedArgNode.name())) {
+                throw new CannotPassNamedArgTwiceError(untypedNamedArgNode.name(), untypedNamedArgNode.source());
             }
 
-            seenNamedArgNames.add(typedArg.name());
+            typedNamedArgs.add(new TypedNamedArgNode(
+                untypedNamedArgNode.name(),
+                typedExpression,
+                untypedNamedArgNode.source()
+            ));
+            seenNamedArgNames.add(untypedNamedArgNode.name());
         }
 
         for (var namedParam : receiverType.namedParams()) {
@@ -134,7 +144,7 @@ public class TypeChecker {
         UntypedAssignNode untyped,
         TypeCheckerContext context
     ) {
-        var expressionResult = typeCheckExpression(untyped.expression(), context);
+        var expressionResult = typeCheckExpression(untyped.expression(), Optional.empty(), context);
         var expression = expressionResult.value();
         context = expressionResult.context();
 
@@ -153,11 +163,11 @@ public class TypeChecker {
         UntypedBinaryOperationNode untyped,
         TypeCheckerContext context
     ) {
-        var leftResult = typeCheckExpression(untyped.left(), context);
+        var leftResult = typeCheckExpression(untyped.left(), Optional.empty(), context);
         var left = leftResult.value();
         context = leftResult.context();
 
-        var rightResult = typeCheckExpression(untyped.right(), context);
+        var rightResult = typeCheckExpression(untyped.right(), Optional.empty(), context);
         var right = rightResult.value();
         context = rightResult.context();
 
@@ -247,7 +257,8 @@ public class TypeChecker {
         UntypedBlockReturnNode untyped,
         TypeCheckerContext context
     ) {
-        return typeCheckExpression(untyped.expression(), context).map(typedExpression -> {
+        // TODO: pass return type as hint if available and applicable
+        return typeCheckExpression(untyped.expression(), Optional.empty(), context).map(typedExpression -> {
             return new TypedBlockReturnNode(
                 typedExpression,
                 untyped.source()
@@ -269,7 +280,7 @@ public class TypeChecker {
         UntypedCallNode untyped,
         TypeCheckerContext context
     ) {
-        var receiverResult = typeCheckExpression(untyped.receiver(), context);
+        var receiverResult = typeCheckExpression(untyped.receiver(), Optional.empty(), context);
         var typedReceiver = receiverResult.value();
         context = receiverResult.context();
 
@@ -349,7 +360,11 @@ public class TypeChecker {
         UntypedConditionalBranchNode conditionalBranch,
         TypeCheckerContext context
     ) {
-        var typedConditionResult = typeCheckExpression(conditionalBranch.condition(), context);
+        var typedConditionResult = typeCheckExpression(
+            conditionalBranch.condition(),
+            Optional.of(Types.BOOL),
+            context
+        );
         var typedCondition = typedConditionResult.value();
         context = typedConditionResult.context();
 
@@ -375,6 +390,7 @@ public class TypeChecker {
 
     public static TypeCheckerResult<? extends TypedExpressionNode> typeCheckExpression(
         UntypedExpressionNode untyped,
+        Optional<Type> hint,
         TypeCheckerContext context
     ) {
         return switch (untyped) {
@@ -391,7 +407,7 @@ public class TypeChecker {
                 typeCheckFieldAccess(untypedFieldAccess, context);
 
             case UntypedFunctionExpressionNode untypedFunctionExpression ->
-                typeCheckFunctionExpression(untypedFunctionExpression, context);
+                typeCheckFunctionExpression(untypedFunctionExpression, hint, context);
 
             case UntypedIfNode untypedIf ->
                 typeCheckIf(untypedIf, context);
@@ -417,7 +433,7 @@ public class TypeChecker {
         UntypedExpressionStatementNode untyped,
         TypeCheckerContext context
     ) {
-        return typeCheckExpression(untyped.expression(), context).map(typedExpression -> {
+        return typeCheckExpression(untyped.expression(), Optional.empty(), context).map(typedExpression -> {
             return new TypedExpressionStatementNode(
                 typedExpression,
                 untyped.source()
@@ -429,7 +445,7 @@ public class TypeChecker {
         UntypedFieldAccessNode untyped,
         TypeCheckerContext context
     ) {
-        var typedReceiverResult = typeCheckExpression(untyped.receiver(), context);
+        var typedReceiverResult = typeCheckExpression(untyped.receiver(), Optional.empty(), context);
         var typedReceiver = typedReceiverResult.value();
         context = typedReceiverResult.context();
 
@@ -644,6 +660,7 @@ public class TypeChecker {
 
     private static TypeCheckerResult<TypedFunctionExpressionNode> typeCheckFunctionExpression(
         UntypedFunctionExpressionNode untyped,
+        Optional<Type> hint,
         TypeCheckerContext context
     ) {
         var typedSignature = typeCheckFunctionSignature(untyped.signature(), context);
@@ -693,7 +710,8 @@ public class TypeChecker {
         UntypedFunctionReturnNode untyped,
         TypeCheckerContext context
     ) {
-        var expressionResult = typeCheckExpression(untyped.expression(), context);
+        // TODO: pass return type as hint if available
+        var expressionResult = typeCheckExpression(untyped.expression(), Optional.empty(), context);
         var expression = expressionResult.value();
         context = expressionResult.context();
 
@@ -1279,7 +1297,8 @@ public class TypeChecker {
         var typedElements = new ArrayList<TypedExpressionNode>();
 
         for (var untypedElement : untyped.elements()) {
-            var typedElementResult = typeCheckExpression(untypedElement, context);
+            // TODO: pass type hint if hint is a tuple
+            var typedElementResult = typeCheckExpression(untypedElement, Optional.empty(), context);
             var typedElement = typedElementResult.value();
             context = typedElementResult.context();
 
@@ -1429,7 +1448,7 @@ public class TypeChecker {
         UntypedWhenNode untyped,
         TypeCheckerContext context
     ) {
-        var typedExpressionResult = typeCheckExpression(untyped.expression(), context);
+        var typedExpressionResult = typeCheckExpression(untyped.expression(), Optional.empty(), context);
         var typedExpression = typedExpressionResult.value();
         context = typedExpressionResult.context();
 
@@ -1530,7 +1549,11 @@ public class TypeChecker {
         UntypedWhileNode untyped,
         TypeCheckerContext context
     ) {
-        var typedConditionResult = typeCheckExpression(untyped.condition(), context);
+        var typedConditionResult = typeCheckExpression(
+            untyped.condition(),
+            Optional.of(Types.BOOL),
+            context
+        );
         var typedCondition = typedConditionResult.value();
         context = typedConditionResult.context();