1 files changed, 20 insertions, 16 deletions
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javatransient0/JavaTransient0Generator.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javatransient0/JavaTransient0Generator.java
index 9bb72a3..00f6b52 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javatransient0/JavaTransient0Generator.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javatransient0/JavaTransient0Generator.java
@@ -472,8 +472,8 @@ public class JavaTransient0Generator implements Generator {
 
             case TypedSumDefinitionNode sumDefinition -> {
                 yield Stream.of(
-                    generateEncodeSumMethod(sumDefinition),
-                    generateDecodeSumMethod(sumDefinition)
+                    generateEncodeSumMethod(sumDefinition, context.typesInfo),
+                    generateDecodeSumMethod(sumDefinition, context.typesInfo)
                 );
             }
         };
@@ -569,29 +569,31 @@ public class JavaTransient0Generator implements Generator {
         );
     }
 
-    private JavaMethodDeclaration generateEncodeSumMethod(TypedSumDefinitionNode sumDefinition) {
+    private JavaMethodDeclaration generateEncodeSumMethod(
+        TypedSumDefinitionNode sumDefinition,
+        TypesInfo typesInfo
+    ) {
         var encodeStatements = List.<JavaBlockStatement>of(new JavaSwitchStatement(
             new JavaRef(VALUE_NAME),
-            IntStream.range(0, sumDefinition.variants().size())
-                .mapToObj(variantIndex -> {
-                    var variant = sumDefinition.variants().get(variantIndex);
-                    var variantVariable = JavaIdentifier.of(VALUE_NAME.value() + "$" + variantIndex);
+            typesInfo.sumVariants(sumDefinition.type()).stream()
+                .map(variant -> {
+                    var variantVariable = JavaIdentifier.of(VALUE_NAME.value() + "$" + variant.tag());
 
                     var statements = new ArrayList<JavaBlockStatement>();
 
                     statements.addAll(generateEncode(
-                        new JavaIntegerLiteral(variantIndex),
+                        new JavaIntegerLiteral(variant.tag()),
                         NativeTypes.INT_32
                     ));
 
                     statements.addAll(generateEncode(
                         new JavaRef(variantVariable),
-                        variant.type().value()
+                        variant.containerType()
                     ));
 
                     return new JavaSwitchRule(
                         new JavaCasePattern(new JavaPattern(
-                            this.javaGenerator.generateTypeRef(variant.type().value()),
+                            this.javaGenerator.generateTypeRef(variant.containerType()),
                             variantVariable
                         )),
                         new JavaBlock(statements)
@@ -606,7 +608,10 @@ public class JavaTransient0Generator implements Generator {
         );
     }
 
-    private JavaMethodDeclaration generateDecodeSumMethod(TypedSumDefinitionNode sumDefinition) {
+    private JavaMethodDeclaration generateDecodeSumMethod(
+        TypedSumDefinitionNode sumDefinition,
+        TypesInfo typesInfo
+    ) {
         var decodeStatements = new ArrayList<JavaBlockStatement>();
 
         var variantTag = JavaIdentifier.of("variantTag");
@@ -615,18 +620,17 @@ public class JavaTransient0Generator implements Generator {
 
         var cases = new ArrayList<JavaSwitchRule>();
 
-        for (var variantIndex = 0; variantIndex < sumDefinition.variants().size(); variantIndex++) {
-            var variantType = sumDefinition.variants().get(variantIndex).type().value();
-            var variantVariable = JavaIdentifier.of(VALUE_NAME.value() + "$" + variantIndex);
+        for (var variant : typesInfo.sumVariants(sumDefinition.type())) {
+            var variantVariable = JavaIdentifier.of(VALUE_NAME.value() + "$" + variant.tag());
 
             var statements = new ArrayList<JavaBlockStatement>();
 
-            statements.addAll(generateDecode(variantVariable, variantType));
+            statements.addAll(generateDecode(variantVariable, variant.containerType()));
 
             statements.add(new JavaReturn(new JavaRef(variantVariable)));
 
             cases.add(new JavaSwitchRule(
-                new JavaCaseConstant(new JavaIntegerLiteral(variantIndex)),
+                new JavaCaseConstant(new JavaIntegerLiteral(variant.tag())),
                 new JavaBlock(statements)
             ));
         }