Remove ProtoSourceInfo and use "_virtual_import" substring for import…

… paths

This reduces heap memory usage, because we can use depsets of Files, instead of depsets of ProtoSourceInfo, which has a new instance of class with an Object array.

The effect of memory saving might be bigger in Bazel, because import_prefixes might be more commonly used.

The drawback is, that we can only strip prefixes using renaming to _virtual_imports subdirectory.

Alternatively we could keep ProtoSourceInfo and implement at least strip_import_prefix without renames. Measurements don't show any improvement in this case.

PiperOrigin-RevId: 537050475
Change-Id: Iff4ed7e0d4e4dfaab79449178c4b1278c1f369d0

src/main/java/com/google/devtools/build/lib/rules/proto/ProtoInfo.java

@@ -21,8 +21,6 @@
 import com.google.devtools.build.lib.collect.nestedset.Depset;
 import com.google.devtools.build.lib.collect.nestedset.Depset.TypeException;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
-import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
-import com.google.devtools.build.lib.collect.nestedset.Order;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.packages.Info;
 import com.google.devtools.build.lib.packages.RuleClass.ConfiguredTargetFactory.RuleErrorException;
@@ -76,16 +74,12 @@ public NestedSet<Artifact> getTransitiveProtoSources() {
 
   /** The {@code .proto} source files in this {@code proto_library}'s {@code srcs}. */
   @VisibleForTesting
-  public ImmutableList<ProtoSource> getDirectSources() throws Exception {
-    ImmutableList.Builder<ProtoSource> directSources = new ImmutableList.Builder<>();
-    for (StarlarkInfo protoSource :
-        Sequence.cast(
+  public ImmutableList<Artifact> getDirectSources() throws Exception {
+    return Sequence.cast(
             value.getValue("_direct_proto_sources", Sequence.class),
-            StarlarkInfo.class,
-            "_direct_proto_sources")) {
-      directSources.add(ProtoSource.create(protoSource));
-    }
-    return directSources.build();
+            Artifact.class,
+            "_direct_proto_sources")
+        .getImmutableList();
   }
 
   /** The proto source files that are used in compiling this {@code proto_library}. */
@@ -127,12 +121,7 @@ public NestedSet<Artifact> getTransitiveDescriptorSets() throws Exception {
    * directly depending on this {@code ProtoInfo}.
    */
   @VisibleForTesting
-  public NestedSet<ProtoSource> getExportedSources() throws Exception {
-    ImmutableList.Builder<ProtoSource> exportedSources = new ImmutableList.Builder<>();
-    for (StarlarkInfo protoSource :
-        value.getValue("_exported_sources", Depset.class).toList(StarlarkInfo.class)) {
-      exportedSources.add(ProtoSource.create(protoSource));
-    }
-    return NestedSetBuilder.wrap(Order.STABLE_ORDER, exportedSources.build());
+  public NestedSet<Artifact> getExportedSources() throws Exception {
+    return value.getValue("_exported_sources", Depset.class).getSet(Artifact.class);
   }
 }

src/main/starlark/builtins_bzl/common/proto/proto_common.bzl

@@ -35,7 +35,7 @@ def _proto_path_flag(path):
     return "--proto_path=%s" % path
 
 def _Iimport_path_equals_fullpath(proto_source):
-    return "-I%s=%s" % (get_import_path(proto_source), proto_source._source_file.path)
+    return "-I%s=%s" % (get_import_path(proto_source), proto_source.path)
 
 def _remove_repo(file):
     """Removes `../repo/` prefix from path, e.g. `../repo/package/path -> package/path`"""
@@ -55,11 +55,12 @@ def get_import_path(proto_source):
     Returns:
       (str) import path
     """
-    proto_path = proto_source._proto_path
-    short_path = _remove_repo(proto_source._source_file)
-    if proto_path and not short_path.startswith(proto_path + "/"):
-        fail("Bad proto_path %s for proto %s" % (proto_path, short_path))
-    return short_path.removeprefix(proto_path + "/")
+    repo_path = _remove_repo(proto_source)
+    index = repo_path.find("_virtual_imports/")
+    if index >= 0:
+        index = repo_path.find("/", index + len("_virtual_imports/"))
+        repo_path = repo_path[index + 1:]
+    return repo_path
 
 def _compile(
         actions,
@@ -151,7 +152,7 @@ def _experimental_filter_sources(proto_info, proto_lang_toolchain_info):
     provided_proto_sources = proto_lang_toolchain_info.provided_proto_sources
     provided_paths = {}
     for src in provided_proto_sources:
-        path = src._source_file.path
+        path = src.path
 
         # For listed protos bundled with the Bazel tools repository, their exec paths start
         # with external/bazel_tools/. This prefix needs to be removed first, because the protos in
@@ -162,7 +163,7 @@ def _experimental_filter_sources(proto_info, proto_lang_toolchain_info):
             provided_paths[path] = None
 
     # Filter proto files
-    proto_files = [src._source_file for src in proto_info._direct_proto_sources]
+    proto_files = proto_info._direct_proto_sources
     excluded = []
     included = []
     for proto_file in proto_files:

src/main/starlark/builtins_bzl/common/proto/proto_info.bzl

@@ -51,7 +51,9 @@ def _create_proto_info(*, srcs, deps, descriptor_set, proto_path = "", workspace
       srcs: ([File]) List of .proto files (possibly under _virtual path)
       deps: ([ProtoInfo]) List of dependencies
       descriptor_set: (File) Descriptor set for this Proto
-      proto_path: (str) Path that should be stripped from files in srcs.
+      proto_path: (str) Path that should be stripped from files in srcs. When
+        stripping is needed, the files should be symlinked into `_virtual_imports/target_name`
+        directory. Only such paths are accepted.
       workspace_root: (str) Set to ctx.workspace_root if this is not the main repository.
       genfiles_dir: (str) Set to ctx.genfiles_dir if _virtual_imports are used.
 
@@ -70,10 +72,15 @@ def _create_proto_info(*, srcs, deps, descriptor_set, proto_path = "", workspace
             fail("srcs parameter expects all files start with %s" % src_prefix)
     if type(descriptor_set) != "File":
         fail("descriptor_set parameter expected to be a File")
-    if "_virtual_imports/" in proto_path and not genfiles_dir:
-        fail("genfiles_dir parameter should be set when _virtual_imports are used")
-
-    direct_proto_sources = [_ProtoSourceInfo(_source_file = src, _proto_path = proto_path) for src in srcs]
+    if proto_path:
+        if "_virtual_imports/" not in proto_path:
+            fail("proto_path needs to contain '_virtual_imports' directory")
+        if proto_path.split("/")[-2] != "_virtual_imports":
+            fail("proto_path needs to be formed like '_virtual_imports/target_name'")
+        if not genfiles_dir:
+            fail("genfiles_dir parameter should be set when _virtual_imports are used")
+
+    direct_proto_sources = srcs
     transitive_proto_sources = depset(
         direct = direct_proto_sources,
         transitive = [dep._transitive_proto_sources for dep in deps],
@@ -170,21 +177,12 @@ ProtoInfo, _ = provider(
         "transitive_imports": """(depset[File]) Deprecated: use `transitive_sources` instead.""",
 
         # Internal fields:
-        "_direct_proto_sources": """(list[ProtoSourceInfo]) The `ProtoSourceInfo`s from the `srcs`
+        "_direct_proto_sources": """(list[File]) The `ProtoSourceInfo`s from the `srcs`
             attribute.""" + _warning,
-        "_transitive_proto_sources": """(depset[ProtoSourceInfo]) The `ProtoSourceInfo`s from this
+        "_transitive_proto_sources": """(depset[File]) The `ProtoSourceInfo`s from this
             rule and all its dependent protocol buffer rules.""" + _warning,
-        "_exported_sources": """(depset[ProtoSourceInfo]) A set of `ProtoSourceInfo`s that may be
+        "_exported_sources": """(depset[File]) A set of `ProtoSourceInfo`s that may be
             imported by another `proto_library` depending on this one.""" + _warning,
     },
     init = _create_proto_info,
 )
-
-_ProtoSourceInfo = provider(
-    doc = "Represents a single `.proto` source file.",
-    fields = {
-        "_source_file": """(File) The `.proto` file. Possibly virtual to handle additional/stripped
-          path prefix.""" + _warning,
-        "_proto_path": "(str) The root of the virtual location." + _warning,
-    },
-)

src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoLibraryTest.java

@@ -518,13 +518,8 @@ public void testExportedStrippedImportPrefixes() throws Exception {
     assertThat(
             Iterables.transform(
                 a.get(ProtoInfo.PROVIDER).getExportedSources().toList(),
-                s -> s.getSourceRoot().getSafePathString()))
-        .containsExactly("a/_virtual_imports/a");
-    assertThat(
-            Iterables.transform(
-                a.get(ProtoInfo.PROVIDER).getExportedSources().toList(),
-                s -> s.getImportPath().getSafePathString()))
-        .containsExactly("a.proto");
+                s -> s.getRootRelativePathString()))
+        .containsExactly("a/_virtual_imports/a/a.proto");
   }
 
   private void testImportPrefixInExternalRepo(boolean siblingRepoLayout) throws Exception {
@@ -554,10 +549,9 @@ private void testImportPrefixInExternalRepo(boolean siblingRepoLayout) throws Ex
 
     ConfiguredTarget target = getConfiguredTarget("@yolo_repo//yolo_pkg:yolo_proto");
     assertThat(
-            Iterables.transform(
-                target.get(ProtoInfo.PROVIDER).getExportedSources().toList(),
-                s -> s.getImportPath().getPathString()))
-        .contains("bazel.build/yolo/yolo_pkg/yolo.proto");
+            Iterables.getOnlyElement(target.get(ProtoInfo.PROVIDER).getExportedSources().toList())
+                .getExecPathString())
+        .endsWith("/_virtual_imports/yolo_proto/bazel.build/yolo/yolo_pkg/yolo.proto");
   }
 
   @Test
