1 /*
2 * Copyright (C) 2015 The Guava Authors
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 package com.google.common.collect;
18
19 import static com.google.common.base.Preconditions.checkArgument;
20 import static com.google.common.base.Preconditions.checkNotNull;
21
22 import com.google.common.annotations.GwtCompatible;
23 import com.google.j2objc.annotations.Weak;
24 import java.util.Comparator;
25 import java.util.Spliterator;
26 import java.util.function.Consumer;
27 import java.util.function.DoubleConsumer;
28 import java.util.function.Function;
29 import java.util.function.IntConsumer;
30 import java.util.function.IntFunction;
31 import java.util.function.LongConsumer;
32 import java.util.function.Predicate;
33 import java.util.stream.IntStream;
34 import org.checkerframework.checker.nullness.qual.Nullable;
35
36 /** Spliterator utilities for {@code common.collect} internals. */
37 @GwtCompatible
38 final class CollectSpliterators {
39 private CollectSpliterators() {}
40
41 static <T> Spliterator<T> indexed(int size, int extraCharacteristics, IntFunction<T> function) {
42 return indexed(size, extraCharacteristics, function, null);
43 }
44
45 static <T> Spliterator<T> indexed(
46 int size,
47 int extraCharacteristics,
48 IntFunction<T> function,
49 Comparator<? super T> comparator) {
50 if (comparator != null) {
51 checkArgument((extraCharacteristics & Spliterator.SORTED) != 0);
52 }
53 class WithCharacteristics implements Spliterator<T> {
54 private final Spliterator.OfInt delegate;
55
56 WithCharacteristics(Spliterator.OfInt delegate) {
57 this.delegate = delegate;
58 }
59
60 @Override
61 public boolean tryAdvance(Consumer<? super T> action) {
62 return delegate.tryAdvance((IntConsumer) i -> action.accept(function.apply(i)));
63 }
64
65 @Override
66 public void forEachRemaining(Consumer<? super T> action) {
67 delegate.forEachRemaining((IntConsumer) i -> action.accept(function.apply(i)));
68 }
69
70 @Override
71 public @Nullable Spliterator<T> trySplit() {
72 Spliterator.OfInt split = delegate.trySplit();
73 return (split == null) ? null : new WithCharacteristics(split);
74 }
75
76 @Override
77 public long estimateSize() {
78 return delegate.estimateSize();
79 }
80
81 @Override
82 public int characteristics() {
83 return Spliterator.ORDERED
84 | Spliterator.SIZED
85 | Spliterator.SUBSIZED
86 | extraCharacteristics;
87 }
88
89 @Override
90 public Comparator<? super T> getComparator() {
91 if (hasCharacteristics(Spliterator.SORTED)) {
92 return comparator;
93 } else {
94 throw new IllegalStateException();
95 }
96 }
97 }
98 return new WithCharacteristics(IntStream.range(0, size).spliterator());
99 }
100
101 /**
102 * Returns a {@code Spliterator} over the elements of {@code fromSpliterator} mapped by {@code
103 * function}.
104 */
105 static <InElementT, OutElementT> Spliterator<OutElementT> map(
106 Spliterator<InElementT> fromSpliterator,
107 Function<? super InElementT, ? extends OutElementT> function) {
108 checkNotNull(fromSpliterator);
109 checkNotNull(function);
110 return new Spliterator<OutElementT>() {
111
112 @Override
113 public boolean tryAdvance(Consumer<? super OutElementT> action) {
114 return fromSpliterator.tryAdvance(
115 fromElement -> action.accept(function.apply(fromElement)));
116 }
117
118 @Override
119 public void forEachRemaining(Consumer<? super OutElementT> action) {
120 fromSpliterator.forEachRemaining(fromElement -> action.accept(function.apply(fromElement)));
121 }
122
123 @Override
124 public Spliterator<OutElementT> trySplit() {
125 Spliterator<InElementT> fromSplit = fromSpliterator.trySplit();
126 return (fromSplit != null) ? map(fromSplit, function) : null;
127 }
128
129 @Override
130 public long estimateSize() {
131 return fromSpliterator.estimateSize();
132 }
133
134 @Override
135 public int characteristics() {
136 return fromSpliterator.characteristics()
137 & ~(Spliterator.DISTINCT | Spliterator.NONNULL | Spliterator.SORTED);
138 }
139 };
140 }
141
142 /** Returns a {@code Spliterator} filtered by the specified predicate. */
143 static <T> Spliterator<T> filter(Spliterator<T> fromSpliterator, Predicate<? super T> predicate) {
144 checkNotNull(fromSpliterator);
145 checkNotNull(predicate);
146 class Splitr implements Spliterator<T>, Consumer<T> {
147 T holder = null;
148
149 @Override
150 public void accept(T t) {
151 this.holder = t;
152 }
153
154 @Override
155 public boolean tryAdvance(Consumer<? super T> action) {
156 while (fromSpliterator.tryAdvance(this)) {
157 try {
158 if (predicate.test(holder)) {
159 action.accept(holder);
160 return true;
161 }
162 } finally {
163 holder = null;
164 }
165 }
166 return false;
167 }
168
169 @Override
170 public Spliterator<T> trySplit() {
171 Spliterator<T> fromSplit = fromSpliterator.trySplit();
172 return (fromSplit == null) ? null : filter(fromSplit, predicate);
173 }
174
175 @Override
176 public long estimateSize() {
177 return fromSpliterator.estimateSize() / 2;
178 }
179
180 @Override
181 public Comparator<? super T> getComparator() {
182 return fromSpliterator.getComparator();
183 }
184
185 @Override
186 public int characteristics() {
187 return fromSpliterator.characteristics()
188 & (Spliterator.DISTINCT
189 | Spliterator.NONNULL
190 | Spliterator.ORDERED
191 | Spliterator.SORTED);
192 }
193 }
194 return new Splitr();
195 }
196
197 /**
198 * Returns a {@code Spliterator} that iterates over the elements of the spliterators generated by
199 * applying {@code function} to the elements of {@code fromSpliterator}.
200 */
201 static <InElementT, OutElementT> Spliterator<OutElementT> flatMap(
202 Spliterator<InElementT> fromSpliterator,
203 Function<? super InElementT, Spliterator<OutElementT>> function,
204 int topCharacteristics,
205 long topSize) {
206 checkArgument(
207 (topCharacteristics & Spliterator.SUBSIZED) == 0,
208 "flatMap does not support SUBSIZED characteristic");
209 checkArgument(
210 (topCharacteristics & Spliterator.SORTED) == 0,
211 "flatMap does not support SORTED characteristic");
212 checkNotNull(fromSpliterator);
213 checkNotNull(function);
214 return new FlatMapSpliteratorOfObject<InElementT, OutElementT>(
215 null, fromSpliterator, function, topCharacteristics, topSize);
216 }
217
218 /**
219 * Returns a {@code Spliterator.OfInt} that iterates over the elements of the spliterators
220 * generated by applying {@code function} to the elements of {@code fromSpliterator}. (If {@code
221 * function} returns {@code null} for an input, it is replaced with an empty stream.)
222 */
223 static <InElementT> Spliterator.OfInt flatMapToInt(
224 Spliterator<InElementT> fromSpliterator,
225 Function<? super InElementT, Spliterator.OfInt> function,
226 int topCharacteristics,
227 long topSize) {
228 checkArgument(
229 (topCharacteristics & Spliterator.SUBSIZED) == 0,
230 "flatMap does not support SUBSIZED characteristic");
231 checkArgument(
232 (topCharacteristics & Spliterator.SORTED) == 0,
233 "flatMap does not support SORTED characteristic");
234 checkNotNull(fromSpliterator);
235 checkNotNull(function);
236 return new FlatMapSpliteratorOfInt<InElementT>(
237 null, fromSpliterator, function, topCharacteristics, topSize);
238 }
239
240 /**
241 * Returns a {@code Spliterator.OfLong} that iterates over the elements of the spliterators
242 * generated by applying {@code function} to the elements of {@code fromSpliterator}. (If {@code
243 * function} returns {@code null} for an input, it is replaced with an empty stream.)
244 */
245 static <InElementT> Spliterator.OfLong flatMapToLong(
246 Spliterator<InElementT> fromSpliterator,
247 Function<? super InElementT, Spliterator.OfLong> function,
248 int topCharacteristics,
249 long topSize) {
250 checkArgument(
251 (topCharacteristics & Spliterator.SUBSIZED) == 0,
252 "flatMap does not support SUBSIZED characteristic");
253 checkArgument(
254 (topCharacteristics & Spliterator.SORTED) == 0,
255 "flatMap does not support SORTED characteristic");
256 checkNotNull(fromSpliterator);
257 checkNotNull(function);
258 return new FlatMapSpliteratorOfLong<InElementT>(
259 null, fromSpliterator, function, topCharacteristics, topSize);
260 }
261
262 /**
263 * Returns a {@code Spliterator.OfDouble} that iterates over the elements of the spliterators
264 * generated by applying {@code function} to the elements of {@code fromSpliterator}. (If {@code
265 * function} returns {@code null} for an input, it is replaced with an empty stream.)
266 */
267 static <InElementT> Spliterator.OfDouble flatMapToDouble(
268 Spliterator<InElementT> fromSpliterator,
269 Function<? super InElementT, Spliterator.OfDouble> function,
270 int topCharacteristics,
271 long topSize) {
272 checkArgument(
273 (topCharacteristics & Spliterator.SUBSIZED) == 0,
274 "flatMap does not support SUBSIZED characteristic");
275 checkArgument(
276 (topCharacteristics & Spliterator.SORTED) == 0,
277 "flatMap does not support SORTED characteristic");
278 checkNotNull(fromSpliterator);
279 checkNotNull(function);
280 return new FlatMapSpliteratorOfDouble<InElementT>(
281 null, fromSpliterator, function, topCharacteristics, topSize);
282 }
283
284 /**
285 * Implements the {@link Stream#flatMap} operation on spliterators.
286 *
287 * @param <InElementT> the element type of the input spliterator
288 * @param <OutElementT> the element type of the output spliterators
289 * @param <OutSpliteratorT> the type of the output spliterators
290 */
291 abstract static class FlatMapSpliterator<
292 InElementT, OutElementT, OutSpliteratorT extends Spliterator<OutElementT>>
293 implements Spliterator<OutElementT> {
294 /** Factory for constructing {@link FlatMapSpliterator} instances. */
295 @FunctionalInterface
296 interface Factory<InElementT, OutSpliteratorT extends Spliterator<?>> {
297 OutSpliteratorT newFlatMapSpliterator(
298 OutSpliteratorT prefix,
299 Spliterator<InElementT> fromSplit,
300 Function<? super InElementT, OutSpliteratorT> function,
301 int splitCharacteristics,
302 long estSplitSize);
303 }
304
305 @Nullable @Weak OutSpliteratorT prefix;
306 final Spliterator<InElementT> from;
307 final Function<? super InElementT, OutSpliteratorT> function;
308 final Factory<InElementT, OutSpliteratorT> factory;
309 int characteristics;
310 long estimatedSize;
311
312 FlatMapSpliterator(
313 OutSpliteratorT prefix,
314 Spliterator<InElementT> from,
315 Function<? super InElementT, OutSpliteratorT> function,
316 Factory<InElementT, OutSpliteratorT> factory,
317 int characteristics,
318 long estimatedSize) {
319 this.prefix = prefix;
320 this.from = from;
321 this.function = function;
322 this.factory = factory;
323 this.characteristics = characteristics;
324 this.estimatedSize = estimatedSize;
325 }
326
327 /*
328 * The tryAdvance and forEachRemaining in FlatMapSpliteratorOfPrimitive are overloads of these
329 * methods, not overrides. They are annotated @Override because they implement methods from
330 * Spliterator.OfPrimitive (and override default implementations from Spliterator.OfPrimitive or
331 * a subtype like Spliterator.OfInt).
332 */
333
334 @Override
335 public final boolean tryAdvance(Consumer<? super OutElementT> action) {
336 while (true) {
337 if (prefix != null && prefix.tryAdvance(action)) {
338 if (estimatedSize != Long.MAX_VALUE) {
339 estimatedSize--;
340 }
341 return true;
342 } else {
343 prefix = null;
344 }
345 if (!from.tryAdvance(fromElement -> prefix = function.apply(fromElement))) {
346 return false;
347 }
348 }
349 }
350
351 @Override
352 public final void forEachRemaining(Consumer<? super OutElementT> action) {
353 if (prefix != null) {
354 prefix.forEachRemaining(action);
355 prefix = null;
356 }
357 from.forEachRemaining(
358 fromElement -> {
359 Spliterator<OutElementT> elements = function.apply(fromElement);
360 if (elements != null) {
361 elements.forEachRemaining(action);
362 }
363 });
364 estimatedSize = 0;
365 }
366
367 @Override
368 public final OutSpliteratorT trySplit() {
369 Spliterator<InElementT> fromSplit = from.trySplit();
370 if (fromSplit != null) {
371 int splitCharacteristics = characteristics & ~Spliterator.SIZED;
372 long estSplitSize = estimateSize();
373 if (estSplitSize < Long.MAX_VALUE) {
374 estSplitSize /= 2;
375 this.estimatedSize -= estSplitSize;
376 this.characteristics = splitCharacteristics;
377 }
378 OutSpliteratorT result =
379 factory.newFlatMapSpliterator(
380 this.prefix, fromSplit, function, splitCharacteristics, estSplitSize);
381 this.prefix = null;
382 return result;
383 } else if (prefix != null) {
384 OutSpliteratorT result = prefix;
385 this.prefix = null;
386 return result;
387 } else {
388 return null;
389 }
390 }
391
392 @Override
393 public final long estimateSize() {
394 if (prefix != null) {
395 estimatedSize = Math.max(estimatedSize, prefix.estimateSize());
396 }
397 return Math.max(estimatedSize, 0);
398 }
399
400 @Override
401 public final int characteristics() {
402 return characteristics;
403 }
404 }
405
406 /**
407 * Implementation of {@link Stream#flatMap} with an object spliterator output type.
408 *
409 * <p>To avoid having this type, we could use {@code FlatMapSpliterator} directly. The main
410 * advantages to having the type are the ability to use its constructor reference below and the
411 * parallelism with the primitive version. In short, it makes its caller ({@code flatMap})
412 * simpler.
413 *
414 * @param <InElementT> the element type of the input spliterator
415 * @param <OutElementT> the element type of the output spliterators
416 */
417 static final class FlatMapSpliteratorOfObject<InElementT, OutElementT>
418 extends FlatMapSpliterator<InElementT, OutElementT, Spliterator<OutElementT>> {
419 FlatMapSpliteratorOfObject(
420 Spliterator<OutElementT> prefix,
421 Spliterator<InElementT> from,
422 Function<? super InElementT, Spliterator<OutElementT>> function,
423 int characteristics,
424 long estimatedSize) {
425 super(
426 prefix, from, function, FlatMapSpliteratorOfObject::new, characteristics, estimatedSize);
427 }
428 }
429
430 /**
431 * Implementation of {@link Stream#flatMap} with a primitive spliterator output type.
432 *
433 * @param <InElementT> the element type of the input spliterator
434 * @param <OutElementT> the (boxed) element type of the output spliterators
435 * @param <OutConsumerT> the specialized consumer type for the primitive output type
436 * @param <OutSpliteratorT> the primitive spliterator type associated with {@code OutElementT}
437 */
438 abstract static class FlatMapSpliteratorOfPrimitive<
439 InElementT,
440 OutElementT,
441 OutConsumerT,
442 OutSpliteratorT extends
443 Spliterator.OfPrimitive<OutElementT, OutConsumerT, OutSpliteratorT>>
444 extends FlatMapSpliterator<InElementT, OutElementT, OutSpliteratorT>
445 implements Spliterator.OfPrimitive<OutElementT, OutConsumerT, OutSpliteratorT> {
446
447 FlatMapSpliteratorOfPrimitive(
448 OutSpliteratorT prefix,
449 Spliterator<InElementT> from,
450 Function<? super InElementT, OutSpliteratorT> function,
451 Factory<InElementT, OutSpliteratorT> factory,
452 int characteristics,
453 long estimatedSize) {
454 super(prefix, from, function, factory, characteristics, estimatedSize);
455 }
456
457 @Override
458 public final boolean tryAdvance(OutConsumerT action) {
459 while (true) {
460 if (prefix != null && prefix.tryAdvance(action)) {
461 if (estimatedSize != Long.MAX_VALUE) {
462 estimatedSize--;
463 }
464 return true;
465 } else {
466 prefix = null;
467 }
468 if (!from.tryAdvance(fromElement -> prefix = function.apply(fromElement))) {
469 return false;
470 }
471 }
472 }
473
474 @Override
475 public final void forEachRemaining(OutConsumerT action) {
476 if (prefix != null) {
477 prefix.forEachRemaining(action);
478 prefix = null;
479 }
480 from.forEachRemaining(
481 fromElement -> {
482 OutSpliteratorT elements = function.apply(fromElement);
483 if (elements != null) {
484 elements.forEachRemaining(action);
485 }
486 });
487 estimatedSize = 0;
488 }
489 }
490
491 /** Implementation of {@link #flatMapToInt}. */
492 static final class FlatMapSpliteratorOfInt<InElementT>
493 extends FlatMapSpliteratorOfPrimitive<InElementT, Integer, IntConsumer, Spliterator.OfInt>
494 implements Spliterator.OfInt {
495 FlatMapSpliteratorOfInt(
496 Spliterator.OfInt prefix,
497 Spliterator<InElementT> from,
498 Function<? super InElementT, Spliterator.OfInt> function,
499 int characteristics,
500 long estimatedSize) {
501 super(prefix, from, function, FlatMapSpliteratorOfInt::new, characteristics, estimatedSize);
502 }
503 }
504
505 /** Implementation of {@link #flatMapToLong}. */
506 static final class FlatMapSpliteratorOfLong<InElementT>
507 extends FlatMapSpliteratorOfPrimitive<InElementT, Long, LongConsumer, Spliterator.OfLong>
508 implements Spliterator.OfLong {
509 FlatMapSpliteratorOfLong(
510 Spliterator.OfLong prefix,
511 Spliterator<InElementT> from,
512 Function<? super InElementT, Spliterator.OfLong> function,
513 int characteristics,
514 long estimatedSize) {
515 super(prefix, from, function, FlatMapSpliteratorOfLong::new, characteristics, estimatedSize);
516 }
517 }
518
519 /** Implementation of {@link #flatMapToDouble}. */
520 static final class FlatMapSpliteratorOfDouble<InElementT>
521 extends FlatMapSpliteratorOfPrimitive<
522 InElementT, Double, DoubleConsumer, Spliterator.OfDouble>
523 implements Spliterator.OfDouble {
524 FlatMapSpliteratorOfDouble(
525 Spliterator.OfDouble prefix,
526 Spliterator<InElementT> from,
527 Function<? super InElementT, Spliterator.OfDouble> function,
528 int characteristics,
529 long estimatedSize) {
530 super(
531 prefix, from, function, FlatMapSpliteratorOfDouble::new, characteristics, estimatedSize);
532 }
533 }
534 }