bazel: handle proto_src_root when dealing with proto includes, generated files and C++ headers

This change completes the handling of proto_src_root when it comes to inclusion of protos, generating the proto files in the right place and adding the generated headers to the include paths.

WANT_LGTM=elenairina
RELNOTES: None.
PiperOrigin-RevId: 208457740

7 files changed, 95 insertions, 46 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CcCompilationHelper.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CcCompilationHelper.java
index 474ea556fc..22ddbbfddf 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CcCompilationHelper.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CcCompilationHelper.java
@@ -51,7 +51,6 @@ import com.google.devtools.build.lib.rules.cpp.CcToolchainVariables.VariablesExt
 import com.google.devtools.build.lib.rules.cpp.CppConfiguration.HeadersCheckingMode;
 import com.google.devtools.build.lib.skylarkbuildapi.cpp.CompilationInfoApi;
 import com.google.devtools.build.lib.syntax.SkylarkNestedSet;
-import com.google.devtools.build.lib.syntax.Type;
 import com.google.devtools.build.lib.util.FileTypeSet;
 import com.google.devtools.build.lib.util.Pair;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
@@ -232,6 +231,8 @@ public final class CcCompilationHelper {
   private boolean generateNoPicAction;
   private boolean generatePicAction;
   private boolean allowCoverageInstrumentation = true;
+  private String stripIncludePrefix = null;
+  private String includePrefix = null;
 
   // TODO(plf): Pull out of class.
   private CcCompilationContext ccCompilationContext;
@@ -705,6 +706,18 @@ public final class CcCompilationHelper {
     return this;
   }
 
+  /** Sets the include prefix to append to the public headers. */
+  public CcCompilationHelper setIncludePrefix(@Nullable String includePrefix) {
+    this.includePrefix = includePrefix;
+    return this;
+  }
+
+  /** Sets the include prefix to remove from the public headers. */
+  public CcCompilationHelper setStripIncludePrefix(@Nullable String stripIncludePrefix) {
+    this.stripIncludePrefix = stripIncludePrefix;
+    return this;
+  }
+
   public void setAllowCoverageInstrumentation(boolean allowCoverageInstrumentation) {
     this.allowCoverageInstrumentation = allowCoverageInstrumentation;
   }
@@ -811,33 +824,23 @@ public final class CcCompilationHelper {
   }
 
   private PublicHeaders computePublicHeaders() {
-    if (!ruleContext.attributes().has("strip_include_prefix", Type.STRING)
-        || !ruleContext.attributes().has("include_prefix", Type.STRING)) {
-      return new PublicHeaders(
-          ImmutableList.copyOf(Iterables.concat(publicHeaders, nonModuleMapHeaders)),
-          ImmutableList.copyOf(publicHeaders),
-          null);
-    }
-
     PathFragment prefix = null;
-    if (ruleContext.attributes().isAttributeValueExplicitlySpecified("include_prefix")) {
-      String prefixAttr = ruleContext.attributes().get("include_prefix", Type.STRING);
-      prefix = PathFragment.create(prefixAttr);
-      if (PathFragment.containsUplevelReferences(prefixAttr)) {
-        ruleContext.attributeError("include_prefix", "should not contain uplevel references");
+    if (includePrefix != null) {
+      prefix = PathFragment.create(includePrefix);
+      if (PathFragment.containsUplevelReferences(includePrefix)) {
+        ruleContext.ruleError("include prefix should not contain uplevel references");
       }
       if (prefix.isAbsolute()) {
-        ruleContext.attributeError("include_prefix", "should be a relative path");
+        ruleContext.ruleError("include prefix should be a relative path");
       }
     }
 
     PathFragment stripPrefix;
-    if (ruleContext.attributes().isAttributeValueExplicitlySpecified("strip_include_prefix")) {
-      String stripPrefixAttr = ruleContext.attributes().get("strip_include_prefix", Type.STRING);
-      if (PathFragment.containsUplevelReferences(stripPrefixAttr)) {
-        ruleContext.attributeError("strip_include_prefix", "should not contain uplevel references");
+    if (stripIncludePrefix != null) {
+      if (PathFragment.containsUplevelReferences(stripIncludePrefix)) {
+        ruleContext.ruleError("strip include prefix should not contain uplevel references");
       }
-      stripPrefix = PathFragment.create(stripPrefixAttr);
+      stripPrefix = PathFragment.create(stripIncludePrefix);
       if (stripPrefix.isAbsolute()) {
         stripPrefix =
             ruleContext
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CcLibrary.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CcLibrary.java
index 4266f667f1..319f7d84f0 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CcLibrary.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CcLibrary.java
@@ -175,6 +175,16 @@ public abstract class CcLibrary implements RuleConfiguredTargetFactory {
       compilationHelper.addPublicTextualHeaders(
           ruleContext.getPrerequisiteArtifacts("textual_hdrs", Mode.TARGET).list());
     }
+    if (ruleContext.getRule().isAttrDefined("include_prefix", Type.STRING)
+        && ruleContext.attributes().isAttributeValueExplicitlySpecified("include_prefix")) {
+      compilationHelper.setIncludePrefix(
+          ruleContext.attributes().get("include_prefix", Type.STRING));
+    }
+    if (ruleContext.getRule().isAttrDefined("strip_include_prefix", Type.STRING)
+        && ruleContext.attributes().isAttributeValueExplicitlySpecified("strip_include_prefix")) {
+      compilationHelper.setStripIncludePrefix(
+          ruleContext.attributes().get("strip_include_prefix", Type.STRING));
+    }
 
     if (common.getLinkopts().contains("-static")) {
       ruleContext.attributeWarning("linkopts", "Using '-static' here won't work. "
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/proto/CcProtoAspect.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/proto/CcProtoAspect.java
index f48e1de8f9..83d0c7d288 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/proto/CcProtoAspect.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/proto/CcProtoAspect.java
@@ -321,6 +321,7 @@ public abstract class CcProtoAspect extends NativeAspectClass implements Configu
     }
 
     private void createProtoCompileAction(SupportData supportData, Collection<Artifact> outputs) {
+      String protoRoot = supportData.getProtoSourceRoot();
       String genfilesPath =
           ruleContext
               .getConfiguration()
@@ -331,6 +332,7 @@ public abstract class CcProtoAspect extends NativeAspectClass implements Configu
                       .getPackageIdentifier()
                       .getRepository()
                       .getPathUnderExecRoot())
+              .getRelative(protoRoot == null ? "" : protoRoot)
               .getPathString();
 
       ImmutableList.Builder<ToolchainInvocation> invocations = ImmutableList.builder();
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 1c8a7d52cb..000aa5ec53 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
@@ -60,6 +60,7 @@ public class BazelProtoLibrary implements RuleConfiguredTargetFactory {
             protosInDirectDeps,
             transitiveImports,
             protoPathFlags,
+            protoSourceRoot,
             directProtoSourceRoots,
             !protoSources.isEmpty());
 
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 ed138d0381..da43b6a2a4 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
@@ -47,6 +47,7 @@ import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.build.lib.util.LazyString;
+import com.google.devtools.build.lib.vfs.PathFragment;
 import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
@@ -570,7 +571,11 @@ public class ProtoCompileActionBuilder {
       @Nullable NestedSet<Artifact> protosInDirectDependencies,
       NestedSet<String> directProtoSourceRoots,
       NestedSet<Artifact> transitiveImports) {
-    commandLine.addAll(VectorArg.of(transitiveImports).mapped(EXPAND_TRANSITIVE_IMPORT_ARG));
+    // For each import, include both the import as well as the import relativized against its
+    // protoSourceRoot. This ensures that protos can reference either the full path or the short
+    // path when including other protos.
+    commandLine.addAll(
+        VectorArg.of(transitiveImports).mapped(new ExpandImportArgsFn(directProtoSourceRoots)));
     if (protosInDirectDependencies != null) {
       if (!protosInDirectDependencies.isEmpty()) {
         commandLine.addAll(
@@ -590,11 +595,31 @@ public class ProtoCompileActionBuilder {
   static final CommandLineItem.MapFn<String> EXPAND_TRANSITIVE_PROTO_PATH_FLAGS =
       (flag, args) -> args.accept("--proto_path=" + flag);
 
-  @AutoCodec @AutoCodec.VisibleForSerialization
-  static final CommandLineItem.MapFn<Artifact> EXPAND_TRANSITIVE_IMPORT_ARG =
-      (artifact, args) ->
-          args.accept(
-              "-I" + getPathIgnoringRepository(artifact) + "=" + artifact.getExecPathString());
+  @AutoCodec
+  @AutoCodec.VisibleForSerialization
+  static final class ExpandImportArgsFn implements CapturingMapFn<Artifact> {
+    private final NestedSet<String> directProtoSourceRoots;
+
+    public ExpandImportArgsFn(NestedSet<String> directProtoSourceRoots) {
+      this.directProtoSourceRoots = directProtoSourceRoots;
+    }
+
+    /**
+     * Generates up to two import flags for each artifact: one for full path (only relative to the
+     * repository root) and one for the path relative to the proto source root (if one exists
+     * corresponding to the artifact).
+     */
+    @Override
+    public void expandToCommandLine(Artifact proto, Consumer<String> args) {
+      for (String directProtoSourceRoot : directProtoSourceRoots) {
+        String path = getPathIgnoringSourceRoot(proto, directProtoSourceRoot);
+        if (path != null) {
+          args.accept("-I" + path + "=" + proto.getExecPathString());
+        }
+      }
+      args.accept("-I" + getPathIgnoringRepository(proto) + "=" + proto.getExecPathString());
+    }
+  }
 
   @AutoCodec
   @AutoCodec.VisibleForSerialization
@@ -608,16 +633,15 @@ public class ProtoCompileActionBuilder {
     @Override
     public void expandToCommandLine(Artifact proto, Consumer<String> args) {
       for (String directProtoSourceRoot : directProtoSourceRoots) {
-        expandToPathIgnoringSourceRoot(proto, directProtoSourceRoot, args);
+        String path = getPathIgnoringSourceRoot(proto, directProtoSourceRoot);
+        if (path != null) {
+          args.accept(path);
+        }
       }
-      EXPAND_TO_PATH_IGNORING_REPOSITORY.expandToCommandLine(proto, args);
+      args.accept(getPathIgnoringRepository(proto));
     }
   }
 
-  @AutoCodec @AutoCodec.VisibleForSerialization
-  static final CommandLineItem.MapFn<Artifact> EXPAND_TO_PATH_IGNORING_REPOSITORY =
-      (artifact, args) -> args.accept(getPathIgnoringRepository(artifact));
-
   /**
    * Gets the artifact's path relative to the root, ignoring the external repository the artifact is
    * at. For example, <code>
@@ -633,27 +657,25 @@ public class ProtoCompileActionBuilder {
         .toString();
   }
 
-  private static void expandToPathIgnoringSourceRoot(
-      Artifact artifact, String directProtoSourceRoot, Consumer<String> args) {
+  /**
+   * Gets the artifact's path relative to the proto source root, ignoring the external repository
+   * the artifact is at. For example, <code>
+   * //a/b/c:d.proto with proto source root a/b --> c/d.proto
+   * {@literal @}foo//a/b/c:d.proto with proto source root a/b --> c/d.proto
+   * </code>
+   */
+  private static String getPathIgnoringSourceRoot(Artifact artifact, String directProtoSourceRoot) {
     // 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);
+      return PathFragment.createAlreadyNormalized(getPathIgnoringRepository(artifact))
+          .relativeTo(directProtoSourceRoot)
+          .toString();
     } catch (IllegalArgumentException exception) {
       // do nothing
     }
+    return null;
   }
 
   /**
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 c3bc3aa553..b3644d4c17 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,
+      String protoSourceRoot,
       NestedSet<String> directProtoSourceRoots,
       boolean hasProtoSources) {
     return new AutoValue_SupportData(
@@ -43,6 +44,7 @@ public abstract class SupportData {
         transitiveImports,
         protosInDirectDeps,
         transitiveProtoPathFlags,
+        protoSourceRoot,
         directProtoSourceRoots,
         hasProtoSources);
   }
@@ -64,6 +66,9 @@ public abstract class SupportData {
    */
   public abstract NestedSet<String> getTransitiveProtoPathFlags();
 
+  /** The {@code proto_source_root} of the current library. */
+  public abstract String getProtoSourceRoot();
+
   /**
    * The {@code proto_source_root}'s collected from the current library and the direct dependencies.
    */
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoCompileActionBuilderTest.java b/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoCompileActionBuilderTest.java
index 1e2311fa84..2044f279ad 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoCompileActionBuilderTest.java
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoCompileActionBuilderTest.java
@@ -83,6 +83,7 @@ public class ProtoCompileActionBuilderTest {
                 artifact("//:dont-care", "import1.proto"),
                 artifact("//:dont-care", "import2.proto")),
             /*transitiveProtoPathFlags=*/ NestedSetBuilder.<String>emptySet(STABLE_ORDER),
+            /*protoSourceRoot=*/ "",
             /*directProtoSourceRoots=*/ NestedSetBuilder.<String>stableOrder().build(),
             /* hasProtoSources= */ true);
 
@@ -122,6 +123,7 @@ public class ProtoCompileActionBuilderTest {
             /* protosInDirectDeps= */ NestedSetBuilder.<Artifact>emptySet(STABLE_ORDER),
             /* transitiveImports= */ NestedSetBuilder.<Artifact>emptySet(STABLE_ORDER),
             /*transitiveProtoPathFlags=*/ NestedSetBuilder.<String>emptySet(STABLE_ORDER),
+            /*protoSourceRoot=*/ "",
             /*directProtoSourceRoots=*/ NestedSetBuilder.<String>stableOrder().build(),
             /* hasProtoSources= */ true);
 
@@ -159,6 +161,7 @@ public class ProtoCompileActionBuilderTest {
                 artifact("//:dont-care", "import1.proto"),
                 artifact("//:dont-care", "import2.proto")),
             /*transitiveProtoPathFlags=*/ NestedSetBuilder.<String>emptySet(STABLE_ORDER),
+            /*protoSourceRoot=*/ "",
             /*directProtoSourceRoots=*/ NestedSetBuilder.<String>stableOrder().build(),
             /* hasProtoSources= */ true);
 
@@ -195,6 +198,7 @@ public class ProtoCompileActionBuilderTest {
             /* protosInDirectDeps= */ NestedSetBuilder.<Artifact>emptySet(STABLE_ORDER),
             NestedSetBuilder.<Artifact>emptySet(STABLE_ORDER),
             /*transitiveProtoPathFlags=*/ NestedSetBuilder.<String>emptySet(STABLE_ORDER),
+            /*protoSourceRoot=*/ "",
             /*directProtoSourceRoots=*/ NestedSetBuilder.<String>stableOrder().build(),
             /* hasProtoSources= */ true);
 
@@ -241,6 +245,7 @@ public class ProtoCompileActionBuilderTest {
             /* protosInDirectDeps= */ NestedSetBuilder.<Artifact>emptySet(STABLE_ORDER),
             NestedSetBuilder.<Artifact>emptySet(STABLE_ORDER),
             /*transitiveProtoPathFlags=*/ NestedSetBuilder.<String>emptySet(STABLE_ORDER),
+            /*protoSourceRoot=*/ "",
             /*directProtoSourceRoots=*/ NestedSetBuilder.<String>stableOrder().build(),
             /* hasProtoSources= */ true);
 
@@ -274,6 +279,7 @@ public class ProtoCompileActionBuilderTest {
             /* protosInDirectDeps= */ NestedSetBuilder.<Artifact>emptySet(STABLE_ORDER),
             NestedSetBuilder.<Artifact>emptySet(STABLE_ORDER),
             /*transitiveProtoPathFlags=*/ NestedSetBuilder.<String>emptySet(STABLE_ORDER),
+            /*protoSourceRoot=*/ "",
             /*directProtoSourceRoots=*/ NestedSetBuilder.<String>stableOrder().build(),
             /* hasProtoSources= */ true);