Add incomplete type-checking for if expressions

10 files changed, 137 insertions, 7 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedBlockNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedBlockNode.java
index b5887627..5c0c8435 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedBlockNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedBlockNode.java
@@ -1,6 +1,8 @@
 package org.shedlang.compiler.ast.typed;
 
 import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.types.Type;
+import org.shedlang.compiler.types.Types;
 import org.shedlang.compiler.util.Lists;
 
 import java.util.List;
@@ -16,4 +18,17 @@ public record TypedBlockNode(
             .map(statement -> statement.behaviour())
             .orElse(TypedFunctionStatementBehaviour.CONTINUE);
     }
+
+    public Type type() {
+        return switch (behaviour()) {
+            case TypedFunctionStatementBehaviour.BlockReturn blockReturn ->
+                blockReturn.type();
+
+            case TypedFunctionStatementBehaviour.Continue continue_ ->
+                Types.UNIT;
+
+            case TypedFunctionStatementBehaviour.Terminate terminate ->
+                Types.BOTTOM;
+        };
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedConditionalBranchNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedConditionalBranchNode.java
new file mode 100644
index 00000000..c811a8ee
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedConditionalBranchNode.java
@@ -0,0 +1,10 @@
+package org.shedlang.compiler.ast.typed;
+
+import org.shedlang.compiler.sources.Source;
+
+public record TypedConditionalBranchNode(
+    TypedExpressionNode condition,
+    TypedBlockNode body,
+    Source source
+) implements TypedNode {
+}
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedExpressionNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedExpressionNode.java
index 963cf8a3..9c9ddbf0 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedExpressionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedExpressionNode.java
@@ -4,6 +4,6 @@ import org.shedlang.compiler.types.Type;
 
 public sealed interface TypedExpressionNode
     extends TypedNode
-    permits TypedBinaryOperationNode, TypedBoolLiteralNode, TypedCallNode, TypedIntLiteralNode, TypedReferenceNode, TypedStringLiteralNode {
+    permits TypedBinaryOperationNode, TypedBoolLiteralNode, TypedCallNode, TypedIfNode, TypedIntLiteralNode, TypedReferenceNode, TypedStringLiteralNode {
     Type type();
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedIfNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedIfNode.java
new file mode 100644
index 00000000..6ee18510
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedIfNode.java
@@ -0,0 +1,14 @@
+package org.shedlang.compiler.ast.typed;
+
+import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.types.Type;
+
+import java.util.List;
+
+public record TypedIfNode(
+    List<TypedConditionalBranchNode> conditionalBranches,
+    TypedBlockNode elseBody,
+    Type type,
+    Source source
+) implements TypedExpressionNode {
+}
diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index 70f38cf6..7365a8b4 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -72,6 +72,9 @@ public class StackCompiler {
             case TypedCallNode typedCallNode ->
                 compileCall(typedCallNode);
 
+            case TypedIfNode typedIfNode ->
+                throw new UnsupportedOperationException("TODO");
+
             case TypedIntLiteralNode typedIntLiteralNode ->
                 compileIntLiteral(typedIntLiteralNode);
 
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 90276186..b6233c04 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -7,12 +7,14 @@ import org.shedlang.compiler.typechecker.errors.*;
 import org.shedlang.compiler.typechecker.typesets.TypeSetTypeLevelValue;
 import org.shedlang.compiler.typechecker.typesets.TypeSets;
 import org.shedlang.compiler.types.*;
+import org.shedlang.compiler.util.P;
 import org.shedlang.compiler.util.PCollectors;
 
 import java.util.*;
 import java.util.stream.Stream;
 
 public class TypeChecker {
+
     private TypeChecker() {
     }
 
@@ -205,11 +207,19 @@ public class TypeChecker {
         UntypedBlockNode untyped,
         TypeCheckerContext context
     ) {
-        return context.mapWithContext(
+        // TODO: get environment and pass to enterBlock()
+        var bodyContext = context.enterBlock(P.map());
+
+        var typedStatementsResult = bodyContext.mapWithContext(
             untyped.statements(),
             TypeChecker::typeCheckFunctionStatement
-        )
-        .map(statements -> new TypedBlockNode(statements, untyped.source()));
+        );
+        var typedStatements = typedStatementsResult.value();
+        bodyContext = typedStatementsResult.context();
+
+        var typed = new TypedBlockNode(typedStatements, untyped.source());
+
+        return context.exit(bodyContext).toResult(typed);
     }
 
     private static TypeCheckerResult<TypedFunctionStatementNode> typeCheckBlockReturn(
@@ -281,6 +291,27 @@ public class TypeChecker {
         ));
     }
 
+    private static TypeCheckerResult<TypedConditionalBranchNode> typeCheckConditionalBranch(
+        UntypedConditionalBranchNode conditionalBranch,
+        TypeCheckerContext context
+    ) {
+        var typedConditionResult = typeCheckExpression(conditionalBranch.condition(), context);
+        var typedCondition = typedConditionResult.value();
+        context = typedConditionResult.context();
+
+        checkExpressionType(typedCondition, Types.BOOL, context);
+
+        var typedBodyResult = typeCheckBlock(conditionalBranch.body(), context);
+        var typedBody = typedBodyResult.value();
+        context = typedBodyResult.context();
+
+        return context.toResult(new TypedConditionalBranchNode(
+            typedCondition,
+            typedBody,
+            conditionalBranch.source()
+        ));
+    }
+
     private static TypedTypeLevelExpressionNode<Effect> typeCheckEffect(
         UntypedTypeLevelExpressionNode untyped,
         TypeCheckerContext context
@@ -303,7 +334,7 @@ public class TypeChecker {
                 typeCheckCall(untypedCall, context);
 
             case UntypedIfNode untypedIf ->
-                throw new UnsupportedOperationException("TODO");
+                typeCheckIf(untypedIf, context);
 
             case UntypedIntLiteralNode untypedIntLiteral ->
                 typeCheckIntLiteral(untypedIntLiteral, context);
@@ -451,6 +482,32 @@ public class TypeChecker {
         };
     }
 
+    private static TypeCheckerResult<? extends TypedExpressionNode> typeCheckIf(
+        UntypedIfNode untyped,
+        TypeCheckerContext context
+    ) {
+        var typedConditionalBranchesResult = context.mapWithContext(
+            untyped.conditionalBranches(),
+            TypeChecker::typeCheckConditionalBranch
+        );
+        var typedConditionalBranches = typedConditionalBranchesResult.value();
+        context = typedConditionalBranchesResult.context();
+
+        var typedElseBodyResult = typeCheckBlock(untyped.elseBody(), context);
+        var typedElseBody = typedElseBodyResult.value();
+        context = typedElseBodyResult.context();
+
+        var typed = new TypedIfNode(
+            typedConditionalBranches,
+            typedElseBody,
+            // TODO: consider type of conditional branches
+            typedElseBody.type(),
+            untyped.source()
+        );
+
+        return context.toResult(typed);
+    }
+
     private static TypeCheckerResult<TypedExpressionNode> typeCheckIntLiteral(
         UntypedIntLiteralNode untyped,
         TypeCheckerContext context
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 151f0a0b..5f95fdc3 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -76,6 +76,10 @@ public class TypeCheckerContext {
         return enter(environment, effect, Optional.of(returnType));
     }
 
+    public TypeCheckerContext enterBlock(PMap<String, Variable> environment) {
+        return enter(environment, this.effect, this.returnType);
+    }
+
     private TypeCheckerContext enter(
         PMap<String, Variable> environment,
         Effect effect,
@@ -89,16 +93,32 @@ public class TypeCheckerContext {
             FreeVariables.builder(),
             this.namespaceTypeInfo
         );
+    }
+
+    public TypeCheckerContext exit(TypeCheckerContext childContext) {
+        // TODO: we need to track free variables, but it's easy to discard the
+        // child context without calling exit. What's the right way of handling
+        // this?  Does using a callback work when we have type phases?
+        var context = this;
+
+        for (var freeVariable : childContext.freeVariables.build()) {
+            context = context.withVariable(freeVariable.name());
+        }
 
+        return context;
     }
 
     public TypeCheckerResult<Optional<Variable>> runtimeLookup(String name) {
         var variable = Optional.ofNullable(environment.get(name));
 
-        var context = boundVariables.contains(name) ? this : withFreeVariable(name);
+        var context = withVariable(name);
         return context.toResult(variable);
     }
 
+    private TypeCheckerContext withVariable(String name) {
+        return boundVariables.contains(name) ? this : withFreeVariable(name);
+    }
+
     public Optional<Variable> compileTimeLookup(String name) {
         return Optional.ofNullable(environment.get(name));
     }
diff --git a/src/main/java/org/shedlang/compiler/types/BottomType.java b/src/main/java/org/shedlang/compiler/types/BottomType.java
new file mode 100644
index 00000000..f975fd56
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/types/BottomType.java
@@ -0,0 +1,10 @@
+package org.shedlang.compiler.types;
+
+public record BottomType() implements Type {
+    public static final BottomType INSTANCE = new BottomType();
+
+    @Override
+    public String describe() {
+        return "Bottom";
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/types/Type.java b/src/main/java/org/shedlang/compiler/types/Type.java
index 72849de1..92948b15 100644
--- a/src/main/java/org/shedlang/compiler/types/Type.java
+++ b/src/main/java/org/shedlang/compiler/types/Type.java
@@ -6,5 +6,5 @@ package org.shedlang.compiler.types;
  * found through auxiliary data.
  */
 public sealed interface Type extends TypeLevelValue
-    permits BoolType, FunctionType, IntType, NamespaceType, StringType, TopType, TypeLevelValueType, UnitType {
+    permits BoolType, BottomType, FunctionType, IntType, NamespaceType, StringType, TopType, TypeLevelValueType, UnitType {
 }
diff --git a/src/main/java/org/shedlang/compiler/types/Types.java b/src/main/java/org/shedlang/compiler/types/Types.java
index 1cf17a0b..f192c98e 100644
--- a/src/main/java/org/shedlang/compiler/types/Types.java
+++ b/src/main/java/org/shedlang/compiler/types/Types.java
@@ -5,6 +5,7 @@ public class Types {
     }
 
     public static final Type BOOL = BoolType.INSTANCE;
+    public static final Type BOTTOM = BottomType.INSTANCE;
     public static final Type INT = IntType.INSTANCE;
     public static final Type STRING = StringType.INSTANCE;
     public static final Type TOP = TopType.INSTANCE;