Make GroupedList.getAllElementsAsIterable return a collection.
When passed to Iterables.size, this is taking the code path that counts all of the iterator's elements even though we already track the number of elements in GroupedList.
Storing the collection view as a field since it is used multiple times per GroupedList on average.
@@ -24,6 +24,7 @@ import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadHostile;
import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
import java.lang.annotation.ElementType;
import java.lang.annotation.Target;
+import java.util.AbstractCollection;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
@@ -45,7 +46,7 @@ import org.checkerframework.framework.qual.SubtypeOf;
* <p>Groups are implemented as lists to minimize memory use. However, {@link #equals} is defined to
* treat groups as unordered.
*/
-public class GroupedList<T> implements Iterable<List<T>> {
+public final class GroupedList<T> implements Iterable<List<T>> {
/**
* Indicates that the annotated element is a compressed {@link GroupedList}, so that it can be
@@ -69,6 +70,8 @@ public class GroupedList<T> implements Iterable<List<T>> {
// Non-final only for #remove.
private List<Object> elements;
+ private final CollectionView collectionView = new CollectionView();
+
public GroupedList() {
// We optimize for small lists.
this.elements = new ArrayList<>(1);
@@ -405,8 +408,27 @@ public class GroupedList<T> implements Iterable<List<T>> {
return first.equals(second) || CompactHashSet.create(first).containsAll(second);
}
+ /**
+ * A grouping-unaware view of a {@code GroupedList} which does not support modifications.
+ *
+ * <p>This is implemented as a {@code Collection} so that calling {@link Iterables#size} on the
+ * return value of {@link #getAllElementsAsIterable} will take constant time.
+ */
+ private final class CollectionView extends AbstractCollection<T> {
+
+ @Override
+ public Iterator<T> iterator() {
+ return new UngroupedIterator();
+ }
+
+ @Override
+ public int size() {
+ return size;
+ }
+ }
+
/** An iterator that loops through every element in each group. */
- private class UngroupedIterator implements Iterator<T> {
+ private final class UngroupedIterator implements Iterator<T> {
private final Iterator<Object> iter = elements.iterator();
int counter = 0;
List<T> currentGroup;
@@ -439,7 +461,7 @@ public class GroupedList<T> implements Iterable<List<T>> {
@ThreadHostile
public Iterable<T> getAllElementsAsIterable() {
- return UngroupedIterator::new;
+ return collectionView;
}
@Override
