Distinguish sum variant from untagged value in java-transient-0

3 files changed, 48 insertions, 8 deletions

diff --git a/examples/11-transient-0/output/java/src/test/java/org/zwobble/example/Transient0Tests.java b/examples/11-transient-0/output/java/src/test/java/org/zwobble/example/Transient0Tests.java
index 610532e..0a87da0 100644
--- a/examples/11-transient-0/output/java/src/test/java/org/zwobble/example/Transient0Tests.java
+++ b/examples/11-transient-0/output/java/src/test/java/org/zwobble/example/Transient0Tests.java
@@ -14,6 +14,7 @@ import java.util.Map;
 import java.util.Optional;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
+import org.zwobble.example.types.HobgoblinTransient0TaggedSharedValue;
 import org.zwobble.example.types.data.EnumWithVariants;
 import org.zwobble.example.types.data.HobgoblinTransient0Data;
 import org.zwobble.example.types.data.StructWithBool;
@@ -290,7 +291,7 @@ public class Transient0Tests {
     }
 
     private interface Encoder<T> {
-        void encode(T value, OutputStream outputStream, Map<Object, Long> sharedValues) throws IOException;
+        void encode(T value, OutputStream outputStream, Map<HobgoblinTransient0TaggedSharedValue, Long> sharedValues) throws IOException;
     }
 
     private interface Decoder<T> {
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 f1d4957..7a27e52 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
@@ -16,9 +16,7 @@ import org.zwobble.json5.reader.Json5ObjectReader;
 
 import java.io.IOException;
 import java.nio.file.Path;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Optional;
+import java.util.*;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
 
@@ -49,10 +47,12 @@ public class JavaTransient0Generator implements Generator {
 
     private final Path sourceRootDirectory;
     private final JavaGenerator javaGenerator;
+    private final Map<Type, Long> typeIds;
 
     public JavaTransient0Generator(Path sourceRootDirectory, JavaGenerator javaGenerator) {
         this.sourceRootDirectory = sourceRootDirectory;
         this.javaGenerator = javaGenerator;
+        this.typeIds = new HashMap<>();
     }
 
     @Override
@@ -62,6 +62,7 @@ public class JavaTransient0Generator implements Generator {
     ) throws IOException {
         var context = new Context(typesInfo);
 
+        generateSupportTypes();
         generateBuiltins();
 
         for (var namespace : namespaces) {
@@ -71,6 +72,44 @@ public class JavaTransient0Generator implements Generator {
         }
     }
 
+    private void generateSupportTypes() throws IOException {
+        JavaGenerator.write(this.sourceRootDirectory, List.of(
+            new JavaCompilationUnit(new JavaRecordDeclaration(
+                this.taggedSharedValueRef(),
+                List.of(
+                    new JavaRecordComponent(JavaTypeRef.LONG, JavaIdentifier.of("typeId")),
+                    new JavaRecordComponent(JavaTypeRef.OBJECT, JavaIdentifier.of("value"))
+                ),
+                List.of(),
+                List.of(),
+                DocComment.EMPTY
+            ))
+        ));
+    }
+
+    private JavaTypeRef taggedSharedValueRef() {
+        var packageName = this.javaGenerator.namespaceToJavaPackageName(NamespaceName.of());
+
+        return JavaTypeRef.topLevel(
+            packageName,
+            JavaIdentifier.of("HobgoblinTransient0TaggedSharedValue")
+        );
+    }
+
+    private JavaExpression taggedSharedValue(JavaExpression value, Type type) {
+        return new JavaNewExpression(
+            this.taggedSharedValueRef(),
+            List.of(new JavaIntegerLiteral(this.typeId(type)), value)
+        );
+    }
+
+    private long typeId(Type type) {
+        var typeId = this.typeIds
+            .computeIfAbsent(type, _ -> (long) this.typeIds.size());
+
+        return typeId;
+    }
+
     private void generateBuiltins() throws IOException {
         var packageName = this.javaGenerator.namespaceToJavaPackageName(NamespaceName.of());
         var transient0JavaName = generateTransient0ClassName("Builtins");
@@ -551,7 +590,7 @@ public class JavaTransient0Generator implements Generator {
                     OUTPUT_STREAM_NAME
                 ),
                 new JavaParam(
-                    JavaTypeRef.map(JavaTypeRef.OBJECT, JavaTypeRef.LONG_BOXED),
+                    JavaTypeRef.map(this.taggedSharedValueRef(), JavaTypeRef.LONG_BOXED),
                     SHARED_VALUES_NAME
                 )
             ),
@@ -912,7 +951,7 @@ public class JavaTransient0Generator implements Generator {
                     Optional.of(new JavaMethodCall(
                         new JavaRef(SHARED_VALUES_NAME),
                         JavaIdentifier.of("get"),
-                        List.of(value)
+                        List.of(taggedSharedValue(value, type))
                     ))
                 ),
                 new JavaIfStatement(
@@ -938,7 +977,7 @@ public class JavaTransient0Generator implements Generator {
                                 new JavaRef(SHARED_VALUES_NAME),
                                 JavaIdentifier.of("put"),
                                 List.of(
-                                    value,
+                                    this.taggedSharedValue(value, type),
                                     new JavaCast(JavaTypeRef.LONG, new JavaMethodCall(
                                         new JavaRef(SHARED_VALUES_NAME),
                                         JavaIdentifier.of("size"),
diff --git a/src/test/java/org/zwobble/hobgoblin/compiler/ExampleTests.java b/src/test/java/org/zwobble/hobgoblin/compiler/ExampleTests.java
index fceb6e3..1c77a83 100644
--- a/src/test/java/org/zwobble/hobgoblin/compiler/ExampleTests.java
+++ b/src/test/java/org/zwobble/hobgoblin/compiler/ExampleTests.java
@@ -52,7 +52,7 @@ public class ExampleTests {
                     var actualBytes = Files.readAllBytes(transient0OutputDir.resolve(name));
                     if (expectedBytes.isEmpty()) {
                         expectedBytes = Optional.of(actualBytes);
-                    } else if (!name.equals("StructWithSharedSumAndVariant")) {
+                    } else {
                         assertThat(Arrays.equals(actualBytes, expectedBytes.get()), equalTo(true));
                     }
                 }