Use Artifact#getGeneratingActionKey to avoid creating Artifact nodes in the graph for all "normal" generated artifacts.

Any generated artifact that does not represent multiple other artifacts (aggregating middleman and tree artifacts are the exceptions) can be looked up directly from the generating action's ActionExecutionValue. The artifact's type gives a conservative decision procedure for this: non-middleman, non-tree generated artifacts are definitely ok. Non-aggregating middlemen and tree artifacts that are not produced by template expansion are also ok, but knowing that requires checking more data, so we continue to call into to ArtifactFunction for those artifacts.

Saves 60/4100 post-execution heap on a medium-sized build, or 1.5% of memory:

$ diffheap.pl ~/blaze/histo{_before,}.txt
 objsize  chg   instances       space KB    class name
------------------------------------------------------
      77   +0        -711        -735 KB    [B
      32   +0      -43594       -1362 KB    com.google.devtools.build.lib.actions.FileArtifactValue$RegularFileArtifactValue
      24   +0      -90829       -2128 KB    java.util.ArrayList
     746 -173          -5       -9230 KB    [Ljava.util.concurrent.ConcurrentHashMap$Node;
      32   +0     -309720       -9678 KB    java.util.concurrent.ConcurrentHashMap$Node
      40   +0     -309641      -12095 KB    com.google.devtools.build.skyframe.InMemoryNodeEntry
      90   +0     -400463      -26764 KB    [Ljava.lang.Object;
------------------------------------------------------
 total change:                 -62092 KB

There are three main external changes in this CL:

(1) Every consumer of Artifact metadata must pass all artifacts through ArtifactSkyKey#key and friends, since generated artifacts' keys are no longer necessarily themselves. This is something of a partial rollback of https://github.com/bazelbuild/bazel/commit/bf4123df23b5f93e572cd920f15afba340f92391 (originally unknown commit): see, for example, https://github.com/bazelbuild/bazel/commit/bf4123df23b5f93e572cd920f15afba340f92391#diff-619984696e738a6f3ccb9b3802ab7d90.

(2) Similarly every such consumer must be prepared for the returned SkyValue for an Artifact to be an ActionExecutionValue, not a FileArtifactValue or TreeArtifactValue. This means that the consumer must iterate over the original list of Artifacts, constructing each key on the fly, not the returned map, since the returned map may not have any indication of the original Artifacts. The construction of some FileArtifactValues on the fly here may slightly increase garbage (although not in the Google-internal or common Bazel case, since we store FileArtifactValues directly in ActionExecutionValue there), but it should be dominated by the memory savings. I'm hoping to clean up ActionExecutionValues in a separate change, so that we don't have this overlapping data, although the garbage issue may still remain.

The main complication in (2) is in ActionExecutionFunction, where we need to construct keys for non-mandatory source artifacts.

(3) Action rewinding no longer needs to invalidate ordinary generated Artifact nodes in the graph. This simplifies the resulting graphs, but can complicate the rewinding logic, since some Artifacts will still have nodes while others won't. Instead of tracking on the basis of artifacts, we now track on the basis of actions and artifacts.

While modifying ActionRewindStrategy, I took the liberty of doing some clean-ups: for instance, after https://github.com/bazelbuild/bazel/commit/efb3f1595ee897484c477168b8da42b67602e10e, the HashMultimap<DerivedArtifact, ActionInput> lostInputsByDepOwners was only using its values for a check that was guaranteed to succeed, so I just made it a set.

PiperOrigin-RevId: 252701678
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 9f8680a..1a8b54f 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
@@ -101,6 +101,11 @@
  *   <li>A 'Fileset' special Artifact. This is a legacy type of Artifact and should not be used by
  *       new rule implementations.
  * </ul>
+ *
+ * <p>While Artifact implements {@link SkyKey} for memory-saving purposes, Skyframe requests
+ * involving artifacts should always go through {@link ArtifactSkyKey#key}, {@link
+ * ArtifactSkyKey#mandatoryKey}, or {@link ArtifactSkyKey#mandatoryKeys}, since ordinary derived
+ * artifacts should not be requested directly from Skyframe.
  */
 @Immutable
 public abstract class Artifact
diff --git a/src/main/java/com/google/devtools/build/lib/actions/ArtifactSkyKey.java b/src/main/java/com/google/devtools/build/lib/actions/ArtifactSkyKey.java
index ef66e83..c3b3080 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/ArtifactSkyKey.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/ArtifactSkyKey.java
@@ -15,6 +15,7 @@
 
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Interner;
+import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.concurrent.BlazeInterners;
 import com.google.devtools.build.lib.concurrent.ThreadSafety;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
@@ -39,12 +40,34 @@
     this.artifact = Preconditions.checkNotNull(sourceArtifact);
   }
 
+  /**
+   * Returns a {@link SkyKey} that, when built, will produce this artifact. For mandatory source
+   * artifacts and generated artifacts that may aggregate other artifacts (middleman, since they may
+   * be aggregating middlemen, and tree), returns the artifact itself. For normal generated
+   * artifacts, returns the key of the generating action. For non-mandatory source artifacts,
+   * returns a custom SkyKey.
+   *
+   * <p>Callers should use this method (or the related ones below) in preference to directly
+   * requesting an {@link Artifact} to be built by Skyframe, since ordinary derived artifacts should
+   * never be directly built by Skyframe.
+   */
   @ThreadSafety.ThreadSafe
   public static SkyKey key(Artifact artifact, boolean isMandatory) {
-    if (isMandatory || !artifact.isSourceArtifact()) {
+    if (artifact.isTreeArtifact() || artifact.isMiddlemanArtifact()) {
       return artifact;
     }
-    return create(((Artifact.SourceArtifact) artifact));
+    if (artifact.isSourceArtifact()) {
+      return isMandatory ? artifact : create((Artifact.SourceArtifact) artifact);
+    }
+    return ((Artifact.DerivedArtifact) artifact).getGeneratingActionKey();
+  }
+
+  public static SkyKey mandatoryKey(Artifact artifact) {
+    return key(artifact, /*isMandatory=*/ true);
+  }
+
+  public static Iterable<SkyKey> mandatoryKeys(Iterable<Artifact> artifacts) {
+    return Iterables.transform(artifacts, ArtifactSkyKey::mandatoryKey);
   }
 
   @AutoCodec.VisibleForSerialization
diff --git a/src/main/java/com/google/devtools/build/lib/buildtool/SkyframeBuilder.java b/src/main/java/com/google/devtools/build/lib/buildtool/SkyframeBuilder.java
index d77ef7f..0a8f4ce 100644
--- a/src/main/java/com/google/devtools/build/lib/buildtool/SkyframeBuilder.java
+++ b/src/main/java/com/google/devtools/build/lib/buildtool/SkyframeBuilder.java
@@ -166,17 +166,6 @@
               options.getOptions(KeepGoingOption.class).keepGoing,
               skyframeExecutor);
 
-      Preconditions.checkState(
-          exitCode != null
-              || result.keyNames().size()
-                  == (artifacts.size()
-                      + targetsToBuild.size()
-                      + aspects.size()
-                      + parallelTests.size()),
-          "Build reported as successful but not all artifacts and targets built: %s, %s",
-          result,
-          artifacts);
-
       if (exitCode != null) {
         exitCodes.add(exitCode.orNull());
       }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java
index f0ed9d2..2d5c9ba 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java
@@ -21,7 +21,6 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Maps;
-import com.google.common.collect.Sets;
 import com.google.devtools.build.lib.actions.Action;
 import com.google.devtools.build.lib.actions.ActionCacheChecker.Token;
 import com.google.devtools.build.lib.actions.ActionCompletionEvent;
@@ -124,6 +123,9 @@
       throws ActionExecutionFunctionException, InterruptedException {
     ActionLookupData actionLookupData = (ActionLookupData) skyKey.argument();
     Action action = getActionForLookupData(env, actionLookupData);
+    if (action == null) {
+      return null;
+    }
     skyframeActionExecutor.noteActionEvaluationStarted(actionLookupData, action);
     if (actionDependsOnBuildId(action)) {
       PrecomputedValue.BUILD_ID.get(env);
@@ -196,10 +198,11 @@
       Preconditions.checkState(!env.valuesMissing(), "%s %s", action, state);
     }
     CheckInputResults checkedInputs = null;
-    Iterable<? extends SkyKey> inputDepKeys =
-        toKeys(
-            state.allInputs.getAllInputs(),
-            action.discoversInputs() ? action.getMandatoryInputs() : null);
+    Iterable<Artifact> allInputs = state.allInputs.getAllInputs();
+    @Nullable
+    ImmutableSet<Artifact> mandatoryInputs =
+        action.discoversInputs() ? ImmutableSet.copyOf(action.getMandatoryInputs()) : null;
+    Iterable<? extends SkyKey> inputDepKeys = toKeys(allInputs, mandatoryInputs);
     // Declare deps on known inputs to action. We do this unconditionally to maintain our
     // invariant of asking for the same deps each build.
     Map<SkyKey, ValueOrException2<MissingInputFileException, ActionExecutionException>> inputDeps =
@@ -208,7 +211,7 @@
     try {
       if (previousExecution == null && !state.hasArtifactData()) {
         // Do we actually need to find our metadata?
-        checkedInputs = checkInputs(env, action, inputDeps);
+        checkedInputs = checkInputs(env, action, inputDeps, allInputs, mandatoryInputs);
       }
     } catch (ActionExecutionException e) {
       // Remove action from state map in case it's there (won't be unless it discovers inputs).
@@ -265,7 +268,8 @@
               skyframeDepsResult,
               actionStartTime);
     } catch (LostInputsActionExecutionException e) {
-      return handleLostInputs(e, actionLookupData, action, actionStartTime, env, inputDeps, state);
+      return handleLostInputs(
+          e, actionLookupData, action, actionStartTime, env, inputDeps, allInputs, state);
     } catch (ActionExecutionException e) {
       // Remove action from state map in case it's there (won't be unless it discovers inputs).
       stateMap.remove(action);
@@ -299,7 +303,7 @@
 
   private boolean declareDepsOnLostDiscoveredInputsIfAny(Environment env, Action action)
       throws InterruptedException, ActionExecutionFunctionException {
-    ImmutableList<Artifact> previouslyLostDiscoveredInputs =
+    ImmutableList<SkyKey> previouslyLostDiscoveredInputs =
         skyframeActionExecutor.getLostDiscoveredInputs(action);
     if (previouslyLostDiscoveredInputs != null) {
       Map<SkyKey, ValueOrException2<MissingInputFileException, ActionExecutionException>>
@@ -341,6 +345,7 @@
       long actionStartTime,
       Environment env,
       Map<SkyKey, ValueOrException2<MissingInputFileException, ActionExecutionException>> inputDeps,
+      Iterable<Artifact> allInputs,
       ContinuationState state)
       throws InterruptedException, ActionExecutionFunctionException {
     // Remove action from state map in case it's there (won't be unless it discovers inputs).
@@ -354,7 +359,16 @@
       Set<ActionInput> lostRunfiles = e.getInputOwners().getRunfilesInputsAndOwners();
       if (!lostRunfiles.isEmpty()) {
         try {
-          runfilesDepOwners = getInputDepOwners(env, action, inputDeps, lostRunfiles);
+          runfilesDepOwners =
+              getInputDepOwners(
+                  env,
+                  action,
+                  inputDeps,
+                  allInputs,
+                  action.discoversInputs()
+                      ? ImmutableSet.copyOf(action.getMandatoryInputs())
+                      : null,
+                  lostRunfiles);
         } catch (ActionExecutionException unexpected) {
           // getInputDepOwners should not be able to throw, because it does the same work as
           // checkInputs, so if getInputDepOwners throws then checkInputs should have thrown, and if
@@ -367,19 +381,24 @@
 
       // Collect the set of direct deps of this action which may be responsible for the lost inputs,
       // some of which may be discovered.
-      ImmutableList<Artifact> lostDiscoveredInputs = ImmutableList.of();
+      ImmutableList<SkyKey> lostDiscoveredInputs = ImmutableList.of();
       Iterable<? extends SkyKey> failedActionDeps;
       if (e.isFromInputDiscovery()) {
-        // Lost inputs found during input discovery are necessarily artifacts. These may not be
-        // direct deps yet, but the next time this Skyframe node is evaluated they will be. See
-        // SkyframeActionExecutor's lostDiscoveredInputsMap.
+        // Lost inputs found during input discovery are necessarily ordinary derived artifacts.
+        // Their keys may not be direct deps yet, but the next time this Skyframe node is evaluated
+        // they will be. See SkyframeActionExecutor's lostDiscoveredInputsMap.
         lostDiscoveredInputs =
             e.getLostInputs().values().stream()
                 .map(i -> (Artifact) i)
+                .map(ArtifactSkyKey::mandatoryKey)
                 .collect(ImmutableList.toImmutableList());
         failedActionDeps = lostDiscoveredInputs;
       } else if (state.discoveredInputs != null) {
-        failedActionDeps = Iterables.concat(inputDeps.keySet(), state.discoveredInputs);
+        failedActionDeps =
+            Iterables.concat(
+                inputDeps.keySet(),
+                Iterables.transform(
+                    state.discoveredInputs, a -> ArtifactSkyKey.key(a, /*isMandatory=*/ false)));
       } else {
         failedActionDeps = inputDeps.keySet();
       }
@@ -421,16 +440,14 @@
     }
   }
 
+  @Nullable
   static Action getActionForLookupData(Environment env, ActionLookupData actionLookupData)
       throws InterruptedException {
-    // Because of the phase boundary separating analysis and execution, all needed
-    // ActionLookupValues must have already been evaluated.
     ActionLookupValue actionLookupValue =
-        Preconditions.checkNotNull(
-            (ActionLookupValue) env.getValue(actionLookupData.getActionLookupKey()),
-            "ActionLookupValue missing: %s",
-            actionLookupData);
-    return actionLookupValue.getAction(actionLookupData.getActionIndex());
+        ArtifactFunction.getActionLookupValue(actionLookupData.getActionLookupKey(), env);
+    return actionLookupValue != null
+        ? actionLookupValue.getAction(actionLookupData.getActionIndex())
+        : null;
   }
 
   /**
@@ -819,18 +836,24 @@
     if (nonMandatoryDiscovered.isEmpty()) {
       return DiscoveredState.NO_DISCOVERED_DATA;
     }
-    for (Map.Entry<SkyKey, SkyValue> entry : nonMandatoryDiscovered.entrySet()) {
-      Artifact input = ArtifactSkyKey.artifact(entry.getKey());
-      if (entry.getValue() instanceof TreeArtifactValue) {
-        TreeArtifactValue treeValue = (TreeArtifactValue) entry.getValue();
+    for (Artifact input : discoveredInputs) {
+      SkyValue retrievedMetadata =
+          nonMandatoryDiscovered.get(ArtifactSkyKey.key(input, /*isMandatory=*/ false));
+      if (retrievedMetadata instanceof TreeArtifactValue) {
+        TreeArtifactValue treeValue = (TreeArtifactValue) retrievedMetadata;
         expandedArtifacts.put(input, ImmutableSet.copyOf(treeValue.getChildren()));
         for (Map.Entry<Artifact.TreeFileArtifact, FileArtifactValue> child :
             treeValue.getChildValues().entrySet()) {
           inputData.putWithNoDepOwner(child.getKey(), child.getValue());
         }
         inputData.putWithNoDepOwner(input, treeValue.getSelfData());
+      } else if (retrievedMetadata instanceof ActionExecutionValue) {
+        inputData.putWithNoDepOwner(
+            input,
+            ArtifactFunction.createSimpleFileArtifactValue(
+                (Artifact.DerivedArtifact) input, (ActionExecutionValue) retrievedMetadata));
       } else {
-        inputData.putWithNoDepOwner(input, (FileArtifactValue) entry.getValue());
+        inputData.putWithNoDepOwner(input, (FileArtifactValue) retrievedMetadata);
       }
     }
     return DiscoveredState.DISCOVERED_DATA;
@@ -868,20 +891,25 @@
   }
 
   private static Iterable<? extends SkyKey> toKeys(
-      Iterable<Artifact> inputs, Iterable<Artifact> mandatoryInputs) {
+      Iterable<Artifact> inputs, @Nullable ImmutableSet<Artifact> mandatoryInputs) {
     if (mandatoryInputs == null) {
       // This is a non inputs-discovering action, so no need to distinguish mandatory from regular
       // inputs.
-      return inputs;
+      return ArtifactSkyKey.mandatoryKeys(inputs);
     }
     Collection<SkyKey> discoveredArtifacts = new HashSet<>();
-    Set<Artifact> mandatory = Sets.newHashSet(mandatoryInputs);
     for (Artifact artifact : inputs) {
-      discoveredArtifacts.add(ArtifactSkyKey.key(artifact, mandatory.contains(artifact)));
+      discoveredArtifacts.add(toKey(artifact, mandatoryInputs));
     }
     return discoveredArtifacts;
   }
 
+  private static SkyKey toKey(Artifact input, @Nullable Set<Artifact> mandatoryInputs) {
+    return ArtifactSkyKey.key(
+        input,
+        !input.isSourceArtifact() || mandatoryInputs == null || mandatoryInputs.contains(input));
+  }
+
   private static class CheckInputResults {
     /** Metadata about Artifacts consumed by this Action. */
     private final ActionInputMap actionInputMap;
@@ -914,9 +942,18 @@
   private static CheckInputResults checkInputs(
       Environment env,
       Action action,
-      Map<SkyKey, ValueOrException2<MissingInputFileException, ActionExecutionException>> inputDeps)
+      Map<SkyKey, ValueOrException2<MissingInputFileException, ActionExecutionException>> inputDeps,
+      Iterable<Artifact> allInputs,
+      ImmutableSet<Artifact> mandatoryInputs)
       throws ActionExecutionException, InterruptedException {
-    return accumulateInputs(env, action, inputDeps, ActionInputMap::new, CheckInputResults::new);
+    return accumulateInputs(
+        env,
+        action,
+        inputDeps,
+        allInputs,
+        mandatoryInputs,
+        ActionInputMap::new,
+        CheckInputResults::new);
   }
 
   /**
@@ -926,12 +963,16 @@
       Environment env,
       Action action,
       Map<SkyKey, ValueOrException2<MissingInputFileException, ActionExecutionException>> inputDeps,
+      Iterable<Artifact> allInputs,
+      ImmutableSet<Artifact> mandatoryInputs,
       Collection<ActionInput> lostInputs)
       throws ActionExecutionException, InterruptedException {
     return accumulateInputs(
         env,
         action,
         inputDeps,
+        allInputs,
+        mandatoryInputs,
         ignoredInputDepsSize -> new ActionInputDepOwnerMap(lostInputs),
         (actionInputMapSink, expandedArtifacts, expandedFilesets) -> actionInputMapSink);
   }
@@ -940,6 +981,8 @@
       Environment env,
       Action action,
       Map<SkyKey, ValueOrException2<MissingInputFileException, ActionExecutionException>> inputDeps,
+      Iterable<Artifact> allInputs,
+      ImmutableSet<Artifact> mandatoryInputs,
       IntFunction<S> actionInputMapSinkFactory,
       AccumulateInputResultsFactory<S, R> accumulateInputResultsFactory)
       throws ActionExecutionException, InterruptedException {
@@ -957,11 +1000,14 @@
     Map<Artifact, ImmutableList<FilesetOutputSymlink>> expandedFilesets = new HashMap<>();
 
     ActionExecutionException firstActionExecutionException = null;
-    for (Map.Entry<SkyKey, ValueOrException2<MissingInputFileException, ActionExecutionException>>
-        depsEntry : inputDeps.entrySet()) {
-      Artifact input = ArtifactSkyKey.artifact(depsEntry.getKey());
+    for (Artifact input : allInputs) {
+      ValueOrException2<MissingInputFileException, ActionExecutionException> valueOrException =
+          inputDeps.get(toKey(input, mandatoryInputs));
+      if (valueOrException == null) {
+        continue;
+      }
       try {
-        SkyValue value = depsEntry.getValue().get();
+        SkyValue value = valueOrException.get();
         if (populateInputData) {
           ActionInputMapHelper.addToMap(
               inputArtifactData, expandedArtifacts, expandedFilesets, input, value, env);
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionInputMapHelper.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionInputMapHelper.java
index d5771fc..be4cc99 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ActionInputMapHelper.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ActionInputMapHelper.java
@@ -84,6 +84,12 @@
     } else if (value instanceof TreeArtifactValue) {
       expandTreeArtifactAndPopulateArtifactData(
           key, (TreeArtifactValue) value, expandedArtifacts, inputMap, /*depOwner=*/ key);
+    } else if (value instanceof ActionExecutionValue) {
+      inputMap.put(
+          key,
+          ArtifactFunction.createSimpleFileArtifactValue(
+              (Artifact.DerivedArtifact) key, (ActionExecutionValue) value),
+          key);
     } else {
       Preconditions.checkState(value instanceof FileArtifactValue);
       inputMap.put(key, (FileArtifactValue) value, /*depOwner=*/ key);
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionRewindStrategy.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionRewindStrategy.java
index caa8388..ba82e90 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ActionRewindStrategy.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ActionRewindStrategy.java
@@ -20,7 +20,6 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ConcurrentHashMultiset;
-import com.google.common.collect.HashMultimap;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -37,6 +36,7 @@
 import com.google.devtools.build.lib.actions.ActionInputDepOwners;
 import com.google.devtools.build.lib.actions.ActionLookupData;
 import com.google.devtools.build.lib.actions.Artifact;
+import com.google.devtools.build.lib.actions.ArtifactSkyKey;
 import com.google.devtools.build.lib.actions.LostInputsExecException;
 import com.google.devtools.build.lib.actions.LostInputsExecException.LostInputsActionExecutionException;
 import com.google.devtools.build.lib.bugreport.BugReport;
@@ -45,8 +45,10 @@
 import com.google.devtools.build.skyframe.SkyFunction.Restart;
 import com.google.devtools.build.skyframe.SkyKey;
 import java.util.ArrayDeque;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.logging.Logger;
@@ -107,22 +109,17 @@
     // failedAction, which the caller of getRewindPlan already knows must be restarted).
     ImmutableList.Builder<Action> additionalActionsToRestart = ImmutableList.builder();
 
-    HashMultimap<Artifact.DerivedArtifact, ActionInput> lostInputsByDepOwners =
+    Set<Artifact.DerivedArtifact> lostArtifacts =
         getLostInputsByDepOwners(
             lostInputs,
             lostInputsException.getInputOwners(),
             runfilesDepOwners,
             ImmutableSet.copyOf(failedActionDeps),
-            failedAction);
+            failedAction,
+            lostInputsException);
 
