View Javadoc
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 java.util.Comparator;
24  import java.util.Spliterator;
25  import java.util.function.Consumer;
26  import java.util.function.Function;
27  import java.util.function.IntConsumer;
28  import java.util.function.IntFunction;
29  import java.util.function.Predicate;
30  import java.util.stream.IntStream;
31  import javax.annotation.Nullable;
32  
33  /**
34   * Spliterator utilities for {@code common.collect} internals.
35   */
36  @GwtCompatible
37  final class CollectSpliterators {
38    private CollectSpliterators() {}
39  
40    static <T> Spliterator<T> indexed(int size, int extraCharacteristics, IntFunction<T> function) {
41      return indexed(size, extraCharacteristics, function, null);
42    }
43  
44    static <T> Spliterator<T> indexed(
45        int size,
46        int extraCharacteristics,
47        IntFunction<T> function,
48        Comparator<? super T> comparator) {
49      if (comparator != null) {
50        checkArgument((extraCharacteristics & (Spliterator.SORTED)) != 0);
51      }
52      class WithCharacteristics implements Spliterator<T> {
53        private final Spliterator.OfInt delegate;
54  
55        WithCharacteristics(Spliterator.OfInt delegate) {
56          this.delegate = delegate;
57        }
58  
59        @Override
60        public boolean tryAdvance(Consumer<? super T> action) {
61          return delegate.tryAdvance((IntConsumer) i -> action.accept(function.apply(i)));
62        }
63  
64        @Override
65        public void forEachRemaining(Consumer<? super T> action) {
66          delegate.forEachRemaining((IntConsumer) i -> action.accept(function.apply(i)));
67        }
68  
69        @Override
70        @Nullable
71        public 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 <F, T> Spliterator<T> map(
106       Spliterator<F> fromSpliterator, Function<? super F, ? extends T> function) {
107     checkNotNull(fromSpliterator);
108     checkNotNull(function);
109     return new Spliterator<T>() {
110 
111       @Override
112       public boolean tryAdvance(Consumer<? super T> action) {
113         return fromSpliterator.tryAdvance(
114             fromElement -> action.accept(function.apply(fromElement)));
115       }
116 
117       @Override
118       public void forEachRemaining(Consumer<? super T> action) {
119         fromSpliterator.forEachRemaining(fromElement -> action.accept(function.apply(fromElement)));
120       }
121 
122       @Override
123       public Spliterator<T> trySplit() {
124         Spliterator<F> fromSplit = fromSpliterator.trySplit();
125         return (fromSplit != null) ? map(fromSplit, function) : null;
126       }
127 
128       @Override
129       public long estimateSize() {
130         return fromSpliterator.estimateSize();
131       }
132 
133       @Override
134       public int characteristics() {
135         return fromSpliterator.characteristics()
136             & ~(Spliterator.DISTINCT | Spliterator.NONNULL | Spliterator.SORTED);
137       }
138     };
139   }
140   
141   /** Returns a {@code Spliterator} filtered by the specified predicate. */
142   static <T> Spliterator<T> filter(Spliterator<T> fromSpliterator, Predicate<? super T> predicate) {
143     checkNotNull(fromSpliterator);
144     checkNotNull(predicate);
145     class Splitr implements Spliterator<T>, Consumer<T> {
146       T holder = null;
147 
148       @Override
149       public void accept(T t) {
150         this.holder = t;
151       }
152 
153       @Override
154       public boolean tryAdvance(Consumer<? super T> action) {
155         while (fromSpliterator.tryAdvance(this)) {
156           try {
157             if (predicate.test(holder)) {
158               action.accept(holder);
159               return true;
160             }
161           } finally {
162             holder = null;
163           }
164         }
165         return false;
166       }
167 
168       @Override
169       public Spliterator<T> trySplit() {
170         Spliterator<T> fromSplit = fromSpliterator.trySplit();
171         return (fromSplit == null) ? null : filter(fromSplit, predicate);
172       }
173 
174       @Override
175       public long estimateSize() {
176         return fromSpliterator.estimateSize() / 2;
177       }
178 
179       @Override
180       public Comparator<? super T> getComparator() {
181         return fromSpliterator.getComparator();
182       }
183 
184       @Override
185       public int characteristics() {
186         return fromSpliterator.characteristics()
187             & (Spliterator.DISTINCT
188                 | Spliterator.NONNULL
189                 | Spliterator.ORDERED
190                 | Spliterator.SORTED);
191       }
192     }
193     return new Splitr();
194   }
195 
196   /**
197    * Returns a {@code Spliterator} that iterates over the elements of the spliterators generated by
198    * applying {@code function} to the elements of {@code fromSpliterator}.
199    */
200   static <F, T> Spliterator<T> flatMap(
201       Spliterator<F> fromSpliterator,
202       Function<? super F, Spliterator<T>> function,
203       int topCharacteristics,
204       long topSize) {
205     checkArgument(
206         (topCharacteristics & Spliterator.SUBSIZED) == 0,
207         "flatMap does not support SUBSIZED characteristic");
208     checkArgument(
209         (topCharacteristics & Spliterator.SORTED) == 0,
210         "flatMap does not support SORTED characteristic");
211     checkNotNull(fromSpliterator);
212     checkNotNull(function);
213     class FlatMapSpliterator implements Spliterator<T> {
214       @Nullable Spliterator<T> prefix;
215       final Spliterator<F> from;
216       int characteristics;
217       long estimatedSize;
218 
219       FlatMapSpliterator(
220           Spliterator<T> prefix, Spliterator<F> from, int characteristics, long estimatedSize) {
221         this.prefix = prefix;
222         this.from = from;
223         this.characteristics = characteristics;
224         this.estimatedSize = estimatedSize;
225       }
226 
227       @Override
228       public boolean tryAdvance(Consumer<? super T> action) {
229         while (true) {
230           if (prefix != null && prefix.tryAdvance(action)) {
231             if (estimatedSize != Long.MAX_VALUE) {
232               estimatedSize--;
233             }
234             return true;
235           } else {
236             prefix = null;
237           }
238           if (!from.tryAdvance(fromElement -> prefix = function.apply(fromElement))) {
239             return false;
240           }
241         }
242       }
243 
244       @Override
245       public void forEachRemaining(Consumer<? super T> action) {
246         if (prefix != null) {
247           prefix.forEachRemaining(action);
248           prefix = null;
249         }
250         from.forEachRemaining(fromElement -> function.apply(fromElement).forEachRemaining(action));
251         estimatedSize = 0;
252       }
253 
254       @Override
255       public Spliterator<T> trySplit() {
256         Spliterator<F> fromSplit = from.trySplit();
257         if (fromSplit != null) {
258           int splitCharacteristics = characteristics & ~Spliterator.SIZED;
259           long estSplitSize = estimateSize();
260           if (estSplitSize < Long.MAX_VALUE) {
261             estSplitSize /= 2;
262             this.estimatedSize -= estSplitSize;
263             this.characteristics = splitCharacteristics;
264           }
265           Spliterator<T> result =
266               new FlatMapSpliterator(this.prefix, fromSplit, splitCharacteristics, estSplitSize);
267           this.prefix = null;
268           return result;
269         } else if (prefix != null) {
270           Spliterator<T> result = prefix;
271           this.prefix = null;
272           return result;
273         } else {
274           return null;
275         }
276       }
277 
278       @Override
279       public long estimateSize() {
280         if (prefix != null) {
281           estimatedSize = Math.max(estimatedSize, prefix.estimateSize());
282         }
283         return Math.max(estimatedSize, 0);
284       }
285 
286       @Override
287       public int characteristics() {
288         return characteristics;
289       }
290     }
291     return new FlatMapSpliterator(null, fromSpliterator, topCharacteristics, topSize);
292   }
293 }