perf: significantly reduce overhead of Neighborhoods

Author: triceo

core/src/main/java/ai/timefold/solver/core/config/solver/PreviewFeature.java

@@ -24,13 +24,7 @@ public enum PreviewFeature {
     /**
      * Unlike other preview features, Neighborhoods are an active research project.
      * It is intended to simplify the creation of custom moves, eventually replacing move selectors.
-     * The component is under heavy development, entirely undocumented, and many key features are yet to be delivered.
-     * Neither the API nor the feature set are complete, and any part can change or be removed at any time.
-     * 
-     * Neighborhoods will eventually stabilize and be promoted from a research project to a true preview feature.
-     * We only expose it now to be able to use it for experimentation and testing.
-     * As such, it is an exception to the rule;
-     * this preview feature is not finished, and it is not yet ready for feedback.
+     * The component is under development, and many key features are yet to be delivered.
      */
     NEIGHBORHOODS
 

core/src/main/java/ai/timefold/solver/core/impl/bavet/common/index/ComparisonIndexer.java

@@ -2,8 +2,10 @@
 
 import java.util.Collections;
 import java.util.Comparator;
-import java.util.List;
+import java.util.Iterator;
+import java.util.Map;
 import java.util.NavigableMap;
+import java.util.NoSuchElementException;
 import java.util.Objects;
 import java.util.TreeMap;
 import java.util.function.Consumer;
@@ -13,6 +15,7 @@
 import ai.timefold.solver.core.impl.util.ListEntry;
 
 import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
 
 @NullMarked
 final class ComparisonIndexer<T, Key_ extends Comparable<Key_>>
@@ -138,7 +141,8 @@ private int sizeManyIndexers(Object compositeKey) {
     public void forEach(Object compositeKey, Consumer<T> tupleConsumer) {
         switch (comparisonMap.size()) {
             case 0 -> {
-                /* Nothing to do. */ }
+                /* Nothing to do. */
+            }
             case 1 -> forEachSingleIndexer(compositeKey, tupleConsumer);
             default -> forEachManyIndexers(compositeKey, tupleConsumer);
         }
@@ -164,47 +168,115 @@ private void forEachManyIndexers(Object compositeKey, Consumer<T> tupleConsumer)
         }
     }
 
+    @Override
+    public Iterator<T> iterator(Object compositeKey) {
+        return switch (comparisonMap.size()) {
+            case 0 -> Collections.emptyIterator();
+            case 1 -> iteratorSingleIndexer(compositeKey);
+            default -> new DefaultIterator(compositeKey);
+        };
+    }
+
+    private Iterator<T> iteratorSingleIndexer(Object compositeKey) {
+        var indexKey = keyRetriever.apply(compositeKey);
+        var entry = comparisonMap.firstEntry();
+        if (boundaryReached(entry.getKey(), indexKey)) {
+            return Collections.emptyIterator();
+        }
+        // Boundary condition not yet reached; include the indexer in the range.
+        return entry.getValue().iterator(compositeKey);
+    }
+
     @Override
     public boolean isEmpty() {
         return comparisonMap.isEmpty();
     }
 
     @Override
