Prohibit tree file artifacts from being stored in both artifactData and treeArtifactData in ActionExecutionValue.

The new method isChildOfDeclaredDirectory is used to distinguish between undeclared tree file artifacts and action template expansion outputs. The former should only be stored as part of the enclosing TreeArtifactValue.

If you have any suggestions for the method name that might be better than isChildOfDeclaredDirectory, please let me know.

There are some additional simplifications that can be made in followups.

RELNOTES: None.
PiperOrigin-RevId: 313466599
diff --git a/src/main/java/com/google/devtools/build/lib/actions/Artifact.java b/src/main/java/com/google/devtools/build/lib/actions/Artifact.java
index 4e99b2f..1349c71 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/Artifact.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/Artifact.java
@@ -610,6 +610,19 @@
   }
 
   /**
+   * Returns {@code true} if this is a {@link TreeFileArtifact} that was created by an action which
+   * declared an output directory, as opposed to an action that was generated by an action template
+   * expansion.
+   *
+   * <p>Such artifacts should always be stored within a {@link
+   * com.google.devtools.build.lib.skyframe.TreeArtifactValue} representing the declared directory
+   * and all children, not individually like other derived artifacts.
+   */
+  public boolean isChildOfDeclaredDirectory() {
+    return false;
+  }
+
+  /**
    * Returns whether the artifact represents a Fileset.
    *
    * <p>if true, this artifact is necessarily a {@link SpecialArtifact} with type {@link
@@ -796,12 +809,14 @@
    * <ol>
    *   <li>Outputs under a directory created by an action using {@code declare_directory}. In this
    *       case, a single action creates both the parent and all of the children. Instances should
-   *       be created by calling {@link #createTreeOutput}.
+   *       be created by calling {@link #createTreeOutput}. {@link #isChildOfDeclaredDirectory} will
+   *       return {@code true}.
    *   <li>Outputs of an action template expansion. In this case, the parent directory is not
    *       actually produced by any action, but rather serves as a placeholder for dependant actions
    *       to declare a dep on during analysis, before the children are known. The children are
    *       created by various actions (from the template expansion). Instances should be created by
-   *       calling {@link #createTemplateExpansionOutput}.
+   *       calling {@link #createTemplateExpansionOutput}. {@link #isChildOfDeclaredDirectory} will
+   *       return {@code false}.
    * </ol>
    */
   @Immutable
@@ -919,6 +934,11 @@
       return parentRelativePath.getPathString();
     }
 
+    @Override
+    public boolean isChildOfDeclaredDirectory() {
+      return !isActionTemplateExpansionKey(getArtifactOwner());
+    }
+
     private static boolean isActionTemplateExpansionKey(ActionLookupKey key) {
       return SkyFunctions.ACTION_TEMPLATE_EXPANSION.equals(key.functionName());
     }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionValue.java
index 7116a0a..7a48ad9 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionValue.java
@@ -25,7 +25,6 @@
 import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact;
 import com.google.devtools.build.lib.actions.ArtifactOwner;
 import com.google.devtools.build.lib.actions.FileArtifactValue;
-import com.google.devtools.build.lib.actions.FileStateType;
 import com.google.devtools.build.lib.actions.FilesetOutputSymlink;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