-    for (Map.Entry<Artifact.DerivedArtifact, Collection<ActionInput>> entry :
-        lostInputsByDepOwners.asMap().entrySet()) {
-      Artifact.DerivedArtifact lostArtifact = entry.getKey();
-      checkIfLostArtifactIsSource(
-          lostArtifact, failedAction, lostInputsException, entry.getValue());
-
-      // Note that this artifact must be restarted.
-      rewindGraph.putEdge(actionLookupData, lostArtifact);
+    for (Artifact.DerivedArtifact lostArtifact : lostArtifacts) {
+      SkyKey artifactKey = ArtifactSkyKey.mandatoryKey(lostArtifact);
 
       Map<ActionLookupData, Action> actionMap = getActionsForLostArtifact(lostArtifact, env);
       if (actionMap == null) {
@@ -132,6 +129,10 @@
       }
       ImmutableList<ActionAndLookupData> newlyVisitedActions =
           addArtifactDepsAndGetNewlyVisitedActions(rewindGraph, lostArtifact, actionMap);
+      // Note that artifactKey must be restarted. We do this after
+      // addArtifactDepsAndGetNewlyVisitedActions so that it can track if actions are already known
+      // to be in the graph.
+      rewindGraph.putEdge(actionLookupData, artifactKey);
       additionalActionsToRestart.addAll(actions(newlyVisitedActions));
       checkActions(newlyVisitedActions, env, rewindGraph, additionalActionsToRestart);
     }
@@ -188,40 +189,28 @@
     }
   }
 
