1 files changed, 36 insertions, 18 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 82ef0cd..42839c9 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
@@ -64,21 +64,7 @@ public class JavaTypesGenerator implements Generator {
     private List<JavaCompilationUnit> generateNamespace(TypedNamespaceNode namespace) {
         var packageParts = namespaceToJavaPackageParts(namespace.namespaceName());
 
-        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()));
-                }
-            }
-        }
+        var context = contextForNamespace(namespace);
 
         var javaCompilationUnits = new ArrayList<JavaCompilationUnit>();
 
@@ -87,7 +73,7 @@ public class JavaTypesGenerator implements Generator {
                 case TypedStructDefinitionNode structDefinition -> {
                     var javaCompilationUnit = new JavaCompilationUnit(
                         packageParts,
-                        generateStructDefinition(structDefinition, javaClassToImplementsType)
+                        generateStructDefinition(structDefinition, context)
                     );
 
                     javaCompilationUnits.add(javaCompilationUnit);
@@ -115,7 +101,7 @@ public class JavaTypesGenerator implements Generator {
 
     private JavaRecordDeclaration generateStructDefinition(
         TypedStructDefinitionNode structDefinition,
-        HashMap<String, List<JavaTypeRef>> javaClassToImplementsType
+        Context context
     ) {
         var components = structDefinition.fields().stream()
             .map(field -> new JavaRecordComponent(
@@ -127,7 +113,7 @@ public class JavaTypesGenerator implements Generator {
         return new JavaRecordDeclaration(
             structDefinition.name(),
             components,
-            javaClassToImplementsType.getOrDefault(structDefinition.name(), List.of()),
+            context.implementsType(structDefinition.type()),
             List.of(
                 new JavaRecordDeclaration(
                     "Builder",
@@ -188,4 +174,36 @@ public class JavaTypesGenerator implements Generator {
                 throw new UnsupportedOperationException("TODO");
         };
     }
+
+    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;
+
+        public Context(HashMap<String, List<JavaTypeRef>> javaClassToImplementsType) {
+            this.javaClassToImplementsType = javaClassToImplementsType;
+        }
+
+        public List<JavaTypeRef> implementsType(StructType type) {
+            return this.javaClassToImplementsType.getOrDefault(type.name(), List.of());
+        }
+    }
 }