Check field presence for sum variants

7 files changed, 182 insertions, 23 deletions

diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/SubtypeIsMissingFieldError.java b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/SubtypeIsMissingFieldError.java
new file mode 100644
index 0000000..ce59593
--- /dev/null
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/SubtypeIsMissingFieldError.java
@@ -0,0 +1,38 @@
+package org.zwobble.hobgoblin.compiler.typechecker;
+
+import org.zwobble.hobgoblin.compiler.errors.SourceError;
+import org.zwobble.hobgoblin.compiler.sources.Source;
+import org.zwobble.hobgoblin.compiler.types.Type;
+
+public class SubtypeIsMissingFieldError extends SourceError {
+    private final Type supertype;
+    private final Type subtype;
+    private final String fieldName;
+
+    public SubtypeIsMissingFieldError(Type supertype, Type subtype, String fieldName, Source source) {
+        super(
+            String.format(
+                "%s is a subtype of %s, but is missing field %s",
+                subtype.describe(),
+                supertype.describe(),
+                fieldName
+            ),
+            source
+        );
+        this.supertype = supertype;
+        this.subtype = subtype;
+        this.fieldName = fieldName;
+    }
+
+    public Type supertype() {
+        return supertype;
+    }
+
+    public Type subtype() {
+        return subtype;
+    }
+
+    public String fieldName() {
+        return fieldName;
+    }
+}
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeChecker.java b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeChecker.java
index c291aaf..c6f8721 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeChecker.java
@@ -2,6 +2,7 @@ package org.zwobble.hobgoblin.compiler.typechecker;
 
 import org.zwobble.hobgoblin.compiler.ast.typed.*;
 import org.zwobble.hobgoblin.compiler.ast.untyped.*;
+import org.zwobble.hobgoblin.compiler.sources.NullSource;
 import org.zwobble.hobgoblin.compiler.sources.Source;
 import org.zwobble.hobgoblin.compiler.types.*;
 
