Add a helper to translate output paths to archived artifact tree and …

…allow

specifying primary output in `SpawnBuilder` test helper.

PiperOrigin-RevId: 345773408

src/main/java/com/google/devtools/build/lib/actions/Artifact.java

@@ -1121,15 +1121,37 @@ private static ArtifactRoot createRootForArchivedArtifact(
         ArtifactRoot treeArtifactRoot,
         PathFragment derivedPathPrefix,
         PathFragment customDerivedTreeRoot) {
-      Path execRoot = getExecRoot(treeArtifactRoot);
+      return ArtifactRoot.asDerivedRoot(
+          getExecRoot(treeArtifactRoot),
+          // e.g. bazel-out/{customDerivedTreeRoot}/k8-fastbuild/bin
+          getExecPathWithinCustomDerivedRoot(
+              derivedPathPrefix, customDerivedTreeRoot, treeArtifactRoot.getExecPath()));
+    }
 
-      // bazel-out/k8-fastbuild/bin -> bazel-out/{customDerivedTreeRoot}/k8-fastbuild/bin
-      PathFragment rootExecPath =
-          derivedPathPrefix
-              .getRelative(customDerivedTreeRoot)
-              .getRelative(treeArtifactRoot.getExecPath().relativeTo(derivedPathPrefix));
+    /**
+     * Returns an exec path within the archived artifacts directory tree corresponding to the
+     * provided one.
+     *
+     * <p>Example: {@code bazel-out/k8-fastbuild/bin ->
+     * bazel-out/{customDerivedTreeRoot}/k8-fastbuild/bin}.
+     */
+    public static PathFragment getExecPathWithinArchivedArtifactsTree(
+        PathFragment derivedPathPrefix, PathFragment execPath) {
+      return getExecPathWithinCustomDerivedRoot(
+          derivedPathPrefix, ARCHIVED_ARTIFACTS_DERIVED_TREE_ROOT, execPath);
+    }
 
-      return ArtifactRoot.asDerivedRoot(execRoot, rootExecPath);
+    /**
+     * Translates provided output {@code execPath} to one under provided derived tree root.
+     *
+     * <p>Example: {@code bazel-out/k8-fastbuild/bin ->
+     * bazel-out/{customDerivedTreeRoot}/k8-fastbuild/bin}.
+     */
+    private static PathFragment getExecPathWithinCustomDerivedRoot(
+        PathFragment derivedPathPrefix, PathFragment customDerivedTreeRoot, PathFragment execPath) {
+      return derivedPathPrefix
+          .getRelative(customDerivedTreeRoot)
+          .getRelative(execPath.relativeTo(derivedPathPrefix));
     }
 
     private static Path getExecRoot(ArtifactRoot artifactRoot) {

src/test/java/com/google/devtools/build/lib/actions/ArtifactTest.java

@@ -610,6 +610,26 @@ public void archivedTreeArtifact_codec_roundTripsArchivedArtifact() throws Excep
         .runTests();
   }
 
+  @Test
+  public void archivedTreeArtifact_getExecPathWithinArchivedArtifactsTree_returnsCorrectPath() {
+    assertThat(
+            ArchivedTreeArtifact.getExecPathWithinArchivedArtifactsTree(
+                PathFragment.create("bazel-out"),
+                PathFragment.create("bazel-out/k8-fastbuild/bin/dir/subdir")))
+        .isEqualTo(
+            PathFragment.create("bazel-out/:archived_tree_artifacts/k8-fastbuild/bin/dir/subdir"));
+  }
+
+  @Test
+  public void archivedTreeArtifact_getExecPathWithinArchivedArtifactsTree_wrongPrefix_fails() {
+    assertThrows(
+        IllegalArgumentException.class,
+        () ->
+            ArchivedTreeArtifact.getExecPathWithinArchivedArtifactsTree(
+                PathFragment.create("wrongPrefix"),
+                PathFragment.create("bazel-out/k8-fastbuild/bin/dir/subdir")));
+  }
+
   private static SpecialArtifact createTreeArtifact(ArtifactRoot root, String relativePath) {
     return createTreeArtifact(root, relativePath, ActionsTestUtil.NULL_ACTION_LOOKUP_DATA);
   }

src/test/java/com/google/devtools/build/lib/exec/util/FakeOwner.java

@@ -14,6 +14,7 @@
 package com.google.devtools.build.lib.exec.util;
 
 import static com.google.common.base.Preconditions.checkNotNull;
+import static com.google.common.base.Preconditions.checkState;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -42,25 +43,34 @@ public class FakeOwner implements ActionExecutionMetadata {
   private final String mnemonic;
   private final String progressMessage;
   @Nullable private final String ownerLabel;
+  @Nullable private final Artifact primaryOutput;
   @Nullable private final PlatformInfo platform;
   private final ImmutableMap<String, String> execProperties;
 
   FakeOwner(
       String mnemonic,
       String progressMessage,
       String ownerLabel,
+      @Nullable Artifact primaryOutput,
       @Nullable PlatformInfo platform,
       ImmutableMap<String, String> execProperties) {
     this.mnemonic = mnemonic;
     this.progressMessage = progressMessage;
     this.ownerLabel = checkNotNull(ownerLabel);
+    this.primaryOutput = primaryOutput;
     this.platform = platform;
     this.execProperties = execProperties;
   }
 
   private FakeOwner(
       String mnemonic, String progressMessage, String ownerLabel, @Nullable PlatformInfo platform) {
-    this(mnemonic, progressMessage, ownerLabel, platform, ImmutableMap.of());
+    this(
+        mnemonic,
+        progressMessage,
+        ownerLabel,
+        /*primaryOutput=*/ null,
+        platform,
+        ImmutableMap.of());
   }
 
   public FakeOwner(String mnemonic, String progressMessage, String ownerLabel) {
@@ -141,7 +151,8 @@ public Artifact getPrimaryInput() {
 
   @Override
   public Artifact getPrimaryOutput() {
-    throw new UnsupportedOperationException();
+    checkState(primaryOutput != null, "primaryOutput not set");
+    return primaryOutput;
   }
 
   @Override

src/test/java/com/google/devtools/build/lib/exec/util/SpawnBuilder.java

@@ -44,6 +44,7 @@ public final class SpawnBuilder {
   private String mnemonic = "Mnemonic";
   private String progressMessage = "progress message";
   private String ownerLabel = "//dummy:label";
+  @Nullable private Artifact ownerPrimaryOutput;
   @Nullable private PlatformInfo platform;
   private final List<String> args;
   private final Map<String, String> environment = new HashMap<>();
@@ -64,7 +65,8 @@ public SpawnBuilder(String... args) {
 
   public Spawn build() {
     ActionExecutionMetadata owner =
-        new FakeOwner(mnemonic, progressMessage, ownerLabel, platform, execProperties);
+        new FakeOwner(
+            mnemonic, progressMessage, ownerLabel, ownerPrimaryOutput, platform, execProperties);
     return new SimpleSpawn(
         owner,
         ImmutableList.copyOf(args),
@@ -98,6 +100,11 @@ public SpawnBuilder withOwnerLabel(String ownerLabel) {
     return this;
   }
 
+  public SpawnBuilder withOwnerPrimaryOutput(Artifact output) {
+    ownerPrimaryOutput = checkNotNull(output);
+    return this;
+  }
+
   public SpawnBuilder withEnvironment(String key, String value) {
     this.environment.put(key, value);
     return this;