-  private void checkIfLostArtifactIsSource(
-      Artifact lostArtifact,
-      Action failedAction,
-      LostInputsActionExecutionException lostInputsException,
-      Collection<ActionInput> associatedLostInputs)
-      throws ActionExecutionException {
-    if (lostArtifact.isSourceArtifact()) {
-      // Rewinding source artifacts is not possible. They should not be losable, but we tolerate
-      // their loss--by failing the build instead of crashing--in case some kind of infrastructure
-      // failure results in their apparent loss.
-      logger.info(
-          String.format(
-              "lostArtifact unexpectedly source. lostArtifact: %s, lostInputs for artifact: %s, "
-                  + "failedAction: %.10000s",
-              lostArtifact, associatedLostInputs, failedAction));
-      // Launder the LostInputs exception as a plain ActionExecutionException so that it may be
-      // processed by SkyframeActionExecutor without short-circuiting.
-      throw new ActionExecutionException(lostInputsException, failedAction, /*catastrophe=*/ false);
-    }
-  }
-
-  private HashMultimap<Artifact.DerivedArtifact, ActionInput> getLostInputsByDepOwners(
+  private static Set<Artifact.DerivedArtifact> getLostInputsByDepOwners(
       ImmutableList<ActionInput> lostInputs,
       LostInputsExecException.InputOwners inputOwners,
       ActionInputDepOwners runfilesDepOwners,
       ImmutableSet<SkyKey> failedActionDeps,
-      Action failedActionForLogging) {
+      Action failedAction,
+      LostInputsActionExecutionException lostInputsException)
+      throws ActionExecutionException {
 
-    HashMultimap<Artifact.DerivedArtifact, ActionInput> lostInputsByDepOwners =
-        HashMultimap.create();
+    Set<Artifact.DerivedArtifact> lostArtifacts = new HashSet<>();
     for (ActionInput lostInput : lostInputs) {
       boolean foundLostInputDepOwner = false;
       Artifact owner = inputOwners.getOwner(lostInput);
 
+      if (owner != null && owner.isSourceArtifact()) {
+        // TODO(mschaller): tighten signatures for InputMappingsSink to make this impossible.
+        BugReport.sendBugReport(
+            new IllegalStateException(
+                "Unexpected source artifact as input owner: " + owner + " " + failedAction));
+        throw new ActionExecutionException(
+            lostInputsException, failedAction, /*catastrophe=*/ false);
+      }
       // Rewinding must invalidate all Skyframe paths from the failed action to the action which
       // generates the lost input. Intermediate nodes not on the shortest path to that action may
       // have values that depend on the output of that action. If these intermediate nodes are not
@@ -231,11 +220,8 @@
       if (owner != null) {
         runfilesTransitiveOwners = runfilesDepOwners.getDepOwners(owner);
         for (Artifact runfilesTransitiveOwner : runfilesTransitiveOwners) {
-          if (failedActionDeps.contains(runfilesTransitiveOwner)) {
-            // The lost input's owning tree artifact or fileset is included in a runfiles middleman
-            // that the action directly depends on.
-            lostInputsByDepOwners.put(
-                (Artifact.DerivedArtifact) runfilesTransitiveOwner, lostInput);
+          if (failedActionDeps.contains(ArtifactSkyKey.mandatoryKey(runfilesTransitiveOwner))) {
+            lostArtifacts.add((Artifact.DerivedArtifact) runfilesTransitiveOwner);
             foundLostInputDepOwner = true;
           }
         }
@@ -243,27 +229,33 @@
 
       Collection<Artifact> runfilesOwners = runfilesDepOwners.getDepOwners(lostInput);
       for (Artifact runfilesOwner : runfilesOwners) {
-        if (failedActionDeps.contains(runfilesOwner)) {
-          lostInputsByDepOwners.put((Artifact.DerivedArtifact) runfilesOwner, lostInput);
+        if (runfilesOwner.isSourceArtifact()) {
+          // TODO(mschaller): tighten signatures for ActionInputMapSink to make this impossible.
+          BugReport.sendBugReport(
+              new IllegalStateException(
+                  "Unexpected source artifact as runfile owner: "
+                      + runfilesOwner
+                      + " "
+                      + failedAction));
+          throw new ActionExecutionException(
+              lostInputsException, failedAction, /*catastrophe=*/ false);
+        }
+        if (failedActionDeps.contains(ArtifactSkyKey.mandatoryKey(runfilesOwner))) {
+          lostArtifacts.add((Artifact.DerivedArtifact) runfilesOwner);
           foundLostInputDepOwner = true;
         }
       }
 
-      if (owner != null && failedActionDeps.contains(owner)) {
+      if (owner != null && failedActionDeps.contains(ArtifactSkyKey.mandatoryKey(owner))) {
         // The lost input is included in a tree artifact or fileset that the action directly depends
         // on.
-        lostInputsByDepOwners.put((Artifact.DerivedArtifact) owner, lostInput);
+        lostArtifacts.add((Artifact.DerivedArtifact) owner);
         foundLostInputDepOwner = true;
       }
 
-      if (failedActionDeps.contains(lostInput)) {
-        Preconditions.checkState(
-            lostInput instanceof Artifact,
-            "unexpected non-artifact lostInput which is a dep of the current action. "
-                + "lostInput: %s, failedAction: %.10000s",
-            lostInput,
-            failedActionForLogging);
-        lostInputsByDepOwners.put((Artifact.DerivedArtifact) lostInput, lostInput);
+      if (lostInput instanceof Artifact
+          && failedActionDeps.contains(ArtifactSkyKey.mandatoryKey((Artifact) lostInput))) {
+        lostArtifacts.add((Artifact.DerivedArtifact) lostInput);
         foundLostInputDepOwner = true;
       }
 
@@ -279,14 +271,14 @@
       //
       // In other cases, such as with bugs, when the action fails enough it will cause a crash in
       // checkIfActionLostInputTooManyTimes. We log that this has occurred.
-      logger.info(
+      logger.warning(
           String.format(
               "lostInput not a dep of the failed action, and can't be associated with such a dep. "
                   + "lostInput: %s, owner: %s, runfilesOwners: %s, runfilesTransitiveOwners:"
                   + " %s, failedAction: %.10000s",
-              lostInput, owner, runfilesOwners, runfilesTransitiveOwners, failedActionForLogging));
+              lostInput, owner, runfilesOwners, runfilesTransitiveOwners, failedAction));
     }
-    return lostInputsByDepOwners;
+    return lostArtifacts;
   }
 
   /**
@@ -306,23 +298,36 @@
       ActionAndLookupData actionAndLookupData = uncheckedActions.removeFirst();
       ActionLookupData actionKey = actionAndLookupData.lookupData();
       Action action = actionAndLookupData.action();
-      HashSet<Artifact.DerivedArtifact> artifactsToCheck = new HashSet<>();
+      ArrayList<Artifact.DerivedArtifact> artifactsToCheck = new ArrayList<>();
+      ArrayList<ActionLookupData> newlyDiscoveredActions = new ArrayList<>();
 
       if (action instanceof SkyframeAwareAction) {
         // This action depends on more than just its input artifact values. We need to also restart
-        // the Skyframe subgraph it depends on, up to and including any artifacts.
-        artifactsToCheck.addAll(
-            addSkyframeAwareDepsAndGetNewlyVisitedArtifacts(
-                rewindGraph, actionKey, (SkyframeAwareAction) action));
+        // the Skyframe subgraph it depends on, up to and including any artifacts, which may
+        // aggregate multiple actions.
+        addSkyframeAwareDepsAndGetNewlyVisitedArtifactsAndActions(
+            rewindGraph,
+            actionKey,
+            (SkyframeAwareAction) action,
+            artifactsToCheck,
+            newlyDiscoveredActions);
       }
 
       if (action.mayInsensitivelyPropagateInputs()) {
         // Restarting this action won't recreate the missing input. We need to also restart this
         // action's non-source inputs and the actions which created those inputs.
-        artifactsToCheck.addAll(
-            addPropagatingActionDepsAndGetNewlyVisitedArtifacts(rewindGraph, actionKey, action));
+        addPropagatingActionDepsAndGetNewlyVisitedArtifactsAndActions(
+            rewindGraph, actionKey, action, artifactsToCheck, newlyDiscoveredActions);
       }
 
+      for (ActionLookupData actionLookupData : newlyDiscoveredActions) {
+        Action additionalAction =
+            Preconditions.checkNotNull(
+                ActionExecutionFunction.getActionForLookupData(env, actionLookupData),
+                actionLookupData);
+        additionalActionsToRestart.add(additionalAction);
+        uncheckedActions.add(ActionAndLookupData.create(actionLookupData, additionalAction));
+      }
       for (Artifact.DerivedArtifact artifact : artifactsToCheck) {
         Map<ActionLookupData, Action> actionMap = getActionsForLostArtifact(artifact, env);
         if (actionMap == null) {
@@ -338,56 +343,67 @@
 
   /**
    * For a {@link SkyframeAwareAction} {@code action} with key {@code actionKey}, add its Skyframe
-   * subgraph to {@code rewindGraph}.
-   *
-   * <p>Returns a list of artifacts which were newly added to the graph.
+   * subgraph to {@code rewindGraph}, any {@link Artifact}s to {@code newlyVisitedArtifacts}, and
+   * any {@link ActionLookupData}s to {@code newlyVisitedActions}.
    */
-  private ImmutableList<Artifact.DerivedArtifact> addSkyframeAwareDepsAndGetNewlyVisitedArtifacts(
-      MutableGraph<SkyKey> rewindGraph, ActionLookupData actionKey, SkyframeAwareAction action) {
+  private static void addSkyframeAwareDepsAndGetNewlyVisitedArtifactsAndActions(
+      MutableGraph<SkyKey> rewindGraph,
+      ActionLookupData actionKey,
+      SkyframeAwareAction<?> action,
+      ArrayList<Artifact.DerivedArtifact> newlyVisitedArtifacts,
+      ArrayList<ActionLookupData> newlyVisitedActions) {
 
     ImmutableGraph<SkyKey> graph = action.getSkyframeDependenciesForRewinding(actionKey);
     if (graph.nodes().isEmpty()) {
-      return ImmutableList.of();
+      return;
     }
     assertSkyframeAwareRewindingGraph(graph, actionKey);
 
-    ImmutableList.Builder<Artifact.DerivedArtifact> newlyVisitedArtifacts = ImmutableList.builder();
     Set<EndpointPair<SkyKey>> edges = graph.edges();
     for (EndpointPair<SkyKey> edge : edges) {
       SkyKey target = edge.target();
       if (target instanceof Artifact && rewindGraph.addNode(target)) {
         newlyVisitedArtifacts.add(((Artifact.DerivedArtifact) target));
       }
+      if (target instanceof ActionLookupData && rewindGraph.addNode(target)) {
+        newlyVisitedActions.add(((ActionLookupData) target));
+      }
       rewindGraph.putEdge(edge.source(), edge.target());
     }
-    return newlyVisitedArtifacts.build();
   }
 
   /**
    * For a propagating {@code action} with key {@code actionKey}, add its generated inputs' keys to
-   * {@code rewindGraph} and add edges from {@code actionKey} to those keys.
-   *
-   * <p>Returns a list of artifacts which were newly added to the graph.
+   * {@code rewindGraph}, add edges from {@code actionKey} to those keys, add any {@link Artifact}s
+   * to {@code newlyVisitedArtifacts}, and add any {@link ActionLookupData}s to {@code
+   * newlyVisitedActions}.
    */
-  private ImmutableList<Artifact.DerivedArtifact>
-      addPropagatingActionDepsAndGetNewlyVisitedArtifacts(
-          MutableGraph<SkyKey> rewindGraph, ActionLookupData actionKey, Action action) {
+  private static void addPropagatingActionDepsAndGetNewlyVisitedArtifactsAndActions(
+      MutableGraph<SkyKey> rewindGraph,
+      ActionLookupData actionKey,
+      Action action,
+      ArrayList<Artifact.DerivedArtifact> newlyVisitedArtifacts,
+      ArrayList<ActionLookupData> newlyVisitedActions) {
 
-    ImmutableList.Builder<Artifact.DerivedArtifact> newlyVisitedArtifacts = ImmutableList.builder();
     for (Artifact input : action.getInputs()) {
       if (input.isSourceArtifact()) {
         continue;
       }
+      SkyKey artifactKey = ArtifactSkyKey.mandatoryKey(input);
       // Restarting all derived inputs of propagating actions is overkill. Preferably, we'd want
       // to only restart the inputs which correspond to the known lost outputs. The information
-      // to do this is probably present in the ActionInputs contained in getRewindPlan's
-      // lostInputsByOwners.
+      // to do this is probably present in the data available to #getRewindPlan.
       //
       // Rewinding is expected to be rare, so refining this may not be necessary.
-      if (rewindGraph.addNode(input)) {
-        newlyVisitedArtifacts.add((Artifact.DerivedArtifact) input);
+      boolean newlyVisited = rewindGraph.addNode(artifactKey);
+      if (newlyVisited) {
+        if (artifactKey instanceof Artifact) {
+          newlyVisitedArtifacts.add((Artifact.DerivedArtifact) artifactKey);
+        } else if (artifactKey instanceof ActionLookupData) {
+          newlyVisitedActions.add((ActionLookupData) artifactKey);
+        }
       }
-      rewindGraph.putEdge(actionKey, input);
+      rewindGraph.putEdge(actionKey, artifactKey);
     }
 
     // Rewinding ignores artifacts returned by Action#getAllowedDerivedInputs because:
@@ -402,8 +418,6 @@
                       + "%.10000s",
                   actionKey, action)));
     }
-
-    return newlyVisitedArtifacts.build();
   }
 
   /**
@@ -420,12 +434,15 @@
 
     ImmutableList.Builder<ActionAndLookupData> newlyVisitedActions =
         ImmutableList.builderWithExpectedSize(actionMap.size());
+    SkyKey artifactKey = ArtifactSkyKey.mandatoryKey(artifact);
     for (Map.Entry<ActionLookupData, Action> actionEntry : actionMap.entrySet()) {
       ActionLookupData actionKey = actionEntry.getKey();
       if (rewindGraph.addNode(actionKey)) {
         newlyVisitedActions.add(ActionAndLookupData.create(actionKey, actionEntry.getValue()));
       }
-      rewindGraph.putEdge(artifact, actionKey);
+      if (!artifactKey.equals(actionKey)) {
+        rewindGraph.putEdge(artifactKey, actionKey);
+      }
     }
     return newlyVisitedActions.build();
   }
@@ -446,7 +463,9 @@
     Map<ActionLookupData, Action> actions =
         Maps.newHashMapWithExpectedSize(actionExecutionDeps.size());
     for (ActionLookupData dep : actionExecutionDeps) {
-      actions.put(dep, ActionExecutionFunction.getActionForLookupData(env, dep));
+      actions.put(
+          dep,
+          Preconditions.checkNotNull(ActionExecutionFunction.getActionForLookupData(env, dep)));
     }
     return actions;
   }
@@ -458,24 +477,26 @@
   @Nullable
   private Set<ActionLookupData> getActionExecutionDeps(
       Artifact.DerivedArtifact lostInput, Environment env) throws InterruptedException {
+    if (!lostInput.isTreeArtifact()) {
+      return ImmutableSet.of(lostInput.getGeneratingActionKey());
+    }
     ArtifactFunction.ArtifactDependencies artifactDependencies =
         ArtifactFunction.ArtifactDependencies.discoverDependencies(lostInput, env);
     if (artifactDependencies == null) {
       return null;
     }
 
-    if (artifactDependencies.isTemplateActionForTreeArtifact()) {
-      ArtifactFunction.ActionTemplateExpansion actionTemplateExpansion =
-          artifactDependencies.getActionTemplateExpansion(env);
-      if (actionTemplateExpansion == null) {
-        return null;
-      }
-      // This ignores the ActionTemplateExpansionKey dependency of the template artifact because we
-      // expect to never need to rewind that.
-      return ImmutableSet.copyOf(actionTemplateExpansion.getExpandedActionExecutionKeys());
+    if (!artifactDependencies.isTemplateActionForTreeArtifact()) {
+      return ImmutableSet.of(lostInput.getGeneratingActionKey());
     }
-
-    return ImmutableSet.of(artifactDependencies.getNontemplateActionExecutionKey());
+    ArtifactFunction.ActionTemplateExpansion actionTemplateExpansion =
+        artifactDependencies.getActionTemplateExpansion(env);
+    if (actionTemplateExpansion == null) {
+      return null;
+    }
+    // This ignores the ActionTemplateExpansionKey dependency of the template artifact because we
+    // expect to never need to rewind that.
+    return ImmutableSet.copyOf(actionTemplateExpansion.getExpandedActionExecutionKeys());
   }
 
   private static void assertSkyframeAwareRewindingGraph(
@@ -505,18 +526,6 @@
 
     for (EndpointPair<SkyKey> edge : graph.edges()) {
       SkyKey target = edge.target();
-
-      // Action rewinding could be changed to support finding other actions in a
-      // SkyframeAwareAction's rewinding graph. For now, fail if any are found, because it's
-      // currently impossible and unsupported.
-      Preconditions.checkArgument(
-          !(target instanceof ActionLookupData),
-          "SkyframeAwareAction's rewinding graph contains non-root action node. graph: %s, "
-              + "rootActionNode: %s, unexpectedActionNode: %s",
-          graph,
-          actionKey,
-          target);
-
       Preconditions.checkArgument(
           !(target instanceof Artifact && ((Artifact) target).isSourceArtifact()),
           "SkyframeAwareAction's rewinding graph contains source artifact. graph: %s, "
@@ -581,8 +590,7 @@
     }
   }
 
-  private static ImmutableList<Action> actions(
-      ImmutableList<ActionAndLookupData> newlyVisitedActions) {
+  private static ImmutableList<Action> actions(List<ActionAndLookupData> newlyVisitedActions) {
     return newlyVisitedActions.stream().map(ActionAndLookupData::action).collect(toImmutableList());
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionTemplateExpansionFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionTemplateExpansionFunction.java
index 41bed76..51fa410 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ActionTemplateExpansionFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ActionTemplateExpansionFunction.java
@@ -28,7 +28,6 @@
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact;
 import com.google.devtools.build.lib.actions.ArtifactPrefixConflictException;
-import com.google.devtools.build.lib.actions.ArtifactSkyKey;
 import com.google.devtools.build.lib.actions.MutableActionGraph.ActionConflictException;
 import com.google.devtools.build.lib.bugreport.BugReport;
 import com.google.devtools.build.lib.events.Event;
@@ -72,9 +71,8 @@
     }
     ActionTemplate<?> actionTemplate = value.getActionTemplate(key.getActionIndex());
 
-    // Requests the TreeArtifactValue object for the input TreeArtifact.
-    SkyKey artifactValueKey = ArtifactSkyKey.key(actionTemplate.getInputTreeArtifact(), true);
-    TreeArtifactValue treeArtifactValue = (TreeArtifactValue) env.getValue(artifactValueKey);
+    TreeArtifactValue treeArtifactValue =
+        (TreeArtifactValue) env.getValue(actionTemplate.getInputTreeArtifact());
 
     // Input TreeArtifact is not ready yet.
     if (env.valuesMissing()) {
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ArtifactFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/ArtifactFunction.java
index b2b4d96..0fad11b 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ArtifactFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ArtifactFunction.java
@@ -13,10 +13,8 @@
 // limitations under the License.
 package com.google.devtools.build.lib.skyframe;
 
-import com.google.common.base.Function;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
-import com.google.common.base.Predicate;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
@@ -38,6 +36,7 @@
 import com.google.devtools.build.lib.actions.FilesetTraversalParams.PackageBoundaryMode;
 import com.google.devtools.build.lib.actions.MissingInputFileException;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.EventHandler;
 import com.google.devtools.build.lib.skyframe.RecursiveFilesystemTraversalFunction.RecursiveFilesystemTraversalException;
@@ -57,9 +56,13 @@
 import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
-/** A builder of values for {@link ArtifactSkyKey} keys. */
+/**
+ * A builder of values for {@link ArtifactSkyKey} keys when the key is not a simple generated
+ * artifact. To save memory, ordinary generated artifacts (non-middleman, non-tree) have their
+ * metadata accessed directly from the corresponding {@link ActionExecutionValue}. This SkyFunction
+ * is therefore only usable for source, middleman, and tree artifacts.
+ */
 class ArtifactFunction implements SkyFunction {
-
   private final Supplier<Boolean> mkdirForTreeArtifacts;
 
   public ArtifactFunction(Supplier<Boolean> mkdirForTreeArtifacts) {
@@ -82,9 +85,10 @@
         throw new ArtifactFunctionException(e, Transience.TRANSIENT);
       }
     }
+    Artifact.DerivedArtifact derivedArtifact = (DerivedArtifact) artifact;
 
     ArtifactDependencies artifactDependencies =
-        ArtifactDependencies.discoverDependencies((Artifact.DerivedArtifact) artifact, env);
+        ArtifactDependencies.discoverDependencies(derivedArtifact, env);
     if (artifactDependencies == null) {
       return null;
     }
@@ -98,30 +102,31 @@
       return createTreeArtifactValueFromActionKey(artifactDependencies, env);
     }
 
-    ActionExecutionValue actionValue =
-        (ActionExecutionValue)
-            env.getValue(artifactDependencies.getNontemplateActionExecutionKey());
+    ActionLookupData generatingActionKey = derivedArtifact.getGeneratingActionKey();
+    ActionExecutionValue actionValue = (ActionExecutionValue) env.getValue(generatingActionKey);
     if (actionValue == null) {
       return null;
     }
 
     if (artifact.isTreeArtifact()) {
-      // We get a request for the whole tree artifact. We can just return the associated
+      // We got a request for the whole tree artifact. We can just return the associated
       // TreeArtifactValue.
       return Preconditions.checkNotNull(actionValue.getTreeArtifactValue(artifact), artifact);
     }
 
-    if (artifact.isMiddlemanArtifact()) {
-      Action action =
-          Preconditions.checkNotNull(
-              artifactDependencies.getAction(), "Null middleman action? %s", artifactDependencies);
-      if (isAggregatingValue(action)) {
-        return createAggregatingValue(
-            artifact, action, actionValue.getArtifactValue(artifact), env);
-      }
+    Preconditions.checkState(artifact.isMiddlemanArtifact(), artifact);
+    Action action =
+        Preconditions.checkNotNull(
+            artifactDependencies.actionLookupValue.getAction(generatingActionKey.getActionIndex()),
+            "Null middleman action? %s",
+            artifactDependencies);
+    FileArtifactValue individualMetadata =
+        Preconditions.checkNotNull(
+            actionValue.getArtifactValue(artifact), "%s %s", artifact, actionValue);
+    if (isAggregatingValue(action)) {
+      return createAggregatingValue(artifact, action, individualMetadata, env);
     }
-
-    return createSimpleFileArtifactValue(artifact, actionValue);
+    return individualMetadata;
   }
 
   private static void mkdirForTreeArtifact(Artifact artifact, Environment env)
@@ -141,10 +146,6 @@
 
   private static TreeArtifactValue createTreeArtifactValueFromActionKey(
       ArtifactDependencies artifactDependencies, Environment env) throws InterruptedException {
-    ActionLookupKey actionLookupKey = artifactDependencies.getActionLookupKey();
-    int actionIndex = artifactDependencies.getActionIndex();
-    Artifact treeArtifact = artifactDependencies.getArtifact();
-
     // Request the list of expanded actions from the ActionTemplate.
     ActionTemplateExpansion actionTemplateExpansion =
         artifactDependencies.getActionTemplateExpansion(env);
@@ -165,47 +166,35 @@
     // Aggregate the ArtifactValues for individual TreeFileArtifacts into a TreeArtifactValue for
     // the parent TreeArtifact.
     ImmutableMap.Builder<TreeFileArtifact, FileArtifactValue> map = ImmutableMap.builder();
-    for (int i = 0; i < expandedActionExecutionKeys.size(); i++) {
-      final ActionExecutionValue actionExecutionValue =
+    for (ActionLookupData actionKey : expandedActionExecutionKeys) {
+      ActionExecutionValue actionExecutionValue =
           (ActionExecutionValue)
               Preconditions.checkNotNull(
-                  expandedActionValueMap.get(expandedActionExecutionKeys.get(i)),
-                  "Missing tree value: %s %s %s %s %s",
-                  treeArtifact,
-                  actionLookupKey,
-                  actionIndex,
+                  expandedActionValueMap.get(actionKey),
+                  "Missing tree value: %s %s %s",
+                  artifactDependencies,
                   expansionValue,
                   expandedActionValueMap);
       Iterable<TreeFileArtifact> treeFileArtifacts =
           Iterables.transform(
               Iterables.filter(
                   actionExecutionValue.getAllFileValues().keySet(),
-                  new Predicate<Artifact>() {
-                    @Override
-                    public boolean apply(Artifact artifact) {
-                      Preconditions.checkState(
-                          artifact.hasParent(),
-                          "No parent: %s %s %s %s %s",
-                          artifact,
-                          treeArtifact,
-                          actionExecutionValue,
-                          actionLookupKey,
-                          actionIndex);
-                      return artifact.getParent().equals(treeArtifact);
-                    }
+                  artifact -> {
+                    Preconditions.checkState(
+                        artifact.hasParent(),
+                        "No parent: %s %s %s",
+                        artifact,
+                        actionExecutionValue,
+                        artifactDependencies);
+                    return artifact.getParent().equals(artifactDependencies.artifact);
                   }),
-              new Function<Artifact, TreeFileArtifact>() {
-                @Override
-                public TreeFileArtifact apply(Artifact artifact) {
-                  return (TreeFileArtifact) artifact;
-                }
-              });
+              artifact -> (TreeFileArtifact) artifact);
 
       Preconditions.checkState(
           !Iterables.isEmpty(treeFileArtifacts),
-          "Action denoted by %s does not output TreeFileArtifact under %s",
-          expandedActionExecutionKeys.get(i),
-          treeArtifact);
+          "Action denoted by %s does not output TreeFileArtifact from %s",
+          actionKey,
+          artifactDependencies);
 
       for (TreeFileArtifact treeFileArtifact : treeFileArtifacts) {
         FileArtifactValue value =
@@ -295,17 +284,18 @@
     return ex;
   }
 
-  // Non-aggregating artifact -- should contain at most one piece of artifact data.
-  // data may be null if and only if artifact is a middleman artifact.
-  private static FileArtifactValue createSimpleFileArtifactValue(
-      Artifact artifact, ActionExecutionValue actionValue) {
+  /**
+   * Create {@link FileArtifactValue} for artifact that must be non-middleman non-tree derived
+   * artifact.
+   */
+  static FileArtifactValue createSimpleFileArtifactValue(
+      Artifact.DerivedArtifact artifact, ActionExecutionValue actionValue) {
+    Preconditions.checkState(!artifact.isMiddlemanArtifact(), "%s %s", artifact, actionValue);
+    Preconditions.checkState(!artifact.isTreeArtifact(), "%s %s", artifact, actionValue);
     FileArtifactValue value = actionValue.getArtifactValue(artifact);
     if (value != null) {
       return value;
     }
-    // Middleman artifacts have no corresponding files, so their ArtifactValues should have already
-    // been constructed during execution of the action.
-    Preconditions.checkState(!artifact.isMiddlemanArtifact(), artifact);
     ArtifactFileMetadata data =
         Preconditions.checkNotNull(actionValue.getData(artifact), "%s %s", artifact, actionValue);
     Preconditions.checkNotNull(
@@ -327,14 +317,26 @@
         ImmutableList.builder();
     ImmutableList.Builder<Pair<Artifact, TreeArtifactValue>> directoryInputsBuilder =
         ImmutableList.builder();
-    for (Map.Entry<SkyKey, SkyValue> entry : env.getValues(action.getInputs()).entrySet()) {
-      Artifact input = ArtifactSkyKey.artifact(entry.getKey());
-      SkyValue inputValue = entry.getValue();
-      if (inputValue == null) {
-        return null;
-      }
+    Iterable<Artifact> inputs = action.getInputs();
+    if (inputs instanceof NestedSet) {
+      // Avoid iterating over nested set twice.
+      inputs = ((NestedSet<Artifact>) inputs).toList();
+    }
+    Map<SkyKey, SkyValue> values = env.getValues(ArtifactSkyKey.mandatoryKeys(inputs));
+    if (env.valuesMissing()) {
+      return null;
+    }
+    for (Artifact input : inputs) {
+      SkyValue inputValue =
+          Preconditions.checkNotNull(values.get(ArtifactSkyKey.mandatoryKey(input)), input);
       if (inputValue instanceof FileArtifactValue) {
         fileInputsBuilder.add(Pair.of(input, (FileArtifactValue) inputValue));
+      } else if (inputValue instanceof ActionExecutionValue) {
+        fileInputsBuilder.add(
+            Pair.of(
+                input,
+                createSimpleFileArtifactValue(
+                    (DerivedArtifact) input, (ActionExecutionValue) inputValue)));
       } else if (inputValue instanceof TreeArtifactValue) {
         directoryInputsBuilder.add(Pair.of(input, (TreeArtifactValue) inputValue));
       } else {
@@ -391,16 +393,12 @@
 
   @Nullable
   static ActionLookupValue getActionLookupValue(
-      SkyKey actionLookupKey, SkyFunction.Environment env, Artifact artifact)
-      throws InterruptedException {
+      ActionLookupKey actionLookupKey, SkyFunction.Environment env) throws InterruptedException {
     ActionLookupValue value = (ActionLookupValue) env.getValue(actionLookupKey);
     if (value == null) {
-      ArtifactOwner artifactOwner = artifact.getArtifactOwner();
       Preconditions.checkState(
-          artifactOwner == CoverageReportValue.COVERAGE_REPORT_KEY,
-          "Not-yet-present artifact owner: %s (%s %s)",
-          artifactOwner,
-          artifact,
+          actionLookupKey == CoverageReportValue.COVERAGE_REPORT_KEY,
+          "Not-yet-present artifact owner: %s",
           actionLookupKey);
       return null;
     }
@@ -448,8 +446,7 @@
 
       ActionLookupData generatingActionKey = derivedArtifact.getGeneratingActionKey();
       ActionLookupValue actionLookupValue =
-          ArtifactFunction.getActionLookupValue(
-              generatingActionKey.getActionLookupKey(), env, derivedArtifact);
+          ArtifactFunction.getActionLookupValue(generatingActionKey.getActionLookupKey(), env);
       if (actionLookupValue == null) {
         return null;
       }
@@ -457,26 +454,9 @@
       return new ArtifactDependencies(derivedArtifact, actionLookupValue);
     }
 
-    Artifact getArtifact() {
-      return artifact;
-    }
-
-    ActionLookupKey getActionLookupKey() {
-      return artifact.getArtifactOwner();
-    }
-
-    int getActionIndex() {
-      return artifact.getGeneratingActionKey().getActionIndex();
-    }
-
     boolean isTemplateActionForTreeArtifact() {
-      return artifact.isTreeArtifact() && actionLookupValue.isActionTemplate(getActionIndex());
-    }
-
-    ActionLookupData getNontemplateActionExecutionKey() {
-      Preconditions.checkState(
-          !isTemplateActionForTreeArtifact(), "Action is unexpectedly template: %s", this);
-      return artifact.getGeneratingActionKey();
+      return artifact.isTreeArtifact()
+          && actionLookupValue.isActionTemplate(artifact.getGeneratingActionKey().getActionIndex());
     }
 
     /**
@@ -491,7 +471,8 @@
       Preconditions.checkState(
           isTemplateActionForTreeArtifact(), "Action is unexpectedly non-template: %s", this);
       ActionTemplateExpansionValue.ActionTemplateExpansionKey key =
-          ActionTemplateExpansionValue.key(getActionLookupKey(), getActionIndex());
+          ActionTemplateExpansionValue.key(
+              artifact.getArtifactOwner(), artifact.getGeneratingActionKey().getActionIndex());
       ActionTemplateExpansionValue value = (ActionTemplateExpansionValue) env.getValue(key);
       if (value == null) {
         return null;
@@ -499,10 +480,6 @@
       return new ActionTemplateExpansion(key, value);
     }
 
-    Action getAction() {
-      return actionLookupValue.getAction(getActionIndex());
-    }
-
     @Override
     public String toString() {
       return MoreObjects.toStringHelper(this)
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/CompletionFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/CompletionFunction.java
index 7f4e8b4..2d5ab17 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/CompletionFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/CompletionFunction.java
@@ -349,9 +349,12 @@
       return null;
     }
 
+    // Avoid iterating over nested set twice.
+    ImmutableList<Artifact> allArtifacts =
+        completor.getAllArtifactsToBuild(value, topLevelContext).getAllArtifacts().toList();
     Map<SkyKey, ValueOrException2<MissingInputFileException, ActionExecutionException>> inputDeps =
         env.getValuesOrThrow(
-            completor.getAllArtifactsToBuild(value, topLevelContext).getAllArtifacts(),
+            ArtifactSkyKey.mandatoryKeys(allArtifacts),
             MissingInputFileException.class,
             ActionExecutionException.class);
 
@@ -363,11 +366,9 @@
     ActionExecutionException firstActionExecutionException = null;
     MissingInputFileException missingInputException = null;
     NestedSetBuilder<Cause> rootCausesBuilder = NestedSetBuilder.stableOrder();
-    for (Map.Entry<SkyKey, ValueOrException2<MissingInputFileException, ActionExecutionException>>
-        depsEntry : inputDeps.entrySet()) {
-      Artifact input = ArtifactSkyKey.artifact(depsEntry.getKey());
+    for (Artifact input : allArtifacts) {
       try {
-        SkyValue artifactValue = depsEntry.getValue().get();
+        SkyValue artifactValue = inputDeps.get(ArtifactSkyKey.mandatoryKey(input)).get();
         if (artifactValue != null) {
           ActionInputMapHelper.addToMap(
               inputMap,
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/RecursiveFilesystemTraversalFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/RecursiveFilesystemTraversalFunction.java
index 49f6870..62f894b 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/RecursiveFilesystemTraversalFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/RecursiveFilesystemTraversalFunction.java
@@ -259,6 +259,11 @@
           fsVal = value;
         } else if (value instanceof TreeArtifactValue) {
           fsVal = value;
+        } else if (value instanceof ActionExecutionValue) {
+          fsVal =
+              Preconditions.checkNotNull(
+                  ArtifactFunction.createSimpleFileArtifactValue(
+                      (Artifact.DerivedArtifact) artifact, (ActionExecutionValue) value));
         } else {
           return NON_EXISTENT_FILE_INFO;
         }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java
index f24127a..8f4f07e 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java
@@ -101,6 +101,7 @@
 import com.google.devtools.build.lib.vfs.Symlinks;
 import com.google.devtools.build.skyframe.SkyFunction;
 import com.google.devtools.build.skyframe.SkyFunction.Environment;
+import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.common.options.OptionsProvider;
 import java.io.FileNotFoundException;
 import java.io.IOException;
@@ -182,7 +183,7 @@
   // the lost input's generating action must be rerun before the failed action tries input discovery
   // again. A previously failed action satisfies that requirement by requesting the deps in this map
   // at the start of its next attempt,
-  private ConcurrentMap<OwnerlessArtifactWrapper, ImmutableList<Artifact>> lostDiscoveredInputsMap;
+  private ConcurrentMap<OwnerlessArtifactWrapper, ImmutableList<SkyKey>> lostDiscoveredInputsMap;
 
   // Errors found when examining all actions in the graph are stored here, so that they can be
   // thrown when execution of the action is requested. This field is set during each call to
@@ -501,11 +502,11 @@
   }
 
   @Nullable
-  ImmutableList<Artifact> getLostDiscoveredInputs(Action action) {
+  ImmutableList<SkyKey> getLostDiscoveredInputs(Action action) {
     return lostDiscoveredInputsMap.get(new OwnerlessArtifactWrapper(action.getPrimaryOutput()));
   }
 
-  void resetFailedActionExecution(Action action, ImmutableList<Artifact> lostDiscoveredInputs) {
+  void resetFailedActionExecution(Action action, ImmutableList<SkyKey> lostDiscoveredInputs) {
     OwnerlessArtifactWrapper ownerlessArtifactWrapper =
         new OwnerlessArtifactWrapper(action.getPrimaryOutput());
     buildActionMap.remove(ownerlessArtifactWrapper);
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
index 5b07b3d..7238381 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
@@ -53,6 +53,7 @@
 import com.google.devtools.build.lib.actions.ArtifactPathResolver;
 import com.google.devtools.build.lib.actions.ArtifactResolver.ArtifactResolverSupplier;
 import com.google.devtools.build.lib.actions.ArtifactRoot;
+import com.google.devtools.build.lib.actions.ArtifactSkyKey;
 import com.google.devtools.build.lib.actions.CommandLineExpansionException;
 import com.google.devtools.build.lib.actions.CompletionContext.PathResolverFactory;
 import com.google.devtools.build.lib.actions.EnvironmentalExecException;
@@ -1564,7 +1565,9 @@
               .setEventHander(reporter)
               .build();
       return buildDriver.evaluate(
-          Iterables.concat(artifactsToBuild, targetKeys, aspectKeys, testKeys), evaluationContext);
+          Iterables.concat(
+              ArtifactSkyKey.mandatoryKeys(artifactsToBuild), targetKeys, aspectKeys, testKeys),
+          evaluationContext);
     } finally {
       progressReceiver.executionProgressReceiver = null;
       // Also releases thread locks.
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/ArtifactFunctionTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/ArtifactFunctionTest.java
index 4ba23c3..7d4f3c7 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/ArtifactFunctionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/ArtifactFunctionTest.java
@@ -152,7 +152,10 @@
     actions.add(action);
     file(input2.getPath(), "contents");
     file(input1.getPath(), "source contents");
-    evaluate(Iterables.toArray(ImmutableSet.of(input2, input1, input2, tree), SkyKey.class));
+    evaluate(
+        Iterables.toArray(
+            ArtifactSkyKey.mandatoryKeys(ImmutableSet.of(input2, input1, input2, tree)),
+            SkyKey.class));
     SkyValue value = evaluateArtifactValue(output);
     ArrayList<Pair<Artifact, ?>> inputs = new ArrayList<>();
     inputs.addAll(((AggregatingArtifactValue) value).getFileArtifacts());
@@ -271,9 +274,7 @@
   public void actionExecutionValueSerialization() throws Exception {
     ActionLookupData dummyData = ActionLookupData.create(ALL_OWNER, 0);
     Artifact.DerivedArtifact artifact1 = createDerivedArtifact("one");
-    artifact1.setGeneratingActionKey(dummyData);
     Artifact.DerivedArtifact artifact2 = createDerivedArtifact("two");
-    artifact2.setGeneratingActionKey(dummyData);
     ArtifactFileMetadata metadata1 =
         ActionMetadataHandler.fileMetadataFromArtifact(artifact1, null, null);
     SpecialArtifact treeArtifact = createDerivedTreeArtifactOnly("tree");
@@ -286,7 +287,6 @@
         TreeArtifactValue.create(
             ImmutableMap.of(treeFileArtifact, FileArtifactValue.create(treeFileArtifact)));
     Artifact.DerivedArtifact artifact3 = createDerivedArtifact("three");
-    artifact3.setGeneratingActionKey(dummyData);
     FilesetOutputSymlink filesetOutputSymlink =
         FilesetOutputSymlink.createForTesting(
             PathFragment.EMPTY_FRAGMENT, PathFragment.EMPTY_FRAGMENT, PathFragment.EMPTY_FRAGMENT);
@@ -328,6 +328,7 @@
         new Artifact.DerivedArtifact(
             ArtifactRoot.asDerivedRoot(root, root.getRelative("out")), execPath, ALL_OWNER);
     actions.add(new DummyAction(ImmutableList.<Artifact>of(), output));
+    output.setGeneratingActionKey(ActionLookupData.create(ALL_OWNER, actions.size() - 1));
     return output;
   }
 
@@ -392,7 +393,12 @@
     if (result.hasError()) {
       throw result.getError().getException();
     }
-    return result.get(key);
+    SkyValue value = result.get(key);
+    if (value instanceof ActionExecutionValue) {
+      return ArtifactFunction.createSimpleFileArtifactValue(
+          (Artifact.DerivedArtifact) artifact, (ActionExecutionValue) value);
+    }
+    return value;
   }
 
   private void setGeneratingActions() throws InterruptedException, ActionConflictException {
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/RecursiveFilesystemTraversalFunctionTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/RecursiveFilesystemTraversalFunctionTest.java
index 1fc6492..87f81ac 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/RecursiveFilesystemTraversalFunctionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/RecursiveFilesystemTraversalFunctionTest.java
@@ -14,6 +14,7 @@
 package com.google.devtools.build.lib.skyframe;
 
 import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth.assertWithMessage;
 import static com.google.devtools.build.lib.actions.FilesetTraversalParams.PackageBoundaryMode.CROSS;
 import static com.google.devtools.build.lib.actions.FilesetTraversalParams.PackageBoundaryMode.DONT_CROSS;
 import static com.google.devtools.build.lib.actions.FilesetTraversalParams.PackageBoundaryMode.REPORT_ERROR;
@@ -28,6 +29,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
+import com.google.devtools.build.lib.actions.ActionLookupData;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact;
 import com.google.devtools.build.lib.actions.Artifact.SpecialArtifactType;
@@ -81,7 +83,10 @@
 import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.UUID;
@@ -104,9 +109,11 @@
   private RecordingDifferencer differencer;
   private AtomicReference<PathPackageLocator> pkgLocator;
   private NonHermeticArtifactFakeFunction artifactFunction;
+  private List<Artifact.DerivedArtifact> artifacts;
 
   @Before
   public final void setUp() {
+    artifacts = new ArrayList<>();
     AnalysisMock analysisMock = AnalysisMock.get();
     pkgLocator =
         new AtomicReference<>(
@@ -175,6 +182,7 @@
     // case of a generated directory, which we have test coverage for.
     skyFunctions.put(Artifact.ARTIFACT, new ArtifactFakeFunction());
     artifactFunction = new NonHermeticArtifactFakeFunction();
+    skyFunctions.put(ActionLookupData.NAME, new ActionFakeFunction());
     skyFunctions.put(NONHERMETIC_ARTIFACT, artifactFunction);
 
     progressReceiver = new RecordingEvaluationProgressReceiver();
@@ -215,10 +223,15 @@
 
   private Artifact derivedArtifact(String path) {
     PathFragment execPath = PathFragment.create("out").getRelative(path);
-    Artifact output =
-        ActionsTestUtil.createArtifactWithExecPath(
-            ArtifactRoot.asDerivedRoot(rootDirectory, rootDirectory.getRelative("out")), execPath);
-    return output;
+    Artifact.DerivedArtifact result =
+        (Artifact.DerivedArtifact)
+            ActionsTestUtil.createArtifactWithExecPath(
+                ArtifactRoot.asDerivedRoot(rootDirectory, rootDirectory.getRelative("out")),
+                execPath);
+    result.setGeneratingActionKey(
+        ActionLookupData.create(ActionsTestUtil.NULL_ARTIFACT_OWNER, artifacts.size()));
+    artifacts.add(result);
+    return result;
   }
 
   private static RootedPath rootedPath(Artifact artifact) {
@@ -338,6 +351,9 @@
       Map<PathFragment, ResolvedFile> nameToActualResolvedFiles,
       ResolvedFile... expectedFilesIgnoringMetadata)
       throws Exception {
+    assertWithMessage("Expected files " + Arrays.toString(expectedFilesIgnoringMetadata))
+        .that(nameToActualResolvedFiles)
+        .hasSize(expectedFilesIgnoringMetadata.length);
     assertEquals(
         "Unequal number of ResolvedFiles in Actual and expected.",
         expectedFilesIgnoringMetadata.length,
@@ -381,7 +397,7 @@
     SkyKey key =
         file.isSourceArtifact()
             ? FileStateValue.key(rootedPath(file))
-            : new NonHermeticArtifactSkyKey(ArtifactSkyKey.key(file, true));
+            : new NonHermeticArtifactSkyKey(file);
     appendToFile(rootedPath(file), key, content);
   }
 
@@ -395,8 +411,7 @@
 
   private void invalidateOutputArtifact(Artifact output) {
     assertThat(output.isSourceArtifact()).isFalse();
-    differencer.invalidate(
-        ImmutableList.of(new NonHermeticArtifactSkyKey(ArtifactSkyKey.key(output, true))));
+    differencer.invalidate(ImmutableList.of(new NonHermeticArtifactSkyKey(output)));
   }
 
   private static final class RecordingEvaluationProgressReceiver
@@ -1036,6 +1051,24 @@
     }
   }
 
+  private class ActionFakeFunction implements SkyFunction {
+    @Nullable
+    @Override
+    public SkyValue compute(SkyKey skyKey, Environment env)
+        throws SkyFunctionException, InterruptedException {
+      return env.getValue(
+          new NonHermeticArtifactSkyKey(
+              Preconditions.checkNotNull(
+                  artifacts.get(((ActionLookupData) skyKey).getActionIndex()), skyKey)));
+    }
+
+    @Nullable
+    @Override
+    public String extractTag(SkyKey skyKey) {
+      return null;
+    }
+  }
+
   private static class NonHermeticArtifactSkyKey extends AbstractSkyKey<SkyKey> {
     private NonHermeticArtifactSkyKey(SkyKey arg) {
       super(arg);
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/TimestampBuilderTestCase.java b/src/test/java/com/google/devtools/build/lib/skyframe/TimestampBuilderTestCase.java
index e128bab..63054cb 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/TimestampBuilderTestCase.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/TimestampBuilderTestCase.java
@@ -159,6 +159,11 @@
 
   protected <T extends ActionAnalysisMetadata> T registerAction(T action) {
     actions.add(action);
+    ActionLookupData actionLookupData =
+        ActionLookupData.create(ACTION_LOOKUP_KEY, actions.size() - 1);
+    for (Artifact output : action.getOutputs()) {
+      ((Artifact.DerivedArtifact) output).setGeneratingActionKey(actionLookupData);
+    }
     return action;
   }
 
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/TreeArtifactBuildTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/TreeArtifactBuildTest.java
index 98c0647..8df59da 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/TreeArtifactBuildTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/TreeArtifactBuildTest.java
@@ -194,11 +194,9 @@
   @Test
   public void testCacheCheckingForTreeArtifactsDoesNotCauseReexecution() throws Exception {
     SpecialArtifact outOne = createTreeArtifact("outputOne");
-    outOne.setGeneratingActionKey(ActionLookupData.create(ACTION_LOOKUP_KEY, 0));
     Button buttonOne = new Button();
 
     SpecialArtifact outTwo = createTreeArtifact("outputTwo");
-    outTwo.setGeneratingActionKey(ActionLookupData.create(ACTION_LOOKUP_KEY, 1));
     Button buttonTwo = new Button();
 
     TouchingTestAction actionOne = new TouchingTestAction(
@@ -416,9 +414,12 @@
     reporter.addHandler(storingEventHandler);
 
     SpecialArtifact outOne = createTreeArtifact("outputOne");
-    outOne.setGeneratingActionKey(ActionLookupData.create(ACTION_LOOKUP_KEY, 0));
-    TreeFileArtifact outOneFileOne = ActionInputHelper.treeFileArtifact(outOne, "out_one_file_one");
-    TreeFileArtifact outOneFileTwo = ActionInputHelper.treeFileArtifact(outOne, "out_one_file_two");
+    TreeFileArtifact outOneFileOne =
+        ActionInputHelper.treeFileArtifactWithNoGeneratingActionSet(
+            outOne, PathFragment.create("out_one_file_one"), ACTION_LOOKUP_KEY);
+    TreeFileArtifact outOneFileTwo =
+        ActionInputHelper.treeFileArtifactWithNoGeneratingActionSet(
+            outOne, PathFragment.create("out_one_file_two"), ACTION_LOOKUP_KEY);
     TreeArtifactTestAction failureOne = new TreeArtifactTestAction(
         Runnables.doNothing(), outOneFileOne, outOneFileTwo) {
       @Override
@@ -435,6 +436,8 @@
     };
 
     registerAction(failureOne);
+    outOneFileOne.setGeneratingActionKey(outOne.getGeneratingActionKey());
+    outOneFileTwo.setGeneratingActionKey(outOne.getGeneratingActionKey());
     assertThrows(BuildFailedException.class, () -> buildArtifact(outOne));
     // not all outputs were created
     List<Event> errors =
@@ -444,9 +447,12 @@
     assertThat(errors.get(1).getMessage()).contains("not all outputs were created or valid");
 
     SpecialArtifact outTwo = createTreeArtifact("outputTwo");
-    outTwo.setGeneratingActionKey(ActionLookupData.create(ACTION_LOOKUP_KEY, 1));
-    TreeFileArtifact outTwoFileOne = ActionInputHelper.treeFileArtifact(outOne, "out_one_file_one");
-    TreeFileArtifact outTwoFileTwo = ActionInputHelper.treeFileArtifact(outOne, "out_one_file_two");
+    TreeFileArtifact outTwoFileOne =
+        ActionInputHelper.treeFileArtifactWithNoGeneratingActionSet(
+            outTwo, PathFragment.create("out_two_file_one"), ACTION_LOOKUP_KEY);
+    TreeFileArtifact outTwoFileTwo =
+        ActionInputHelper.treeFileArtifactWithNoGeneratingActionSet(
+            outTwo, PathFragment.create("out_two_file_two"), ACTION_LOOKUP_KEY);
     TreeArtifactTestAction failureTwo = new TreeArtifactTestAction(
         Runnables.doNothing(), outTwoFileOne, outTwoFileTwo) {
       @Override
@@ -465,6 +471,8 @@
     };
 
     registerAction(failureTwo);
+    outTwoFileOne.setGeneratingActionKey(outTwo.getGeneratingActionKey());
+    outTwoFileTwo.setGeneratingActionKey(outTwo.getGeneratingActionKey());
     storingEventHandler.clear();
     assertThrows(BuildFailedException.class, () -> buildArtifact(outTwo));
     errors =
@@ -766,12 +774,15 @@
   public void testExpandedActionsBuildInActionTemplate() throws Throwable {
     // artifact1 is a tree artifact generated by a TouchingTestAction.
     SpecialArtifact artifact1 = createTreeArtifact("treeArtifact1");
-    artifact1.setGeneratingActionKey(ActionLookupData.create(ACTION_LOOKUP_KEY, 0));
-    TreeFileArtifact treeFileArtifactA = ActionInputHelper.treeFileArtifact(
-        artifact1, PathFragment.create("child1"));
-    TreeFileArtifact treeFileArtifactB = ActionInputHelper.treeFileArtifact(
-        artifact1, PathFragment.create("child2"));
+    TreeFileArtifact treeFileArtifactA =
+        ActionInputHelper.treeFileArtifactWithNoGeneratingActionSet(
+            artifact1, PathFragment.create("child1"), ACTION_LOOKUP_KEY);
+    TreeFileArtifact treeFileArtifactB =
+        ActionInputHelper.treeFileArtifactWithNoGeneratingActionSet(
+            artifact1, PathFragment.create("child2"), ACTION_LOOKUP_KEY);
     registerAction(new TouchingTestAction(treeFileArtifactA, treeFileArtifactB));
+    treeFileArtifactA.setGeneratingActionKey(artifact1.getGeneratingActionKey());
+    treeFileArtifactB.setGeneratingActionKey(artifact1.getGeneratingActionKey());
 
     // artifact2 is a tree artifact generated by an action template.
     SpecialArtifact artifact2 = createTreeArtifact("treeArtifact2");
@@ -807,12 +818,15 @@
 
     // artifact1 is a tree artifact generated by a TouchingTestAction.
     SpecialArtifact artifact1 = createTreeArtifact("treeArtifact1");
-    artifact1.setGeneratingActionKey(ActionLookupData.create(ACTION_LOOKUP_KEY, 0));
-    TreeFileArtifact treeFileArtifactA = ActionInputHelper.treeFileArtifact(
-        artifact1, PathFragment.create("child1"));
-    TreeFileArtifact treeFileArtifactB = ActionInputHelper.treeFileArtifact(
-        artifact1, PathFragment.create("child2"));
+    TreeFileArtifact treeFileArtifactA =
+        ActionInputHelper.treeFileArtifactWithNoGeneratingActionSet(
+            artifact1, PathFragment.create("child1"), ACTION_LOOKUP_KEY);
+    TreeFileArtifact treeFileArtifactB =
+        ActionInputHelper.treeFileArtifactWithNoGeneratingActionSet(
+            artifact1, PathFragment.create("child2"), ACTION_LOOKUP_KEY);
     registerAction(new TouchingTestAction(treeFileArtifactA, treeFileArtifactB));
+    treeFileArtifactA.setGeneratingActionKey(artifact1.getGeneratingActionKey());
+    treeFileArtifactB.setGeneratingActionKey(artifact1.getGeneratingActionKey());
 
     // artifact2 is a tree artifact generated by an action template.
     SpecialArtifact artifact2 = createTreeArtifact("treeArtifact2");
@@ -852,12 +866,15 @@
 
     // artifact1 is a tree artifact generated by a TouchingTestAction.
     SpecialArtifact artifact1 = createTreeArtifact("treeArtifact1");
-    artifact1.setGeneratingActionKey(ActionLookupData.create(ACTION_LOOKUP_KEY, 0));
-    TreeFileArtifact treeFileArtifactA = ActionInputHelper.treeFileArtifact(
-        artifact1, PathFragment.create("child1"));
-    TreeFileArtifact treeFileArtifactB = ActionInputHelper.treeFileArtifact(
-        artifact1, PathFragment.create("child2"));
+    TreeFileArtifact treeFileArtifactA =
+        ActionInputHelper.treeFileArtifactWithNoGeneratingActionSet(
+            artifact1, PathFragment.create("child1"), ACTION_LOOKUP_KEY);
+    TreeFileArtifact treeFileArtifactB =
+        ActionInputHelper.treeFileArtifactWithNoGeneratingActionSet(
+            artifact1, PathFragment.create("child2"), ACTION_LOOKUP_KEY);
     registerAction(new TouchingTestAction(treeFileArtifactA, treeFileArtifactB));
+    treeFileArtifactA.setGeneratingActionKey(artifact1.getGeneratingActionKey());
+    treeFileArtifactB.setGeneratingActionKey(artifact1.getGeneratingActionKey());
 
     // artifact2 is a tree artifact generated by an action template.
     SpecialArtifact artifact2 = createTreeArtifact("treeArtifact2");
@@ -900,12 +917,15 @@
 
     // artifact1 is a tree artifact generated by a TouchingTestAction.
     SpecialArtifact artifact1 = createTreeArtifact("treeArtifact1");
-    artifact1.setGeneratingActionKey(ActionLookupData.create(ACTION_LOOKUP_KEY, 0));
-    TreeFileArtifact treeFileArtifactA = ActionInputHelper.treeFileArtifact(
-        artifact1, PathFragment.create("child1"));
-    TreeFileArtifact treeFileArtifactB = ActionInputHelper.treeFileArtifact(
-        artifact1, PathFragment.create("child2"));
+    TreeFileArtifact treeFileArtifactA =
+        ActionInputHelper.treeFileArtifactWithNoGeneratingActionSet(
+            artifact1, PathFragment.create("child1"), ACTION_LOOKUP_KEY);
+    TreeFileArtifact treeFileArtifactB =
+        ActionInputHelper.treeFileArtifactWithNoGeneratingActionSet(
+            artifact1, PathFragment.create("child2"), ACTION_LOOKUP_KEY);
     registerAction(new TouchingTestAction(treeFileArtifactA, treeFileArtifactB));
+    treeFileArtifactA.setGeneratingActionKey(artifact1.getGeneratingActionKey());
+    treeFileArtifactB.setGeneratingActionKey(artifact1.getGeneratingActionKey());
 
     // artifact2 is a tree artifact generated by an action template.
     SpecialArtifact artifact2 = createTreeArtifact("treeArtifact2");
@@ -1117,7 +1137,8 @@
             @Nullable
             @Override
             public TreeFileArtifact apply(String s) {
-              return ActionInputHelper.treeFileArtifact(parent, s);
+              return ActionInputHelper.treeFileArtifactWithNoGeneratingActionSet(
+                  parent, PathFragment.create(s), parent.getArtifactOwner());
             }
           });
     }