@@ -599,10 +593,9 @@ private void testImportPrefixAndStripInExternalRepo(boolean siblingRepoLayout) t
     ConfiguredTarget target =
         getConfiguredTarget("@yolo_repo//yolo_pkg_to_be_stripped/yolo_pkg:yolo_proto");
     assertThat(
-            Iterables.transform(
-                target.get(ProtoInfo.PROVIDER).getExportedSources().toList(),
-                s -> s.getImportPath().getPathString()))
-        .contains("bazel.build/yolo/yolo_pkg/yolo.proto");
+            Iterables.getOnlyElement(target.get(ProtoInfo.PROVIDER).getExportedSources().toList())
+                .getExecPathString())
+        .endsWith("/_virtual_imports/yolo_proto/bazel.build/yolo/yolo_pkg/yolo.proto");
   }
 
   @Test
@@ -643,10 +636,9 @@ private void testStripImportPrefixInExternalRepo(boolean siblingRepoLayout) thro
     ConfiguredTarget target =
         getConfiguredTarget("@yolo_repo//yolo_pkg_to_be_stripped/yolo_pkg:yolo_proto");
     assertThat(
-            Iterables.transform(
-                target.get(ProtoInfo.PROVIDER).getExportedSources().toList(),
-                s -> s.getImportPath().getPathString()))
-        .contains("yolo_pkg/yolo.proto");
+            Iterables.getOnlyElement(target.get(ProtoInfo.PROVIDER).getExportedSources().toList())
+                .getExecPathString())
+        .endsWith("/_virtual_imports/yolo_proto/yolo_pkg/yolo.proto");
   }
 
   @Test
