Improve artifact->owner label accounting in two ways. First, don't do the full mapping unless requested. This gets rid of any performance issue for the vast majority of builds. Second, if requested, use a custom data structure so that we don't have to create a full HashSet for artifacts whose only owning labels are their own owner labels.

PiperOrigin-RevId: 206610370
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/AnalysisResult.java b/src/main/java/com/google/devtools/build/lib/analysis/AnalysisResult.java
index ff8dbe4..1fc7048 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/AnalysisResult.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/AnalysisResult.java
@@ -16,12 +16,9 @@
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.SetMultimap;
 import com.google.devtools.build.lib.actions.ActionGraph;
-import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.PackageRoots;
 import com.google.devtools.build.lib.analysis.config.BuildConfigurationCollection;
-import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.skyframe.AspectValue;
 import java.util.Collection;
 import java.util.List;
@@ -37,7 +34,7 @@
   private final ImmutableSet<ConfiguredTarget> targetsToSkip;
   @Nullable private final String error;
   private final ActionGraph actionGraph;
-  private final SetMultimap<Artifact, Label> topLevelArtifactsToOwnerLabels;
+  private final ArtifactsToOwnerLabels topLevelArtifactsToOwnerLabels;
   private final ImmutableSet<ConfiguredTarget> parallelTests;
   private final ImmutableSet<ConfiguredTarget> exclusiveTests;
   @Nullable private final TopLevelArtifactContext topLevelContext;
@@ -54,7 +51,7 @@
       Collection<ConfiguredTarget> targetsToSkip,
       @Nullable String error,
       ActionGraph actionGraph,
-      SetMultimap<Artifact, Label> topLevelArtifactsToOwnerLabels,
+      ArtifactsToOwnerLabels topLevelArtifactsToOwnerLabels,
       Collection<ConfiguredTarget> parallelTests,
       Collection<ConfiguredTarget> exclusiveTests,
       TopLevelArtifactContext topLevelContext,
@@ -122,7 +119,7 @@
     return targetsToSkip;
   }
 