@@ -183,6 +184,17 @@ public class TypeChecker {
 
         context.defineStructType(structType, typeCheckedFieldDefinitions.fields());
 
+        for (var sumType : context.variantOf(structType)) {
+            // TODO: better error if sum type fields aren't defined
+            var sumTypeFields = context.fieldsOf(sumType).orElseThrow();
+            typeCheckVariantType(
+                sumType,
+                sumTypeFields,
+                structType,
+                typeCheckedFieldDefinitions.fields
+            );
+        }
+
         return new TypedStructDefinitionNode(
             structType,
             typeCheckedFieldDefinitions.typedFieldNodes(),
@@ -233,8 +245,7 @@ public class TypeChecker {
         var sumType = (SumType) lookupMetaType(untyped.name(), untyped.source(), context);
 
         var typedVariants = new ArrayList<TypedSumVariantDefinitionNode>();
-        // TODO: should be StructType?
-        var variantTypes = new ArrayList<Type>();
+        var variantTypes = new ArrayList<StructType>();
         for (var untypedVariant : untyped.variants()) {
             var variantType = typeCheckMetaType(untypedVariant.type(), context);
             var typedVariant = new TypedSumVariantDefinitionNode(
@@ -242,12 +253,25 @@ public class TypeChecker {
                 untyped.source()
             );
             typedVariants.add(typedVariant);
-            variantTypes.add(variantType.value());
+            // TODO: handle not struct type
+            variantTypes.add((StructType) variantType.value());
         }
 
         var typeCheckedFieldDefinitions = typeCheckFieldDefinitions(untyped.fields(), context);
 
-        context.defineSumType(sumType, variantTypes);
+        context.defineSumType(sumType, variantTypes, typeCheckedFieldDefinitions.fields);
+
+        for (var variantType : variantTypes) {
+            var variantTypeFields = context.fieldsOf(variantType);
+            if (variantTypeFields.isPresent()) {
+                typeCheckVariantType(
+                    sumType,
+                    typeCheckedFieldDefinitions.fields,
+                    variantType,
+                    variantTypeFields.get()
+                );
+            }
+        }
 
         return new TypedSumDefinitionNode(
             sumType,
@@ -274,6 +298,24 @@ public class TypeChecker {
         }
     }
 
+    private static void typeCheckVariantType(
+        SumType sumType,
+        List<Field> sumTypeFields,
+        StructType variantType,
+        List<Field> variantTypeFields
+    ) {
+        for (var sumTypeField : sumTypeFields) {
+            var variantTypeField = variantTypeFields.stream()
+                .filter(field -> field.name().equals(sumTypeField.name()))
+                .findFirst();
+
+            if (variantTypeField.isEmpty()) {
+                // TODO: set source appropriately.
+                throw new SubtypeIsMissingFieldError(sumType, variantType, sumTypeField.name(), NullSource.INSTANCE);
+            }
+        }
+    }
+
     static TypedTypeLevelExpressionNode<? extends TypeLevelValue> typeCheckTypeLevelExpression(
         UntypedTypeLevelExpressionNode untyped,
         TypeCheckerNamespaceContext context
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerGlobalContext.java b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerGlobalContext.java
index 3e7665c..ca2d731 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerGlobalContext.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerGlobalContext.java
@@ -8,7 +8,7 @@ import java.util.*;
 public class TypeCheckerGlobalContext {
     private final Map<String, Type> nativeTypes = new HashMap<>();
     private final Map<EnumType, List<EnumVariant>> enumVariants = new HashMap<>();
-    private final Map<StructType, List<Field>> structFieldsOf = new HashMap<>();
+    private final Map<Type, List<Field>> fieldsOf = new HashMap<>();
     private final Map<NamespaceName, Fields> namespaceFieldsOf = new HashMap<>();
     private final ManyToMany<SumType, Type> sumVariants = new ManyToMany<>();
 
@@ -22,7 +22,7 @@ public class TypeCheckerGlobalContext {
     public TypesInfo toTypesInfo() {
         return new TypesInfo(
             this.enumVariants,
-            this.structFieldsOf,
+            this.fieldsOf,
             this.sumVariants
         );
     }
@@ -40,7 +40,7 @@ public class TypeCheckerGlobalContext {
             namespaceName,
             this.nativeTypes,
             this.enumVariants,
-            this.structFieldsOf,
+            this.fieldsOf,
             this.sumVariants
         );
     }
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerNamespaceContext.java b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerNamespaceContext.java
index 906cdfa..0df7e16 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerNamespaceContext.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerNamespaceContext.java
@@ -13,7 +13,7 @@ public class TypeCheckerNamespaceContext {
         NamespaceName namespaceName,
         Map<String, Type> nativeTypes,
         Map<EnumType, List<EnumVariant>> enumVariants,
-        Map<StructType, List<Field>> fieldsOf,
+        Map<Type, List<Field>> fieldsOf,
         ManyToMany<SumType, Type> variants
     ) {
         return new TypeCheckerNamespaceContext(
@@ -28,14 +28,14 @@ public class TypeCheckerNamespaceContext {
     private final NamespaceName namespaceName;
     private final Map<String, Type> variables;
     private final Map<EnumType, List<EnumVariant>> enumVariants;
-    private final Map<StructType, List<Field>> fieldsOf;
+    private final Map<Type, List<Field>> fieldsOf;
     private final ManyToMany<SumType, Type> sumVariants;
 
     private TypeCheckerNamespaceContext(
         NamespaceName namespaceName,
         Map<String, Type> variables,
         Map<EnumType, List<EnumVariant>> enumVariants,
-        Map<StructType, List<Field>> fieldsOf,
+        Map<Type, List<Field>> fieldsOf,
         ManyToMany<SumType, Type> sumVariants
     ) {
         this.namespaceName = namespaceName;
@@ -66,9 +66,26 @@ public class TypeCheckerNamespaceContext {
         this.fieldsOf.put(structType, fields);
     }
 
-    public void defineSumType(SumType sumType, List<Type> variantTypes) {
+    public Optional<List<Field>> fieldsOf(StructType type) {
+        return Optional.ofNullable(this.fieldsOf.get(type));
+    }
+
+    public void defineSumType(
+        SumType sumType,
+        List<StructType> variantTypes,
+        List<Field> fields
+    ) {
         for (var variantType : variantTypes) {
             this.sumVariants.add(sumType, variantType);
         }
+        this.fieldsOf.put(sumType, fields);
+    }
+
+    public List<SumType> variantOf(StructType variantType) {
+        return this.sumVariants.rightToLeft(variantType);
+    }
+
+    public Optional<List<Field>> fieldsOf(SumType type) {
+        return Optional.ofNullable(this.fieldsOf.get(type));
     }
 }
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypesInfo.java b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypesInfo.java
index 0fb4f81..63d21b6 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypesInfo.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypesInfo.java
@@ -9,16 +9,16 @@ import java.util.Optional;
 
 public class TypesInfo {
     private final Map<EnumType, List<EnumVariant>> enumVariants;
-    private final Map<StructType, List<Field>> structFieldsOf;
+    private final Map<Type, List<Field>> fieldsOf;
     private final ManyToMany<SumType, Type> sumVariants;
 
     public TypesInfo(
         Map<EnumType, List<EnumVariant>> enumVariants,
-        Map<StructType, List<Field>> structFieldsOf,
+        Map<Type, List<Field>> fieldsOf,
         ManyToMany<SumType, Type> sumVariants
     ) {
         this.enumVariants = enumVariants;
-        this.structFieldsOf = structFieldsOf;
+        this.fieldsOf = fieldsOf;
         this.sumVariants = sumVariants;
     }
 
@@ -28,7 +28,7 @@ public class TypesInfo {
 
     public List<Field> fieldsOf(StructType structType) {
         // TODO: handle error better
-        return Optional.ofNullable(this.structFieldsOf.get(structType)).orElseThrow();
+        return Optional.ofNullable(this.fieldsOf.get(structType)).orElseThrow();
     }
 
     public List<SumType> variantOf(StructType variantType) {
diff --git a/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerStructDefinitionTests.java b/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerStructDefinitionTests.java
index f45b215..81bf9b8 100644
--- a/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerStructDefinitionTests.java
+++ b/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerStructDefinitionTests.java
@@ -5,17 +5,12 @@ import org.zwobble.hobgoblin.compiler.ast.typed.TypedNamespaceNode;
 import org.zwobble.hobgoblin.compiler.ast.typed.TypedStructDefinitionNode;
 import org.zwobble.hobgoblin.compiler.ast.typed.TypedStructFieldDefinitionNode;
 import org.zwobble.hobgoblin.compiler.ast.typed.TypedTypeLevelExpressionNode;
-import org.zwobble.hobgoblin.compiler.ast.untyped.UntypedArb;
-import org.zwobble.hobgoblin.compiler.ast.untyped.UntypedNamespaceNode;
-import org.zwobble.hobgoblin.compiler.ast.untyped.UntypedStructDefinitionNode;
-import org.zwobble.hobgoblin.compiler.ast.untyped.UntypedStructFieldDefinitionNode;
-import org.zwobble.hobgoblin.compiler.types.Field;
-import org.zwobble.hobgoblin.compiler.types.NamespaceName;
-import org.zwobble.hobgoblin.compiler.types.SimpleNativeType;
-import org.zwobble.hobgoblin.compiler.types.StructType;
+import org.zwobble.hobgoblin.compiler.ast.untyped.*;
+import org.zwobble.hobgoblin.compiler.types.*;
 
 import java.util.List;
 
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.zwobble.hobgoblin.compiler.typechecker.TypeCheckerTesting.typeCheckNamespaceStatement;
 import static org.zwobble.precisely.AssertThat.assertThat;
 import static org.zwobble.precisely.Matchers.*;
diff --git a/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerSumDefinitionTests.java b/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerSumDefinitionTests.java
index 57c40f5..36a4647 100644
--- a/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerSumDefinitionTests.java
+++ b/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerSumDefinitionTests.java
@@ -10,6 +10,7 @@ import org.zwobble.hobgoblin.compiler.types.*;
 
 import java.util.List;
 
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.zwobble.hobgoblin.compiler.typechecker.TypeCheckerTesting.typeCheckNamespaceStatement;
 import static org.zwobble.precisely.AssertThat.assertThat;
 import static org.zwobble.precisely.Matchers.*;
@@ -130,4 +131,70 @@ public class TypeCheckerSumDefinitionTests {
             )
         ));
     }
+
+    @Test
+    public void givenVariantTypeIsDefinedFirstWhenVariantTypeIsMissingFieldThenErrorIsThrown() {
+        var int32Type = SimpleNativeType.builtin("Int32");
+        var untyped = UntypedSumDefinitionNode.arbitrary()
+            .withName("Shape")
+            .withVariants(List.of(
+                UntypedSumVariantDefinitionNode.arbitrary()
+                    .withType(UntypedTypeLevelReferenceNode.arbitrary().withName("Square"))
+                    .build()
+            ))
+            .withFields(List.of(
+                UntypedStructFieldDefinitionNode.arbitrary()
+                    .withName("area")
+                    .withType(UntypedTypeLevelReferenceNode.arbitrary().withName("Int32"))
+                    .build()
+            ))
+            .build();
+        var globalContext = TypeCheckerGlobalContext.initial();
+        globalContext.addNativeType(int32Type);
+        var namespaceName = NamespaceName.of("a", "b");
+        var namespaceContext = globalContext
+            .enterNamespace(namespaceName);
+        var variantType = new StructType(namespaceName, "Square");
+        namespaceContext.declare("Square", new TypeLevelValueType(variantType));
+        namespaceContext.defineStructType(variantType, List.of());
+
+        var error = assertThrows(
+            SubtypeIsMissingFieldError.class,
+            () -> typeCheckNamespaceStatement(untyped, namespaceContext)
+        );
+
+        assertThat(error.supertype(), equalTo(new SumType(namespaceName, "Shape")));
+        assertThat(error.subtype(), equalTo(variantType));
+        assertThat(error.fieldName(), equalTo("area"));
+    }
+
+    @Test
+    public void givenSumTypeIsDefinedFirstWhenVariantTypeIsMissingFieldThenErrorIsThrown() {
+        var int32Type = SimpleNativeType.builtin("Int32");
+        var untyped = UntypedStructDefinitionNode.arbitrary()
+            .withName("Square")
+            .build();
+        var globalContext = TypeCheckerGlobalContext.initial();
+        globalContext.addNativeType(int32Type);
+        var namespaceName = NamespaceName.of("a", "b");
+        var namespaceContext = globalContext
+            .enterNamespace(namespaceName);
+        var sumType = new SumType(namespaceName, "Shape");
+        var variantType = new StructType(namespaceName, "Square");
+        namespaceContext.declare("Sum", new TypeLevelValueType(sumType));
+        namespaceContext.defineSumType(
+            sumType,
+            List.of(variantType),
+            List.of(new Field("area", int32Type))
+        );
+
+        var error = assertThrows(
+            SubtypeIsMissingFieldError.class,
+            () -> typeCheckNamespaceStatement(untyped, namespaceContext)
+        );
+
+        assertThat(error.supertype(), equalTo(sumType));
+        assertThat(error.subtype(), equalTo(variantType));
+        assertThat(error.fieldName(), equalTo("area"));
+    }
 }