@@ -687,10 +679,9 @@ private void testRelativeStripImportPrefixInExternalRepo(boolean siblingRepoLayo
 
     ConfiguredTarget target = getConfiguredTarget("@yolo_repo//:yolo_proto");
     assertThat(
-            Iterables.transform(
-                target.get(ProtoInfo.PROVIDER).getExportedSources().toList(),
-                s -> s.getImportPath().getPathString()))
-        .contains("yolo_pkg/yolo.proto");
+            Iterables.getOnlyElement(target.get(ProtoInfo.PROVIDER).getExportedSources().toList())
+                .getExecPathString())
+        .endsWith("/_virtual_imports/yolo_proto/yolo_pkg/yolo.proto");
   }
 
   @Test
@@ -1115,17 +1106,7 @@ public void testProtoLibrary() throws Exception {
     scratch.file("x/BUILD", "proto_library(name='foo', srcs=['a.proto', 'b.proto', 'c.proto'])");
 
     ProtoInfo provider = getConfiguredTarget("//x:foo").get(ProtoInfo.PROVIDER);
-    assertThat(
-            Iterables.transform(
-                provider.getDirectSources(), s -> s.getSourceFile().getExecPath().getPathString()))
-        .containsExactly("x/a.proto", "x/b.proto", "x/c.proto");
-    assertThat(
-            Iterables.transform(
-                provider.getDirectSources(), s -> s.getSourceRoot().getSafePathString()))
-        .containsExactly(".", ".", ".");
-    assertThat(
-            Iterables.transform(
-                provider.getDirectSources(), s -> s.getImportPath().getPathString()))
+    assertThat(Iterables.transform(provider.getDirectSources(), s -> s.getExecPathString()))
         .containsExactly("x/a.proto", "x/b.proto", "x/c.proto");
   }
 
