Automatically map paths in Args#map_each

When path mapping is enabled, `File` objects accessed in a
user-supplied callback passed to `Args#map_each` automatically have
their paths mapped.

Automatic rewriting is preferable to e.g. passing a rewriter object into
the callback: All paths emitted into command lines must be rewritten
with path mapping as inputs and outputs are staged at the mapped
locations, so the user would need to manually map all paths - there is
no choice. As an added benefit, the automatic rewriting ensures that
existing rules relying on `map_each` do not need to be modified to use
path mapping.

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

@@ -28,6 +28,7 @@
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Streams;
+import com.google.devtools.build.docgen.annot.DocCategory;
 import com.google.devtools.build.lib.actions.ArtifactRoot.RootType;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.LabelConstants;
@@ -42,6 +43,7 @@
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec.VisibleForSerialization;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.SerializationConstant;
 import com.google.devtools.build.lib.starlarkbuildapi.FileApi;
+import com.google.devtools.build.lib.starlarkbuildapi.FileRootApi;
 import com.google.devtools.build.lib.util.FileType;
 import com.google.devtools.build.lib.util.FileTypeSet;
 import com.google.devtools.build.lib.util.HashCodes;
@@ -61,9 +63,11 @@
 import java.util.Map;
 import java.util.function.UnaryOperator;
 import javax.annotation.Nullable;
+import net.starlark.java.annot.StarlarkBuiltin;
 import net.starlark.java.eval.EvalException;
 import net.starlark.java.eval.Printer;
 import net.starlark.java.eval.Starlark;
+import net.starlark.java.eval.StarlarkSemantics;
 
 /**
  * An Artifact represents a file used by the build system, whether it's a source file or a derived
@@ -581,17 +585,23 @@ public SpecialArtifact getParent() {
   /**
    * Returns the directory name of this artifact, similar to dirname(1).
    *
-   * <p> The directory name is always a relative path to the execution directory.
+   * <p>The directory name is always a relative path to the execution directory.
    */
-  @Override
   public final String getDirname() {
+    return getDirname(execPath);
+  }
+
+  @Override
+  public final String getDirnameForStarlark(StarlarkSemantics semantics) {
+    return getDirname(PathMapper.loadFrom(semantics).map(execPath));
+  }
+
+  private static String getDirname(PathFragment execPath) {
     PathFragment parent = execPath.getParentDirectory();
     return (parent == null) ? "/" : parent.getSafePathString();
   }
 
