11package ai .timefold .solver .core .impl .bavet .common .index ;
22
3- import ai .timefold .solver .core .impl .util .ElementAwareList ;
4- import org .jspecify .annotations .NullMarked ;
5- import org .jspecify .annotations .Nullable ;
6-
73import java .util .ArrayList ;
84import java .util .BitSet ;
95import java .util .Collections ;
106import java .util .List ;
117import java .util .Objects ;
128import java .util .function .Consumer ;
9+ import java .util .function .Predicate ;
10+
11+ import ai .timefold .solver .core .impl .util .ElementAwareList ;
12+
13+ import org .jspecify .annotations .NullMarked ;
14+ import org .jspecify .annotations .Nullable ;
1315
1416/**
1517 * {@link ArrayList}-backed set which allows to {@link #remove(Object)} an element
@@ -72,16 +74,24 @@ public void add(T element) {
7274 var actualElementList = getElementList ();
7375 if (gapCount > 0 ) {
7476 var gapIndex = gaps .nextSetBit (0 );
75- actualElementList .set (gapIndex , element );
76- elementPositionTracker .setPosition (element , gapIndex );
77- gaps .clear (gapIndex );
78- gapCount --;
77+ putElementIntoGap (actualElementList , element , gapIndex );
7978 } else {
8079 actualElementList .add (element );
8180 elementPositionTracker .setPosition (element , actualElementList .size () - 1 );
8281 }
8382 }
8483
84+ private void putElementIntoGap (List <T > elementList , T element , int gap ) {
85+ setElementList (elementList , element , gap );
86+ gaps .clear (gap );
87+ gapCount --;
88+ }
89+
90+ private void setElementList (List <T > elementList , T element , int position ) {
91+ elementList .set (position , element );
92+ elementPositionTracker .setPosition (element , position );
93+ }
94+
8595 /**
8696 * Removes the first occurrence of the specified element from this collection, if it is present.
8797 * Will use identity comparison to check for presence;
@@ -109,6 +119,7 @@ private boolean innerRemove(T element) {
109119 if (insertionPosition == actualElementList .size () - 1 ) {
110120 // The element was the last one added; we can simply remove it.
111121 actualElementList .remove (insertionPosition );
122+ removeTailGap (actualElementList );
112123 } else {
113124 actualElementList .set (insertionPosition , null );
114125 gaps .set (insertionPosition );
@@ -125,18 +136,37 @@ public int size() {
125136 * Performs the given action for each element of the collection
126137 * until all elements have been processed.
127138 *
128- * @param tupleConsumer the action to be performed for each element
139+ * @param elementConsumer the action to be performed for each element
129140 */
130- public void forEach (Consumer <T > tupleConsumer ) {
131- if (elementList == null ) {
141+ public void forEach (Consumer <T > elementConsumer ) {
142+ if (isEmpty () ) {
132143 return ;
133144 }
145+ findFirst (element -> {
146+ elementConsumer .accept (element );
147+ return false ; // Iterate until the end.
148+ });
149+ }
150+
151+ public @ Nullable T findFirst (Predicate <T > elementPredicate ) {
152+ if (isEmpty ()) {
153+ return null ;
154+ }
134155 for (var i = 0 ; i < elementList .size (); i ++) {
135156 var element = elementList .get (i );
136- if (element != null ) {
137- tupleConsumer .accept (element );
157+ if (element == null ) {
158+ var nonGap = removeTailGap (elementList );
159+ if (i >= nonGap ) {
160+ return null ;
161+ }
162+ element = elementList .remove (nonGap );
163+ putElementIntoGap (elementList , element , i );
164+ }
165+ if (elementPredicate .test (element )) {
166+ return element ;
138167 }
139168 }
169+ return null ;
140170 }
141171
142172 public boolean isEmpty () {
@@ -153,9 +183,8 @@ public List<T> asList() {
153183 if (elementList == null ) {
154184 return Collections .emptyList ();
155185 }
156- var actualElementList = getElementList ();
157- defrag (actualElementList );
158- return actualElementList ;
186+ defrag (elementList );
187+ return elementList ;
159188 }
160189
161190 private void defrag (List <T > actualElementList ) {
@@ -164,26 +193,37 @@ private void defrag(List<T> actualElementList) {
164193 }
165194 var gap = gaps .nextSetBit (0 );
166195 while (gap >= 0 ) {
167- var lastNonGapIndex = findNonGapFromEnd (actualElementList );
196+ var lastNonGapIndex = removeTailGap (actualElementList );
168197 if (lastNonGapIndex < 0 || gap >= lastNonGapIndex ) {
169198 break ;
170199 }
171200 var lastElement = actualElementList .remove (lastNonGapIndex );
172- actualElementList .set (gap , lastElement );
173- elementPositionTracker .setPosition (lastElement , gap );
201+ setElementList (actualElementList , lastElement , gap );
174202 gap = gaps .nextSetBit (gap + 1 );
175203 }
204+ resetGaps ();
205+ }
206+
207+ private void resetGaps () {
176208 gaps .clear ();
177209 gapCount = 0 ;
178210 }
179211
180- private int findNonGapFromEnd (List <T > actualElementList ) {
212+ private int removeTailGap (List <T > actualElementList ) {
181213 var end = actualElementList .size () - 1 ;
182214 var lastNonGap = gaps .previousClearBit (end );
183- for (var i = end ; i > lastNonGap ; i --) {
184- actualElementList .remove (i );
215+ if (lastNonGap < 0 ) {
216+ actualElementList .clear ();
217+ resetGaps ();
218+ return -1 ;
219+ } else {
220+ for (var i = end ; i > lastNonGap ; i --) {
221+ actualElementList .remove (i );
222+ }
223+ gaps .clear (lastNonGap , end + 1 );
224+ gapCount = gaps .cardinality ();
225+ return lastNonGap ;
185226 }
186- return lastNonGap ;
187227 }
188228
189229}
0 commit comments