Implement shared values properly in java-transient-0

3 files changed, 156 insertions, 12 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 0110de3..a482612 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
@@ -5,6 +5,7 @@ import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -194,37 +195,44 @@ public class Transient0Tests {
 
     @Test
     public void structWithShared() throws IOException {
-        var value = new StructWithShared(List.of(
-            new StructWithInt32(10, 25)
-        ));
+        var a = new StructWithInt32(10, 25);
+        var b = new StructWithInt32(42, 47);
 
-        assertRoundTripEncoding(
+        var value = new StructWithShared(List.of(a, a, b, a, b));
+
+        var decodedValue = assertRoundTripEncoding(
             value,
             HobgoblinTransient0Data::encodeStructWithShared,
             HobgoblinTransient0Data::decodeStructWithShared
         );
+        Assertions.assertSame(decodedValue.inner().get(0), decodedValue.inner().get(1));
+        Assertions.assertSame(decodedValue.inner().get(0), decodedValue.inner().get(3));
+        Assertions.assertSame(decodedValue.inner().get(2), decodedValue.inner().get(4));
     }
 
-    private <T> void assertRoundTripEncoding(
+    private <T> T assertRoundTripEncoding(
         T value,
         Encoder<T> encoder,
         Decoder<T> decoder
     ) throws IOException {
         var outputStream = new ByteArrayOutputStream();
+        // TODO: should this use IdentityHashMap? Leave it up the caller?
         encoder.encode(value, outputStream, new HashMap<>());
         var bytes = outputStream.toByteArray();
 
         var inputStream = new ByteArrayInputStream(bytes);
-        var decodedValue = decoder.decode(inputStream, new HashMap<>());
+        var decodedValue = decoder.decode(inputStream, new ArrayList<>());
 
         Assertions.assertEquals(value, decodedValue);
+
+        return decodedValue;
     }
 
     private interface Encoder<T> {
-        void encode(T value, OutputStream outputStream, Map<String, Object> sharedValues) throws IOException;
+        void encode(T value, OutputStream outputStream, Map<Object, Long> sharedValues) throws IOException;
     }
 
     private interface Decoder<T> {
-        T decode(InputStream inputStream, Map<String, Object> sharedValues) throws IOException;
+        T decode(InputStream inputStream, List<Object> sharedValues) throws IOException;
     }
 }
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 a775773..ca654e3 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
@@ -11,6 +11,7 @@ import org.zwobble.hobgoblin.compiler.output.lang.java.ast.*;
 import org.zwobble.hobgoblin.compiler.typechecker.TypesInfo;
 import org.zwobble.hobgoblin.compiler.types.*;
 import org.zwobble.hobgoblin.compiler.util.Casing;
+import org.zwobble.hobgoblin.compiler.util.Lists;
 import org.zwobble.json5.reader.Json5ObjectReader;
 
 import java.io.IOException;
@@ -550,7 +551,7 @@ public class JavaTransient0Generator implements Generator {
                     OUTPUT_STREAM_NAME
                 ),
                 new JavaParam(
-                    JavaTypeRef.map(JavaTypeRef.STRING, JavaTypeRef.OBJECT),
+                    JavaTypeRef.map(JavaTypeRef.OBJECT, JavaTypeRef.LONG_BOXED),
                     SHARED_VALUES_NAME
                 )
             ),
@@ -576,7 +577,7 @@ public class JavaTransient0Generator implements Generator {
                     INPUT_STREAM_NAME
                 ),
                 new JavaParam(
-                    JavaTypeRef.map(JavaTypeRef.STRING, JavaTypeRef.OBJECT),
+                    JavaTypeRef.list(JavaTypeRef.OBJECT),
                     SHARED_VALUES_NAME
                 )
             ),
