[Reland] Accept proto paths relative to proto_source_root as direct dependencies.

This is a reland of https://github.com/bazelbuild/bazel/commit/5deca4cf88f5568771f2c836a9b8c693b88bd749.

This will make protoc see as direct dependencies the .proto files that were included using the proto_source_root flag.

Until now, Bazel passed to protoc the direct dependencies of a target as the path relative to the WORKSPACE, which made it fail when a shorter path, relative to the package was used.

Progress on #4544.

RELNOTES: None.
PiperOrigin-RevId: 203808292

5 files changed, 128 insertions, 18 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/BazelProtoLibrary.java b/src/main/java/com/google/devtools/build/lib/rules/proto/BazelProtoLibrary.java
index e3687be773..1c8a7d52cb 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/BazelProtoLibrary.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/BazelProtoLibrary.java
@@ -45,10 +45,14 @@ public class BazelProtoLibrary implements RuleConfiguredTargetFactory {
 
     NestedSet<Artifact> transitiveImports =
         ProtoCommon.collectTransitiveImports(ruleContext, protoSources);
-    NestedSet<String> protoPathFlags = ProtoCommon.collectTransitiveProtoPathFlags(ruleContext);
-
     NestedSet<Artifact> protosInDirectDeps = ProtoCommon.computeProtosInDirectDeps(ruleContext);
 
+    String protoSourceRoot = ProtoCommon.getProtoSourceRoot(ruleContext);
+    NestedSet<String> directProtoSourceRoots =
+        ProtoCommon.getProtoSourceRootsOfDirectDependencies(ruleContext, protoSourceRoot);
+    NestedSet<String> protoPathFlags =
+        ProtoCommon.collectTransitiveProtoPathFlags(ruleContext, protoSourceRoot);
+
     final SupportData supportData =
         SupportData.create(
             /* nonWeakDepsPredicate= */ Predicates.<TransitiveInfoCollection>alwaysTrue(),
@@ -56,6 +60,7 @@ public class BazelProtoLibrary implements RuleConfiguredTargetFactory {
             protosInDirectDeps,
             transitiveImports,
             protoPathFlags,
+            directProtoSourceRoots,
             !protoSources.isEmpty());
 
     Artifact descriptorSetOutput =
@@ -75,7 +80,8 @@ public class BazelProtoLibrary implements RuleConfiguredTargetFactory {
         descriptorSetOutput,
         /* allowServices= */ true,
         dependenciesDescriptorSets,
-        protoPathFlags);
+        protoPathFlags,
+        directProtoSourceRoots);
 
     Runfiles dataRunfiles =
         ProtoCommon.createDataRunfilesProvider(transitiveImports, ruleContext)
@@ -91,7 +97,8 @@ public class BazelProtoLibrary implements RuleConfiguredTargetFactory {
             checkDepsProtoSources,
             descriptorSetOutput,
             transitiveDescriptorSetOutput,
-            protoPathFlags);
+            protoPathFlags,
+            protoSourceRoot);
 
     return new RuleConfiguredTargetBuilder(ruleContext)
         .setFilesToBuild(NestedSetBuilder.create(STABLE_ORDER, descriptorSetOutput))
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoCommon.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoCommon.java
index a86bdda6dc..f3cc3d6751 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoCommon.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoCommon.java
@@ -100,28 +100,66 @@ public class ProtoCommon {
   }
 
   /**
-   * Returns all proto source roots in this lib and in its transitive dependencies.
+   * Returns all proto source roots in this lib ({@code currentProtoSourceRoot}) and in its
+   * transitive dependencies.
    *
-   * Build will fail if the {@code proto_source_root} of the current lib is different than the
-   * package name.
+   * <p>Assumes {@code currentProtoSourceRoot} is the same as the package name.
    */