-  public SetMultimap<Artifact, Label> getTopLevelArtifactsToOwnerLabels() {
+  public ArtifactsToOwnerLabels getTopLevelArtifactsToOwnerLabels() {
     return topLevelArtifactsToOwnerLabels;
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/ArtifactsToOwnerLabels.java b/src/main/java/com/google/devtools/build/lib/analysis/ArtifactsToOwnerLabels.java
new file mode 100644
index 0000000..e8cc545
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/analysis/ArtifactsToOwnerLabels.java
@@ -0,0 +1,116 @@
+// Copyright 2018 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.lib.analysis;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.HashMultimap;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.SetMultimap;
+import com.google.common.collect.Sets;
+import com.google.devtools.build.lib.actions.Artifact;
+import com.google.devtools.build.lib.cmdline.Label;
+import java.util.HashSet;
+import java.util.Set;
+
+/**
+ * Morally a {@link SetMultimap} from artifacts to the labels that own them, which may not be just
+ * {@link Artifact#getOwnerLabel} depending on the build. Optimizes for artifacts that are owned by
+ * their {@link Artifact#getOwnerLabel} by storing them separately.
+ */
+public class ArtifactsToOwnerLabels {
+  private final SetMultimap<Artifact, Label> artifactToMultipleOrDifferentOwnerLabels;
+  private final Set<Artifact> artifactsOwnedOnlyByTheirLabels;
+
+  private ArtifactsToOwnerLabels(
+      SetMultimap<Artifact, Label> artifactToMultipleOrDifferentOwnerLabels,
+      Set<Artifact> artifactsOwnedOnlyByTheirLabels) {
+    this.artifactToMultipleOrDifferentOwnerLabels = artifactToMultipleOrDifferentOwnerLabels;
+    this.artifactsOwnedOnlyByTheirLabels = artifactsOwnedOnlyByTheirLabels;
+  }
+
+  public Set<Label> getOwners(Artifact artifact) {
+    if (artifactsOwnedOnlyByTheirLabels.contains(artifact)) {
+      Preconditions.checkState(
+          !artifactToMultipleOrDifferentOwnerLabels.containsKey(artifact),
+          "Artifact %s incorrectly in multiple categories",
+          artifact);
+      Label ownerLabel = artifact.getOwnerLabel();
+      if (ownerLabel == null) {
+        return ImmutableSet.of();
+      }
+      return ImmutableSet.of(ownerLabel);
+    }
+    return Preconditions.checkNotNull(
+        artifactToMultipleOrDifferentOwnerLabels.get(artifact), artifact);
+  }
+
+  public Set<Artifact> getArtifacts() {
+    return Sets.union(
+        artifactsOwnedOnlyByTheirLabels, artifactToMultipleOrDifferentOwnerLabels.keySet());
+  }
+
+  public Builder toBuilder() {
+    return new Builder(
+        HashMultimap.create(artifactToMultipleOrDifferentOwnerLabels),
+        new HashSet<>(artifactsOwnedOnlyByTheirLabels));
+  }
+
+  static class Builder {
+    private final SetMultimap<Artifact, Label> artifactToMultipleOrDifferentOwnerLabels;
+    private final Set<Artifact> artifactsOwnedOnlyByTheirLabels;
+
+    Builder() {
+      this(HashMultimap.create(), new HashSet<>());
+    }
+
+    private Builder(
+        SetMultimap<Artifact, Label> artifactToMultipleOrDifferentOwnerLabels,
+        Set<Artifact> artifactsOwnedOnlyByTheirLabels) {
+      this.artifactToMultipleOrDifferentOwnerLabels = artifactToMultipleOrDifferentOwnerLabels;
+      this.artifactsOwnedOnlyByTheirLabels = artifactsOwnedOnlyByTheirLabels;
+    }
+
+    public Builder addArtifact(Artifact artifact) {
+      if (artifactToMultipleOrDifferentOwnerLabels.containsKey(artifact)) {
+        Label ownerLabel = artifact.getOwnerLabel();
+        if (ownerLabel != null) {
+          artifactToMultipleOrDifferentOwnerLabels.put(artifact, artifact.getOwnerLabel());
+        }
+      } else {
+        artifactsOwnedOnlyByTheirLabels.add(artifact);
+      }
+      return this;
+    }
+
+    public Builder addArtifact(Artifact artifact, Label label) {
+      Preconditions.checkNotNull(label, artifact);
+      if (label.equals(artifact.getOwnerLabel())) {
+        addArtifact(artifact);
+      } else {
+        artifactToMultipleOrDifferentOwnerLabels.put(artifact, label);
+        if (artifactsOwnedOnlyByTheirLabels.remove(artifact)) {
+          // Redoing this call now that we have a mismatched label will force addition into the
+          // multimap.
+          addArtifact(artifact);
+        }
+      }
+      return this;
+    }
+
+    ArtifactsToOwnerLabels build() {
+      return new ArtifactsToOwnerLabels(
+          artifactToMultipleOrDifferentOwnerLabels, artifactsOwnedOnlyByTheirLabels);
+    }
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/BuildView.java b/src/main/java/com/google/devtools/build/lib/analysis/BuildView.java
index ec049cd..9995de0 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/BuildView.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/BuildView.java
@@ -18,13 +18,11 @@
 import com.google.common.base.Function;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ArrayListMultimap;
-import com.google.common.collect.HashMultimap;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Multimap;
-import com.google.common.collect.SetMultimap;
 import com.google.common.collect.Sets;
 import com.google.common.eventbus.EventBus;
 import com.google.devtools.build.lib.actions.ActionAnalysisMetadata;
@@ -377,7 +375,8 @@
       allTargetsToTest = filterTestsByTargets(configuredTargets, testsToRun);
     }
 
-    SetMultimap<Artifact, Label> topLevelArtifactsToOwnerLabels = HashMultimap.create();
+    ArtifactsToOwnerLabels.Builder artifactsToOwnerLabelsBuilder =
+        new ArtifactsToOwnerLabels.Builder();
     Set<ConfiguredTarget> parallelTests = new HashSet<>();
     Set<ConfiguredTarget> exclusiveTests = new HashSet<>();
 
@@ -385,15 +384,15 @@
     Collection<Artifact> buildInfoArtifacts =
         skyframeExecutor.getWorkspaceStatusArtifacts(eventHandler);
     Preconditions.checkState(buildInfoArtifacts.size() == 2, buildInfoArtifacts);
-    addArtifactsWithNoOwner(buildInfoArtifacts, topLevelArtifactsToOwnerLabels);
+    buildInfoArtifacts.forEach(artifactsToOwnerLabelsBuilder::addArtifact);
 
     // Extra actions
     addExtraActionsIfRequested(
-        viewOptions, configuredTargets, aspects, topLevelArtifactsToOwnerLabels, eventHandler);
+        viewOptions, configuredTargets, aspects, artifactsToOwnerLabelsBuilder, eventHandler);
 
     // Coverage
     NestedSet<Artifact> baselineCoverageArtifacts =
-        getBaselineCoverageArtifacts(configuredTargets, topLevelArtifactsToOwnerLabels);
+        getBaselineCoverageArtifacts(configuredTargets, artifactsToOwnerLabelsBuilder);
     if (coverageReportActionFactory != null) {
       CoverageReportActionsWrapper actionsWrapper;
       actionsWrapper =
@@ -408,8 +407,7 @@
       if (actionsWrapper != null) {
         ImmutableList<ActionAnalysisMetadata> actions = actionsWrapper.getActions();
         skyframeExecutor.injectCoverageReportData(actions);
-        addArtifactsWithNoOwner(
-            actionsWrapper.getCoverageOutputs(), topLevelArtifactsToOwnerLabels);
+        actionsWrapper.getCoverageOutputs().forEach(artifactsToOwnerLabelsBuilder::addArtifact);
       }
     }
 
@@ -453,7 +451,7 @@
         targetsToSkip,
         error,
         actionGraph,
-        topLevelArtifactsToOwnerLabels,
+        artifactsToOwnerLabelsBuilder.build(),
         parallelTests,
         exclusiveTests,
         topLevelOptions,
@@ -462,11 +460,6 @@
         topLevelTargetsWithConfigs);
   }
 