-  /**
-   * Returns the base file name of this artifact, similar to basename(1).
-   */
+  /** Returns the base file name of this artifact, similar to basename(1). */
   @Override
   public final String getFilename() {
     return execPath.getBaseName();
@@ -648,16 +658,37 @@ public final Label getOwner() {
    * package-path entries (for source Artifacts), or one of the bin, genfiles or includes dirs (for
    * derived Artifacts). It will always be an ancestor of getPath().
    */
-  @Override
   public final ArtifactRoot getRoot() {
     return root;
   }
 
+  @Override
+  public final FileRootApi getRootForStarlark(StarlarkSemantics semantics) {
+    // It would *not* be correct to just apply PathMapper#map to the exec path of the root: The
+    // root part of the mapped exec path of this artifact may depend on its complete exec path as
+    // well as on e.g. the digest of the artifact.
+    PathFragment mappedExecPath = PathMapper.loadFrom(semantics).map(execPath);
+    if (mappedExecPath.equals(execPath)) {
+      return root;
+    }
+    // PathMapper#map never changes the root-relative part of the exec path, so we can remove that
+    // suffix to get the mapped root part.
+    int rootRelativeSegmentCount = execPath.segmentCount() - root.getExecPath().segmentCount();
+    PathFragment mappedRootExecPath =
+        mappedExecPath.subFragment(0, mappedExecPath.segmentCount() - rootRelativeSegmentCount);
+    return new MappedArtifactRoot(mappedRootExecPath);
+  }
+
   @Override
   public final PathFragment getExecPath() {
     return execPath;
   }
 
+  @Override
+  public String getExecPathStringForStarlark(StarlarkSemantics semantics) {
+    return PathMapper.loadFrom(semantics).getMappedExecPathString(this);
+  }
+
   /**
    * Returns the relative path to this artifact relative to its root. It makes no guarantees as to
    * the semantic meaning or the completeness of the returned path value. In other words, no
@@ -1641,4 +1672,61 @@ public String toString() {
       return MoreObjects.toStringHelper(this).add("artifact", artifact.toDebugString()).toString();
     }
   }
+
+  /** A {@link FileRootApi} obtained by applying a {@link PathMapper} to an {@link ArtifactRoot}. */
+  @StarlarkBuiltin(
+      name = "mapped_root",
+      category = DocCategory.BUILTIN,
+      doc = "A root for files that have been subject to path mapping")
+  private static final class MappedArtifactRoot
+      implements FileRootApi, Comparable<MappedArtifactRoot> {
+    private final PathFragment mappedRootExecPath;
+
+    public MappedArtifactRoot(PathFragment mappedRootExecPath) {
+      this.mappedRootExecPath = mappedRootExecPath;
+    }
+
+    @Override
+    public String getExecPathString() {
+      return mappedRootExecPath.getPathString();
+    }
+
+    @Override
+    public int compareTo(MappedArtifactRoot otherRoot) {
+      return mappedRootExecPath.compareTo(otherRoot.mappedRootExecPath);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+      // Per the contract of PathMapper#map, mapped roots never have exec paths that are equal to
+      // exec paths of non-mapped roots, that is, of instances of ArtifactRoot. Thus, it is correct
+      // for both equals implementations to return false if the other object is not an instance of
+      // the respective class.
+      if (!(obj instanceof MappedArtifactRoot)) {
+        return false;
+      }
+      MappedArtifactRoot other = (MappedArtifactRoot) obj;
+      return mappedRootExecPath.equals(other.mappedRootExecPath);
+    }
+
+    @Override
+    public int hashCode() {
+      return mappedRootExecPath.hashCode();
+    }
+
+    @Override
+    public String toString() {
+      return mappedRootExecPath + " [mapped]";
+    }
+
+    @Override
+    public void repr(Printer printer) {
+      printer.append("<mapped root>");
+    }
+
+    @Override
+    public boolean isImmutable() {
+      return true;
+    }
+  }
 }

src/main/java/com/google/devtools/build/lib/actions/BUILD

@@ -304,6 +304,7 @@ java_library(
         "ArtifactResolver.java",
         "ArtifactRoot.java",
         "Artifacts.java",
+        "PathMapper.java",
     ],
     deps = [
         ":action_lookup_data",
@@ -313,6 +314,7 @@ java_library(
         ":fileset_output_symlink",
         ":package_roots",
         ":path_strippable",
+        "//src/main/java/com/google/devtools/build/docgen/annot",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/collect/nestedset",
         "//src/main/java/com/google/devtools/build/lib/concurrent",
@@ -330,6 +332,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//src/main/java/com/google/devtools/build/skyframe:execution_phase_skykey",
         "//src/main/java/com/google/devtools/build/skyframe:skyframe-objects",
+        "//src/main/java/net/starlark/java/annot",
         "//src/main/java/net/starlark/java/eval",
         "//src/main/protobuf:failure_details_java_proto",
         "//third_party:guava",

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

@@ -17,30 +17,92 @@
 import com.google.devtools.build.lib.actions.CommandLineItem.ExceptionlessMapFn;
 import com.google.devtools.build.lib.actions.CommandLineItem.MapFn;
 import com.google.devtools.build.lib.vfs.PathFragment;
+import com.google.errorprone.annotations.CheckReturnValue;
 import java.util.List;
+import java.util.Objects;
 import javax.annotation.Nullable;
+import net.starlark.java.eval.StarlarkSemantics;
 
 /**
  * Support for mapping config parts of exec paths in an action's command line as well as when
  * staging its inputs and outputs for execution, with the aim of making the resulting {@link Spawn}
  * more cacheable.
  *
+ * <p>Implementations must be pure functions of the set of inputs (including paths and potentially
+ * content) of a given action.
+ *
  * <p>Actions that want to support path mapping should use {@link
  * com.google.devtools.build.lib.analysis.actions.PathMappers}.
  *
  * <p>An example of an implementing class is {@link
  * com.google.devtools.build.lib.analysis.actions.StrippingPathMapper}, which removes the config
  * part (e.g. "k8-fastbuild") from exec paths to allow for cross-configuration cache hits.
  */
-public interface PathMapper {
-  /** Returns the exec path with the path mapping applied. */
-  PathFragment map(PathFragment execPath);
+public abstract class PathMapper {
+  /**
+   * Returns the exec path with the path mapping applied.
+   *
+   * <p>An exec path typically consists of a root part (such as <code>"bazel-out/k8-opt/bin"</code>
+   * or <code>""</code>) and a root-relative part (such as <code>"path/to/pkg/file"</code>).
+   *
+   * <p>Overrides must satisfy the following properties:
+   *
+   * <ul>
+   *   <li>The root-relative part of the path must not be modified.
+   *   <li>If the path is modified, the new root part must be different from any possible valid root
+   *       part of an unmapped path.
+   * </ul>
+   */
+  public abstract PathFragment map(PathFragment execPath);
 
   /** Returns the exec path of the input with the path mapping applied. */
-  default String getMappedExecPathString(ActionInput artifact) {
+  public String getMappedExecPathString(ActionInput artifact) {
     return map(artifact.getExecPath()).getPathString();
   }
 
+  /** A {@link PathMapper} that doesn't change paths. */
+  public static final PathMapper NOOP =
+      new PathMapper() {
+        @Override
+        public PathFragment map(PathFragment execPath) {
+          return execPath;
+        }
+      };
+
+  private static final StarlarkSemantics.Key<PathMapper> SEMANTICS_KEY =
+      new StarlarkSemantics.Key<>("path_mapper", PathMapper.NOOP);
+
+  /**
+   * Retrieve the {@link PathMapper} instance stored in the given {@link StarlarkSemantics} via
+   * {@link #storeIn(StarlarkSemantics)}.
+   */
+  public static PathMapper loadFrom(StarlarkSemantics semantics) {
+    return semantics.get(SEMANTICS_KEY);
+  }
+
+  /**
+   * Creates a new {@link StarlarkSemantics} instance which causes all Starlark threads using it to
+   * automatically apply this {@link PathMapper} to all struct fields of {@link
+   * com.google.devtools.build.lib.starlarkbuildapi.FileApi}.
+   *
+   * <p>This is meant to be used when evaluating user-defined callbacks to Starlark variants of
+   * custom command lines that are evaluated during the execution phase.
+   *
+   * <p>Since any unmapped path appearing in a command line will prevent cross-configuration cache
+   * hits, this mapping is applied automatically instead of requiring users to explicitly map all
+   * paths themselves. As an added benefit, this allows actions to opt into path mapping without
+   * actual changes to their command line code.
+   */
+  @CheckReturnValue
+  public StarlarkSemantics storeIn(StarlarkSemantics semantics) {
+    // Since PathMapper#equals returns true for all instances of the same class, every non-noop
+    // instance is different from the default value for the key and thus persisted. Furthermore,
+    // since this causes the resulting semantics to compare equal if they PathMapper class
+    // agrees, there is only a single cache entry for it in Starlark's CallUtils cache for Starlark
+    // methods.
+    return semantics.toBuilder().set(SEMANTICS_KEY, this).build();
+  }
+
   /**
    * We don't yet have a Starlark API for mapping paths in command lines. Simple Starlark calls like
    * {@code args.add(arg_name, file_path} are automatically handled. But calls that involve custom
@@ -49,7 +111,7 @@ default String getMappedExecPathString(ActionInput artifact) {
    * <p>This method allows implementations to hard-code support for specific command line entries
    * for specific Starlark actions.
    */
-  default List<String> mapCustomStarlarkArgs(List<String> args) {
+  public List<String> mapCustomStarlarkArgs(List<String> args) {
     return args;
   }
 
@@ -66,7 +128,7 @@ default List<String> mapCustomStarlarkArgs(List<String> args) {
    *
    * <p>By default, this method returns {@link MapFn#DEFAULT}.
    */
-  default ExceptionlessMapFn<Object> getMapFn(@Nullable String previousFlag) {
+  public ExceptionlessMapFn<Object> getMapFn(@Nullable String previousFlag) {
     return MapFn.DEFAULT;
   }
 
@@ -76,10 +138,24 @@ default ExceptionlessMapFn<Object> getMapFn(@Nullable String previousFlag) {
    * <p>Can be used by actions to skip additional work that isn't needed if path mapping is not
    * enabled.
    */
-  default boolean isNoop() {
+  public boolean isNoop() {
     return this == NOOP;
   }
 
-  /** A {@link PathMapper} that doesn't change paths. */
-  PathMapper NOOP = execPath -> execPath;
+  @Override
+  public final boolean equals(Object obj) {
+    if (obj == null || obj.getClass() != getClass()) {
+      return false;
+    }
+    // We store PathMapper instances in StarlarkSemantics to allow mapping of File struct fields
+    // such as "path" in map_each callbacks of Args. Since this code is only executed during the
+    // execution phase, we do not want two different non-noop PathMapper instances to result in
+    // distinct StarlarkSemantics instances for caching purposes.
+    return isNoop() == ((PathMapper) obj).isNoop();
+  }
+
+  @Override
+  public final int hashCode() {
+    return Objects.hashCode(isNoop());
+  }
 }