1 files changed, 39 insertions, 4 deletions
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 f3cc3d6751..cb9dbb6818 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
@@ -140,19 +140,19 @@ public class ProtoCommon {
 
   /**
    * Returns a set of the {@code proto_source_root} collected from the current library and the
-   * direct dependencies.
+   * specified attribute.
    *
    * <p>Assumes {@code currentProtoSourceRoot} is the same as the package name.
    */
-  public static NestedSet<String> getProtoSourceRootsOfDirectDependencies(
-      RuleContext ruleContext, String currentProtoSourceRoot) {
+  private static NestedSet<String> getProtoSourceRootsOfAttribute(
+      RuleContext ruleContext, String currentProtoSourceRoot, String attributeName) {
     NestedSetBuilder<String> protoSourceRoots = NestedSetBuilder.stableOrder();
     if (currentProtoSourceRoot != null && !currentProtoSourceRoot.isEmpty()) {
       protoSourceRoots.add(currentProtoSourceRoot);
     }
 
     for (ProtoSourcesProvider provider :
-        ruleContext.getPrerequisites("deps", Mode.TARGET, ProtoSourcesProvider.class)) {
+        ruleContext.getPrerequisites(attributeName, Mode.TARGET, ProtoSourcesProvider.class)) {
       String protoSourceRoot = provider.getProtoSourceRoot();
       if (protoSourceRoot != null && !protoSourceRoot.isEmpty()) {
         protoSourceRoots.add(provider.getProtoSourceRoot());
@@ -162,6 +162,28 @@ public class ProtoCommon {
     return protoSourceRoots.build();
   }
 
+  /**
+   * 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) {
+    return getProtoSourceRootsOfAttribute(ruleContext, currentProtoSourceRoot, "deps");
+  }
+
+  /**
+   * Returns a set of the {@code proto_source_root} collected from the current library and the
+   * exported dependencies.
+   *
+   * <p>Assumes {@code currentProtoSourceRoot} is the same as the package name.
+   */
+  public static NestedSet<String> getProtoSourceRootsOfExportedDependencies(
+      RuleContext ruleContext, String currentProtoSourceRoot) {
+    return getProtoSourceRootsOfAttribute(ruleContext, currentProtoSourceRoot, "exports");
+  }
+
   private static void checkProtoSourceRootIsTheSameAsPackage(
       String protoSourceRoot, RuleContext ruleContext) {
     if (!ruleContext.getLabel().getPackageName().equals(protoSourceRoot)) {
@@ -272,6 +294,19 @@ public class ProtoCommon {
   }
 
   /**
+   * Returns the .proto files that are the direct srcs of the exported dependencies of this rule.
+   */
+  @Nullable
+  public static NestedSet<Artifact> computeProtosInExportedDeps(RuleContext ruleContext) {
+    NestedSetBuilder<Artifact> result = NestedSetBuilder.stableOrder();
+    for (ProtoSupportDataProvider provider :
+        ruleContext.getPrerequisites("exports", TARGET, ProtoSupportDataProvider.class)) {
+      result.addTransitive(provider.getSupportData().getProtosInDirectDeps());
+    }
+    return result.build();
+  }
+
+  /**
    * Decides whether this proto_library should check for strict proto deps.
    *
    * <p>Takes into account command-line flags, package-level attributes and rule attributes.