-  private static void addArtifactsWithNoOwner(
-      Collection<Artifact> artifacts, SetMultimap<Artifact, Label> topLevelArtifactsToOwnerLabels) {
-    artifacts.forEach((a) -> topLevelArtifactsToOwnerLabels.put(a, a.getOwnerLabel()));
-  }
-
   @Nullable
   public static String createErrorMessage(
       TargetPatternPhaseValue loadingResult,
@@ -483,7 +476,7 @@
 
   private static NestedSet<Artifact> getBaselineCoverageArtifacts(
       Collection<ConfiguredTarget> configuredTargets,
-      SetMultimap<Artifact, Label> topLevelArtifactsToOwnerLabels) {
+      ArtifactsToOwnerLabels.Builder topLevelArtifactsToOwnerLabels) {
     NestedSetBuilder<Artifact> baselineCoverageArtifacts = NestedSetBuilder.stableOrder();
     for (ConfiguredTarget target : configuredTargets) {
       InstrumentedFilesProvider provider = target.getProvider(InstrumentedFilesProvider.class);
@@ -503,7 +496,7 @@
       AnalysisOptions viewOptions,
       Collection<ConfiguredTarget> configuredTargets,
       Collection<AspectValue> aspects,
-      SetMultimap<Artifact, Label> artifactsToTopLevelLabelsMap,
+      ArtifactsToOwnerLabels.Builder artifactsToTopLevelLabelsMap,
       ExtendedEventHandler eventHandler) {
     RegexFilter filter = viewOptions.extraActionFilter;
     for (ConfiguredTarget target : configuredTargets) {
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/TopLevelArtifactHelper.java b/src/main/java/com/google/devtools/build/lib/analysis/TopLevelArtifactHelper.java
index 853ecb9..ff15303 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/TopLevelArtifactHelper.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/TopLevelArtifactHelper.java
@@ -17,16 +17,16 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.collect.HashMultimap;
-import com.google.common.collect.SetMultimap;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.analysis.test.TestProvider;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
+import com.google.devtools.build.lib.profiler.AutoProfiler;
 import com.google.devtools.build.lib.skyframe.AspectValue;
 import com.google.devtools.build.lib.util.RegexFilter;
+import java.util.logging.Logger;
 import javax.annotation.Nullable;
 
 /**
@@ -34,6 +34,7 @@
  * extra top-level artifacts into the build.
  */
 public final class TopLevelArtifactHelper {
+  private static Logger logger = Logger.getLogger(TopLevelArtifactHelper.class.getName());
 
   /** Set of {@link Artifact}s in an output group. */
   @Immutable
@@ -119,46 +120,49 @@
   }
 
   @VisibleForTesting
-  public static SetMultimap<Artifact, Label> makeTopLevelArtifactsToOwnerLabels(
+  public static ArtifactsToOwnerLabels makeTopLevelArtifactsToOwnerLabels(
       AnalysisResult analysisResult, Iterable<AspectValue> aspects) {
-    SetMultimap<Artifact, Label> topLevelArtifactsToOwnerLabels =
-        HashMultimap.create(analysisResult.getTopLevelArtifactsToOwnerLabels());
+    try (AutoProfiler ignored = AutoProfiler.logged("assigning owner labels", logger, 10)) {
+
+      ArtifactsToOwnerLabels.Builder artifactsToOwnerLabelsBuilder =
+          analysisResult.getTopLevelArtifactsToOwnerLabels().toBuilder();
     TopLevelArtifactContext artifactContext = analysisResult.getTopLevelContext();
     for (ConfiguredTarget target : analysisResult.getTargetsToBuild()) {
-      addArtifactsWithOwnerLabel(
-          getAllArtifactsToBuild(target, artifactContext).getAllArtifacts(),
-          null,
-          target.getLabel(),
-          topLevelArtifactsToOwnerLabels);
+        addArtifactsWithOwnerLabel(
+            getAllArtifactsToBuild(target, artifactContext).getAllArtifacts(),
+            null,
+            target.getLabel(),
+            artifactsToOwnerLabelsBuilder);
     }
     for (AspectValue aspect : aspects) {
-      addArtifactsWithOwnerLabel(
-          getAllArtifactsToBuild(aspect, artifactContext).getAllArtifacts(),
-          null,
-          aspect.getLabel(),
-          topLevelArtifactsToOwnerLabels);
+        addArtifactsWithOwnerLabel(
+            getAllArtifactsToBuild(aspect, artifactContext).getAllArtifacts(),
+            null,
+            aspect.getLabel(),
+            artifactsToOwnerLabelsBuilder);
     }
     if (analysisResult.getTargetsToTest() != null) {
       for (ConfiguredTarget target : analysisResult.getTargetsToTest()) {
-        addArtifactsWithOwnerLabel(
-            TestProvider.getTestStatusArtifacts(target),
-            null,
-            target.getLabel(),
-            topLevelArtifactsToOwnerLabels);
+          addArtifactsWithOwnerLabel(
+              TestProvider.getTestStatusArtifacts(target),
+              null,
+              target.getLabel(),
+              artifactsToOwnerLabelsBuilder);
       }
     }
-    // TODO(dslomov): Artifacts to test from aspects?
-    return topLevelArtifactsToOwnerLabels;
+      // TODO(dslomov): Artifacts to test from aspects?
+      return artifactsToOwnerLabelsBuilder.build();
+    }
   }
 
-  public static void addArtifactsWithOwnerLabel(
+  static void addArtifactsWithOwnerLabel(
       Iterable<Artifact> artifacts,
       @Nullable RegexFilter filter,
       Label ownerLabel,
-      SetMultimap<Artifact, Label> artifactToTopLevelLabels) {
+      ArtifactsToOwnerLabels.Builder artifactsToOwnerLabelsBuilder) {
     for (Artifact artifact : artifacts) {
       if (filter == null || filter.isIncluded(artifact.getOwnerLabel().toString())) {
-        artifactToTopLevelLabels.put(artifact, ownerLabel);
+        artifactsToOwnerLabelsBuilder.addArtifact(artifact, ownerLabel);
       }
     }
   }
diff --git a/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java b/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java
index 44a5fd7..9a31023 100644
--- a/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java
+++ b/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java
@@ -18,15 +18,14 @@
 import com.google.common.base.Preconditions;
 import com.google.common.base.Predicate;
 import com.google.common.base.Stopwatch;
+import com.google.common.base.Suppliers;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.SetMultimap;
 import com.google.devtools.build.lib.actions.Action;
 import com.google.devtools.build.lib.actions.ActionCacheChecker;
 import com.google.devtools.build.lib.actions.ActionGraph;
 import com.google.devtools.build.lib.actions.ActionInputPrefetcher;
-import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.ArtifactFactory;
 import com.google.devtools.build.lib.actions.BuildFailedException;
 import com.google.devtools.build.lib.actions.ExecException;
@@ -49,7 +48,6 @@
 import com.google.devtools.build.lib.analysis.config.BuildConfiguration;
 import com.google.devtools.build.lib.buildtool.buildevent.ExecutionPhaseCompleteEvent;
 import com.google.devtools.build.lib.buildtool.buildevent.ExecutionStartingEvent;
-import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.PackageIdentifier;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.EventHandler;
@@ -311,16 +309,17 @@
     }
     boolean buildCompleted = false;
     try {
-      SetMultimap<Artifact, Label> topLevelArtifactsToOwnerLabels =
-          TopLevelArtifactHelper.makeTopLevelArtifactsToOwnerLabels(analysisResult, aspects);
       for (ActionContextProvider actionContextProvider : actionContextProviders) {
         try (SilentCloseable c =
             Profiler.instance().profile(actionContextProvider + ".executionPhaseStarting")) {
-          actionContextProvider.executionPhaseStarting(actionGraph, topLevelArtifactsToOwnerLabels);
+          actionContextProvider.executionPhaseStarting(
+              actionGraph,
+              Suppliers.memoize(
+                  () ->
+                      TopLevelArtifactHelper.makeTopLevelArtifactsToOwnerLabels(
+                          analysisResult, aspects)));
         }
       }
-      // Don't retain memory.
-      topLevelArtifactsToOwnerLabels = null;
       executor.executionPhaseStarting();
       skyframeExecutor.drainChangedFiles();
 
@@ -341,7 +340,7 @@
 
       builder.buildArtifacts(
           env.getReporter(),
-          analysisResult.getTopLevelArtifactsToOwnerLabels().keySet(),
+          analysisResult.getTopLevelArtifactsToOwnerLabels().getArtifacts(),
           analysisResult.getParallelTests(),
           analysisResult.getExclusiveTests(),
           analysisResult.getTargetsToBuild(),
diff --git a/src/main/java/com/google/devtools/build/lib/exec/ActionContextProvider.java b/src/main/java/com/google/devtools/build/lib/exec/ActionContextProvider.java
index 909f36a..d4730a8 100644
--- a/src/main/java/com/google/devtools/build/lib/exec/ActionContextProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/exec/ActionContextProvider.java
@@ -13,13 +13,12 @@
 // limitations under the License.
 package com.google.devtools.build.lib.exec;
 
-import com.google.common.collect.SetMultimap;
+import com.google.common.base.Supplier;
 import com.google.devtools.build.lib.actions.ActionContext;
 import com.google.devtools.build.lib.actions.ActionGraph;
-import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.ExecutorInitException;
 import com.google.devtools.build.lib.actions.MetadataProvider;
-import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.analysis.ArtifactsToOwnerLabels;
 
 /**
  * An object that provides execution strategies to {@link BlazeExecutor}.
@@ -52,7 +51,7 @@
 
   /** Called when the execution phase is started. */
   public void executionPhaseStarting(
-      ActionGraph actionGraph, SetMultimap<Artifact, Label> topLevelArtifactsToOwnerLabels)
+      ActionGraph actionGraph, Supplier<ArtifactsToOwnerLabels> topLevelArtifactsToOwnerLabels)
       throws ExecutorInitException, InterruptedException {}
 
   /**
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/ArtifactsToOwnerLabelsTest.java b/src/test/java/com/google/devtools/build/lib/analysis/ArtifactsToOwnerLabelsTest.java
new file mode 100644
index 0000000..6f3bebc
--- /dev/null
+++ b/src/test/java/com/google/devtools/build/lib/analysis/ArtifactsToOwnerLabelsTest.java
@@ -0,0 +1,65 @@
+// Copyright 2018 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.analysis;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.devtools.build.lib.actions.Artifact;
+import com.google.devtools.build.lib.cmdline.Label;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
+
+/** Tests for {@link ArtifactsToOwnerLabels}. */
+@RunWith(JUnit4.class)
+public class ArtifactsToOwnerLabelsTest {
+  @Test
+  public void smoke() {
+    Artifact label1Artifact = Mockito.mock(Artifact.class);
+    Label label1 = Label.parseAbsoluteUnchecked("//label:one");
+    Mockito.when(label1Artifact.getOwnerLabel()).thenReturn(label1);
+    ArtifactsToOwnerLabels underTest =
+        new ArtifactsToOwnerLabels.Builder().addArtifact(label1Artifact).build();
+    assertThat(underTest.getOwners(label1Artifact)).containsExactly(label1);
+    underTest = new ArtifactsToOwnerLabels.Builder().addArtifact(label1Artifact, label1).build();
+    assertThat(underTest.getOwners(label1Artifact)).containsExactly(label1);
+    underTest =
+        new ArtifactsToOwnerLabels.Builder()
+            .addArtifact(label1Artifact)
+            .addArtifact(label1Artifact, label1)
+            .build();
+    assertThat(underTest.getOwners(label1Artifact)).containsExactly(label1);
+    underTest =
+        new ArtifactsToOwnerLabels.Builder()
+            .addArtifact(label1Artifact, label1)
+            .addArtifact(label1Artifact)
+            .build();
+    assertThat(underTest.getOwners(label1Artifact)).containsExactly(label1);
+    Label label2 = Label.parseAbsoluteUnchecked("//label:two");
+    underTest =
+        new ArtifactsToOwnerLabels.Builder()
+            .addArtifact(label1Artifact, label2)
+            .addArtifact(label1Artifact)
+            .build();
+    assertThat(underTest.getOwners(label1Artifact)).containsExactly(label1, label2);
+    underTest =
+        new ArtifactsToOwnerLabels.Builder()
+            .addArtifact(label1Artifact)
+            .addArtifact(label1Artifact, label2)
+            .build();
+    assertThat(underTest.getOwners(label1Artifact)).containsExactly(label1, label2);
+  }
+}
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/AspectTest.java b/src/test/java/com/google/devtools/build/lib/analysis/AspectTest.java
index 69bbade..954ac20 100644
--- a/src/test/java/com/google/devtools/build/lib/analysis/AspectTest.java
+++ b/src/test/java/com/google/devtools/build/lib/analysis/AspectTest.java
@@ -614,7 +614,7 @@
 
       // Get owners of all extra-action artifacts.
       List<Label> extraArtifactOwners = new ArrayList<>();
-      for (Artifact artifact : analysisResult.getTopLevelArtifactsToOwnerLabels().keySet()) {
+      for (Artifact artifact : analysisResult.getTopLevelArtifactsToOwnerLabels().getArtifacts()) {
         if (artifact.getRootRelativePathString().endsWith(".xa")) {
           extraArtifactOwners.add(artifact.getOwnerLabel());
         }
@@ -628,7 +628,7 @@
 
       // Get owners of all extra-action artifacts.
       List<Label> extraArtifactOwners = new ArrayList<>();
-      for (Artifact artifact : analysisResult.getTopLevelArtifactsToOwnerLabels().keySet()) {
+      for (Artifact artifact : analysisResult.getTopLevelArtifactsToOwnerLabels().getArtifacts()) {
         if (artifact.getRootRelativePathString().endsWith(".xa")) {
           extraArtifactOwners.add(artifact.getOwnerLabel());
         }
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/BuildViewTest.java b/src/test/java/com/google/devtools/build/lib/analysis/BuildViewTest.java
index 91d7438..48e4baf 100644
--- a/src/test/java/com/google/devtools/build/lib/analysis/BuildViewTest.java
+++ b/src/test/java/com/google/devtools/build/lib/analysis/BuildViewTest.java
@@ -1271,7 +1271,7 @@
     AnalysisResult analysisResult = update("//x:a");
 
     List<String> owners = new ArrayList<>();
-    for (Artifact artifact : analysisResult.getTopLevelArtifactsToOwnerLabels().keySet()) {
+    for (Artifact artifact : analysisResult.getTopLevelArtifactsToOwnerLabels().getArtifacts()) {
       if ("xa".equals(artifact.getExtension())) {
         owners.add(artifact.getOwnerLabel().toString());
       }
diff --git a/src/test/java/com/google/devtools/build/lib/skylark/SkylarkDefinedAspectsTest.java b/src/test/java/com/google/devtools/build/lib/skylark/SkylarkDefinedAspectsTest.java
index 7de706c..4f09a47 100644
--- a/src/test/java/com/google/devtools/build/lib/skylark/SkylarkDefinedAspectsTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skylark/SkylarkDefinedAspectsTest.java
@@ -1790,7 +1790,8 @@
         "//test:xxx");
     assertThat(
             Iterables.transform(
-                analysisResult.getTopLevelArtifactsToOwnerLabels().keySet(), Artifact::getFilename))
+                analysisResult.getTopLevelArtifactsToOwnerLabels().getArtifacts(),
+                Artifact::getFilename))
         .contains("file.xa");
   }