Support encoding/decoding of parameterized records (#1005)

jyemin · web-flow · commit ba39a2c34598 · 2022-09-30T11:06:10.000-04:00
JAVA-4740
diff --git a/bson-record-codec/src/main/org/bson/codecs/record/RecordCodec.java b/bson-record-codec/src/main/org/bson/codecs/record/RecordCodec.java
@@ -38,6 +38,8 @@
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.ParameterizedType;
 import java.lang.reflect.RecordComponent;
+import java.lang.reflect.Type;
+import java.lang.reflect.TypeVariable;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -48,9 +50,10 @@
 import static java.lang.String.format;
 import static org.bson.assertions.Assertions.notNull;
 
-final class RecordCodec<T extends Record> implements Codec<T> {
+final class RecordCodec<T extends Record> implements Codec<T>, Parameterizable {
     private static final Logger LOGGER = Loggers.getLogger("RecordCodec");
     private final Class<T> clazz;
+    private final boolean requiresParameterization;
     private final Constructor<?> canonicalConstructor;
     private final List<ComponentModel> componentModels;
     private final ComponentModel componentModelForId;
@@ -62,10 +65,11 @@ private static final class ComponentModel {
         private final int index;
         private final String fieldName;
 
-        private ComponentModel(final RecordComponent component, final CodecRegistry codecRegistry, final int index) {
+        private ComponentModel(final List<Type> typeParameters, final RecordComponent component, final CodecRegistry codecRegistry,
+                final int index) {
             validateAnnotations(component, index);
             this.component = component;
-            this.codec = computeCodec(component, codecRegistry);
+            this.codec = computeCodec(typeParameters, component, codecRegistry);
             this.index = index;
             this.fieldName = computeFieldName(component);
         }
@@ -83,11 +87,13 @@ Object getValue(final Record record) throws InvocationTargetException, IllegalAc
         }
 
         @SuppressWarnings("deprecation")
-        private static Codec<?> computeCodec(final RecordComponent component, final CodecRegistry codecRegistry) {
-            var codec = codecRegistry.get(toWrapper(component.getType()));
+        private static Codec<?> computeCodec(final List<Type> typeParameters, final RecordComponent component,
+                final CodecRegistry codecRegistry) {
+            var codec = codecRegistry.get(toWrapper(resolveComponentType(typeParameters, component)));
             if (codec instanceof Parameterizable parameterizableCodec
                     && component.getGenericType() instanceof ParameterizedType parameterizedType) {
-                codec = parameterizableCodec.parameterize(codecRegistry, Arrays.asList(parameterizedType.getActualTypeArguments()));
+                codec = parameterizableCodec.parameterize(codecRegistry,
+                        resolveActualTypeArguments(typeParameters, component.getDeclaringRecord(), parameterizedType));
             }
             BsonType bsonRepresentationType = null;
 
@@ -109,6 +115,36 @@ private static Codec<?> computeCodec(final RecordComponent component, final Code
             return codec;
         }
 
+        private static Class<?> resolveComponentType(final List<Type> typeParameters, final RecordComponent component) {
+            Type resolvedType = resolveType(component.getGenericType(), typeParameters, component.getDeclaringRecord());
+            return resolvedType instanceof Class<?> clazz ? clazz : component.getType();
+        }
+
+        private static List<Type> resolveActualTypeArguments(final List<Type> typeParameters, final Class<?> recordClass,
+                final ParameterizedType parameterizedType) {
+            return Arrays.stream(parameterizedType.getActualTypeArguments())
+                    .map(type -> resolveType(type, typeParameters, recordClass))
+                    .toList();
+        }
+
+        private static Type resolveType(final Type type, final List<Type> typeParameters, final Class<?> recordClass) {
+            return type instanceof TypeVariable<?> typeVariable
+                    ? typeParameters.get(getIndexOfTypeParameter(typeVariable.getName(), recordClass))
+                    : type;
+        }
+
+        // Get
+        private static int getIndexOfTypeParameter(final String typeParameterName, final Class<?> recordClass) {
+            var typeParameters = recordClass.getTypeParameters();
+            for (int i = 0; i < typeParameters.length; i++) {
+                if (typeParameters[i].getName().equals(typeParameterName)) {
+                    return i;
+                }
+            }
+            throw new CodecConfigurationException(String.format("Could not find type parameter on record %s with name %s",
+                    recordClass.getName(), typeParameterName));
+        }
+
         @SuppressWarnings("deprecation")
         private static String computeFieldName(final RecordComponent component) {
             if (component.isAnnotationPresent(BsonId.class)) {
@@ -218,16 +254,47 @@ private static <T extends Annotation> void validateAnnotationOnlyOnField(final R
 
     RecordCodec(final Class<T> clazz, final CodecRegistry codecRegistry) {
         this.clazz = notNull("class", clazz);
+        if (clazz.getTypeParameters().length > 0) {
+            requiresParameterization = true;
+            canonicalConstructor = null;
+            componentModels = null;
+            fieldNameToComponentModel = null;
+            componentModelForId = null;
+        } else {
+            requiresParameterization = false;
+            canonicalConstructor = notNull("canonicalConstructor", getCanonicalConstructor(clazz));
+            componentModels = getComponentModels(clazz, codecRegistry, List.of());
+            fieldNameToComponentModel = componentModels.stream()
+                    .collect(Collectors.toMap(ComponentModel::getFieldName, Function.identity()));
+            componentModelForId = getComponentModelForId(clazz, componentModels);
+        }
+    }
+
+    RecordCodec(final Class<T> clazz, final CodecRegistry codecRegistry, final List<Type> types) {
+        if (types.size() != clazz.getTypeParameters().length) {
+            throw new CodecConfigurationException("Unexpected number of type parameters for record class " + clazz);
+        }
+        this.clazz = notNull("class", clazz);
+        requiresParameterization = false;
         canonicalConstructor = notNull("canonicalConstructor", getCanonicalConstructor(clazz));
-        componentModels = getComponentModels(clazz, codecRegistry);
+        componentModels = getComponentModels(clazz, codecRegistry, types);
         fieldNameToComponentModel = componentModels.stream()
                 .collect(Collectors.toMap(ComponentModel::getFieldName, Function.identity()));
         componentModelForId = getComponentModelForId(clazz, componentModels);
     }
 
+    @Override
+    public Codec<?> parameterize(final CodecRegistry codecRegistry, final List<Type> types) {
+        return new RecordCodec<>(clazz, codecRegistry, types);
+    }
+
     @SuppressWarnings("unchecked")
     @Override
     public T decode(final BsonReader reader, final DecoderContext decoderContext) {
+        if (requiresParameterization) {
+            throw new CodecConfigurationException("Can not decode to a record with type parameters that has not been parameterized");
+        }
+
         reader.readStartDocument();
 
         Object[] constructorArguments = new Object[componentModels.size()];
@@ -254,6 +321,10 @@ public T decode(final BsonReader reader, final DecoderContext decoderContext) {
 
     @Override
     public void encode(final BsonWriter writer, final T record, final EncoderContext encoderContext) {
+        if (requiresParameterization) {
+            throw new CodecConfigurationException("Can not decode to a record with type parameters that has not been parameterized");
+        }
+
         writer.writeStartDocument();
         if (componentModelForId != null) {
             writeComponent(writer, record, componentModelForId);
@@ -287,11 +358,12 @@ private void writeComponent(final BsonWriter writer, final T record, final Compo
         }
     }
 
-    private static <T> List<ComponentModel> getComponentModels(final Class<T> clazz, final CodecRegistry codecRegistry) {
+    private static <T> List<ComponentModel> getComponentModels(final Class<T> clazz, final CodecRegistry codecRegistry,
+            final List<Type> typeParameters) {
         var recordComponents = clazz.getRecordComponents();
         var componentModels = new ArrayList<ComponentModel>(recordComponents.length);
         for (int i = 0; i < recordComponents.length; i++) {
-            componentModels.add(new ComponentModel(recordComponents[i], codecRegistry, i));
+            componentModels.add(new ComponentModel(typeParameters, recordComponents[i], codecRegistry, i));
         }
         return componentModels;
     }
diff --git a/bson-record-codec/src/test/unit/org/bson/codecs/record/RecordCodecTest.java b/bson-record-codec/src/test/unit/org/bson/codecs/record/RecordCodecTest.java
@@ -20,13 +20,15 @@
 import org.bson.BsonDocument;
 import org.bson.BsonDocumentReader;
 import org.bson.BsonDocumentWriter;
+import org.bson.BsonDouble;
 import org.bson.BsonInt32;
 import org.bson.BsonObjectId;
 import org.bson.BsonString;
 import org.bson.codecs.DecoderContext;
 import org.bson.codecs.EncoderContext;
 import org.bson.codecs.configuration.CodecConfigurationException;
 import org.bson.codecs.record.samples.TestRecordEmbedded;
+import org.bson.codecs.record.samples.TestRecordParameterized;
 import org.bson.codecs.record.samples.TestRecordWithDeprecatedAnnotations;
 import org.bson.codecs.record.samples.TestRecordWithIllegalBsonCreatorOnConstructor;
 import org.bson.codecs.record.samples.TestRecordWithIllegalBsonCreatorOnMethod;
@@ -44,6 +46,9 @@
 import org.bson.codecs.record.samples.TestRecordWithListOfRecords;
 import org.bson.codecs.record.samples.TestRecordWithMapOfListOfRecords;
 import org.bson.codecs.record.samples.TestRecordWithMapOfRecords;
+import org.bson.codecs.record.samples.TestRecordWithNestedParameterized;
+import org.bson.codecs.record.samples.TestRecordWithNestedParameterizedRecord;
+import org.bson.codecs.record.samples.TestRecordWithParameterizedRecord;
 import org.bson.codecs.record.samples.TestRecordWithPojoAnnotations;
 import org.bson.conversions.Bson;
 import org.bson.types.ObjectId;
@@ -227,6 +232,71 @@ public void testRecordWithNestedMapOfListRecords() {
         assertEquals(testRecord, decoded);
     }
 
+    @Test
+    public void testRecordWithNestedParameterizedRecord() {
+        var codec = new RecordCodec<>(TestRecordWithParameterizedRecord.class,
+                fromProviders(new RecordCodecProvider(), Bson.DEFAULT_CODEC_REGISTRY));
+        var identifier = new ObjectId();
+        var testRecord = new TestRecordWithParameterizedRecord(identifier,
+                new TestRecordParameterized<>(42.0, List.of(new TestRecordEmbedded("embedded"))));
+
+        var document = new BsonDocument();
+        var writer = new BsonDocumentWriter(document);
+
+        // when
+        codec.encode(writer, testRecord, EncoderContext.builder().build());
+
+        // then
+        assertEquals(
+                new BsonDocument("_id", new BsonObjectId(identifier))
+                        .append("parameterizedRecord",
+                                new BsonDocument("number", new BsonDouble(42.0))
+                                        .append("parameterizedList",
+                                                new BsonArray(List.of(new BsonDocument("name", new BsonString("embedded")))))),
+                document);
+        assertEquals("_id", document.getFirstKey());
+
+        // when
+        var decoded = codec.decode(new BsonDocumentReader(document), DecoderContext.builder().build());
+
+        // then
+        assertEquals(testRecord, decoded);
+    }
+
+    @Test
+    public void testRecordWithNestedParameterizedRecordWithDifferentlyOrderedTypeParameters() {
+        var codec = new RecordCodec<>(TestRecordWithNestedParameterizedRecord.class,
+                fromProviders(new RecordCodecProvider(), Bson.DEFAULT_CODEC_REGISTRY));
+        var identifier = new ObjectId();
+        var testRecord = new TestRecordWithNestedParameterizedRecord(identifier,
+                new TestRecordWithNestedParameterized<>(
+                        new TestRecordParameterized<>(42.0, List.of(new TestRecordEmbedded("p"))),
+                        "o"));
+
+        var document = new BsonDocument();
+        var writer = new BsonDocumentWriter(document);
+
+        // when
+        codec.encode(writer, testRecord, EncoderContext.builder().build());
+
+        // then
+        assertEquals(
+                new BsonDocument("_id", new BsonObjectId(identifier))
+                        .append("nestedParameterized",
+                                new BsonDocument("parameterizedRecord",
+                                        new BsonDocument("number", new BsonDouble(42.0))
+                                                .append("parameterizedList",
+                                                        new BsonArray(List.of(new BsonDocument("name", new BsonString("p"))))))
+                                        .append("other", new BsonString("o"))),
+                document);
+
+        // when
+        var decoded = codec.decode(new BsonDocumentReader(document), DecoderContext.builder().build());
+
+        // then
+        assertEquals(testRecord, decoded);
+    }
+
     @Test
     public void testRecordWithNulls() {
         var codec = new RecordCodec<>(TestRecordWithDeprecatedAnnotations.class, Bson.DEFAULT_CODEC_REGISTRY);
diff --git a/bson-record-codec/src/test/unit/org/bson/codecs/record/samples/TestRecordParameterized.java b/bson-record-codec/src/test/unit/org/bson/codecs/record/samples/TestRecordParameterized.java
@@ -0,0 +1,22 @@
+/*
+ * Copyright 2008-present MongoDB, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.bson.codecs.record.samples;
+
+import java.util.List;
+
+public record TestRecordParameterized<N extends Number, T>(N number, List<T> parameterizedList) {
+}
diff --git a/bson-record-codec/src/test/unit/org/bson/codecs/record/samples/TestRecordWithNestedParameterized.java b/bson-record-codec/src/test/unit/org/bson/codecs/record/samples/TestRecordWithNestedParameterized.java
@@ -0,0 +1,22 @@
+/*
+ * Copyright 2008-present MongoDB, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.bson.codecs.record.samples;
+
+public record TestRecordWithNestedParameterized<A, B, C extends Number>(
+        TestRecordParameterized<C, A> parameterizedRecord,
+        B other) {
+}
diff --git a/bson-record-codec/src/test/unit/org/bson/codecs/record/samples/TestRecordWithNestedParameterizedRecord.java b/bson-record-codec/src/test/unit/org/bson/codecs/record/samples/TestRecordWithNestedParameterizedRecord.java
@@ -0,0 +1,25 @@
+/*
+ * Copyright 2008-present MongoDB, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.bson.codecs.record.samples;
+
+import org.bson.codecs.pojo.annotations.BsonId;
+import org.bson.types.ObjectId;
+
+public record TestRecordWithNestedParameterizedRecord(
+        @BsonId ObjectId id,
+        TestRecordWithNestedParameterized<TestRecordEmbedded, String, Double> nestedParameterized) {
+}
diff --git a/bson-record-codec/src/test/unit/org/bson/codecs/record/samples/TestRecordWithParameterizedRecord.java b/bson-record-codec/src/test/unit/org/bson/codecs/record/samples/TestRecordWithParameterizedRecord.java
@@ -0,0 +1,24 @@
+/*
+ * Copyright 2008-present MongoDB, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.bson.codecs.record.samples;
+
+import org.bson.codecs.pojo.annotations.BsonId;
+import org.bson.types.ObjectId;
+
+public record TestRecordWithParameterizedRecord(@BsonId ObjectId id,
+                                                TestRecordParameterized<Double, TestRecordEmbedded> parameterizedRecord) {
+}
