Add "proto_source_root" flag to proto_library.

Fixes #4544.

RELNOTES: Add "proto_source_root" flag to proto_library.
PiperOrigin-RevId: 185997723

6 files changed, 80 insertions, 6 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 15e0e5f82b..d1cd914016 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
@@ -44,6 +44,7 @@ public class BazelProtoLibrary implements RuleConfiguredTargetFactory {
 
     NestedSet<Artifact> transitiveImports =
         ProtoCommon.collectTransitiveImports(ruleContext, protoSources);
+    NestedSet<String> protoPathFlags = ProtoCommon.collectTransitiveProtoPathFlags(ruleContext);
 
     NestedSet<Artifact> protosInDirectDeps = ProtoCommon.computeProtosInDirectDeps(ruleContext);
 
@@ -53,6 +54,7 @@ public class BazelProtoLibrary implements RuleConfiguredTargetFactory {
             protoSources,
             protosInDirectDeps,
             transitiveImports,
+            protoPathFlags,
             !protoSources.isEmpty());
 
     Artifact descriptorSetOutput =
@@ -71,7 +73,8 @@ public class BazelProtoLibrary implements RuleConfiguredTargetFactory {
         protosInDirectDeps,
         descriptorSetOutput,
         true /* allowServices */,
-        dependenciesDescriptorSets);
+        dependenciesDescriptorSets,
+        protoPathFlags);
 
     Runfiles dataRunfiles =
         ProtoCommon.createDataRunfilesProvider(transitiveImports, ruleContext)
@@ -86,7 +89,8 @@ public class BazelProtoLibrary implements RuleConfiguredTargetFactory {
             protoSources,
             checkDepsProtoSources,
             descriptorSetOutput,
-            transitiveDescriptorSetOutput);
+            transitiveDescriptorSetOutput,
+            protoPathFlags);
 
     return new RuleConfiguredTargetBuilder(ruleContext)
         .setFilesToBuild(NestedSetBuilder.create(STABLE_ORDER, descriptorSetOutput))
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/BazelProtoLibraryRule.java b/src/main/java/com/google/devtools/build/lib/rules/proto/BazelProtoLibraryRule.java
index 557cf5fa15..3eab22ee77 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/BazelProtoLibraryRule.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/BazelProtoLibraryRule.java
@@ -17,6 +17,7 @@ package com.google.devtools.build.lib.rules.proto;
 import static com.google.devtools.build.lib.packages.Attribute.attr;
 import static com.google.devtools.build.lib.packages.BuildType.LABEL;
 import static com.google.devtools.build.lib.packages.BuildType.LABEL_LIST;
+import static com.google.devtools.build.lib.syntax.Type.STRING;
 
 import com.google.devtools.build.lib.analysis.BaseRuleClasses;
 import com.google.devtools.build.lib.analysis.RuleDefinition;
@@ -72,6 +73,11 @@ public final class BazelProtoLibraryRule implements RuleDefinition {
             attr("srcs", LABEL_LIST)
                 .direct_compile_time_input()
                 .allowedFileTypes(FileType.of(".proto"), FileType.of(".protodevel")))
