Nested generics support for Structs

airlift · Dec 24, 2018 · 0d7b886 · 0d7b886
1 parent 800c823
commit 0d7b886
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 3 deletions.
diff --git a/.../src/main/java/io/airlift/drift/codec/internal/compiler/ThriftCodecByteCodeGenerator.java b/.../src/main/java/io/airlift/drift/codec/internal/compiler/ThriftCodecByteCodeGenerator.java
@@ -63,6 +63,7 @@
 import java.lang.reflect.Type;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -84,6 +85,7 @@
 import static io.airlift.bytecode.BytecodeUtils.dumpBytecodeTree;
 import static io.airlift.bytecode.ClassGenerator.classGenerator;
 import static io.airlift.bytecode.Parameter.arg;
+import static io.airlift.bytecode.ParameterizedType.getPathName;
 import static io.airlift.bytecode.ParameterizedType.type;
 import static io.airlift.bytecode.ParameterizedType.typeFromPathName;
 import static io.airlift.bytecode.control.SwitchStatement.switchBuilder;
@@ -111,6 +113,7 @@
 import static io.airlift.drift.codec.metadata.FieldKind.THRIFT_FIELD;
 import static io.airlift.drift.codec.metadata.FieldKind.THRIFT_UNION_ID;
 import static java.lang.String.format;
+import static java.util.stream.Collectors.joining;
 
 @NotThreadSafe
 public class ThriftCodecByteCodeGenerator<T>
@@ -1038,8 +1041,16 @@ private static boolean needsCodec(ThriftFieldMetadata fieldMetadata)
 
     private static ParameterizedType toCodecType(ThriftStructMetadata metadata)
     {
-        String className = type(metadata.getStructClass()).getClassName();
-        return typeFromPathName(PACKAGE + "/" + className + "Codec");
+        String className = getPathName(metadata.getStructClass()) + "Codec";
+
+        Type type = metadata.getStructType();
+        if (type instanceof java.lang.reflect.ParameterizedType) {
+            className += Arrays.stream(((java.lang.reflect.ParameterizedType) type).getActualTypeArguments())
+                    .map(arg -> arg.getTypeName().replaceAll("[^a-zA-Z0-9]+", "_"))
+                    .collect(joining("$", "$$", ""));
+        }
+
+        return typeFromPathName(PACKAGE + "/" + className);
     }
 
     private static boolean isOptionalWrapper(Type javaType)

diff --git a/drift-codec/src/main/java/io/airlift/drift/codec/metadata/ThriftType.java b/drift-codec/src/main/java/io/airlift/drift/codec/metadata/ThriftType.java
@@ -210,7 +210,7 @@ private ThriftType(ThriftStructMetadata structMetadata)
         requireNonNull(structMetadata, "structMetadata is null");
 
         this.protocolType = ThriftProtocolType.STRUCT;
-        this.javaType = structMetadata.getStructClass();
+        this.javaType = structMetadata.getStructType();
         keyTypeReference = null;
         valueTypeReference = null;
         this.structMetadata = structMetadata;

diff --git a/drift-codec/src/test/java/io/airlift/drift/codec/AbstractThriftCodecManagerTest.java b/drift-codec/src/test/java/io/airlift/drift/codec/AbstractThriftCodecManagerTest.java
@@ -57,6 +57,8 @@
 import java.lang.reflect.Type;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.OptionalDouble;
 import java.util.OptionalInt;
@@ -467,6 +469,21 @@ public void testBeanGeneric()
         bean.setGenericProperty("genericValue");
 
         testRoundTripSerialize(new TypeToken<GenericThriftStructBean<String>>() {}, bean);
+
+        GenericThriftStructBean<Long> beanForLong = new GenericThriftStructBean<>();
+        beanForLong.setGenericProperty(123L);
+
+        testRoundTripSerialize(new TypeToken<GenericThriftStructBean<Long>>() {}, beanForLong);
+
+        GenericThriftStructBean<List<String>> beanForList = new GenericThriftStructBean<>();
+        beanForList.setGenericProperty(ImmutableList.of("abc", "xyz"));
+
+        testRoundTripSerialize(new TypeToken<GenericThriftStructBean<List<String>>>() {}, beanForList);
+
+        GenericThriftStructBean<Map<String, List<String>>> beanForMap = new GenericThriftStructBean<>();
+        beanForMap.setGenericProperty(ImmutableMap.of("test", ImmutableList.of("abc", "xyz")));
+
+        testRoundTripSerialize(new TypeToken<GenericThriftStructBean<Map<String, List<String>>>>() {}, beanForMap);
     }
 
     @Test