@@ -586,12 +587,14 @@ public class JavaTransient0Generator implements Generator {
     }
 
     private List<JavaBlockStatement> generateEncode(JavaExpression value, Type type) {
-        return switch (this.javaGenerator.collapseType(type)) {
+        return switch (type) {
             case ConstructedNativeType constructedNativeType -> {
                 if (constructedNativeType.constructor().equals(NativeTypes.LIST)) {
                     yield generateEncodeList(value, constructedNativeType.args().getFirst());
                 } else if (constructedNativeType.constructor().equals(NativeTypes.OPTION)) {
                     yield generateEncodeOption(value, constructedNativeType.args().getFirst());
+                } else if (constructedNativeType.constructor().equals(NativeTypes.SHARED)) {
+                    yield generateEncodeShared(value, constructedNativeType.args().getFirst());
                 } else {
                     throw new UnsupportedOperationException("TODO");
                 }
@@ -658,12 +661,14 @@ public class JavaTransient0Generator implements Generator {
     }
 
     private List<JavaBlockStatement> generateDecode(JavaIdentifier target, Type type) {
-        return switch (this.javaGenerator.collapseType(type)) {
+        return switch (type) {
             case ConstructedNativeType constructedNativeType -> {
                 if (constructedNativeType.constructor().equals(NativeTypes.LIST)) {
                     yield generateDecodeList(target, constructedNativeType.args().getFirst());
                 } else if (constructedNativeType.constructor().equals(NativeTypes.OPTION)) {
                     yield generateDecodeOption(target, constructedNativeType.args().getFirst());
+                } else if (constructedNativeType.constructor().equals(NativeTypes.SHARED)) {
+                    yield generateDecodeShared(target, constructedNativeType.args().getFirst());
                 } else {
                     throw new UnsupportedOperationException("TODO");
                 }
@@ -893,6 +898,119 @@ public class JavaTransient0Generator implements Generator {
         return statements;
     }
 
+    private List<JavaBlockStatement> generateEncodeShared(
+        JavaExpression value,
+        Type type
+    ) {
+        var existingId = JavaIdentifier.of("existingId");
+
+        return List.of(
+            new JavaLocalVariableDeclaration(
+                Optional.empty(),
+                existingId,
+                Optional.of(new JavaMethodCall(
+                    new JavaRef(SHARED_VALUES_NAME),
+                    JavaIdentifier.of("get"),
+                    List.of(value)
+                ))
+            ),
+            new JavaIfStatement(
+                new JavaBinaryOperation(
+                    JavaBinaryOperator.NOT_EQUAL_TO,
+                    new JavaRef(existingId),
+                    new JavaNullLiteral()
+                ),
+                new JavaBlock(
+                    generateEncode(new JavaRef(existingId), NativeTypes.INT_64)
+                ),
+                new JavaBlock(Lists.concat(List.of(
+                    generateEncode(
+                        new JavaMethodCall(
+                            new JavaRef(SHARED_VALUES_NAME),
+                            JavaIdentifier.of("size"),
+                            List.of()
+                        ),
+                        NativeTypes.INT_64
+                    ),
+                    List.of(
+                        new JavaExpressionStatement(new JavaMethodCall(
+                            new JavaRef(SHARED_VALUES_NAME),
+                            JavaIdentifier.of("put"),
+                            List.of(
+                                value,
+                                new JavaCast(JavaTypeRef.LONG, new JavaMethodCall(
+                                    new JavaRef(SHARED_VALUES_NAME),
+                                    JavaIdentifier.of("size"),
+                                    List.of()
+                                ))
+                            )
+                        ))
+                    ),
+                    generateEncode(value, type)
+                )))
+            )
+        );
+    }
+
+    private List<JavaBlockStatement> generateDecodeShared(
+        JavaIdentifier target,
+        Type type
+    ) {
+        var id = JavaIdentifier.of("id");
+        var decodedValue = JavaIdentifier.of("decodedValue");
+
+        return Lists.concat(List.of(
+            generateDecode(id, NativeTypes.INT_64),
+            List.of(
+                new JavaLocalVariableDeclaration(
+                    Optional.of(this.javaGenerator.generateTypeRef(type)),
+                    target,
+                    Optional.empty()
+                ),
+                new JavaIfStatement(
+                    new JavaBinaryOperation(
+                        JavaBinaryOperator.LESS_THAN,
+                        new JavaRef(id),
+                        new JavaMethodCall(
+                            new JavaRef(SHARED_VALUES_NAME),
+                            JavaIdentifier.of("size"),
+                            List.of()
+                        )
+                    ),
+                    new JavaBlock(List.of(
+                        new JavaExpressionStatement(new JavaBinaryOperation(
+                            JavaBinaryOperator.ASSIGN,
+                            new JavaRef(target),
+                            new JavaCast(
+                                this.javaGenerator.generateTypeRef(type),
+                                new JavaMethodCall(
+                                    new JavaRef(SHARED_VALUES_NAME),
+                                    JavaIdentifier.of("get"),
+                                    List.of(new JavaCast(JavaTypeRef.INT, new JavaRef(id)))
+                                )
+                            )
+                        ))
+                    )),
+                    new JavaBlock(Lists.concat(List.of(
+                        generateDecode(decodedValue, type),
+                        List.of(
+                            new JavaExpressionStatement(new JavaMethodCall(
+                                new JavaRef(SHARED_VALUES_NAME),
+                                JavaIdentifier.of("add"),
+                                List.of(new JavaRef(decodedValue))
+                            )),
+                            new JavaExpressionStatement(new JavaBinaryOperation(
+                                JavaBinaryOperator.ASSIGN,
+                                new JavaRef(target),
+                                new JavaRef(decodedValue)
+                            ))
+                        )
+                    )))
+                )
+            )
+        ));
+    }
+
     private static JavaMethodCall generateOutputStreamWrite(JavaExpression byteToWrite) {
         return new JavaMethodCall(
             new JavaRef(OUTPUT_STREAM_NAME),
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/util/Lists.java b/src/main/java/org/zwobble/hobgoblin/compiler/util/Lists.java
new file mode 100644
index 0000000..31bbd42
--- /dev/null
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/util/Lists.java
@@ -0,0 +1,18 @@
+package org.zwobble.hobgoblin.compiler.util;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class Lists {
+    private Lists() {}
+
+    public static <T> List<T> concat(List<List<T>> lists) {
+        var result = new ArrayList<T>();
+
+        for (var list : lists) {
+            result.addAll(list);
+        }
+
+        return result;
+    }
+}