1 files changed, 14 insertions, 31 deletions
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javatypes/JavaTypesGenerator.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javatypes/JavaTypesGenerator.java
index 7e785f1..2efd1ec 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javatypes/JavaTypesGenerator.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javatypes/JavaTypesGenerator.java
@@ -20,7 +20,6 @@ import java.io.File;
 import java.io.IOException;
 import java.nio.file.Path;
 import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.List;
 
 public class JavaTypesGenerator implements Generator {
@@ -45,8 +44,10 @@ public class JavaTypesGenerator implements Generator {
 
     @Override
     public void generate(List<TypedNamespaceNode> namespaces, TypesInfo typesInfo) throws IOException {
+        var context = new Context(typesInfo);
+
         for (var namespace : namespaces) {
-            var javaCompilationUnits = generateNamespace(namespace);
+            var javaCompilationUnits = generateNamespace(namespace, context);
 
             for (var javaCompilationUnit : javaCompilationUnits) {
                 var javaCompilationUnitPath = this.sourceRootDirectory
@@ -62,11 +63,9 @@ public class JavaTypesGenerator implements Generator {
         }
     }
 
-    private List<JavaCompilationUnit> generateNamespace(TypedNamespaceNode namespace) {
+    private List<JavaCompilationUnit> generateNamespace(TypedNamespaceNode namespace, Context context) {
         var packageParts = namespaceToJavaPackageParts(namespace.namespaceName());
 
-        var context = contextForNamespace(namespace);
-
         var javaCompilationUnits = new ArrayList<JavaCompilationUnit>();
 
         for (var statement : namespace.body()) {
@@ -111,10 +110,14 @@ public class JavaTypesGenerator implements Generator {
             ))
             .toList();
 
+        var implementsTypes = context.variantOf(structDefinition.type()).stream()
+            .map(this::generateTypeRef)
+            .toList();
+
         return new JavaRecordDeclaration(
             structDefinition.name(),
             components,
-            context.implementsType(structDefinition.type()),
+            implementsTypes,
             List.of(
                 new JavaRecordDeclaration(
                     "Builder",
@@ -176,35 +179,15 @@ public class JavaTypesGenerator implements Generator {
         };
     }
 
-    private Context contextForNamespace(JavaTypesGenerator this, TypedNamespaceNode namespace) {
-        var javaClassToImplementsType = new HashMap<String, List<JavaTypeRef>>();
-        for (var statement : namespace.body()) {
-            if (statement instanceof TypedSumDefinitionNode sumDefinition) {
-                for (var variant : sumDefinition.variants()) {
-                    // TODO: handle not struct types
-                    var variantTypeName = ((StructType) variant.type().value()).name();
-                    javaClassToImplementsType.putIfAbsent(
-                        variantTypeName,
-                        new ArrayList<>()
-                    );
-                    javaClassToImplementsType.get(variantTypeName)
-                        .add(generateTypeRef(sumDefinition.type()));
-                }
-            }
-        }
-
-        return new Context(javaClassToImplementsType);
-    }
-
     private static class Context {
-        private final HashMap<String, List<JavaTypeRef>> javaClassToImplementsType;
+        private final TypesInfo typesInfo;
 
-        public Context(HashMap<String, List<JavaTypeRef>> javaClassToImplementsType) {
-            this.javaClassToImplementsType = javaClassToImplementsType;
+        public Context(TypesInfo typesInfo) {
+            this.typesInfo = typesInfo;
         }
 
-        public List<JavaTypeRef> implementsType(StructType type) {
-            return this.javaClassToImplementsType.getOrDefault(type.name(), List.of());
+        public List<SumType> variantOf(StructType type) {
+            return this.typesInfo.variantOf(type);
         }
     }
 }