Some optimizations in Type#flatten (used under the covers by AggregatingAttributeMapper#visitLabels):

-Rename Type#flatten to Type#extractLabels.
-Change the return type of Type#extractLabels from Collection to Iterable. This way we don't need to create and concatenate large lists.
-Add an internal-only Type#containsLabels so this way ListType and DictType can have efficient implementations of Type#extractLabels.

Note that AggregatingAttributeMapper#visitLabels is called multiple times on the same  in several different places during the lifetime of a non-incremental Blaze invocation (e.g. during Package loading, during transitive target visitation, etc)

--
MOS_MIGRATED_REVID=131311698

5 files changed, 145 insertions, 51 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/packages/AbstractAttributeMapper.java b/src/main/java/com/google/devtools/build/lib/packages/AbstractAttributeMapper.java
index d89e278ab7..c4c47b4b7a 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/AbstractAttributeMapper.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/AbstractAttributeMapper.java
@@ -245,9 +245,9 @@ public abstract class AbstractAttributeMapper implements AttributeMap {
     return attribute != null && attribute.getType() == type;
   }
 
-  protected static Iterable<Label> extractLabels(Type type, Object value) {
+  protected static Iterable<Label> extractLabels(Type<?> type, Object value) {
     return value == null
         ? ImmutableList.<Label>of()
-        : Iterables.filter(type.flatten(value), Label.class);
+        : Iterables.filter(type.extractLabels(value), Label.class);
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/packages/BuildType.java b/src/main/java/com/google/devtools/build/lib/packages/BuildType.java
index bbef95be91..3a8c69c677 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/BuildType.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/BuildType.java
@@ -97,8 +97,13 @@ public final class BuildType {
     }
 
     @Override
-    public Collection<Object> flatten(Object value) {
-      return NOT_COMPOSITE_TYPE;
+    protected boolean containsLabels() {
+      return false;
+    }
+
+    @Override
+    public Collection<Object> extractLabels(Object value) {
+      return NO_LABELS;
     }
 
     @Override
@@ -191,7 +196,12 @@ public final class BuildType {
     }
 
     @Override
-    public Collection<? extends Object> flatten(Object value) {
+    protected boolean containsLabels() {
+      return true;
+    }
+
+    @Override
+    public Collection<? extends Object> extractLabels(Object value) {
       return cast(value).getLabels();
     }
   }
@@ -208,7 +218,12 @@ public final class BuildType {
     }
 
     @Override
-    public Collection<Label> flatten(Object value) {
+    protected boolean containsLabels() {
+      return true;
+    }
+
+    @Override
+    public Iterable<Label> extractLabels(Object value) {
       return ImmutableList.of(cast(value));
     }
 
@@ -259,8 +274,13 @@ public final class BuildType {
     }
 
     @Override
-    public Collection<Object> flatten(Object value) {
-      return NOT_COMPOSITE_TYPE;
+    protected boolean containsLabels() {
+      return false;
+    }
+
+    @Override
+    public Collection<Object> extractLabels(Object value) {
+      return NO_LABELS;
     }
 
     @Override
@@ -299,8 +319,13 @@ public final class BuildType {
     }
 
     @Override
-    public Collection<Object> flatten(Object what) {
-      return NOT_COMPOSITE_TYPE;
+    protected boolean containsLabels() {
+      return false;
+    }
+
+    @Override
+    public Collection<Object> extractLabels(Object what) {
+      return NO_LABELS;
     }
 
     @Override
@@ -326,7 +351,12 @@ public final class BuildType {
     }
 
     @Override
-    public Collection<Label> flatten(Object value) {
+    protected boolean containsLabels() {
+      return true;
+    }
+
+    @Override
+    public Collection<Label> extractLabels(Object value) {
       return ImmutableList.of(cast(value));
     }
 
@@ -587,8 +617,13 @@ public final class BuildType {
     }
 
     @Override
-    public Collection<Object> flatten(Object value) {
-      return NOT_COMPOSITE_TYPE;
+    protected boolean containsLabels() {
+      return false;
+    }
+
+    @Override
+    public Collection<Object> extractLabels(Object value) {
+      return NO_LABELS;
     }
 
     @Override
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/Type.java b/src/main/java/com/google/devtools/build/lib/syntax/Type.java
index 49d0b74c3b..a8987e8eca 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/Type.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/Type.java
@@ -14,6 +14,7 @@
 
 package com.google.devtools.build.lib.syntax;
 
+import com.google.common.base.Function;
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -52,6 +53,14 @@ import javax.annotation.Nullable;
  */
 public abstract class Type<T> {
 
+  private final Function<Object, Iterable<? extends Object>> flattenFunction =
+      new Function<Object, Iterable<? extends Object>>() {
+        @Override
+        public Iterable<? extends Object> apply(Object value) {
+          return extractLabels(value);
+        }
+      };
+
   protected Type() {}
 
   /**
@@ -126,7 +135,13 @@ public abstract class Type<T> {
   public abstract T getDefaultValue();
 
   /**
-   * Flatten the an instance of the type if the type is a composite one.
+   * Returns whether there exists an {@code x} such that {@code extractLabels(x)} will return a
+   * non-{@code NO_LABELS} value.
+   */
+  protected abstract boolean containsLabels();
+
+  /**
+   * Extracts all the labels from the given instance of the type.
    *
    * <p>This is used to support reliable label visitation in
    * {@link com.google.devtools.build.lib.packages.AbstractAttributeMapper#visitLabels}. To preserve
@@ -134,12 +149,12 @@ public abstract class Type<T> {
    * words, be careful about defining default instances in base types that get auto-inherited by
    * their children. Keep all definitions as explicit as possible.
    */
-  public abstract Collection<? extends Object> flatten(Object value);
+  public abstract Iterable<? extends Object> extractLabels(Object value);
 
   /**
-   * {@link #flatten} return value for types that don't contain labels.
+   * {@link #extractLabels} return value for types that don't contain labels.
    */
-  protected static final Collection<Object> NOT_COMPOSITE_TYPE = ImmutableList.of();
+  protected static final ImmutableList<Object> NO_LABELS = ImmutableList.of();
 
   /**
    * Implementation of concatenation for this type (e.g. "val1 + val2"). Returns null to
@@ -265,8 +280,13 @@ public abstract class Type<T> {
     }
 
     @Override
-    public Collection<Object> flatten(Object value) {
-      return NOT_COMPOSITE_TYPE;
+    protected boolean containsLabels() {
+      return false;
+    }
+
+    @Override
+    public Iterable<? extends Object> extractLabels(Object value) {
+      return NO_LABELS;
     }
 
     @Override
@@ -292,8 +312,13 @@ public abstract class Type<T> {
     }
 
     @Override
-    public Collection<Object> flatten(Object value) {
-      return NOT_COMPOSITE_TYPE;
+    protected boolean containsLabels() {
+      return false;
+    }
+
+    @Override
+    public Collection<Object> extractLabels(Object value) {
+      return NO_LABELS;
     }
 
     @Override
@@ -332,8 +357,13 @@ public abstract class Type<T> {
     }
 
     @Override
-    public Collection<Object> flatten(Object value) {
-      return NOT_COMPOSITE_TYPE;
+    protected boolean containsLabels() {
+      return false;
+    }
+
+    @Override
+    public Iterable<Object> extractLabels(Object value) {
+      return NO_LABELS;
     }
 
     @Override
@@ -383,8 +413,13 @@ public abstract class Type<T> {
     }
 
     @Override
-    public Collection<Object> flatten(Object value) {
-      return NOT_COMPOSITE_TYPE;
+    protected boolean containsLabels() {
+      return false;
+    }
+
+    @Override
+    public Collection<Object> extractLabels(Object value) {
+      return NO_LABELS;
     }
 
     @Override
@@ -429,6 +464,24 @@ public abstract class Type<T> {
 
     private final Map<KeyT, ValueT> empty = ImmutableMap.of();
 
+    private final Function<
+        Map.Entry<KeyT, ValueT>, Iterable<? extends Object>> mapEntryFlattenFunction =
+        new Function<Map.Entry<KeyT, ValueT>, Iterable<? extends Object>>() {
+          @Override
+          public Iterable<? extends Object> apply(Entry<KeyT, ValueT> entry) {
+            Iterable<? extends Object> flattenedKeys = keyType.extractLabels(entry.getKey());
+            Iterable<? extends Object> flattenedValues = valueType.extractLabels(entry.getValue());
+            if (keyType.containsLabels() && valueType.containsLabels()) {
+              return Iterables.concat(flattenedKeys, flattenedValues);
+            } else if (keyType.containsLabels()) {
+              return flattenedKeys;
+            } else if (valueType.containsLabels()) {
+              return flattenedValues;
+            }
+            throw new IllegalStateException(this.toString());
+          }
+        };
+
     public static <KEY, VALUE> DictType<KEY, VALUE> create(
         Type<KEY> keyType, Type<VALUE> valueType) {
       return new DictType<>(keyType, valueType);
@@ -482,13 +535,15 @@ public abstract class Type<T> {
     }
 
     @Override
-    public Collection<Object> flatten(Object value) {
-      ImmutableList.Builder<Object> result = ImmutableList.builder();
-      for (Map.Entry<KeyT, ValueT> entry : cast(value).entrySet()) {
-        result.addAll(keyType.flatten(entry.getKey()));
-        result.addAll(valueType.flatten(entry.getValue()));
-      }
-      return result.build();
+    protected boolean containsLabels() {
+      return keyType.containsLabels() || valueType.containsLabels();
+    }
+
+    @Override
+    public Iterable<Object> extractLabels(Object value) {
+      return containsLabels()
+          ? Iterables.concat(Iterables.transform(cast(value).entrySet(), mapEntryFlattenFunction))
+          : NO_LABELS;
     }
   }
 
@@ -524,12 +579,15 @@ public abstract class Type<T> {
     }
 
     @Override
-    public Collection<Object> flatten(Object value) {
-      ImmutableList.Builder<Object> labels = ImmutableList.builder();
-      for (ElemT entry : cast(value)) {
-        labels.addAll(elemType.flatten(entry));
-      }
-      return labels.build();
+    protected boolean containsLabels() {
+      return elemType.containsLabels();
+    }
+
+    @Override
+    public Iterable<? extends Object> extractLabels(Object value) {
+      return containsLabels()
+          ? Iterables.concat(Iterables.transform(cast(value), elemType.flattenFunction))
+          : NO_LABELS;
     }
 
     @Override
diff --git a/src/test/java/com/google/devtools/build/lib/packages/BuildTypeTest.java b/src/test/java/com/google/devtools/build/lib/packages/BuildTypeTest.java
index 089430dfdf..dfd4f28cef 100644
--- a/src/test/java/com/google/devtools/build/lib/packages/BuildTypeTest.java
+++ b/src/test/java/com/google/devtools/build/lib/packages/BuildTypeTest.java
@@ -66,7 +66,7 @@ public class BuildTypeTest {
             /* symlinkBehavior */ null,
             /* stripPrefix */ null);
     assertEquals(input, BuildType.FILESET_ENTRY.convert(input, null, currentRule));
-    assertThat(BuildType.FILESET_ENTRY.flatten(input)).containsExactly(entryLabel);
+    assertThat(BuildType.FILESET_ENTRY.extractLabels(input)).containsExactly(entryLabel);
   }
 
   @Test
@@ -90,7 +90,8 @@ public class BuildTypeTest {
             /* symlinkBehavior */ null,
             /* stripPrefix */ null));
     assertEquals(input, BuildType.FILESET_ENTRY_LIST.convert(input, null, currentRule));
-    assertThat(BuildType.FILESET_ENTRY_LIST.flatten(input)).containsExactly(entry1Label, entry2Label);
+    assertThat(BuildType.FILESET_ENTRY_LIST.extractLabels(input)).containsExactly(
+        entry1Label, entry2Label);
   }
 
   /**
diff --git a/src/test/java/com/google/devtools/build/lib/syntax/TypeTest.java b/src/test/java/com/google/devtools/build/lib/syntax/TypeTest.java
index d85bd79d9e..0714c83362 100644
--- a/src/test/java/com/google/devtools/build/lib/syntax/TypeTest.java
+++ b/src/test/java/com/google/devtools/build/lib/syntax/TypeTest.java
@@ -59,7 +59,7 @@ public class TypeTest {
   public void testInteger() throws Exception {
     Object x = 3;
     assertEquals(x, Type.INTEGER.convert(x, null));
-    assertThat(Type.INTEGER.flatten(x)).isEmpty();
+    assertThat(Type.INTEGER.extractLabels(x)).isEmpty();
   }
 
   @Test
@@ -91,7 +91,7 @@ public class TypeTest {
   public void testString() throws Exception {
     Object s = "foo";
     assertEquals(s, Type.STRING.convert(s, null));
-    assertThat(Type.STRING.flatten(s)).isEmpty();
+    assertThat(Type.STRING.extractLabels(s)).isEmpty();
   }
 
   @Test
@@ -114,7 +114,7 @@ public class TypeTest {
     assertTrue(Type.BOOLEAN.convert(myTrue, null));
     assertFalse(Type.BOOLEAN.convert(false, null));
     assertFalse(Type.BOOLEAN.convert(myFalse, null));
-    assertThat(Type.BOOLEAN.flatten(myTrue)).isEmpty();
+    assertThat(Type.BOOLEAN.extractLabels(myTrue)).isEmpty();
   }
 
   @Test
@@ -151,7 +151,7 @@ public class TypeTest {
     assertEquals(TriState.YES, BuildType.TRISTATE.convert(TriState.YES, null));
     assertEquals(TriState.NO, BuildType.TRISTATE.convert(TriState.NO, null));
     assertEquals(TriState.AUTO, BuildType.TRISTATE.convert(TriState.AUTO, null));
-    assertThat(BuildType.TRISTATE.flatten(TriState.YES)).isEmpty();
+    assertThat(BuildType.TRISTATE.extractLabels(TriState.YES)).isEmpty();
   }
 
   @Test
@@ -227,7 +227,7 @@ public class TypeTest {
     Label label = Label
         .parseAbsolute("//foo:bar");
     assertEquals(label, BuildType.LABEL.convert("//foo:bar", null, currentRule));
-    assertThat(BuildType.LABEL.flatten(label)).containsExactly(label);
+    assertThat(BuildType.LABEL.extractLabels(label)).containsExactly(label);
   }
 
   @Test
@@ -235,7 +235,7 @@ public class TypeTest {
     Label label = Label
         .parseAbsolute("//foo:bar");
     assertEquals(label, BuildType.NODEP_LABEL.convert("//foo:bar", null, currentRule));
-    assertThat(BuildType.NODEP_LABEL.flatten(label)).containsExactly(label);
+    assertThat(BuildType.NODEP_LABEL.extractLabels(label)).containsExactly(label);
   }
 
   @Test
@@ -279,7 +279,7 @@ public class TypeTest {
         Type.STRING_LIST.convert(input, null);
     assertEquals(input, converted);
     assertNotSame(input, converted);
-    assertThat(Type.STRING_LIST.flatten(input)).isEmpty();
+    assertThat(Type.STRING_LIST.extractLabels(input)).isEmpty();
   }
 
   @Test
@@ -289,7 +289,7 @@ public class TypeTest {
     Map<String, String> converted = Type.STRING_DICT.convert(input, null);
     assertEquals(input, converted);
     assertNotSame(input, converted);
-    assertThat(Type.STRING_DICT.flatten(converted)).isEmpty();
+    assertThat(Type.STRING_DICT.extractLabels(converted)).isEmpty();
   }
 
   @Test
@@ -336,7 +336,7 @@ public class TypeTest {
                     Label.parseAbsolute("//quux:wiz"));
     assertEquals(expected, converted);
     assertNotSame(expected, converted);
-    assertThat(BuildType.LABEL_LIST.flatten(converted)).containsExactlyElementsIn(expected);
+    assertThat(BuildType.LABEL_LIST.extractLabels(converted)).containsExactlyElementsIn(expected);
   }
 
   @Test
@@ -385,7 +385,7 @@ public class TypeTest {
             "wiz", Arrays.asList("bang"));
     assertEquals(expected, converted);
     assertNotSame(expected, converted);
-    assertThat(Type.STRING_LIST_DICT.flatten(converted)).isEmpty();
+    assertThat(Type.STRING_LIST_DICT.extractLabels(converted)).isEmpty();
   }
 
   @Test
@@ -438,7 +438,7 @@ public class TypeTest {
             "wiz", "bang");
     assertEquals(expected, converted);
     assertNotSame(expected, converted);
-    assertThat(Type.STRING_DICT_UNARY.flatten(converted)).isEmpty();
+    assertThat(Type.STRING_DICT_UNARY.extractLabels(converted)).isEmpty();
   }
 
   @Test