Skip to content

Commit 259364b

Browse files
moschecbuescher
authored andcommitted
Cleanup / move some serverless code to ES (elastic#112360)
1 parent 7305cc3 commit 259364b

File tree

2 files changed

+76
-0
lines changed

2 files changed

+76
-0
lines changed

server/src/main/java/org/elasticsearch/common/collect/Iterators.java

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,11 @@
1010

1111
import org.elasticsearch.core.Nullable;
1212

13+
import java.util.ArrayList;
1314
import java.util.Arrays;
1415
import java.util.Collections;
1516
import java.util.Iterator;
17+
import java.util.List;
1618
import java.util.NoSuchElementException;
1719
import java.util.Objects;
1820
import java.util.function.BiFunction;
@@ -233,6 +235,57 @@ public T next() {
233235
}
234236
}
235237

238+
/**
239+
* Returns an iterator that yields at most the first {@code n} elements of the provided {@code input} iterator.
240+
*/
241+
public static <T> Iterator<T> limit(Iterator<? extends T> input, int n) {
242+
assert n >= 0 : "negative limit";
243+
if (n > 0 && input.hasNext()) {
244+
return new LimitIterator<>(input, n);
245+
} else {
246+
return Collections.emptyIterator();
247+
}
248+
}
249+
250+
private static final class LimitIterator<T> implements Iterator<T> {
251+
private final Iterator<? extends T> input;
252+
private final int limit;
253+
private int current;
254+
255+
LimitIterator(Iterator<? extends T> input, int limit) {
256+
this.input = input;
257+
this.limit = limit;
258+
}
259+
260+
@Override
261+
public boolean hasNext() {
262+
return current < limit && input.hasNext();
263+
}
264+
265+
@Override
266+
public T next() {
267+
if (current >= limit) {
268+
throw new NoSuchElementException();
269+
}
270+
++current;
271+
return input.next();
272+
}
273+
}
274+
275+
/**
276+
* Returns a list containing the elements of the provided {@code iterator}.
277+
*/
278+
public static <T> List<T> toList(Iterator<T> iterator) {
279+
if (iterator.hasNext()) {
280+
var list = new ArrayList<T>();
281+
while (iterator.hasNext()) {
282+
list.add(iterator.next());
283+
}
284+
return Collections.unmodifiableList(list);
285+
}
286+
return Collections.emptyList();
287+
}
288+
236289
public static <T, U> Iterator<U> flatMap(Iterator<? extends T> input, Function<T, Iterator<? extends U>> fn) {
237290
while (input.hasNext()) {
238291
final var value = fn.apply(input.next());

server/src/test/java/org/elasticsearch/common/collect/IteratorsTests.java

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,10 @@
2828
import java.util.function.ToIntFunction;
2929
import java.util.stream.IntStream;
3030

31+
import static org.hamcrest.Matchers.contains;
32+
import static org.hamcrest.Matchers.empty;
33+
import static org.hamcrest.Matchers.is;
34+
3135
public class IteratorsTests extends ESTestCase {
3236
public void testConcatentation() {
3337
List<Integer> threeTwoOne = Arrays.asList(3, 2, 1);
@@ -242,6 +246,25 @@ public void testFilter() {
242246
}
243247
}
244248

249+
public void testLimit() {
250+
var result = Iterators.limit(Collections.emptyIterator(), 10);
251+
assertThat(result.hasNext(), is(false));
252+
assertThat(Iterators.toList(result), is(empty()));
253+
254+
var values = List.of(1, 2, 3);
255+
result = Iterators.limit(values.iterator(), 10);
256+
assertThat(result.hasNext(), is(true));
257+
assertThat(Iterators.toList(result), contains(1, 2, 3));
258+
259+
result = Iterators.limit(values.iterator(), 2);
260+
assertThat(result.hasNext(), is(true));
261+
assertThat(Iterators.toList(result), contains(1, 2));
262+
263+
result = Iterators.limit(values.iterator(), 0);
264+
assertThat(result.hasNext(), is(false));
265+
assertThat(Iterators.toList(result), is(empty()));
266+
}
267+
245268
public void testFailFast() {
246269
final var array = randomIntegerArray();
247270
assertEmptyIterator(Iterators.failFast(Iterators.forArray(array), () -> true));

0 commit comments

Comments
 (0)