PiperOrigin-RevId: 203687257

4 files changed, 101 insertions, 15 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/actions/ParamFileInfo.java b/src/main/java/com/google/devtools/build/lib/actions/ParamFileInfo.java
index 27f4d4597d..171a073450 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/ParamFileInfo.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/ParamFileInfo.java
@@ -17,6 +17,8 @@ package com.google.devtools.build.lib.actions;
 import static java.nio.charset.StandardCharsets.ISO_8859_1;
 
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.actions.ParameterFile.ParameterFileType;
 import java.nio.charset.Charset;
 import java.util.Objects;
@@ -32,13 +34,24 @@ public final class ParamFileInfo {
   private final Charset charset;
   private final String flagFormatString;
   private final boolean always;
+  private final Iterable<Artifact> inputs;
 
   private ParamFileInfo(
-      ParameterFileType fileType, Charset charset, String flagFormatString, boolean always) {
+      ParameterFileType fileType,
+      Charset charset,
+      String flagFormatString,
+      Iterable<Artifact> inputs,
+      boolean always) {
     this.fileType = Preconditions.checkNotNull(fileType);
     this.charset = Preconditions.checkNotNull(charset);
     this.flagFormatString = Preconditions.checkNotNull(flagFormatString);
     this.always = always;
+    this.inputs = Iterables.transform(inputs, input -> {
+      Preconditions.checkArgument(input.isTreeArtifact(),
+          input.getExecPath() + " is not tree artifacts, "
+              + "it should not be put into ParamFileInfo's inputs");
+      return input;
+    });
   }
 
   /**
@@ -65,6 +78,11 @@ public final class ParamFileInfo {
     return always;
   }
 
+  /** Returns the tree artifacts that are needed as input for {@code ParameterFileWriteAction} */
+  public Iterable<Artifact> getInputs() {
+    return inputs;
+  }
+
   @Override
   public int hashCode() {
     return Objects.hash(charset, flagFormatString, fileType, always);
@@ -95,6 +113,7 @@ public final class ParamFileInfo {
     private Charset charset = ISO_8859_1;
     private String flagFormatString = "@%s";
     private boolean always;
+    private Iterable<Artifact> inputs = ImmutableList.<Artifact>of();
 
     private Builder(ParameterFileType fileType) {
       this.fileType = fileType;
@@ -123,8 +142,13 @@ public final class ParamFileInfo {
       return this;
     }
 
+    public Builder setInputs(Iterable<Artifact> inputs) {
+      this.inputs = inputs;
+      return this;
+    }
+
     public ParamFileInfo build() {
-      return new ParamFileInfo(fileType, charset, flagFormatString, always);
+      return new ParamFileInfo(fileType, charset, flagFormatString, inputs, always);
     }
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/CustomCommandLine.java b/src/main/java/com/google/devtools/build/lib/analysis/actions/CustomCommandLine.java
index 770d16aef7..0582a93a3f 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/actions/CustomCommandLine.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/actions/CustomCommandLine.java
@@ -50,6 +50,7 @@ import java.util.Set;
 import java.util.TreeSet;
 import java.util.UUID;
 import java.util.function.Consumer;
+import java.util.function.Function;
 import javax.annotation.Nullable;
 
 /** A customizable, serializable class for building memory efficient command lines. */
@@ -617,16 +618,27 @@ public final class CustomCommandLine extends CommandLine {
   }
 
   @AutoCodec
-  static final class ExpandedTreeArtifactExecPathsArg extends TreeArtifactExpansionArgvFragment {
+  static final class ExpandedTreeArtifactArg extends TreeArtifactExpansionArgvFragment {
     private final Artifact treeArtifact;
     private static final UUID TREE_UUID = UUID.fromString("13b7626b-c77d-4a30-ad56-ff08c06b1cee");
+    private final Function<Artifact, Iterable<String>> expandFunction;
 
     @AutoCodec.Instantiator
     @VisibleForSerialization
-    ExpandedTreeArtifactExecPathsArg(Artifact treeArtifact) {
+    ExpandedTreeArtifactArg(Artifact treeArtifact) {
       Preconditions.checkArgument(
           treeArtifact.isTreeArtifact(), "%s is not a TreeArtifact", treeArtifact);
       this.treeArtifact = treeArtifact;
+      this.expandFunction = artifact -> ImmutableList.of(artifact.getExecPathString());
+    }
+
+    @VisibleForSerialization
+    ExpandedTreeArtifactArg(
+        Artifact treeArtifact, Function<Artifact, Iterable<String>> expandFunction) {
+      Preconditions.checkArgument(
+          treeArtifact.isTreeArtifact(), "%s is not a TreeArtifact", treeArtifact);
+      this.treeArtifact = treeArtifact;
+      this.expandFunction = expandFunction;
     }
 
     @Override
@@ -635,14 +647,16 @@ public final class CustomCommandLine extends CommandLine {
       artifactExpander.expand(treeArtifact, expandedArtifacts);
 
       for (Artifact expandedArtifact : expandedArtifacts) {
-        builder.add(expandedArtifact.getExecPathString());
+        for (String commandLine : expandFunction.apply(expandedArtifact)) {
+          builder.add(commandLine);
+        }
       }
     }
 
     @Override
     public String describe() {
       return String.format(
-          "ExpandedTreeArtifactExecPathsArg{ treeArtifact: %s}", treeArtifact.getExecPathString());
+          "ExpandedTreeArtifactArg{ treeArtifact: %s}", treeArtifact.getExecPathString());
     }
 
     @Override
@@ -1040,7 +1054,31 @@ public final class CustomCommandLine extends CommandLine {
      */
     public Builder addExpandedTreeArtifactExecPaths(Artifact treeArtifact) {
       Preconditions.checkNotNull(treeArtifact);
-      arguments.add(new ExpandedTreeArtifactExecPathsArg(treeArtifact));
+      arguments.add(new ExpandedTreeArtifactArg(treeArtifact));
+      return this;
+    }
+
+    public Builder addExpandedTreeArtifactExecPaths(String arg, Artifact treeArtifact) {
+      Preconditions.checkNotNull(arg);
+      Preconditions.checkNotNull(treeArtifact);
+      arguments.add(
+          new ExpandedTreeArtifactArg(
+              treeArtifact, artifact -> ImmutableList.of(arg, artifact.getExecPathString())));
+      return this;
+    }
+
+    /**
+     * Adds the arguments for all {@link TreeFileArtifact}s under
+     * {@code treeArtifact}, one argument per file. Using {@code expandingFunction} to expand each
+     * {@link TreeFileArtifact} to expected argument.
+     *
+     * @param treeArtifact the TreeArtifact containing the {@link TreeFileArtifact}s to add.
+     * @param expandFunction the function to generate the argument for each{@link TreeFileArtifact}.
+     */
+    public Builder addExpandedTreeArtifact(
+        Artifact treeArtifact, Function<Artifact, Iterable<String>> expandFunction) {
+      Preconditions.checkNotNull(treeArtifact);
+      arguments.add(new ExpandedTreeArtifactArg(treeArtifact, expandFunction));
       return this;
     }
 
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java b/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java
index e2cbcc89bd..74509858d9 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java
@@ -726,6 +726,7 @@ public class SpawnAction extends AbstractAction implements ExecutionInfoSpecifie
           ParameterFileWriteAction paramFileWriteAction =
               new ParameterFileWriteAction(
                   owner,
+                  paramFileInfo.getInputs(),
                   paramFile,
                   commandLine,
                   paramFileInfo.getFileType(),
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/actions/ParamFileWriteActionTest.java b/src/test/java/com/google/devtools/build/lib/analysis/actions/ParamFileWriteActionTest.java
index 11ced4c8b6..59364a3328 100644
--- a/src/test/java/com/google/devtools/build/lib/analysis/actions/ParamFileWriteActionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/analysis/actions/ParamFileWriteActionTest.java
@@ -76,9 +76,9 @@ public class ParamFileWriteActionTest extends BuildViewTestCase {
 
   @Test
   public void testInputs() {
-    Action action = createParameterFileWriteAction(
-        ImmutableList.of(treeArtifact),
-        createTreeArtifactExpansionCommandLine());
+    Action action =
+        createParameterFileWriteAction(
+            ImmutableList.of(treeArtifact), createTreeArtifactExpansionCommandLineDefault());
     assertThat(Artifact.toExecPaths(action.getInputs()))
         .containsExactly("out/artifact/myTreeFileArtifact");
   }
@@ -95,10 +95,10 @@ public class ParamFileWriteActionTest extends BuildViewTestCase {
   }
 
   @Test
-  public void testWriteCommandLineWithTreeArtifactExpansion() throws Exception {
-    Action action = createParameterFileWriteAction(
-        ImmutableList.of(treeArtifact),
-        createTreeArtifactExpansionCommandLine());
+  public void testWriteCommandLineWithTreeArtifactExpansionDefault() throws Exception {
+    Action action =
+        createParameterFileWriteAction(
+            ImmutableList.of(treeArtifact), createTreeArtifactExpansionCommandLineDefault());
     ActionExecutionContext context = actionExecutionContext();
     ActionResult actionResult = action.execute(context);
     assertThat(actionResult.spawnResults()).isEmpty();
@@ -110,6 +110,22 @@ public class ParamFileWriteActionTest extends BuildViewTestCase {
                 + "out/artifact/myTreeFileArtifact/artifacts/treeFileArtifact2");
   }
 
+  @Test
+  public void testWriteCommandLineWithTreeArtifactExpansionExpandedFunction() throws Exception {
+    Action action =
+        createParameterFileWriteAction(
+            ImmutableList.of(treeArtifact),
+            createTreeArtifactExpansionCommandLineExpandedFunction());
+    ActionExecutionContext context = actionExecutionContext();
+    ActionResult actionResult = action.execute(context);
+    assertThat(actionResult.spawnResults()).isEmpty();
+    String content = new String(FileSystemUtils.readContentAsLatin1(outputArtifact.getPath()));
+    assertThat(content.trim())
+        .isEqualTo(
+            "--flag1=out/artifact/myTreeFileArtifact/artifacts/treeFileArtifact1\n"
+                + "--flag1=out/artifact/myTreeFileArtifact/artifacts/treeFileArtifact2");
+  }
+
   private SpecialArtifact createTreeArtifact(String rootRelativePath) {
     PathFragment relpath = PathFragment.create(rootRelativePath);
     return new SpecialArtifact(
@@ -145,13 +161,20 @@ public class ParamFileWriteActionTest extends BuildViewTestCase {
         .build();
   }
 
-  private CommandLine createTreeArtifactExpansionCommandLine() {
+  private CommandLine createTreeArtifactExpansionCommandLineDefault() {
     return CustomCommandLine.builder()
         .add("--flag1")
         .addExpandedTreeArtifactExecPaths(treeArtifact)
         .build();
   }
 
+  private CommandLine createTreeArtifactExpansionCommandLineExpandedFunction() {
+    return CustomCommandLine.builder()
+        .addExpandedTreeArtifact(
+            treeArtifact, artifact -> ImmutableList.of("--flag1=" + artifact.getExecPath()))
+        .build();
+  }
+
   private ActionExecutionContext actionExecutionContext() throws Exception {
     final Iterable<TreeFileArtifact> treeFileArtifacts = ImmutableList.of(
         createTreeFileArtifact(treeArtifact, "artifacts/treeFileArtifact1"),