clean up discriminator mapping

fixes #2637
MorphiaOrg · Nov 6, 2023 · 574c1ba · 574c1ba
1 parent df312f0
commit 574c1ba
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 18 deletions.
diff --git a/core/src/main/java/dev/morphia/mapping/DiscriminatorLookup.java b/core/src/main/java/dev/morphia/mapping/DiscriminatorLookup.java
@@ -27,8 +27,6 @@
 import dev.morphia.mapping.codec.pojo.EntityModel;
 import dev.morphia.sofia.Sofia;
 
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
-
 import org.bson.codecs.configuration.CodecConfigurationException;
 
 import static java.lang.String.format;
@@ -48,7 +46,6 @@ public final class DiscriminatorLookup {
      * Creates a new lookup
      *
      */
-    @SuppressFBWarnings("EI_EXPOSE_REP2")
     public DiscriminatorLookup() {
         this.classLoader = Thread.currentThread().getContextClassLoader();
 
@@ -85,8 +82,6 @@ public Class<?> lookup(String discriminator) {
 
         if (clazz == null) {
             throw new CodecConfigurationException(format("A class could not be found for the discriminator: '%s'.", discriminator));
-        } else {
-            discriminatorClassMap.put(discriminator, clazz);
         }
         return clazz;
     }

diff --git a/core/src/main/java/dev/morphia/mapping/Mapper.java b/core/src/main/java/dev/morphia/mapping/Mapper.java
@@ -72,7 +72,7 @@ public class Mapper {
     /**
      * Set of classes that registered by this mapper
      */
-    private final Map<Class, EntityModel> mappedEntities = new ConcurrentHashMap<>();
+    private final Map<String, EntityModel> mappedEntities = new ConcurrentHashMap<>();
     private final ConcurrentHashMap<String, Set<EntityModel>> mappedEntitiesByCollection = new ConcurrentHashMap<>();
     private final List<EntityListener<?>> listeners = new ArrayList<>();
     private final MorphiaConfig config;
@@ -117,7 +117,7 @@ private EntityModel clone(@Nullable EntityModel original) {
         }
         EntityModel superClone = clone(original.superClass);
         if (superClone == null || superClone.getSubtype(original.getType()) == null) {
-            EntityModel copy = documentNewModel(new EntityModel(original));
+            EntityModel copy = register(new EntityModel(original), false);
 
             Set<EntityModel> subtypes = original.subtypes.stream().map(subtype -> {
                 EntityModel clonedSubtype = clone(subtype);
@@ -260,7 +260,7 @@ public EntityModel getEntityModel(Class type) {
         if (actual == null && MorphiaProxy.class.equals(type)) {
             throw new NotMappableException(type);
         }
-        EntityModel model = mappedEntities.get(actual);
+        EntityModel model = mappedEntities.get(actual.getName());
 
         if (model == null) {
             if (!isMappable(actual)) {
@@ -418,7 +418,7 @@ public <T> boolean isMappable(Class<T> type) {
      * @return true if the Class has been mapped
      */
     public boolean isMapped(Class c) {
-        return mappedEntities.containsKey(c);
+        return mappedEntities.containsKey(c.getName());
     }
 
     /**
@@ -570,7 +570,10 @@ public EntityModel register(EntityModel entityModel) {
 
     private EntityModel register(EntityModel entityModel, boolean validate) {
 
-        documentNewModel(entityModel);
+        discriminatorLookup.addModel(entityModel);
+        mappedEntities.put(entityModel.getType().getName(), entityModel);
+        mappedEntitiesByCollection.computeIfAbsent(entityModel.getCollectionName(), s -> new CopyOnWriteArraySet<>())
+                .add(entityModel);
         EntityModel superClass = entityModel.getSuperClass();
         if (superClass != null) {
             superClass.addSubtype(entityModel);
@@ -584,14 +587,6 @@ private EntityModel register(EntityModel entityModel, boolean validate) {
         return entityModel;
     }
 
-    private EntityModel documentNewModel(EntityModel entityModel) {
-        discriminatorLookup.addModel(entityModel);
-        mappedEntities.put(entityModel.getType(), entityModel);
-        mappedEntitiesByCollection.computeIfAbsent(entityModel.getCollectionName(), s -> new CopyOnWriteArraySet<>())
-                .add(entityModel);
-        return entityModel;
-    }
-
     private List<Class> getClasses(ClassLoader loader, String packageName)
             throws ClassNotFoundException {
         final Set<Class> classes = new HashSet<>();