Per-codec CodecPlan; skip the reader mark for codecs that can't fail part-way (plan items 3.0 + 3.1)

core/src/main/java/dev/morphia/mapping/codec/pojo/CodecPlan.java

@@ -0,0 +1,136 @@
+package dev.morphia.mapping.codec.pojo;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+
+import com.mongodb.lang.Nullable;
+
+import dev.morphia.annotations.internal.MorphiaInternal;
+import dev.morphia.mapping.codec.EnumCodec;
+import dev.morphia.mapping.codec.MorphiaDateCodec;
+import dev.morphia.mapping.codec.MorphiaLocalDateTimeCodec;
+
+import org.bson.BsonType;
+import org.bson.codecs.BigDecimalCodec;
+import org.bson.codecs.BooleanCodec;
+import org.bson.codecs.Codec;
+import org.bson.codecs.DateCodec;
+import org.bson.codecs.Decimal128Codec;
+import org.bson.codecs.DoubleCodec;
+import org.bson.codecs.IntegerCodec;
+import org.bson.codecs.LongCodec;
+import org.bson.codecs.ObjectIdCodec;
+import org.bson.codecs.StringCodec;
+
+import static java.util.Map.entry;
+
+/**
+ * The per-codec state a {@link MorphiaCodec} reads while encoding and decoding, built once when the codec is created.
+ * Properties are addressed by their {@link PropertyModel#slot() slot}, their position in their entity model, so a
+ * lookup is an array index.
+ * <p>
+ * It holds what belongs to one codec rather than to the shared mapping model: the codecs resolved for this datastore
+ * and the order properties are encoded in. Later work adds more per-slot data here.
+ * </p>
+ *
+ * @hidden
+ * @morphia.internal
+ */
+@MorphiaInternal
+public final class CodecPlan {
+    private final EntityModel model;
+    private final PropertyModel[] properties;
+    private final Codec<? super Object>[] codecs;
+    private final PropertyModel[] encodeOrder;
+    // per slot, a bitmask of the BSON types (by ordinal) the slot's codec reads without being able to fail part-way
+    private final int[] directBsonTypes;
+
+    /**
+     * Codecs that read the listed BSON types directly: they either read the value and return it, or throw before
+     * reading anything. A property whose codec reads the current type directly can be decoded without a reader mark,
+     * since there's never a partly consumed value to rewind. Keyed on the exact codec class, so subclasses and user
+     * codecs keep the mark. Types a codec reads and then may reject, such as an {@code IntegerCodec} reading an INT64
+     * that doesn't fit, aren't listed.
+     */
+    static final Map<Class<?>, Integer> DIRECT_BSON_TYPES = Map.ofEntries(
+            entry(StringCodec.class, mask(BsonType.STRING, BsonType.SYMBOL)),
+            entry(ObjectIdCodec.class, mask(BsonType.OBJECT_ID)),
+            entry(IntegerCodec.class, mask(BsonType.INT32)),
+            entry(LongCodec.class, mask(BsonType.INT32, BsonType.INT64)),
+            entry(DoubleCodec.class, mask(BsonType.DOUBLE, BsonType.INT32)),
+            entry(BooleanCodec.class, mask(BsonType.BOOLEAN)),
+            entry(Decimal128Codec.class, mask(BsonType.DECIMAL128)),
+            entry(BigDecimalCodec.class, mask(BsonType.DECIMAL128)),
+            entry(DateCodec.class, mask(BsonType.DATE_TIME)),
+            entry(EnumCodec.class, mask(BsonType.STRING)),
+            entry(MorphiaLocalDateTimeCodec.class, mask(BsonType.DATE_TIME, BsonType.INT64)),
+            entry(MorphiaDateCodec.class, mask(BsonType.DATE_TIME)));
+
+    @SuppressWarnings("unchecked")
+    CodecPlan(EntityModel model) {
+        this.model = model;
+        List<PropertyModel> list = model.getProperties();
+        properties = list.toArray(new PropertyModel[0]);
+        codecs = new Codec[properties.length];
+        directBsonTypes = new int[properties.length];
+        for (int slot = 0; slot < properties.length; slot++) {
+            PropertyModel property = properties[slot];
+            if (property.slot() != slot) {
+                throw new IllegalStateException("%s is in slot %d of %s but records slot %d".formatted(property.getName(), slot,
+                        model.getType().getName(), property.slot()));
+            }
+            codecs[slot] = (Codec<? super Object>) property.getCodec();
+            directBsonTypes[slot] = codecs[slot] != null ? DIRECT_BSON_TYPES.getOrDefault(codecs[slot].getClass(), 0) : 0;
+        }
+        PropertyModel idProperty = model.getIdProperty();
+        encodeOrder = Arrays.stream(properties)
+                .filter(property -> property != idProperty)
+                .toArray(PropertyModel[]::new);
+    }
+
+    /**
+     * @return the entity model this plan is for
+     */
+    public EntityModel model() {
+        return model;
+    }
+
+    /**
+     * @param property a property of this plan's model
+     * @return the codec for the property, or null if none was resolved
+     */
+    @Nullable
+    public Codec<? super Object> codec(PropertyModel property) {
+        int slot = property.slot();
+        if (slot < 0 || slot >= properties.length || properties[slot] != property) {
+            throw new IllegalArgumentException("%s isn't a property of %s".formatted(property, model.getType().getName()));
+        }
+        return codecs[slot];
+    }
+
+    /**
+     * @param property a property of this plan's model
+     * @param type     the BSON type about to be read for it
+     * @return true if the property's codec reads that type without being able to fail part-way, so decoding it needs no
+     *         reader mark to fall back to
+     */
+    public boolean readsDirectly(PropertyModel property, BsonType type) {
+        return (directBsonTypes[property.slot()] & (1 << type.ordinal())) != 0;
+    }
+
+    private static int mask(BsonType... types) {
+        int mask = 0;
+        for (BsonType type : types) {
+            mask |= 1 << type.ordinal();
+        }
+        return mask;
+    }
+
+    /**
+     * @return the properties to encode after the id, in order. Callers must not modify the array.
+     */
+    PropertyModel[] encodeOrder() {
+        return encodeOrder;
+    }
+}

