Replace TreeMap with EntityIdMap to reduce comparisons count

repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/EntityIdMap.java

@@ -0,0 +1,70 @@
+package tech.ydb.yoj.repository.test.inmemory;
+
+import tech.ydb.yoj.repository.db.Entity;
+
+import java.util.AbstractMap;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * TreeMap<T, E> - is very expensive due to very complicated comparison login
+ */
+class EntityIdMap<T extends Entity<T>, E> extends AbstractMap<Entity.Id<T>, E> {
+    private final Comparator<Entity.Id<T>> comparator;
+    private Map<Entity.Id<T>, E> entries = new LinkedHashMap<>(); // probably not the best here but very simple
+
+    public EntityIdMap(Comparator<Entity.Id<T>> comparator) {
+        this.comparator = comparator;
+    }
+
+    public EntityIdMap(EntityIdMap<T, E> entityLines) {
+        this(entityLines.comparator);
+        entries = new LinkedHashMap<>(entityLines);
+    }
+
+    @Override
+    public Set<Entry<Entity.Id<T>, E>> entrySet() {
+        return entries.entrySet();
+    }
+
+    @Override
+    public E get(Object key) {
+        return entries.get(key);
+    }
+
+    @Override
+    public E put(Entity.Id<T> key, E value) {
+        List<Entity.Id<T>> temp = new ArrayList<>(entries.keySet());
+        int index = Collections.binarySearch(temp, key, comparator);
+        Map<Entity.Id<T>, E> result = new LinkedHashMap<>();
+        E oldValue = null;
+        int i = 0;
+        for (Entry<Entity.Id<T>, E> entry : entries.entrySet()) {
+            if (index < 0 && -index - 1 == i) {
+                result.put(key, value);
+                result.put(entry.getKey(), entry.getValue());
+            } else if (index == i) {
+                oldValue = entry.getValue();
+                result.put(key, value);
+            } else {
+                result.put(entry.getKey(), entry.getValue());
+            }
+            i++;
+        }
+        if (-index - 1 == result.size()) {
+            result.put(key, value); // add last (item i.e first item or item with max value to existing list)
+        }
+        entries = result;
+        return oldValue;
+    }
+
+    @Override
+    public E remove(Object key) {
+        return entries.remove(key);
+    }
+}

repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryDataShard.java

@@ -16,16 +16,15 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.TreeMap;
 
 final class InMemoryDataShard<T extends Entity<T>> {
     private final Class<T> type;
     private final EntitySchema<T> schema;
-    private final TreeMap<Entity.Id<T>, InMemoryEntityLine> entityLines;
+    private final Map<Entity.Id<T>, InMemoryEntityLine> entityLines;
     private final Map<Long, Set<Entity.Id<T>>> uncommited = new HashMap<>();
 
     private InMemoryDataShard(
-            Class<T> type, EntitySchema<T> schema, TreeMap<Entity.Id<T>, InMemoryEntityLine> entityLines
+            Class<T> type, EntitySchema<T> schema, Map<Entity.Id<T>, InMemoryEntityLine> entityLines
     ) {
         this.type = type;
         this.schema = schema;
@@ -36,15 +35,13 @@ public InMemoryDataShard(Class<T> type) {
         this(type, EntitySchema.of(type), createEmptyLines(type));
     }
 
-    private static <T extends Entity<T>> TreeMap<Entity.Id<T>, InMemoryEntityLine> createEmptyLines(Class<T> type) {
-        return new TreeMap<>(EntityIdSchema.getIdComparator(type));
+    private static <T extends Entity<T>> Map<Entity.Id<T>, InMemoryEntityLine> createEmptyLines(Class<T> type) {
+        return new EntityIdMap<>(EntityIdSchema.getIdComparator(type));
     }
 
     public synchronized InMemoryDataShard<T> createSnapshot() {
-        TreeMap<Entity.Id<T>, InMemoryEntityLine> snapshotLines = createEmptyLines(type);
-        for (Map.Entry<Entity.Id<T>, InMemoryEntityLine> entry : entityLines.entrySet()) {
-            snapshotLines.put(entry.getKey(), entry.getValue().createSnapshot());
-        }
+        Map<Entity.Id<T>, InMemoryEntityLine> snapshotLines = new EntityIdMap<>((EntityIdMap<T, InMemoryEntityLine>)entityLines);
+        snapshotLines.replaceAll((k, v) -> v.createSnapshot());
         return new InMemoryDataShard<>(type, schema, snapshotLines);
     }