From 98c4b3e61abe39179f472db86f5ba19404df1110 Mon Sep 17 00:00:00 2001
From: praveenkrishna <praveenkrishna.d>
Date: Mon, 17 Dec 2018 18:41:45 +0530
Subject: [PATCH] Nested generics support for Structs

---
 .../compiler/ThriftCodecByteCodeGenerator.java  | 15 +++++++++++++--
 .../drift/codec/metadata/ThriftType.java        |  2 +-
 .../codec/AbstractThriftCodecManagerTest.java   | 17 +++++++++++++++++
 3 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/drift-codec/src/main/java/io/airlift/drift/codec/internal/compiler/ThriftCodecByteCodeGenerator.java b/drift-codec/src/main/java/io/airlift/drift/codec/internal/compiler/ThriftCodecByteCodeGenerator.java
index e907d810b..78778d0ca 100644
--- a/drift-codec/src/main/java/io/airlift/drift/codec/internal/compiler/ThriftCodecByteCodeGenerator.java
+++ b/drift-codec/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
index 809c68476..ecfeb4b6e 100644
--- 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
index b01b3aeb6..39de25481 100644
--- 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