+        /* <!-- #BLAZE_RULE(proto_library).ATTRIBUTE(proto_source_root) -->
+        Directory containing .proto files. If set, it must be equal to the package name. If not set,
+        the source root will be the workspace directory (default).
+        <!-- #END_BLAZE_RULE.ATTRIBUTE --> */
+        .add(attr("proto_source_root", STRING))
         .advertiseProvider(ProtoSourcesProvider.class, ProtoSupportDataProvider.class)
         .build();
   }
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 a09a9afa62..374b534cd9 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
@@ -30,6 +30,7 @@ import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.packages.BuildType;
+import com.google.devtools.build.lib.syntax.Type;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import javax.annotation.Nullable;
@@ -99,6 +100,44 @@ public class ProtoCommon {
   }
 
   /**
+   * Returns all proto source roots in this lib and in its transitive dependencies, each prefixed
+   * by {@code --proto_path}.
+   *
+   * Build will fail if the {@code proto_source_root} of the current lib is different than the
+   * package name.
+   */
+  public static NestedSet<String> collectTransitiveProtoPathFlags(RuleContext ruleContext) {
+    NestedSetBuilder<String> protoPathFlags = NestedSetBuilder.stableOrder();
+
+    // first add the protoSourceRoot of the current target, if any
+    String protoSourceRoot =
+        ruleContext.attributes().get("proto_source_root", Type.STRING);
+    if (protoSourceRoot != null && !protoSourceRoot.isEmpty()) {
+      checkProtoSourceRootIsTheSameAsPackage(protoSourceRoot, ruleContext);
+      protoPathFlags.add("--proto_path=" + protoSourceRoot);
+    }
+
+    for (ProtoSourcesProvider provider : ruleContext.getPrerequisites(
+            "deps", Mode.TARGET, ProtoSourcesProvider.class)) {
+      protoPathFlags.addTransitive(provider.getTransitiveProtoPathFlags());
+    }
+
+    return protoPathFlags.build();
+  }
+
+  private static void checkProtoSourceRootIsTheSameAsPackage(
+      String protoSourceRoot, RuleContext ruleContext) {
+    if (!ruleContext.getLabel().getPackageName().equals(protoSourceRoot)) {
+      ruleContext.attributeError(
+          "proto_source_root",
+          "proto_source_root must be the same as the package name ("
+              + ruleContext.getLabel().getPackageName() + ")."
+              + " not '" + protoSourceRoot + "'."
+      );
+    }
+  }
+
+  /**
    * Check that .proto files in sources are from the same package. This is done to avoid clashes
    * with the generated sources.
    */
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 639b07a032..a78fcd3e3f 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
@@ -359,7 +359,8 @@ public class ProtoCompileActionBuilder {
       NestedSet<Artifact> protosInDirectDeps,
       Artifact output,
       boolean allowServices,
-      NestedSet<Artifact> transitiveDescriptorSets) {
+      NestedSet<Artifact> transitiveDescriptorSets,
+      NestedSet<String> protoSourceRoots) {
     if (protosToCompile.isEmpty()) {
       ruleContext.registerAction(
           FileWriteAction.createEmptyWithInputs(
@@ -374,6 +375,7 @@ public class ProtoCompileActionBuilder {
             protosToCompile,
             transitiveSources,
             protosInDirectDeps,
+            protoSourceRoots,
             ruleContext.getLabel(),
             ImmutableList.of(output),
             "Descriptor Set",
@@ -420,6 +422,7 @@ public class ProtoCompileActionBuilder {
       Iterable<Artifact> protosToCompile,
       NestedSet<Artifact> transitiveSources,
       NestedSet<Artifact> protosInDirectDeps,
+      NestedSet<String> protoSourceRoots,
       Label ruleLabel,
       Iterable<Artifact> outputs,
       String flavorName,
@@ -431,6 +434,7 @@ public class ProtoCompileActionBuilder {
             protosToCompile,
             transitiveSources,
             protosInDirectDeps,
+            protoSourceRoots,
             ruleLabel,
             outputs,
             flavorName,
@@ -447,6 +451,7 @@ public class ProtoCompileActionBuilder {
       Iterable<Artifact> protosToCompile,
       NestedSet<Artifact> transitiveSources,
       @Nullable NestedSet<Artifact> protosInDirectDeps,
+      NestedSet<String> protoSourceRoots,
       Label ruleLabel,
       Iterable<Artifact> outputs,
       String flavorName,
@@ -481,6 +486,7 @@ public class ProtoCompileActionBuilder {
                 toolchainInvocations,
                 protosToCompile,
                 transitiveSources,
+                protoSourceRoots,
                 areDepsStrict(ruleContext) ? protosInDirectDeps : null,
                 ruleLabel,
                 allowServices,
@@ -517,6 +523,7 @@ public class ProtoCompileActionBuilder {
       List<ToolchainInvocation> toolchainInvocations,
       Iterable<Artifact> protosToCompile,
       NestedSet<Artifact> transitiveSources,
+      NestedSet<String> transitiveProtoPathFlags,
       @Nullable NestedSet<Artifact> protosInDirectDeps,
       Label ruleLabel,
       boolean allowServices,
@@ -570,6 +577,8 @@ public class ProtoCompileActionBuilder {
       cmdLine.add("--disallow_services");
     }
 
+    cmdLine.addAll(transitiveProtoPathFlags);
+
     return cmdLine.build();
   }
 
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 1207ae0231..bf2ed12b3b 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
@@ -41,14 +41,16 @@ public abstract class ProtoSourcesProvider implements TransitiveInfoProvider {
       ImmutableList<Artifact> protoSources,
       NestedSet<Artifact> checkDepsProtoSources,
       Artifact directDescriptorSet,
-      NestedSet<Artifact> transitiveDescriptorSets) {
+      NestedSet<Artifact> transitiveDescriptorSets,
+      NestedSet<String> protoPathFlags) {
     return new AutoValue_ProtoSourcesProvider(
         transitiveImports,
         transitiveProtoSources,
         protoSources,
         checkDepsProtoSources,
         directDescriptorSet,
-        transitiveDescriptorSets);
+        transitiveDescriptorSets,
+        protoPathFlags);
   }
 
   /**
@@ -127,5 +129,11 @@ public abstract class ProtoSourcesProvider implements TransitiveInfoProvider {
   )
   public abstract NestedSet<Artifact> transitiveDescriptorSets();
 
+  /**
+   * Directories of .proto sources collected from the transitive closure, each prefixed with
+   * {@code --proto_path}. These flags will be passed to {@code protoc} in the specified oreder.
+   */
+  public abstract NestedSet<String> getTransitiveProtoPathFlags();
+
   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 da35d13c34..9e02d190d5 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
@@ -33,9 +33,11 @@ public abstract class SupportData {
       ImmutableList<Artifact> protoSources,
       NestedSet<Artifact> protosInDirectDeps,
       NestedSet<Artifact> transitiveImports,
+      NestedSet<String> transitiveProtoPathFlags,
       boolean hasProtoSources) {
     return new AutoValue_SupportData(
-        nonWeakDepsPredicate, protoSources, transitiveImports, protosInDirectDeps, hasProtoSources);
+        nonWeakDepsPredicate, protoSources, transitiveImports, protosInDirectDeps,
+        transitiveProtoPathFlags, hasProtoSources);
   }
 
   public abstract Predicate<TransitiveInfoCollection> getNonWeakDepsPredicate();
@@ -49,6 +51,12 @@ public abstract class SupportData {
    */
   public abstract NestedSet<Artifact> getProtosInDirectDeps();
 
+  /**
+   * Directories of .proto sources collected from the transitive closure, each prefixed with
+   * {@code --proto_path}. These flags will be passed to {@code protoc} in the specified oreder.
+   */
+  public abstract NestedSet<String> getTransitiveProtoPathFlags();
+
   public abstract boolean hasProtoSources();
 
   SupportData() {}