-  public static NestedSet<String> collectTransitiveProtoPathFlags(RuleContext ruleContext) {
+  public static NestedSet<String> collectTransitiveProtoPathFlags(
+      RuleContext ruleContext, String currentProtoSourceRoot) {
     NestedSetBuilder<String> protoPath = NestedSetBuilder.stableOrder();
 
     // first add the protoSourceRoot of the current target, if any
+    if (currentProtoSourceRoot != null && !currentProtoSourceRoot.isEmpty()) {
+      protoPath.add(currentProtoSourceRoot);
+    }
+
+    for (ProtoSourcesProvider provider :
+        ruleContext.getPrerequisites("deps", Mode.TARGET, ProtoSourcesProvider.class)) {
+      protoPath.addTransitive(provider.getTransitiveProtoPathFlags());
+    }
+
+    return protoPath.build();
+  }
+
+  /**
+   * Returns the {@code proto_source_root} of the current library or null if none is specified.
+   *
+   * <p>Build will fail if the {@code proto_source_root} of the current library is different than
+   * the package name.
+   */
+  @Nullable
+  public static String getProtoSourceRoot(RuleContext ruleContext) {
     String protoSourceRoot =
         ruleContext.attributes().get("proto_source_root", Type.STRING);
     if (protoSourceRoot != null && !protoSourceRoot.isEmpty()) {
       checkProtoSourceRootIsTheSameAsPackage(protoSourceRoot, ruleContext);
-      protoPath.add(protoSourceRoot);
     }
+    return protoSourceRoot;
+  }
 
-    for (ProtoSourcesProvider provider : ruleContext.getPrerequisites(
-            "deps", Mode.TARGET, ProtoSourcesProvider.class)) {
-      protoPath.addTransitive(provider.getTransitiveProtoPathFlags());
+  /**
+   * Returns a set of the {@code proto_source_root} collected from the current library and the
+   * direct dependencies.
+   *
+   * <p>Assumes {@code currentProtoSourceRoot} is the same as the package name.
+   */
+  public static NestedSet<String> getProtoSourceRootsOfDirectDependencies(
+      RuleContext ruleContext, String currentProtoSourceRoot) {
+    NestedSetBuilder<String> protoSourceRoots = NestedSetBuilder.stableOrder();
+    if (currentProtoSourceRoot != null && !currentProtoSourceRoot.isEmpty()) {
+      protoSourceRoots.add(currentProtoSourceRoot);
     }
 
-    return protoPath.build();
+    for (ProtoSourcesProvider provider :
+        ruleContext.getPrerequisites("deps", Mode.TARGET, ProtoSourcesProvider.class)) {
+      String protoSourceRoot = provider.getProtoSourceRoot();
+      if (protoSourceRoot != null && !protoSourceRoot.isEmpty()) {
+        protoSourceRoots.add(provider.getProtoSourceRoot());
+      }
+    }
+
+    return protoSourceRoots.build();
   }
 
   private static void checkProtoSourceRootIsTheSameAsPackage(
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoCompileActionBuilder.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoCompileActionBuilder.java
index 808c2a309a..ed138d0381 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoCompileActionBuilder.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoCompileActionBuilder.java
@@ -28,6 +28,7 @@ import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.actions.Action;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.CommandLineItem;
+import com.google.devtools.build.lib.actions.CommandLineItem.CapturingMapFn;
 import com.google.devtools.build.lib.actions.ParamFileInfo;
 import com.google.devtools.build.lib.actions.ParameterFile.ParameterFileType;
 import com.google.devtools.build.lib.actions.ResourceSet;
@@ -50,6 +51,7 @@ import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Consumer;
 import javax.annotation.Nullable;
 
 /** Constructs actions to run the protocol compiler to generate sources from .proto files. */
@@ -296,6 +298,7 @@ public class ProtoCompileActionBuilder {
     addIncludeMapArguments(
         result,
         areDepsStrict ? supportData.getProtosInDirectDeps() : null,
+        supportData.getDirectProtoSourceRoots(),
         supportData.getTransitiveImports());
 
     if (areDepsStrict) {
@@ -331,7 +334,8 @@ public class ProtoCompileActionBuilder {
       Artifact output,
       boolean allowServices,
       NestedSet<Artifact> transitiveDescriptorSets,
-      NestedSet<String> protoSourceRoots) {
+      NestedSet<String> protoSourceRoots,
+      NestedSet<String> directProtoSourceRoots) {
     if (protosToCompile.isEmpty()) {
       ruleContext.registerAction(
           FileWriteAction.createEmptyWithInputs(
@@ -347,6 +351,7 @@ public class ProtoCompileActionBuilder {
             transitiveSources,
             protosInDirectDeps,
             protoSourceRoots,
+            directProtoSourceRoots,
             ruleContext.getLabel(),
             ImmutableList.of(output),
             "Descriptor Set",
@@ -394,6 +399,7 @@ public class ProtoCompileActionBuilder {
       NestedSet<Artifact> transitiveSources,
       NestedSet<Artifact> protosInDirectDeps,
       NestedSet<String> protoSourceRoots,
+      NestedSet<String> directProtoSourceRoots,
       Label ruleLabel,
       Iterable<Artifact> outputs,
       String flavorName,
@@ -406,6 +412,7 @@ public class ProtoCompileActionBuilder {
             transitiveSources,
             protosInDirectDeps,
             protoSourceRoots,
+            directProtoSourceRoots,
             ruleLabel,
             outputs,
             flavorName,
@@ -423,6 +430,7 @@ public class ProtoCompileActionBuilder {
       NestedSet<Artifact> transitiveSources,
       @Nullable NestedSet<Artifact> protosInDirectDeps,
       NestedSet<String> protoSourceRoots,
+      NestedSet<String> directProtoSourceRoots,
       Label ruleLabel,
       Iterable<Artifact> outputs,
       String flavorName,
@@ -458,6 +466,7 @@ public class ProtoCompileActionBuilder {
                 protosToCompile,
                 transitiveSources,
                 protoSourceRoots,
+                directProtoSourceRoots,
                 areDepsStrict(ruleContext) ? protosInDirectDeps : null,
                 ruleLabel,
                 allowServices,
@@ -495,6 +504,7 @@ public class ProtoCompileActionBuilder {
       Iterable<Artifact> protosToCompile,
       NestedSet<Artifact> transitiveSources,
       NestedSet<String> transitiveProtoPathFlags,
+      NestedSet<String> directProtoSourceRoots,
       @Nullable NestedSet<Artifact> protosInDirectDeps,
       Label ruleLabel,
       boolean allowServices,
@@ -537,7 +547,7 @@ public class ProtoCompileActionBuilder {
     cmdLine.addAll(protocOpts);
 
     // Add include maps
-    addIncludeMapArguments(cmdLine, protosInDirectDeps, transitiveSources);
+    addIncludeMapArguments(cmdLine, protosInDirectDeps, directProtoSourceRoots, transitiveSources);
 
     if (protosInDirectDeps != null) {
       cmdLine.addFormatted(STRICT_DEPS_FLAG_TEMPLATE, ruleLabel);
@@ -558,6 +568,7 @@ public class ProtoCompileActionBuilder {
   static void addIncludeMapArguments(
       CustomCommandLine.Builder commandLine,
       @Nullable NestedSet<Artifact> protosInDirectDependencies,
+      NestedSet<String> directProtoSourceRoots,
       NestedSet<Artifact> transitiveImports) {
     commandLine.addAll(VectorArg.of(transitiveImports).mapped(EXPAND_TRANSITIVE_IMPORT_ARG));
     if (protosInDirectDependencies != null) {
@@ -566,7 +577,8 @@ public class ProtoCompileActionBuilder {
             "--direct_dependencies",
             VectorArg.join(":")
                 .each(protosInDirectDependencies)
-                .mapped(EXPAND_TO_PATH_IGNORING_REPOSITORY));
+                .mapped(new ExpandToPathFn(directProtoSourceRoots)));
+
       } else {
         // The proto compiler requires an empty list to turn on strict deps checking
         commandLine.add("--direct_dependencies=");
@@ -584,6 +596,24 @@ public class ProtoCompileActionBuilder {
           args.accept(
               "-I" + getPathIgnoringRepository(artifact) + "=" + artifact.getExecPathString());
 
+  @AutoCodec
+  @AutoCodec.VisibleForSerialization
+  static final class ExpandToPathFn implements CapturingMapFn<Artifact> {
+    private final NestedSet<String> directProtoSourceRoots;
+
+    public ExpandToPathFn(NestedSet<String> directProtoSourceRoots) {
+      this.directProtoSourceRoots = directProtoSourceRoots;
+    }
+
+    @Override
+    public void expandToCommandLine(Artifact proto, Consumer<String> args) {
+      for (String directProtoSourceRoot : directProtoSourceRoots) {
+        expandToPathIgnoringSourceRoot(proto, directProtoSourceRoot, args);
+      }
+      EXPAND_TO_PATH_IGNORING_REPOSITORY.expandToCommandLine(proto, args);
+    }
+  }
+
   @AutoCodec @AutoCodec.VisibleForSerialization
   static final CommandLineItem.MapFn<Artifact> EXPAND_TO_PATH_IGNORING_REPOSITORY =
       (artifact, args) -> args.accept(getPathIgnoringRepository(artifact));
@@ -603,6 +633,29 @@ public class ProtoCompileActionBuilder {
         .toString();
   }
 
+  private static void expandToPathIgnoringSourceRoot(
+      Artifact artifact, String directProtoSourceRoot, Consumer<String> args) {
+    // TODO(bazel-team): IAE is caught here because every artifact is relativized against every
+    // directProtoSourceRoot. Instead of catching the exception, a check should be performed
+    // to see if the artifact has the root as a substring before relativizing.
+    try {
+      String relativePath =
+          artifact
+              .getRootRelativePath()
+              .relativeTo(
+                  artifact
+                      .getOwnerLabel()
+                      .getPackageIdentifier()
+                      .getRepository()
+                      .getPathUnderExecRoot())
+              .relativeTo(directProtoSourceRoot)
+              .toString();
+      args.accept(relativePath);
+    } catch (IllegalArgumentException exception) {
+      // do nothing
+    }
+  }
+
   /**
    * Describes a toolchain and the value to replace for a $(OUT) that might appear in its
    * commandLine() (e.g., "bazel-out/foo.srcjar").
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoSourcesProvider.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoSourcesProvider.java
index 26075cf07f..d976b4d121 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoSourcesProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoSourcesProvider.java
@@ -44,7 +44,8 @@ public abstract class ProtoSourcesProvider
       NestedSet<Artifact> checkDepsProtoSources,
       Artifact directDescriptorSet,
       NestedSet<Artifact> transitiveDescriptorSets,
-      NestedSet<String> transitiveProtoPathFlags) {
+      NestedSet<String> transitiveProtoPathFlags,
+      String protoSourceRoot) {
     return new AutoValue_ProtoSourcesProvider(
         transitiveImports,
         transitiveProtoSources,
@@ -52,7 +53,8 @@ public abstract class ProtoSourcesProvider
         checkDepsProtoSources,
         directDescriptorSet,
         transitiveDescriptorSets,
-        transitiveProtoPathFlags);
+        transitiveProtoPathFlags,
+        protoSourceRoot);
   }
 
   /**
@@ -108,5 +110,8 @@ public abstract class ProtoSourcesProvider
   @Override
   public abstract NestedSet<String> getTransitiveProtoPathFlags();
 
+  /** The {@code proto_source_root} of the current library. */
+  public abstract String getProtoSourceRoot();
+
   ProtoSourcesProvider() {}
 }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/SupportData.java b/src/main/java/com/google/devtools/build/lib/rules/proto/SupportData.java
index d21a4d9d60..c3bc3aa553 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/SupportData.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/SupportData.java
@@ -35,6 +35,7 @@ public abstract class SupportData {
       NestedSet<Artifact> protosInDirectDeps,
       NestedSet<Artifact> transitiveImports,
       NestedSet<String> transitiveProtoPathFlags,
+      NestedSet<String> directProtoSourceRoots,
       boolean hasProtoSources) {
     return new AutoValue_SupportData(
         nonWeakDepsPredicate,
@@ -42,6 +43,7 @@ public abstract class SupportData {
         transitiveImports,
         protosInDirectDeps,
         transitiveProtoPathFlags,
+        directProtoSourceRoots,
         hasProtoSources);
   }
 
@@ -62,6 +64,11 @@ public abstract class SupportData {
    */
   public abstract NestedSet<String> getTransitiveProtoPathFlags();
 
+  /**
+   * The {@code proto_source_root}'s collected from the current library and the direct dependencies.
+   */
+  public abstract NestedSet<String> getDirectProtoSourceRoots();
+
   public abstract boolean hasProtoSources();
 
   SupportData() {}