@@ -63,9 +62,13 @@
       @Nullable ImmutableList<FilesetOutputSymlink> outputSymlinks,
       @Nullable NestedSet<Artifact> discoveredModules) {
     for (Map.Entry<Artifact, FileArtifactValue> entry : artifactData.entrySet()) {
-      if (entry.getValue().getType() == FileStateType.REGULAR_FILE) {
-        Preconditions.checkArgument(
-            entry.getValue().getDigest() != null, "missing digest for %s", entry.getKey());
+      Preconditions.checkArgument(
+          !entry.getKey().isChildOfDeclaredDirectory(),
+          "%s should only be stored in a TreeArtifactValue",
+          entry.getKey());
+      if (entry.getValue().getType().isFile()) {
+        Preconditions.checkNotNull(
+            entry.getValue().getDigest(), "missing digest for %s", entry.getKey());
       }
     }
 
@@ -75,9 +78,9 @@
         // We should only have RegularFileValue instances in here, but apparently tree artifacts
         // sometimes store their own root directory in here. Sad.
         // https://github.com/bazelbuild/bazel/issues/9058
-        if (file.getValue().getType() == FileStateType.REGULAR_FILE) {
-          Preconditions.checkArgument(
-              file.getValue().getDigest() != null,
+        if (file.getValue().getType().isFile()) {
+          Preconditions.checkNotNull(
+              file.getValue().getDigest(),
               "missing digest for file %s in tree artifact %s",
               file.getKey(),
               tree.getKey());
@@ -161,19 +164,20 @@
   }
 
   /**
-   * @return The map from {@link Artifact}s to the corresponding {@link
-   *     com.google.devtools.build.lib.actions.FileValue}s that would be returned by {@link
-   *     #getArtifactValue}. Primarily needed by {@link FilesystemValueChecker}, also called by
-   *     {@link ArtifactFunction} when aggregating a {@link TreeArtifactValue}.
+   * Returns a map containing all artifacts output by the action, except for tree artifacts which
+   * are accesible via {@link #getAllTreeArtifactValues}.
+   *
+   * <p>Primarily needed by {@link FilesystemValueChecker}. Also called by {@link ArtifactFunction}
+   * when aggregating a {@link TreeArtifactValue} out of action template expansion outputs.
    */
-  Map<Artifact, FileArtifactValue> getAllFileValues() {
+  ImmutableMap<Artifact, FileArtifactValue> getAllFileValues() {
     return artifactData;
   }
 
   /**
-   * @return The map from {@link Artifact}s to the corresponding {@link TreeArtifactValue}s that
-   *     would be returned by {@link #getTreeArtifactValue}. Should only be needed by {@link
-   *     FilesystemValueChecker}.
+   * Returns a map containing all tree artifacts output by the action.
+   *
+   * <p>Should only be needed by {@link FilesystemValueChecker}.
    */
   ImmutableMap<Artifact, TreeArtifactValue> getAllTreeArtifactValues() {
     return treeArtifactData;
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/OutputStore.java b/src/main/java/com/google/devtools/build/lib/skyframe/OutputStore.java
index b230929..5d75eb9 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/OutputStore.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/OutputStore.java
@@ -13,11 +13,14 @@
 // limitations under the License.
 package com.google.devtools.build.lib.skyframe;
 
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Sets;
 import com.google.devtools.build.lib.actions.Artifact;
+import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact;
 import com.google.devtools.build.lib.actions.FileArtifactValue;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
@@ -26,9 +29,13 @@
 /**
  * Storage layer for data associated with outputs of an action.
  *
- * <p>Data is mainly stored in two maps, {@link #artifactData} and {@link #treeArtifactData}, both
- * of which are keyed on an {@link Artifact}. For each of these maps, this class exposes standard
- * methods such as {@code get}, {@code put}, {@code add}, and {@code getAll}.
+ * <p>Data stored in {@link #artifactData} and {@link #treeArtifactData} will be passed along to the
+ * final {@link ActionExecutionValue}.
+ *
+ * <p>Tree file artifacts which should be stored in a {@link TreeArtifactValue} (according to {@link
+ * Artifact#isChildOfDeclaredDirectory}) are temporarily cached in {@link #treeFileCache}, but it is
+ * expected that the final {@link TreeArtifactValue} will eventually be added via {@link
+ * #putTreeArtifactData}.
  *
  * <p>This implementation aggressively stores all data. Subclasses may override mutating methods to
  * avoid storing unnecessary data.
@@ -38,18 +45,23 @@
 
   private final ConcurrentMap<Artifact, FileArtifactValue> artifactData = new ConcurrentHashMap<>();
 
-  private final ConcurrentMap<Artifact, TreeArtifactValue> treeArtifactData =
+  private final ConcurrentMap<SpecialArtifact, TreeArtifactValue> treeArtifactData =
+      new ConcurrentHashMap<>();
+
+  // The keys in this map are all TreeFileArtifact, but the declared type is Artifact to make it
+  // interchangeable with artifactData syntactically.
+  private final ConcurrentMap<Artifact, FileArtifactValue> treeFileCache =
       new ConcurrentHashMap<>();
 
   private final Set<Artifact> injectedFiles = Sets.newConcurrentHashSet();
 
   @Nullable
   final FileArtifactValue getArtifactData(Artifact artifact) {
-    return artifactData.get(artifact);
+    return mapFor(artifact).get(artifact);
   }
 
   void putArtifactData(Artifact artifact, FileArtifactValue value) {
-    artifactData.put(artifact, value);
+    mapFor(artifact).put(artifact, value);
   }
 
   final ImmutableMap<Artifact, FileArtifactValue> getAllArtifactData() {
@@ -61,8 +73,9 @@
     return treeArtifactData.get(artifact);
   }
 
-  void putTreeArtifactData(Artifact artifact, TreeArtifactValue value) {
-    treeArtifactData.put(artifact, value);
+  final void putTreeArtifactData(SpecialArtifact treeArtifact, TreeArtifactValue value) {
+    Preconditions.checkArgument(treeArtifact.isTreeArtifact(), "%s is not a tree artifact");
+    treeArtifactData.put(treeArtifact, value);
   }
 
   /**
@@ -75,7 +88,7 @@
 
   final void injectOutputData(Artifact output, FileArtifactValue artifactValue) {
     injectedFiles.add(output);
-    artifactData.put(output, artifactValue);
+    mapFor(output).put(output, artifactValue);
   }
 
   /** Returns a set that tracks which Artifacts have had metadata injected. */
@@ -87,13 +100,26 @@
   final void clear() {
     artifactData.clear();
     treeArtifactData.clear();
+    treeFileCache.clear();
     injectedFiles.clear();
   }
 
-  /** Clears data about a specific Artifact from this store. */
+  /**
+   * Clears data about a specific artifact from this store.
+   *
+   * <p>If a tree artifact parent is given, it will be cleared from {@link #treeArtifactData} but
+   * its children will remain in {@link #treeFileCache} if present. If a tree artifact child is
+   * given, it will only be removed from {@link #treeFileCache}.
+   */
   final void remove(Artifact artifact) {
-    artifactData.remove(artifact);
-    treeArtifactData.remove(artifact);
+    mapFor(artifact).remove(artifact);
+    if (artifact.isTreeArtifact()) {
+      treeArtifactData.remove(artifact);
+    }
     injectedFiles.remove(artifact);
   }
+
+  private Map<Artifact, FileArtifactValue> mapFor(Artifact artifact) {
+    return artifact.isChildOfDeclaredDirectory() ? treeFileCache : artifactData;
+  }
 }
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/ActionMetadataHandlerTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/ActionMetadataHandlerTest.java
index 82cd4d7..91b15b5 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/ActionMetadataHandlerTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/ActionMetadataHandlerTest.java
@@ -400,6 +400,48 @@
   }
 
   @Test
+  public void injectRemoteTreeFileArtifactMetadata() throws Exception {
+    scratch.file("/output/bin/foo/bar/child1", "child1");
+    scratch.file("/output/bin/foo/bar/child2", "child2");
+    SpecialArtifact treeArtifact =
+        ActionsTestUtil.createTreeArtifactWithGeneratingAction(
+            outputRoot, PathFragment.create("bin/foo/bar"));
+    TreeFileArtifact child1 = TreeFileArtifact.createTreeOutput(treeArtifact, "child1");
+    TreeFileArtifact child2 = TreeFileArtifact.createTreeOutput(treeArtifact, "child2");
+    assertThat(child1.getPath().exists()).isTrue();
+    assertThat(child2.getPath().exists()).isTrue();
+
+    OutputStore store = new OutputStore();
+    ActionMetadataHandler handler =
+        new ActionMetadataHandler(
+            /*inputArtifactData=*/ new ActionInputMap(1),
+            ImmutableMap.of(),
+            /*missingArtifactsAllowed=*/ false,
+            /*outputs=*/ ImmutableList.of(treeArtifact),
+            /*tsgm=*/ null,
+            ArtifactPathResolver.IDENTITY,
+            store,
+            outputRoot.getRoot().asPath());
+    handler.discardOutputMetadata();
+
+    RemoteFileArtifactValue child1Value = new RemoteFileArtifactValue(new byte[] {1, 2, 3}, 5, 1);
+    RemoteFileArtifactValue child2Value = new RemoteFileArtifactValue(new byte[] {4, 5, 6}, 10, 1);
+
+    handler.injectRemoteFile(child1, child1Value);
+    handler.injectRemoteFile(child2, child2Value);
+
+    FileArtifactValue treeMetadata = handler.getMetadata(treeArtifact);
+    FileArtifactValue child1Metadata = handler.getMetadata(child1);
+    FileArtifactValue child2Metadata = handler.getMetadata(child2);
+    TreeArtifactValue tree = store.getTreeArtifactData(treeArtifact);
+
+    assertThat(tree.getMetadata()).isEqualTo(treeMetadata);
+    assertThat(tree.getChildValues())
+        .containsExactly(child1, child1Metadata, child2, child2Metadata);
+    assertThat(store.getAllArtifactData()).isEmpty(); // All data should be in treeArtifactData.
+  }
+
+  @Test
   public void injectRemoteTreeArtifactMetadata() throws Exception {
     PathFragment path = PathFragment.create("bin/dir");
     SpecialArtifact treeArtifact =
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/TreeArtifactMetadataTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/TreeArtifactMetadataTest.java
index ef80a1e..c158fc9 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/TreeArtifactMetadataTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/TreeArtifactMetadataTest.java
@@ -70,8 +70,7 @@
 import org.junit.runners.JUnit4;
 
 /**
- * Test the behavior of ActionMetadataHandler and ArtifactFunction
- * with respect to TreeArtifacts.
+ * Test the behavior of ActionMetadataHandler and ArtifactFunction with respect to TreeArtifacts.
  */
 @RunWith(JUnit4.class)
 public class TreeArtifactMetadataTest extends ArtifactFunctionTestCase {
@@ -84,9 +83,8 @@
     delegateActionExecutionFunction = new TreeArtifactExecutionFunction();
   }
 
-  private TreeArtifactValue evaluateTreeArtifact(Artifact treeArtifact,
-      Iterable<PathFragment> children)
-    throws Exception {
+  private TreeArtifactValue evaluateTreeArtifact(
+      Artifact treeArtifact, Iterable<PathFragment> children) throws Exception {
     testTreeArtifactContents = ImmutableList.copyOf(children);
     for (PathFragment child : children) {
       file(treeArtifact.getPath().getRelative(child), child.toString());
@@ -155,13 +153,14 @@
         ImmutableList.of(PathFragment.create("one"), PathFragment.create("two"));
     TreeArtifactValue valueOne = evaluateTreeArtifact(treeArtifact, children);
     MemoizingEvaluator evaluator = driver.getGraphForTesting();
-    evaluator.delete(new Predicate<SkyKey>() {
-      @Override
-      public boolean apply(SkyKey key) {
-        // Delete action execution node to force our artifacts to be re-evaluated.
-        return actions.contains(key.argument());
-      }
-    });
+    evaluator.delete(
+        new Predicate<SkyKey>() {
+          @Override
+          public boolean apply(SkyKey key) {
+            // Delete action execution node to force our artifacts to be re-evaluated.
+            return actions.contains(key.argument());
+          }
+        });
     TreeArtifactValue valueTwo = evaluateTreeArtifact(treeArtifact, children);
     assertThat(valueOne.getDigest()).isNotSameInstanceAs(valueTwo.getDigest());
     assertThat(valueOne).isEqualTo(valueTwo);
@@ -197,8 +196,8 @@
   }
 
   /**
-   * Tests that ArtifactFunction rethrows transitive {@link IOException}s as
-   * {@link MissingInputFileException}s.
+   * Tests that ArtifactFunction rethrows transitive {@link IOException}s as {@link
+   * MissingInputFileException}s.
    */
   @Test
   public void testIOExceptionEndToEnd() throws Throwable {
@@ -276,7 +275,6 @@
     @Override
     public SkyValue compute(SkyKey skyKey, Environment env)
         throws SkyFunctionException, InterruptedException {
-      Map<Artifact, FileArtifactValue> fileData = new HashMap<>();
       Map<TreeFileArtifact, FileArtifactValue> treeArtifactData = new HashMap<>();
       ActionLookupData actionLookupData = (ActionLookupData) skyKey.argument();
       ActionLookupValue actionLookupValue =
@@ -295,7 +293,6 @@
           FileArtifactValue withDigest =
               FileArtifactValue.createFromInjectedDigest(
                   noDigest, path.getDigest(), !output.isConstantMetadata());
-          fileData.put(suboutput, withDigest);
           treeArtifactData.put(suboutput, withDigest);
         } catch (IOException e) {
           throw new SkyFunctionException(e, Transience.TRANSIENT) {};
@@ -305,7 +302,7 @@
       TreeArtifactValue treeArtifactValue = TreeArtifactValue.create(treeArtifactData);
 
       return ActionExecutionValue.create(
-          fileData,
+          /*artifactData=*/ ImmutableMap.of(),
           ImmutableMap.of(output, treeArtifactValue),
           /*outputSymlinks=*/ null,
           /*discoveredModules=*/ null,