JAVA-2144: Optimize BsonDocumentCodec via use of BsonTypeCodecMap

Author: jyemin

bson/src/main/org/bson/codecs/BsonDocumentCodec.java

@@ -30,6 +30,7 @@
 import java.util.List;
 import java.util.Map;
 
+import static org.bson.codecs.BsonValueCodecProvider.getBsonTypeClassMap;
 import static org.bson.codecs.configuration.CodecRegistries.fromProviders;
 
 /**
@@ -42,12 +43,13 @@ public class BsonDocumentCodec implements CollectibleCodec<BsonDocument> {
     private static final CodecRegistry DEFAULT_REGISTRY = fromProviders(new BsonValueCodecProvider());
 
     private final CodecRegistry codecRegistry;
+    private final BsonTypeCodecMap bsonTypeCodecMap;
 
     /**
      * Creates a new instance with a default codec registry that uses the {@link BsonValueCodecProvider}.
      */
     public BsonDocumentCodec() {
-        codecRegistry = DEFAULT_REGISTRY;
+        this(DEFAULT_REGISTRY);
     }
 
     /**
@@ -60,6 +62,7 @@ public BsonDocumentCodec(final CodecRegistry codecRegistry) {
             throw new IllegalArgumentException("Codec registry can not be null");
         }
         this.codecRegistry = codecRegistry;
+        this.bsonTypeCodecMap = new BsonTypeCodecMap(getBsonTypeClassMap(), codecRegistry);
     }
 
     /**
@@ -95,7 +98,7 @@ public BsonDocument decode(final BsonReader reader, final DecoderContext decoder
      * @return the non-null value read from the reader
      */
     protected BsonValue readValue(final BsonReader reader, final DecoderContext decoderContext) {
-        return codecRegistry.get(BsonValueCodecProvider.getClassForBsonType(reader.getCurrentBsonType())).decode(reader, decoderContext);
+        return (BsonValue) bsonTypeCodecMap.get(reader.getCurrentBsonType()).decode(reader, decoderContext);
     }
 
     @Override

bson/src/main/org/bson/codecs/BsonValueCodecProvider.java

@@ -43,7 +43,6 @@
 import org.bson.codecs.configuration.CodecProvider;
 import org.bson.codecs.configuration.CodecRegistry;
 
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -53,7 +52,7 @@
  * @since 3.0
  */
 public class BsonValueCodecProvider implements CodecProvider {
-    private static final Map<BsonType, Class<? extends BsonValue>> DEFAULT_BSON_TYPE_CLASS_MAP;
+    private static final BsonTypeClassMap DEFAULT_BSON_TYPE_CLASS_MAP;
 
     private final Map<Class<?>, Codec<?>> codecs = new HashMap<Class<?>, Codec<?>>();
 
@@ -69,8 +68,19 @@ public BsonValueCodecProvider() {
      * @param bsonType the BsonType
      * @return the class associated with the given type
      */
+    @SuppressWarnings("unchecked")
     public static Class<? extends BsonValue> getClassForBsonType(final BsonType bsonType) {
-        return DEFAULT_BSON_TYPE_CLASS_MAP.get(bsonType);
+        return (Class<? extends BsonValue>) DEFAULT_BSON_TYPE_CLASS_MAP.get(bsonType);
+    }
+
+    /**
+     * Gets the BsonTypeClassMap used by this provider.
+     *
+     * @return the non-null BsonTypeClassMap
+     * @since 3.3
+     */
+    public static BsonTypeClassMap getBsonTypeClassMap() {
+        return DEFAULT_BSON_TYPE_CLASS_MAP;
     }
 
     @Override
@@ -132,7 +142,7 @@ private <T extends BsonValue> void addCodec(final Codec<T> codec) {
     }
 
     static {
-        Map<BsonType, Class<? extends BsonValue>> map = new HashMap<BsonType, Class<? extends BsonValue>>();
+        Map<BsonType, Class<?>> map = new HashMap<BsonType, Class<?>>();
 
         map.put(BsonType.NULL, BsonNull.class);
         map.put(BsonType.ARRAY, BsonArray.class);
@@ -155,6 +165,6 @@ private <T extends BsonValue> void addCodec(final Codec<T> codec) {
         map.put(BsonType.TIMESTAMP, BsonTimestamp.class);
         map.put(BsonType.UNDEFINED, BsonUndefined.class);
 
-        DEFAULT_BSON_TYPE_CLASS_MAP = Collections.unmodifiableMap(map);
+        DEFAULT_BSON_TYPE_CLASS_MAP = new BsonTypeClassMap(map);
     }
 }