@@ -1147,18 +1128,8 @@ public void testProtoLibraryWithVirtualProtoSourceRoot() throws Exception {
 
     String genfiles = getTargetConfiguration().getGenfilesFragment(RepositoryName.MAIN).toString();
     ProtoInfo provider = getConfiguredTarget("//x:foo").get(ProtoInfo.PROVIDER);
-    assertThat(
-            Iterables.transform(
-                provider.getDirectSources(), s -> s.getSourceFile().getExecPath().getPathString()))
+    assertThat(Iterables.transform(provider.getDirectSources(), s -> s.getExecPathString()))
         .containsExactly(genfiles + "/x/_virtual_imports/foo/foo/x/a.proto");
-    assertThat(
-            Iterables.transform(
-                provider.getDirectSources(), s -> s.getSourceRoot().getSafePathString()))
-        .containsExactly("x/_virtual_imports/foo");
-    assertThat(
-            Iterables.transform(
-                provider.getDirectSources(), s -> s.getImportPath().getPathString()))
-        .containsExactly("foo/x/a.proto");
   }
 
   @Test
@@ -1174,18 +1145,8 @@ public void testProtoLibraryWithGeneratedSources_Blaze() throws Exception {
 
     String genfiles = getTargetConfiguration().getGenfilesFragment(RepositoryName.MAIN).toString();
     ProtoInfo provider = getConfiguredTarget("//x:foo").get(ProtoInfo.PROVIDER);
-    assertThat(
-            Iterables.transform(
-                provider.getDirectSources(), s -> s.getSourceFile().getExecPath().getPathString()))
+    assertThat(Iterables.transform(provider.getDirectSources(), s -> s.getExecPathString()))
         .containsExactly(genfiles + "/x/generated.proto");
-    assertThat(
-            Iterables.transform(
-                provider.getDirectSources(), s -> s.getSourceRoot().getSafePathString()))
-        .containsExactly(".");
-    assertThat(
-            Iterables.transform(
-                provider.getDirectSources(), s -> s.getImportPath().getPathString()))
-        .containsExactly("x/generated.proto");
   }
 
   @Test
@@ -1201,17 +1162,7 @@ public void testProtoLibraryWithMixedSources_Blaze() throws Exception {
 
     String genfiles = getTargetConfiguration().getGenfilesFragment(RepositoryName.MAIN).toString();
     ProtoInfo provider = getConfiguredTarget("//x:foo").get(ProtoInfo.PROVIDER);
-    assertThat(
-            Iterables.transform(
-                provider.getDirectSources(), s -> s.getSourceFile().getExecPath().getPathString()))
+    assertThat(Iterables.transform(provider.getDirectSources(), s -> s.getExecPathString()))
         .containsExactly(genfiles + "/x/generated.proto", "x/a.proto");
-    assertThat(
-            Iterables.transform(
-                provider.getDirectSources(), s -> s.getSourceRoot().getSafePathString()))
-        .containsExactly(".", ".");
-    assertThat(
-            Iterables.transform(
-                provider.getDirectSources(), s -> s.getImportPath().getPathString()))
-        .containsExactly("x/generated.proto", "x/a.proto");
   }
 }