core/src/main/java/dev/morphia/mapping/codec/pojo/EntityDecoder.java

@@ -64,7 +64,7 @@ public T decode(BsonReader reader, DecoderContext decoderContext) {
 
             boolean decoded = false;
             try {
-                decodeProperties(reader, decoderContext, instanceCreator, classModel);
+                decodeProperties(reader, decoderContext, instanceCreator, morphiaCodec.getPlan());
                 decoded = true;
             } finally {
                 if (prereadId != null) {
@@ -99,35 +99,79 @@ public T decode(BsonReader reader, DecoderContext decoderContext) {
 
     protected void decodeModel(BsonReader reader, DecoderContext decoderContext,
             MorphiaInstanceCreator instanceCreator, @Nullable PropertyModel model) {
+        decodeModel(reader, decoderContext, instanceCreator, morphiaCodec.getPlan(), model);
+    }
 
-        if (model != null) {
-            final BsonReaderMark mark = reader.getMark();
-            try {
-                if (reader.getCurrentBsonType() == BsonType.NULL) {
-                    reader.readNull();
-                } else {
-                    Object value = decoderContext.decodeWithChildContext(model.getCodec(), reader);
-                    instanceCreator.set(value, model);
-                }
-            } catch (BsonInvalidOperationException e) {
-                mark.reset();
-                final Object value = morphiaCodec.getRegistry().get(Object.class).decode(reader, decoderContext);
-                instanceCreator.set(morphiaCodec.getConversions().convert(value, model.getTypeData().getType()), model);
-            }
-        } else {
+    /**
+     * Decodes a property with the codec the given plan holds for it.
+     *
+     * @param reader          the reader
+     * @param decoderContext  the context
+     * @param instanceCreator the creator to give the value to
+     * @param plan            the plan of the codec for the model the property belongs to
+     * @param model           the property, or null to skip the value
+     * @hidden
+     * @morphia.internal
+     */
+    @MorphiaInternal
+    protected void decodeModel(BsonReader reader, DecoderContext decoderContext,
+            MorphiaInstanceCreator instanceCreator, CodecPlan plan, @Nullable PropertyModel model) {
+
+        if (model == null) {
             reader.skipValue();
+            return;
+        }
+        BsonType type = reader.getCurrentBsonType();
+        if (type == BsonType.NULL) {
+            reader.readNull();
+            return;
+        }
+        if (plan.readsDirectly(model, type)) {
+            // the codec can't fail part-way on this type, so there's nothing to rewind and no mark to take
+            instanceCreator.set(decoderContext.decodeWithChildContext(plan.codec(model), reader), model);
+            return;
+        }
+        final BsonReaderMark mark = reader.getMark();
+        try {
+            Object value = decoderContext.decodeWithChildContext(plan.codec(model), reader);
+            instanceCreator.set(value, model);
+        } catch (BsonInvalidOperationException e) {
+            mark.reset();
+            final Object value = morphiaCodec.getRegistry().get(Object.class).decode(reader, decoderContext);
+            instanceCreator.set(morphiaCodec.getConversions().convert(value, model.getTypeData().getType()), model);
         }
     }
 
     protected void decodeProperties(BsonReader reader, DecoderContext decoderContext,
             MorphiaInstanceCreator instanceCreator, EntityModel classModel) {
+        CodecPlan plan = morphiaCodec.getPlan();
+        if (plan.model() != classModel) {
+            plan = ((MorphiaCodec<?>) morphiaCodec.getRegistry().get(classModel.getType())).getPlan();
+        }
+        decodeProperties(reader, decoderContext, instanceCreator, plan);
+    }
+
+    /**
+     * Decodes a document's properties with the codecs a plan holds.
+     *
+     * @param reader          the reader
+     * @param decoderContext  the context
+     * @param instanceCreator the creator to give the values to
+     * @param plan            the plan of the codec for the model being decoded
+     * @hidden
+     * @morphia.internal
+     */
+    @MorphiaInternal
+    protected void decodeProperties(BsonReader reader, DecoderContext decoderContext,
+            MorphiaInstanceCreator instanceCreator, CodecPlan plan) {
+        EntityModel classModel = plan.model();
         reader.readStartDocument();
         while (reader.readBsonType() != BsonType.END_OF_DOCUMENT) {
             String name = reader.readName();
             if (classModel.useDiscriminator() && classModel.discriminatorKey().equals(name)) {
                 reader.readString();
             } else {
-                decodeModel(reader, decoderContext, instanceCreator, classModel.getProperty(name));
+                decodeModel(reader, decoderContext, instanceCreator, plan, classModel.getProperty(name));
             }
         }
         reader.readEndDocument();

core/src/main/java/dev/morphia/mapping/codec/pojo/EntityEncoder.java

@@ -8,7 +8,6 @@
 import dev.morphia.annotations.internal.MorphiaInternal;
 
 import org.bson.BsonWriter;
-import org.bson.codecs.Codec;
 import org.bson.codecs.EncoderContext;
 import org.bson.codecs.IdGenerator;
 import org.bson.codecs.ObjectIdGenerator;
@@ -52,10 +51,7 @@ public void encode(BsonWriter writer, T value, EncoderContext encoderContext) {
                     encodeDiscriminator(writer, model);
                 }
 
-                for (PropertyModel propertyModel : model.properties()) {
-                    if (propertyModel.equals(idModel)) {
-                        continue;
-                    }
+                for (PropertyModel propertyModel : morphiaCodec.getPlan().encodeOrder()) {
                     encodeProperty(writer, propertyModel, value, encoderContext);
                 }
             });
@@ -138,8 +134,7 @@ protected void writeValue(BsonWriter writer, EncoderContext encoderContext, Prop
         if (value == null) {
             writer.writeNull();
         } else {
-            Codec<? super Object> codec = (Codec<? super Object>) model.getCodec();
-            encoderContext.encodeWithChildContext(codec, writer, value);
+            encoderContext.encodeWithChildContext(morphiaCodec.getPlan().codec(model), writer, value);
         }
     }
 }

core/src/main/java/dev/morphia/mapping/codec/pojo/EntityModel.java

@@ -61,7 +61,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import static dev.morphia.mapping.PropertyDiscovery.FIELDS;
 import static java.lang.String.format;
 import static java.util.Arrays.stream;
 import static java.util.Collections.emptyList;
@@ -170,6 +170,7 @@
         added &= propertyModelsByMappedName.put(property.getMappedName(), property) == null;
 
         if (added) {
+            property.slot(propertyModelsByName.size() - 1);
             if (property.hasAnnotation(Id.class)) {
                 idProperty = property;
             } else if (property.hasAnnotation(Version.class)) {

core/src/main/java/dev/morphia/mapping/codec/pojo/LifecycleDecoder.java

@@ -37,6 +37,7 @@ public LifecycleDecoder(MorphiaCodec<T> codec) {
     public T decode(BsonReader reader, DecoderContext decoderContext) {
         Document document = getMorphiaCodec().getRegistry().get(Document.class).decode(reader, decoderContext);
         EntityModel model = getMorphiaCodec().getEntityModel();
+        CodecPlan plan = getMorphiaCodec().getPlan();
         if (model.useDiscriminator()) {
             String discriminator = document.getString(model.discriminatorKey());
             if (discriminator != null) {
@@ -45,6 +46,7 @@ public T decode(BsonReader reader, DecoderContext decoderContext) {
                 Codec<?> codec = getMorphiaCodec().getRegistry().get(discriminatorClass);
                 if (codec instanceof MorphiaCodec) {
                     model = ((MorphiaCodec<?>) codec).getEntityModel();
+                    plan = ((MorphiaCodec<?>) codec).getPlan();
                 } else {
                     throw new CodecConfigurationException(format("Non-entity class used as discriminator: '%s'.", discriminator));
                 }
@@ -70,7 +72,7 @@ public T decode(BsonReader reader, DecoderContext decoderContext) {
         try {
             model.callLifecycleMethods(PreLoad.class, entity, document, getMorphiaCodec().getDatastore());
             decodeProperties(new DocumentReader(document, getMorphiaCodec().getConversions()), decoderContext, instanceCreator,
-                    model);
+                    plan);
             model.callLifecycleMethods(PostLoad.class, entity, document, getMorphiaCodec().getDatastore());
             decoded = true;
         } finally {

core/src/main/java/dev/morphia/mapping/codec/pojo/MorphiaCodec.java

@@ -47,6 +47,7 @@ public class MorphiaCodec<T> implements CollectibleCodec<T> {
     private final PropertyCodecRegistry propertyCodecRegistry;
     private final DiscriminatorLookup discriminatorLookup;
     private final Conversions conversions;
+    private final CodecPlan plan;
     private EntityEncoder<T> encoder;
     private EntityDecoder<T> decoder;
     private MorphiaDatastore datastore;
@@ -74,6 +75,7 @@ public MorphiaCodec(MorphiaDatastore datastore, EntityModel model,
         this.propertyCodecRegistry = new PropertyCodecRegistryImpl(this, registry, propertyCodecProviders);
         idProperty = model.getIdProperty();
         specializePropertyCodecs();
+        plan = new CodecPlan(model);
     }
 
     @Override
@@ -173,6 +175,16 @@ public EntityModel getEntityModel() {
         return entityModel;
     }
 
+    /**
+     * @return the plan this codec encodes and decodes with
+     * @hidden
+     * @morphia.internal
+     */
+    @MorphiaInternal
+    public CodecPlan getPlan() {
+        return plan;
+    }
+
     /**
      * @return the codec registry
      */

core/src/main/java/dev/morphia/mapping/codec/pojo/PropertyModel.java

@@ -67,6 +67,8 @@ public class PropertyModel {
     private final List<String> loadNames = new ArrayList<>();
     private final EntityModel entityModel;
     private Codec<? super Object> codec;
+    // this property's position in its entity model, which addresses it in a CodecPlan; assigned when it's added
+    private int slot = -1;
     private Class<?> normalizedType;
     private Conversions conversions;
 
@@ -425,4 +427,15 @@ void codec(Codec<? super Object> codec) {
         this.codec = codec;
     }
 
+    /**
+     * @return this property's position in its entity model, or -1 if it hasn't been added to one
+     */
+    int slot() {
+        return slot;
+    }
+
+    void slot(int slot) {
+        this.slot = slot;
+    }
+
 }