-    public List<? extends ListEntry<T>> asList(Object compositeKey) {
+    @Nullable
+    public ListEntry<T> get(Object compositeKey, int index) {
         return switch (comparisonMap.size()) {
-            case 0 -> Collections.emptyList();
-            case 1 -> asListSingleIndexer(compositeKey);
-            default -> asListManyIndexers(compositeKey);
+            case 0 -> null;
+            case 1 -> getSingleIndexer(compositeKey, index);
+            default -> getManyIndexers(compositeKey, index);
         };
     }
 
-    private List<? extends ListEntry<T>> asListSingleIndexer(Object compositeKey) {
+    private @Nullable ListEntry<T> getSingleIndexer(Object compositeKey, int index) {
         var indexKey = keyRetriever.apply(compositeKey);
         var entry = comparisonMap.firstEntry();
-        return boundaryReached(entry.getKey(), indexKey) ? Collections.emptyList() : entry.getValue().asList(compositeKey);
+        return boundaryReached(entry.getKey(), indexKey) ? null : entry.getValue().get(compositeKey, index);
     }
 
-    @SuppressWarnings("unchecked")
-    private List<? extends ListEntry<T>> asListManyIndexers(Object compositeKey) {
-        // The index backend's asList() may take a while to build.
-        // At the same time, the elements in these lists will be accessed randomly.
-        // Therefore we build this abstraction to avoid building unnecessary lists that would never get accessed.
-        var result = new ComposingList<ListEntry<T>>();
+    private @Nullable ListEntry<T> getManyIndexers(Object compositeKey, int index) {
+        var seenCount = 0;
         var indexKey = keyRetriever.apply(compositeKey);
         for (var entry : comparisonMap.entrySet()) {
             if (boundaryReached(entry.getKey(), indexKey)) {
-                return result;
+                return null;
             } else { // Boundary condition not yet reached; include the indexer in the range.
                 var value = entry.getValue();
-                result.addSubList(() -> (List<ListEntry<T>>) value.asList(compositeKey), value.size(compositeKey));
+                var size = value.size(compositeKey);
+                if (seenCount + size > index) {
+                    return value.get(compositeKey, index - seenCount);
+                }
+                seenCount += size;
             }
         }
-        return result;
+        return null;
     }
 
     @Override
     public String toString() {
         return "size = " + comparisonMap.size();
     }
 
+    private final class DefaultIterator implements Iterator<T> {
+
+        private final Object compositeKey;
+        private final Key_ indexKey;
+        private final Iterator<Map.Entry<Key_, Indexer<T>>> indexerIterator = comparisonMap.entrySet().iterator();
+        private @Nullable Iterator<T> downstreamIterator = null;
+        private @Nullable T next = null;
+
+        public DefaultIterator(Object compositeKey) {
+            this.compositeKey = compositeKey;
+            this.indexKey = keyRetriever.apply(compositeKey);
+        }
+
+        @Override
+        public boolean hasNext() {
+            if (next != null) {
+                return true;
+            }
+            if (downstreamIterator != null && downstreamIterator.hasNext()) {
+                next = downstreamIterator.next();
+                return true;
+            }
+            while (indexerIterator.hasNext()) {
+                var entry = indexerIterator.next();
+                if (boundaryReached(entry.getKey(), indexKey)) {
+                    return false;
+                }
+                // Boundary condition not yet reached; include the indexer in the range.
+                downstreamIterator = entry.getValue().iterator(compositeKey);
+                if (downstreamIterator.hasNext()) {
+                    next = downstreamIterator.next();
+                    return true;
+                }
+            }
+            return false;
+        }
+
+        @Override
+        public T next() {
+            if (!hasNext()) {
+                throw new NoSuchElementException();
+            }
+            var result = next;
+            next = null;
+            return result;
+        }
+    }
+
 }

core/src/main/java/ai/timefold/solver/core/impl/bavet/common/index/EqualsIndexer.java

@@ -2,7 +2,7 @@
 
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.List;
+import java.util.Iterator;
 import java.util.Map;
 import java.util.Objects;
 import java.util.function.Consumer;
@@ -11,6 +11,7 @@
 import ai.timefold.solver.core.impl.util.ListEntry;
 
 import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
 
 @NullMarked
 final class EqualsIndexer<T, Key_> implements Indexer<T> {
@@ -108,18 +109,29 @@ public void forEach(Object compositeKey, Consumer<T> tupleConsumer) {
     }
 
     @Override
-    public boolean isEmpty() {
-        return downstreamIndexerMap.isEmpty();
+    public Iterator<T> iterator(Object compositeKey) {
+        Key_ indexKey = keyRetriever.apply(compositeKey);
+        Indexer<T> downstreamIndexer = downstreamIndexerMap.get(indexKey);
+        if (downstreamIndexer == null) {
+            return Collections.emptyIterator();
+        }
+        return downstreamIndexer.iterator(compositeKey);
     }
 
     @Override
-    public List<? extends ListEntry<T>> asList(Object compositeKey) {
+    @Nullable
+    public ListEntry<T> get(Object compositeKey, int index) {
         Key_ indexKey = keyRetriever.apply(compositeKey);
         Indexer<T> downstreamIndexer = downstreamIndexerMap.get(indexKey);
         if (downstreamIndexer == null) {
-            return Collections.emptyList();
+            return null;
         }
-        return downstreamIndexer.asList(compositeKey);
+        return downstreamIndexer.get(compositeKey, index);
+    }
+
+    @Override
+    public boolean isEmpty() {
+        return downstreamIndexerMap.isEmpty();
     }
 
     @Override

core/src/main/java/ai/timefold/solver/core/impl/bavet/common/index/Indexer.java

@@ -1,12 +1,13 @@
 package ai.timefold.solver.core.impl.bavet.common.index;
 
-import java.util.List;
+import java.util.Iterator;
 import java.util.function.Consumer;
 
 import ai.timefold.solver.core.impl.bavet.common.tuple.TupleState;
 import ai.timefold.solver.core.impl.util.ListEntry;
 
 import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
 
 /**
  * An indexer for entity or fact {@code X},
@@ -35,18 +36,11 @@ public sealed interface Indexer<T>
 
     void forEach(Object compositeKey, Consumer<T> tupleConsumer);
 
-    boolean isEmpty();
+    Iterator<T> iterator(Object compositeKey);
+
+    @Nullable
+    ListEntry<T> get(Object compositeKey, int index);
 
-    /**
-     * Returns all entries for the given composite key as a list.
-     * The index must not be modified while iterating over the returned list.
-     * If the index is modified, a new instance of this list must be retrieved;
-     * the previous instance is no longer valid and its behavior is undefined.
-     * 
-     * @param compositeKey the composite key
-     * @return all entries for a given composite key;
-     *         the caller must not modify the list
-     */
-    List<? extends ListEntry<T>> asList(Object compositeKey);
+    boolean isEmpty();
 
 }

core/src/main/java/ai/timefold/solver/core/impl/bavet/common/index/IndexerBackend.java

@@ -1,9 +1,5 @@
 package ai.timefold.solver.core.impl.bavet.common.index;
 
-import java.util.List;
-
-import ai.timefold.solver.core.impl.util.ListEntry;
-
 import org.jspecify.annotations.NullMarked;
 
 /**
@@ -18,10 +14,4 @@ public sealed interface IndexerBackend<T>
         extends Indexer<T>
         permits RandomAccessIndexerBackend, LinkedListIndexerBackend {
 
-    @Override
-    default List<? extends ListEntry<T>> asList(Object compositeKey) {
-        throw new UnsupportedOperationException("Indexer backend (%s) does not support random access."
-                .formatted(this.getClass().getSimpleName()));
-    }
-
 }

core/src/main/java/ai/timefold/solver/core/impl/bavet/common/index/LinkedListIndexerBackend.java

@@ -1,5 +1,6 @@
 package ai.timefold.solver.core.impl.bavet.common.index;
 
+import java.util.Iterator;
 import java.util.function.Consumer;
 
 import ai.timefold.solver.core.impl.util.ElementAwareLinkedList;
@@ -37,6 +38,16 @@ public void forEach(Object compositeKey, Consumer<T> tupleConsumer) {
         tupleList.forEach(tupleConsumer);
     }
 
+    @Override
+    public Iterator<T> iterator(Object compositeKey) {
+        return tupleList.iterator();
+    }
+
+    @Override
+    public ListEntry<T> get(Object compositeKey, int index) {
+        throw new UnsupportedOperationException(); // Random access uses a different backend.
+    }
+
     @Override
     public boolean isEmpty() {
         return tupleList.size() == 0;

core/src/main/java/ai/timefold/solver/core/impl/bavet/common/index/RandomAccessIndexerBackend.java

@@ -1,6 +1,6 @@
 package ai.timefold.solver.core.impl.bavet.common.index;
 
-import java.util.List;
+import java.util.Iterator;
 import java.util.function.Consumer;
 
 import ai.timefold.solver.core.impl.util.ElementAwareArrayList;
@@ -41,13 +41,18 @@ public void forEach(Object compositeKey, Consumer<T> tupleConsumer) {
     }
 
     @Override
-    public boolean isEmpty() {
-        return tupleList.isEmpty();
+    public Iterator<T> iterator(Object compositeKey) {
+        return tupleList.iterator();
     }
 
     @Override
-    public List<ElementAwareArrayList.Entry<T>> asList(Object compositeKey) {
-        return tupleList.asList();
+    public ListEntry<T> get(Object compositeKey, int index) {
+        return tupleList.get(index);
+    }
+
+    @Override
+    public boolean isEmpty() {
+        return tupleList.isEmpty();
     }
 
     @Override

core/src/main/java/ai/timefold/solver/core/impl/localsearch/DefaultLocalSearchPhaseFactory.java

@@ -72,7 +72,7 @@ private LocalSearchDecider<Solution_> buildDecider(HeuristicConfigPolicy<Solutio
             PhaseTermination<Solution_> phaseTermination) {
         var neighborhoodsEnabled = phaseConfigPolicy.isPreviewFeatureEnabled(PreviewFeature.NEIGHBORHOODS);
         var neighborhoodProviderClass = phaseConfig.<Solution_> getNeighborhoodProviderClass();
-        if (phaseConfigPolicy.isPreviewFeatureEnabled(PreviewFeature.NEIGHBORHOODS)) {
+        if (neighborhoodsEnabled) {
             if (neighborhoodProviderClass == null) {
                 // Neighborhoods are enabled, but no provider was specified: use the default one.
                 neighborhoodProviderClass = (Class) DefaultNeighborhoodProvider.class;