Refactor entity association and caching logic (#1281)

Author: nakamura-to

doma-core/src/main/java/org/seasar/doma/jdbc/criteria/command/CacheKey.java renamed to doma-core/src/main/java/org/seasar/doma/jdbc/EntityId.java

@@ -13,13 +13,15 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.seasar.doma.jdbc.criteria.command;
+package org.seasar.doma.jdbc;
 
 import java.util.List;
+import java.util.Objects;
 import org.seasar.doma.jdbc.entity.EntityType;
 
-public record CacheKey(EntityType<?> entityType, List<?> items) {
-  static CacheKey of(EntityKey key) {
-    return new CacheKey(key.getEntityMetamodel().asType(), key.getItems());
+public record EntityId(EntityType<?> entityType, List<?> items) {
+  public EntityId {
+    Objects.requireNonNull(entityType);
+    Objects.requireNonNull(items);
   }
 }

doma-core/src/main/java/org/seasar/doma/jdbc/aggregate/AggregateCommand.java

@@ -25,6 +25,7 @@
 import java.util.stream.Stream;
 import org.seasar.doma.internal.util.Combinations;
 import org.seasar.doma.internal.util.Pair;
+import org.seasar.doma.jdbc.EntityId;
 import org.seasar.doma.jdbc.command.Command;
 import org.seasar.doma.jdbc.command.SelectCommand;
 import org.seasar.doma.jdbc.entity.EntityType;
@@ -40,52 +41,48 @@
  */
 public class AggregateCommand<RESULT, ENTITY> implements Command<RESULT> {
   private final SelectQuery query;
-  private final EntityType<ENTITY> entityType;
+  private final EntityType<ENTITY> rootEntityType;
   private final StreamReducer<RESULT, ENTITY> streamReducer;
   private final AggregateStrategyType aggregateStrategyType;
 
   public AggregateCommand(
       SelectQuery query,
-      EntityType<ENTITY> entityType,
+      EntityType<ENTITY> rootEntityType,
       StreamReducer<RESULT, ENTITY> streamReducer,
       AggregateStrategyType aggregateStrategyType) {
     this.query = Objects.requireNonNull(query);
-    this.entityType = Objects.requireNonNull(entityType);
+    this.rootEntityType = Objects.requireNonNull(rootEntityType);
     this.streamReducer = Objects.requireNonNull(streamReducer);
     this.aggregateStrategyType = Objects.requireNonNull(aggregateStrategyType);
   }
 
   @Override
   public RESULT execute() {
-    Set<AggregateEntityCacheKey> rootEntityKeys = new LinkedHashSet<>();
-    Map<AggregateEntityCacheKey, Object> entityCache = new HashMap<>();
-    SelectCommand<List<AggregateEntityPool>> command =
+    Set<EntityId> rootEntityIds = new LinkedHashSet<>();
+    Map<EntityId, Object> entityCache = new HashMap<>();
+    Combinations<AssociationEntityKey> combinations = new Combinations<>();
+
+    SelectCommand<List<AssociationEntityPool>> command =
         new SelectCommand<>(
             query,
-            new AggregateEntityPoolIterationHandler(
-                entityType,
+            new AssociationEntityPoolIterationHandler(
+                rootEntityType,
                 aggregateStrategyType,
                 query.isResultMappingEnsured(),
-                rootEntityKeys,
+                rootEntityIds,
                 entityCache));
-    List<AggregateEntityPool> entityPools = command.execute();
-
-    Combinations<AggregateEntityKey> combinations = new Combinations<>();
-    for (AggregateEntityPool entityPool : entityPools) {
-      Map<AggregatePathKey, AggregateEntityPoolEntry> associationCandidate = new HashMap<>();
-      for (AggregateEntityPoolEntry entry : entityPool) {
-        associationCandidate.put(entry.entityKey().pathKey(), entry);
-      }
-      associate(entityCache, combinations, associationCandidate);
+    List<AssociationEntityPool> entityPools = command.execute();
+    for (AssociationEntityPool entityPool : entityPools) {
+      associate(entityCache, combinations, entityPool);
     }
 
     @SuppressWarnings("unchecked")
     Stream<ENTITY> stream =
         (Stream<ENTITY>)
-            rootEntityKeys.stream()
+            rootEntityIds.stream()
                 .map(entityCache::get)
                 .filter(Objects::nonNull)
-                .filter(entityType.getEntityClass()::isInstance);
+                .filter(rootEntityType.getEntityClass()::isInstance);
 
     return streamReducer.reduce(stream);
   }
@@ -95,19 +92,19 @@ public RESULT execute() {
    * of linkage rules and updates the cache with newly associated entities.
    */
   private void associate(
-      Map<AggregateEntityCacheKey, Object> entityCache,
-      Combinations<AggregateEntityKey> combinations,
-      Map<AggregatePathKey, AggregateEntityPoolEntry> associationCandidate) {
+      Map<EntityId, Object> entityCache,
+      Combinations<AssociationEntityKey> combinations,
+      AssociationEntityPool entityPool) {
 
     for (AssociationLinkerType<?, ?> linkerType :
         aggregateStrategyType.getAssociationLinkerTypes()) {
-      AggregateEntityPoolEntry source = associationCandidate.get(linkerType.getSourcePathKey());
-      AggregateEntityPoolEntry target = associationCandidate.get(linkerType.getTargetPathKey());
+      AssociationEntityPoolEntry source = entityPool.get(linkerType.getSourcePathKey());
+      AssociationEntityPoolEntry target = entityPool.get(linkerType.getTargetPathKey());
       if (source == null || target == null) {
         continue;
       }
 
-      Pair<AggregateEntityKey, AggregateEntityKey> keyPair =
+      Pair<AssociationEntityKey, AssociationEntityKey> keyPair =
           new Pair<>(source.entityKey(), target.entityKey());
       if (combinations.contains(keyPair)) {
         continue;
@@ -117,12 +114,10 @@ private void associate(
       @SuppressWarnings("unchecked")
       BiFunction<Object, Object, Object> linker =
           (BiFunction<Object, Object, Object>) linkerType.getLinker();
-      Object entity = linker.apply(source.entity(), target.entity());
-      if (entity != null) {
-        AggregateEntityCacheKey cacheKey = AggregateEntityCacheKey.of(source.entityKey());
-        entityCache.replace(cacheKey, entity);
-        associationCandidate.replace(
-            source.pathKey(), new AggregateEntityPoolEntry(source.entityKey(), entity));
+      Object newEntity = linker.apply(source.entity(), target.entity());
+      if (newEntity != null && newEntity != source.entity()) {
+        entityCache.replace(source.entityId(), newEntity);
+        entityPool.replace(new AssociationEntityPoolEntry(source.entityKey(), newEntity));
       }
     }
   }

doma-core/src/main/java/org/seasar/doma/jdbc/aggregate/AggregateEntityKey.java renamed to doma-core/src/main/java/org/seasar/doma/jdbc/aggregate/AssociationEntityKey.java

@@ -17,19 +17,15 @@
 
 import java.util.List;
 import java.util.Objects;
-import org.seasar.doma.jdbc.entity.EntityType;
+import org.seasar.doma.jdbc.EntityId;
 
-public record AggregateEntityKey(AggregatePathKey pathKey, List<?> items) {
-  public AggregateEntityKey {
+public record AssociationEntityKey(AssociationPathKey pathKey, List<?> items) {
+  public AssociationEntityKey {
     Objects.requireNonNull(pathKey);
     Objects.requireNonNull(items);
   }
 
-  public String propertyPath() {
-    return pathKey.propertyPath();
-  }
-
-  public EntityType<?> entityType() {
-    return pathKey.entityType();
+  public EntityId entityId() {
+    return new EntityId(pathKey.entityType(), items);
   }
 }

doma-core/src/main/java/org/seasar/doma/jdbc/aggregate/AggregateEntityPool.java renamed to doma-core/src/main/java/org/seasar/doma/jdbc/aggregate/AssociationEntityPool.java

@@ -15,22 +15,26 @@
  */
 package org.seasar.doma.jdbc.aggregate;
 
-import java.util.ArrayList;
-import java.util.Iterator;
-import java.util.List;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Objects;
 
-public class AggregateEntityPool implements Iterable<AggregateEntityPoolEntry> {
-  private final List<AggregateEntityPoolEntry> entities = new ArrayList<>();
+public class AssociationEntityPool {
+  private final Map<AssociationPathKey, AssociationEntityPoolEntry> entries = new HashMap<>();
 
-  public void add(AggregateEntityPoolEntry entry) {
+  public void add(AssociationEntityPoolEntry entry) {
     Objects.requireNonNull(entry);
-    entities.add(entry);
+    entries.put(entry.pathKey(), entry);
   }
 
-  @SuppressWarnings("NullableProblems")
-  @Override
-  public Iterator<AggregateEntityPoolEntry> iterator() {
-    return entities.iterator();
+  public AssociationEntityPoolEntry get(AssociationPathKey pathKey) {
+    Objects.requireNonNull(pathKey);
+    return entries.get(pathKey);
+  }
+
+  public void replace(AssociationEntityPoolEntry entry) {
+    Objects.requireNonNull(entry);
+    Objects.requireNonNull(entry);
+    entries.replace(entry.pathKey(), entry);
   }
 }

doma-core/src/main/java/org/seasar/doma/jdbc/aggregate/AggregateEntityPoolEntry.java renamed to doma-core/src/main/java/org/seasar/doma/jdbc/aggregate/AssociationEntityPoolEntry.java

@@ -16,14 +16,19 @@
 package org.seasar.doma.jdbc.aggregate;
 
 import java.util.Objects;
+import org.seasar.doma.jdbc.EntityId;
 
-public record AggregateEntityPoolEntry(AggregateEntityKey entityKey, Object entity) {
-  public AggregateEntityPoolEntry {
+public record AssociationEntityPoolEntry(AssociationEntityKey entityKey, Object entity) {
+  public AssociationEntityPoolEntry {
     Objects.requireNonNull(entityKey);
     Objects.requireNonNull(entity);
   }
 
-  public AggregatePathKey pathKey() {
+  public AssociationPathKey pathKey() {
     return entityKey.pathKey();
   }
+
+  public EntityId entityId() {
+    return entityKey.entityId();
+  }
 }

doma-core/src/main/java/org/seasar/doma/jdbc/aggregate/AggregateEntityPoolIterationHandler.java renamed to doma-core/src/main/java/org/seasar/doma/jdbc/aggregate/AssociationEntityPoolIterationHandler.java

@@ -21,45 +21,46 @@
 import java.util.Set;
 import org.seasar.doma.internal.jdbc.command.AbstractIterationHandler;
 import org.seasar.doma.internal.jdbc.command.ResultListCallback;
+import org.seasar.doma.jdbc.EntityId;
 import org.seasar.doma.jdbc.ObjectProvider;
 import org.seasar.doma.jdbc.entity.EntityType;
 import org.seasar.doma.jdbc.query.SelectQuery;
 
 /**
- * Handles the iteration of {@link AggregateEntityPool} results from a query execution and manages
+ * Handles the iteration of {@link AssociationEntityPool} results from a query execution and manages
  * the creation and configuration of object providers used to process the results.
  */
-public class AggregateEntityPoolIterationHandler
-    extends AbstractIterationHandler<AggregateEntityPool, List<AggregateEntityPool>> {
+public class AssociationEntityPoolIterationHandler
+    extends AbstractIterationHandler<AssociationEntityPool, List<AssociationEntityPool>> {
 
-  private final EntityType<?> entityType;
+  private final EntityType<?> rootEntityType;
   private final AggregateStrategyType aggregateStrategyType;
   private final boolean resultMappingEnsured;
-  private final Set<AggregateEntityCacheKey> rootEntityKeys;
-  private final Map<AggregateEntityCacheKey, Object> entityCache;
+  private final Set<EntityId> rootEntityIds;
+  private final Map<EntityId, Object> entityCache;
 
-  public AggregateEntityPoolIterationHandler(
-      EntityType<?> entityType,
+  public AssociationEntityPoolIterationHandler(
+      EntityType<?> rootEntityType,
       AggregateStrategyType aggregateStrategyType,
       boolean resultMappingEnsured,
-      Set<AggregateEntityCacheKey> rootEntityKeys,
-      Map<AggregateEntityCacheKey, Object> entityCache) {
+      Set<EntityId> rootEntityIds,
+      Map<EntityId, Object> entityCache) {
     super(new ResultListCallback<>());
-    this.entityType = Objects.requireNonNull(entityType);
+    this.rootEntityType = Objects.requireNonNull(rootEntityType);
     this.aggregateStrategyType = Objects.requireNonNull(aggregateStrategyType);
     this.resultMappingEnsured = resultMappingEnsured;
-    this.rootEntityKeys = Objects.requireNonNull(rootEntityKeys);
+    this.rootEntityIds = Objects.requireNonNull(rootEntityIds);
     this.entityCache = Objects.requireNonNull(entityCache);
   }
 
   @Override
-  protected ObjectProvider<AggregateEntityPool> createObjectProvider(SelectQuery query) {
-    return new AggregateEntityPoolProvider(
-        entityType,
+  protected ObjectProvider<AssociationEntityPool> createObjectProvider(SelectQuery query) {
+    return new AssociationEntityPoolProvider(
+        rootEntityType,
         aggregateStrategyType,
         query,
         resultMappingEnsured,
-        rootEntityKeys,
+        rootEntityIds,
         entityCache);
   }
 }