From d9b403e7642720264eefe021cffa48674a00284b Mon Sep 17 00:00:00 2001
From: ilist <ilist@google.com>
Date: Thu, 11 Nov 2021 07:26:15 -0800
Subject: [PATCH 01/24] Put protoc label to a constant.

Renamed StrictProtoDepsViolationMessage to ProtoConstants and added protoc label there.

PiperOrigin-RevId: 409142099
---
 .../devtools/build/lib/rules/proto/BUILD      |  1 +
 .../proto/ProtoCompileActionBuilder.java      |  2 +-
 .../lib/rules/proto/ProtoConfiguration.java   |  8 +++++-
 ...lationMessage.java => ProtoConstants.java} | 28 +++++++++++--------
 4 files changed, 26 insertions(+), 13 deletions(-)
 rename src/main/java/com/google/devtools/build/lib/rules/proto/{StrictProtoDepsViolationMessage.java => ProtoConstants.java} (51%)

diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/BUILD b/src/main/java/com/google/devtools/build/lib/rules/proto/BUILD
index 6c45232c43407d..25e76419a7a5de 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/BUILD
@@ -38,6 +38,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/analysis:rule_definition_environment",
         "//src/main/java/com/google/devtools/build/lib/analysis:transitive_info_collection",
         "//src/main/java/com/google/devtools/build/lib/analysis:transitive_info_provider",
+        "//src/main/java/com/google/devtools/build/lib/analysis/starlark/annotations",
         "//src/main/java/com/google/devtools/build/lib/analysis/stringtemplate",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/collect/nestedset",
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 7f4b2daa3ec174..3b27faea755aa9 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
@@ -53,7 +53,7 @@
 public class ProtoCompileActionBuilder {
   @VisibleForTesting
   public static final String STRICT_DEPS_FLAG_TEMPLATE =
-      "--direct_dependencies_violation_msg=" + StrictProtoDepsViolationMessage.MESSAGE;
+      "--direct_dependencies_violation_msg=" + ProtoConstants.STRICT_PROTO_DEPS_VIOLATION_MESSAGE;
 
   private static final String MNEMONIC = "GenProto";
 
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoConfiguration.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoConfiguration.java
index 722b6ab41b1692..e626684e8df74b 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoConfiguration.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoConfiguration.java
@@ -21,6 +21,7 @@
 import com.google.devtools.build.lib.analysis.config.Fragment;
 import com.google.devtools.build.lib.analysis.config.FragmentOptions;
 import com.google.devtools.build.lib.analysis.config.RequiresOptions;
+import com.google.devtools.build.lib.analysis.starlark.annotations.StarlarkConfigurationField;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.starlarkbuildapi.ProtoConfigurationApi;
@@ -38,6 +39,7 @@
 // configuration fragment in aspect definitions.
 @RequiresOptions(options = {ProtoConfiguration.Options.class})
 public class ProtoConfiguration extends Fragment implements ProtoConfigurationApi {
+
   /** Command line options. */
   public static class Options extends FragmentOptions {
     @Option(
@@ -81,7 +83,7 @@ public static class Options extends FragmentOptions {
 
     @Option(
         name = "proto_compiler",
-        defaultValue = "@com_google_protobuf//:protoc",
+        defaultValue = ProtoConstants.DEFAULT_PROTOC_LABEL,
         converter = CoreOptionConverters.LabelConverter.class,
         documentationCategory = OptionDocumentationCategory.UNCATEGORIZED,
         effectTags = {OptionEffectTag.AFFECTS_OUTPUTS, OptionEffectTag.LOADING_AND_ANALYSIS},
@@ -220,6 +222,10 @@ public boolean runExperimentalProtoExtraActions() {
     return options.experimentalProtoExtraActions;
   }
 
+  @StarlarkConfigurationField(
+      name = "proto_compiler",
+      doc = "Label for the proto compiler.",
+      defaultLabel = ProtoConstants.DEFAULT_PROTOC_LABEL)
   public Label protoCompiler() {
     return options.protoCompiler;
   }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/StrictProtoDepsViolationMessage.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoConstants.java
similarity index 51%
rename from src/main/java/com/google/devtools/build/lib/rules/proto/StrictProtoDepsViolationMessage.java
rename to src/main/java/com/google/devtools/build/lib/rules/proto/ProtoConstants.java
index 6d6d03844d1b64..419756cae76a92 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/StrictProtoDepsViolationMessage.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoConstants.java
@@ -14,16 +14,22 @@
 
 package com.google.devtools.build.lib.rules.proto;
 
-/**
- * This class is used in ProtoCompileActionBuilder to generate an error message that's displayed
- * when a strict proto deps violation occurs.
- *
- * <p>%1$s is replaced with the label of the proto_library rule that's currently being built.
- *
- * <p>%%s is replaced with the literal "%s", which is passed to the proto-compiler, which replaces
- * it with the .proto file that violates strict proto deps.
- */
-public class StrictProtoDepsViolationMessage {
-  static final String MESSAGE =
+/** Constants used in Proto rules. */
+public final class ProtoConstants {
+  /** Default label for proto compiler. */
+  static final String DEFAULT_PROTOC_LABEL =  "@com_google_protobuf//:protoc";
+
+  /**
+   * This constant is used in ProtoCompileActionBuilder to generate an error message that's
+   * displayed when a strict proto deps violation occurs.
+   *
+   * <p>%1$s is replaced with the label of the proto_library rule that's currently being built.
+   *
+   * <p>%%s is replaced with the literal "%s", which is passed to the proto-compiler, which replaces
+   * it with the .proto file that violates strict proto deps.
+   */
+  static final String STRICT_PROTO_DEPS_VIOLATION_MESSAGE =
       "%%s is imported, but %1$s doesn't directly depend on a proto_library that 'srcs' it.";
+
+  private ProtoConstants() {}
 }

From 16b62f5df549e163fde1f99dae77784a1c4a0d1a Mon Sep 17 00:00:00 2001
From: ilist <ilist@google.com>
Date: Thu, 25 Nov 2021 08:12:20 -0800
Subject: [PATCH 02/24] Remove proto_lang_toolchain rule's $(PLUGIN_OUT)
 placeholder and simplify ProtoCompileActionBuilder.

`$(OUT)` placeholder is replaced with `%s` so it can be directly used in addFormatted. This simplifies construction of proto compile action and makes it possible for Starlark version to have the same performance.

github shows no uses of the placeholder: https://github.com/search?q=%22%24%28PLUGIN_OUT%29%22&type=Repositories

PiperOrigin-RevId: 412286743
---
 .../build/lib/rules/proto/ProtoCommon.java    |  4 +---
 .../proto/ProtoCompileActionBuilder.java      | 20 +++++++++----------
 .../lib/rules/proto/ProtoLangToolchain.java   | 11 +++++++++-
 .../proto/ProtoLangToolchainProvider.java     | 18 ++++++++++++-----
 .../rules/proto/ProtoLangToolchainRule.java   |  4 +---
 .../proto/ProtoCompileActionBuilderTest.java  | 10 +++++-----
 .../rules/proto/ProtoLangToolchainTest.java   |  2 +-
 7 files changed, 40 insertions(+), 29 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 ced81605610e6c..0d468dd69d4a75 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
@@ -37,9 +37,7 @@
 import com.google.devtools.build.lib.vfs.PathFragment;
 import javax.annotation.Nullable;
 
-/**
- * Utility functions for proto_library and proto aspect implementations.
- */
+/** Utility functions for proto_library and proto aspect implementations. */
 public class ProtoCommon {
   private ProtoCommon() {
     throw new UnsupportedOperationException();
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 3b27faea755aa9..5d7f63613fffa9 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
@@ -21,7 +21,6 @@
 import com.google.common.base.Preconditions;
 import com.google.common.base.Supplier;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.actions.AbstractAction;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.CommandLineItem;
@@ -362,7 +361,7 @@ private static ToolchainInvocation createDescriptorSetToolchain(
             // output size to become quadratic, so don't.
             // A rule that concatenates the artifacts from ctx.deps.proto.transitive_descriptor_sets
             // provides similar results.
-            "--descriptor_set_out=$(OUT)",
+            "--descriptor_set_out=%s",
             /* pluginExecutable= */ null,
             /* runtime= */ null,
             /* providedProtoSources= */ ImmutableList.of()),
@@ -488,8 +487,6 @@ public static boolean arePublicImportsStrict(RuleContext ruleContext) {
    * <ul>
    *   <li>Each toolchain contributes a command-line, formatted from its commandLine() method.
    *   <li>$(OUT) is replaced with the outReplacement field of ToolchainInvocation.
-   *   <li>$(PLUGIN_out) is replaced with PLUGIN_<key>_out where 'key' is the key of
-   *       toolchainInvocations. The key thus allows multiple plugins in one command-line.
    *   <li>If a toolchain's {@code plugin()} is non-null, we point at it by emitting
    *       --plugin=protoc-gen-PLUGIN_<key>=<location of plugin>.
    * </ul>
@@ -533,14 +530,15 @@ static CustomCommandLine createCommandLineFromToolchains(
 
       ProtoLangToolchainProvider toolchain = invocation.toolchain;
 
+      final String formatString = toolchain.outReplacementFormatFlag();
+      final CharSequence outReplacement = invocation.outReplacement;
       cmdLine.addLazyString(
-          new OnDemandCommandLineExpansion(
-              toolchain.commandLine(),
-              ImmutableMap.of(
-                  "OUT",
-                  invocation.outReplacement,
-                  "PLUGIN_OUT",
-                  String.format("PLUGIN_%s_out", invocation.name))));
+          new OnDemandString() {
+            @Override
+            public String toString() {
+              return String.format(formatString, outReplacement);
+            }
+          });
 
       if (toolchain.pluginExecutable() != null) {
         cmdLine.addFormatted(
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java
index fbc331d35be35d..db7ad8fa74a636 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java
@@ -39,10 +39,19 @@ public ConfiguredTarget create(RuleContext ruleContext)
       providedProtoSources.addTransitive(protoInfo.getTransitiveSources());
     }
 
+    String flag = ruleContext.attributes().get("command_line", Type.STRING);
+
+    if (flag.contains("$(PLUGIN_OUT)")) {
+      ruleContext.attributeError("command_line", "Placeholder '$(PLUGIN_OUT)' is not supported.");
+      return null;
+    }
+
+    flag = flag.replace("$(OUT)", "%s");
+
     return new RuleConfiguredTargetBuilder(ruleContext)
         .addProvider(
             ProtoLangToolchainProvider.create(
-                ruleContext.attributes().get("command_line", Type.STRING),
+                flag,
                 ruleContext.getPrerequisite("plugin", FilesToRunProvider.class),
                 ruleContext.getPrerequisite("runtime"),
                 // We intentionally flatten the NestedSet here.
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
index 0a6bd07462e738..1567eb5550ada5 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
@@ -34,7 +34,7 @@
 @AutoValue
 @AutoCodec
 public abstract class ProtoLangToolchainProvider implements TransitiveInfoProvider {
-  public abstract String commandLine();
+  public abstract String outReplacementFormatFlag();
 
   @Nullable
   public abstract FilesToRunProvider pluginExecutable();
@@ -63,17 +63,21 @@ public NestedSet<Artifact> blacklistedProtos() {
 
   @AutoCodec.Instantiator
   public static ProtoLangToolchainProvider createForDeserialization(
-      String commandLine,
+      String outReplacementFormatFlag,
       FilesToRunProvider pluginExecutable,
       TransitiveInfoCollection runtime,
       ImmutableList<ProtoSource> providedProtoSources,
       NestedSet<Artifact> blacklistedProtos) {
     return new AutoValue_ProtoLangToolchainProvider(
-        commandLine, pluginExecutable, runtime, providedProtoSources, blacklistedProtos);
+        outReplacementFormatFlag,
+        pluginExecutable,
+        runtime,
+        providedProtoSources,
+        blacklistedProtos);
   }
 
   public static ProtoLangToolchainProvider create(
-      String commandLine,
+      String outReplacementFormatFlag,
       FilesToRunProvider pluginExecutable,
       TransitiveInfoCollection runtime,
       ImmutableList<ProtoSource> providedProtoSources) {
@@ -82,6 +86,10 @@ public static ProtoLangToolchainProvider create(
       blacklistedProtos.add(protoSource.getOriginalSourceFile());
     }
     return new AutoValue_ProtoLangToolchainProvider(
-        commandLine, pluginExecutable, runtime, providedProtoSources, blacklistedProtos.build());
+        outReplacementFormatFlag,
+        pluginExecutable,
+        runtime,
+        providedProtoSources,
+        blacklistedProtos.build());
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java
index b05895452a5716..4d040c86356744 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java
@@ -39,8 +39,6 @@ public RuleClass build(RuleClass.Builder builder, RuleDefinitionEnvironment envi
           <li><code>$(OUT)</code> is LANG_proto_library-specific. The rules are expected to define
               how they interpret this variable. For Java, for example, $(OUT) will be replaced with
               the src-jar filename to create.</li>
-          <li><code>$(PLUGIN_out)</code> will be substituted to work with a
-              `--plugin=protoc-gen-PLUGIN` command line.</li>
         </ul>
         <!-- #END_BLAZE_RULE.ATTRIBUTE --> */
         .add(attr("command_line", Type.STRING).mandatory())
@@ -122,7 +120,7 @@ Specifies how a LANG_proto_library rule (e.g., <code>java_proto_library</code>)
 <pre class="code">
 proto_lang_toolchain(
     name = "javalite_toolchain",
-    command_line = "--$(PLUGIN_OUT)=shared,immutable:$(OUT)",
+    command_line = "--javalite_out=shared,immutable:$(OUT)",
     plugin = ":javalite_plugin",
     runtime = ":protobuf_lite",
 )
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 ce9edd7450bf05..e82a7f6eebf246 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
@@ -104,14 +104,14 @@ public void commandLine_basic() throws Exception {
 
     ProtoLangToolchainProvider toolchainNoPlugin =
         ProtoLangToolchainProvider.create(
-            "--java_out=param1,param2:$(OUT)",
+            "--java_out=param1,param2:%s",
             /* pluginExecutable= */ null,
             /* runtime= */ mock(TransitiveInfoCollection.class),
             /* providedProtoSources= */ ImmutableList.of());
 
     ProtoLangToolchainProvider toolchainWithPlugin =
         ProtoLangToolchainProvider.create(
-            "--$(PLUGIN_OUT)=param3,param4:$(OUT)",
+            "--PLUGIN_pluginName_out=param3,param4:%s",
             plugin,
             /* runtime= */ mock(TransitiveInfoCollection.class),
             /* providedProtoSources= */ ImmutableList.of());
@@ -174,7 +174,7 @@ public void commandline_derivedArtifact() throws Exception {
   public void commandLine_strictDeps() throws Exception {
     ProtoLangToolchainProvider toolchain =
         ProtoLangToolchainProvider.create(
-            "--java_out=param1,param2:$(OUT)",
+            "--java_out=param1,param2:%s",
             /* pluginExecutable= */ null,
             /* runtime= */ mock(TransitiveInfoCollection.class),
             /* providedProtoSources= */ ImmutableList.of());
@@ -212,7 +212,7 @@ public void commandLine_strictDeps() throws Exception {
   public void commandLine_exports() throws Exception {
     ProtoLangToolchainProvider toolchain =
         ProtoLangToolchainProvider.create(
-            "--java_out=param1,param2:$(OUT)",
+            "--java_out=param1,param2:%s",
             /* pluginExecutable= */ null,
             /* runtime= */ mock(TransitiveInfoCollection.class),
             /* providedProtoSources= */ ImmutableList.of());
@@ -282,7 +282,7 @@ public String toString() {
 
     ProtoLangToolchainProvider toolchain =
         ProtoLangToolchainProvider.create(
-            "--java_out=param1,param2:$(OUT)",
+            "--java_out=param1,param2:%s",
             /* pluginExecutable= */ null,
             /* runtime= */ mock(TransitiveInfoCollection.class),
             /* providedProtoSources= */ ImmutableList.of());
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java b/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
index 55c49ea1d8689f..0fc6dcb69ad5f5 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
@@ -40,7 +40,7 @@ public void setUp() throws Exception {
   }
 
   private void validateProtoLangToolchain(ProtoLangToolchainProvider toolchain) throws Exception {
-    assertThat(toolchain.commandLine()).isEqualTo("cmd-line");
+    assertThat(toolchain.outReplacementFormatFlag()).isEqualTo("cmd-line");
     assertThat(toolchain.pluginExecutable().getExecutable().getRootRelativePathString())
         .isEqualTo("third_party/x/plugin");
 

From 4158e5259ef5615dad4e354dd501b470381310eb Mon Sep 17 00:00:00 2001
From: ilist <ilist@google.com>
Date: Thu, 25 Nov 2021 08:15:25 -0800
Subject: [PATCH 03/24] Add plugin_format_flag attribute to
 ProtoLangToolchainRule

This makes it possible to pair it with command_line attribute which could contain dependent data,
for example:

```
proto_lang_toolchain(
    name = "j2objc_proto_toolchain",
    blacklisted_protos = [],
    command_line = "--PLUGIN_j2objc_out=file_dir_mapping,generate_class_mappings:$(OUT)",
+   plugin_format_flag = "--plugin=protoc-gen-PLUGIN_j2objc=%s",
    plugin = "//third_party/java/j2objc:proto_plugin",
    runtime = "//third_party/java/j2objc:proto_runtime",
    visibility = ["//visibility:public"],
)
```

It also retains reference to the flag on proto_lang_toolchain rule and so doesn't cause a memory regression when proto_lang_libraries are Starlarkyfied.

PiperOrigin-RevId: 412287195
---
 .../build/lib/rules/proto/ProtoCompileActionBuilder.java   | 1 +
 .../devtools/build/lib/rules/proto/ProtoLangToolchain.java | 1 +
 .../build/lib/rules/proto/ProtoLangToolchainProvider.java  | 7 +++++++
 .../build/lib/rules/proto/ProtoLangToolchainRule.java      | 7 +++++++
 .../lib/rules/proto/ProtoCompileActionBuilderTest.java     | 7 +++++++
 .../build/lib/rules/proto/ProtoLangToolchainTest.java      | 4 ++++
 6 files changed, 27 insertions(+)

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 5d7f63613fffa9..fe8301b4bd3454 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
@@ -362,6 +362,7 @@ private static ToolchainInvocation createDescriptorSetToolchain(
             // A rule that concatenates the artifacts from ctx.deps.proto.transitive_descriptor_sets
             // provides similar results.
             "--descriptor_set_out=%s",
+            /* pluginFormatFlag = */ null,
             /* pluginExecutable= */ null,
             /* runtime= */ null,
             /* providedProtoSources= */ ImmutableList.of()),
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java
index db7ad8fa74a636..4ab6296920f70b 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java
@@ -52,6 +52,7 @@ public ConfiguredTarget create(RuleContext ruleContext)
         .addProvider(
             ProtoLangToolchainProvider.create(
                 flag,
+                ruleContext.attributes().get("plugin_format_flag", Type.STRING),
                 ruleContext.getPrerequisite("plugin", FilesToRunProvider.class),
                 ruleContext.getPrerequisite("runtime"),
                 // We intentionally flatten the NestedSet here.
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
index 1567eb5550ada5..3fb68a361bcf0d 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
@@ -36,6 +36,9 @@
 public abstract class ProtoLangToolchainProvider implements TransitiveInfoProvider {
   public abstract String outReplacementFormatFlag();
 
+  @Nullable
+  public abstract String pluginFormatFlag();
+
   @Nullable
   public abstract FilesToRunProvider pluginExecutable();
 
@@ -64,12 +67,14 @@ public NestedSet<Artifact> blacklistedProtos() {
   @AutoCodec.Instantiator
   public static ProtoLangToolchainProvider createForDeserialization(
       String outReplacementFormatFlag,
+      String pluginFormatFlag,
       FilesToRunProvider pluginExecutable,
       TransitiveInfoCollection runtime,
       ImmutableList<ProtoSource> providedProtoSources,
       NestedSet<Artifact> blacklistedProtos) {
     return new AutoValue_ProtoLangToolchainProvider(
         outReplacementFormatFlag,
+        pluginFormatFlag,
         pluginExecutable,
         runtime,
         providedProtoSources,
@@ -78,6 +83,7 @@ public static ProtoLangToolchainProvider createForDeserialization(
 
   public static ProtoLangToolchainProvider create(
       String outReplacementFormatFlag,
+      String pluginFormatFlag,
       FilesToRunProvider pluginExecutable,
       TransitiveInfoCollection runtime,
       ImmutableList<ProtoSource> providedProtoSources) {
@@ -87,6 +93,7 @@ public static ProtoLangToolchainProvider create(
     }
     return new AutoValue_ProtoLangToolchainProvider(
         outReplacementFormatFlag,
+        pluginFormatFlag,
         pluginExecutable,
         runtime,
         providedProtoSources,
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java
index 4d040c86356744..324d5e90ce2e09 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java
@@ -43,6 +43,13 @@ public RuleClass build(RuleClass.Builder builder, RuleDefinitionEnvironment envi
         <!-- #END_BLAZE_RULE.ATTRIBUTE --> */
         .add(attr("command_line", Type.STRING).mandatory())
 
+        /* <!-- #BLAZE_RULE(proto_lang_toolchain).ATTRIBUTE(plugin_format_flag) -->
+        If provided, this value will be passed to proto-compiler to use the plugin. The value must
+        contain a single %s which is replaced with plugin executable.
+        <code>--plugin=protoc-gen-PLUGIN=<executable>.</code>
+        <!-- #END_BLAZE_RULE.ATTRIBUTE --> */
+        .add(attr("plugin_format_flag", Type.STRING))
+
         /* <!-- #BLAZE_RULE(proto_lang_toolchain).ATTRIBUTE(plugin) -->
         If provided, will be made available to the action that calls the proto-compiler, and will be
         passed to the proto-compiler:
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 e82a7f6eebf246..4dff20c5ae1f22 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
@@ -105,6 +105,7 @@ public void commandLine_basic() throws Exception {
     ProtoLangToolchainProvider toolchainNoPlugin =
         ProtoLangToolchainProvider.create(
             "--java_out=param1,param2:%s",
+            /* pluginFormatFlag= */ null,
             /* pluginExecutable= */ null,
             /* runtime= */ mock(TransitiveInfoCollection.class),
             /* providedProtoSources= */ ImmutableList.of());
@@ -112,6 +113,7 @@ public void commandLine_basic() throws Exception {
     ProtoLangToolchainProvider toolchainWithPlugin =
         ProtoLangToolchainProvider.create(
             "--PLUGIN_pluginName_out=param3,param4:%s",
+            /* pluginFormatFlag= */ "--plugin=protoc-gen-PLUGIN_pluginName=%s",
             plugin,
             /* runtime= */ mock(TransitiveInfoCollection.class),
             /* providedProtoSources= */ ImmutableList.of());
@@ -175,6 +177,7 @@ public void commandLine_strictDeps() throws Exception {
     ProtoLangToolchainProvider toolchain =
         ProtoLangToolchainProvider.create(
             "--java_out=param1,param2:%s",
+            /* pluginFormatFlag= */ null,
             /* pluginExecutable= */ null,
             /* runtime= */ mock(TransitiveInfoCollection.class),
             /* providedProtoSources= */ ImmutableList.of());
@@ -213,6 +216,7 @@ public void commandLine_exports() throws Exception {
     ProtoLangToolchainProvider toolchain =
         ProtoLangToolchainProvider.create(
             "--java_out=param1,param2:%s",
+            /* pluginFormatFlag= */ null,
             /* pluginExecutable= */ null,
             /* runtime= */ mock(TransitiveInfoCollection.class),
             /* providedProtoSources= */ ImmutableList.of());
@@ -283,6 +287,7 @@ public String toString() {
     ProtoLangToolchainProvider toolchain =
         ProtoLangToolchainProvider.create(
             "--java_out=param1,param2:%s",
+            /* pluginFormatFlag= */ null,
             /* pluginExecutable= */ null,
             /* runtime= */ mock(TransitiveInfoCollection.class),
             /* providedProtoSources= */ ImmutableList.of());
@@ -317,6 +322,7 @@ public void exceptionIfSameName() throws Exception {
     ProtoLangToolchainProvider toolchain1 =
         ProtoLangToolchainProvider.create(
             "dontcare",
+            /* pluginFormatFlag= */ null,
             /* pluginExecutable= */ null,
             /* runtime= */ mock(TransitiveInfoCollection.class),
             /* providedProtoSources= */ ImmutableList.of());
@@ -324,6 +330,7 @@ public void exceptionIfSameName() throws Exception {
     ProtoLangToolchainProvider toolchain2 =
         ProtoLangToolchainProvider.create(
             "dontcare",
+            /* pluginFormatFlag= */ null,
             /* pluginExecutable= */ null,
             /* runtime= */ mock(TransitiveInfoCollection.class),
             /* providedProtoSources= */ ImmutableList.of());
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java b/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
index 0fc6dcb69ad5f5..d6eb3a8124caca 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
@@ -41,6 +41,7 @@ public void setUp() throws Exception {
 
   private void validateProtoLangToolchain(ProtoLangToolchainProvider toolchain) throws Exception {
     assertThat(toolchain.outReplacementFormatFlag()).isEqualTo("cmd-line");
+    assertThat(toolchain.pluginFormatFlag()).isEqualTo("--plugin=%s");
     assertThat(toolchain.pluginExecutable().getExecutable().getRootRelativePathString())
         .isEqualTo("third_party/x/plugin");
 
@@ -73,6 +74,7 @@ public void protoToolchain() throws Exception {
         "proto_lang_toolchain(",
         "    name = 'toolchain',",
         "    command_line = 'cmd-line',",
+        "    plugin_format_flag = '--plugin=%s',",
         "    plugin = '//third_party/x:plugin',",
         "    runtime = '//third_party/x:runtime',",
         "    blacklisted_protos = ['//third_party/x:denied']",
@@ -101,6 +103,7 @@ public void protoToolchainBlacklistProtoLibraries() throws Exception {
         "proto_lang_toolchain(",
         "    name = 'toolchain',",
         "    command_line = 'cmd-line',",
+        "    plugin_format_flag = '--plugin=%s',",
         "    plugin = '//third_party/x:plugin',",
         "    runtime = '//third_party/x:runtime',",
         "    blacklisted_protos = ['//third_party/x:descriptors', '//third_party/x:any']",
@@ -129,6 +132,7 @@ public void protoToolchainBlacklistTransitiveProtos() throws Exception {
         "proto_lang_toolchain(",
         "    name = 'toolchain',",
         "    command_line = 'cmd-line',",
+        "    plugin_format_flag = '--plugin=%s',",
         "    plugin = '//third_party/x:plugin',",
         "    runtime = '//third_party/x:runtime',",
         "    blacklisted_protos = ['//third_party/x:any']",

From 4099a4300494322d310c925657885745c83d6e79 Mon Sep 17 00:00:00 2001
From: ilist <ilist@google.com>
Date: Mon, 28 Mar 2022 03:35:30 -0700
Subject: [PATCH 04/24] Proxy proto compiler and proto opts over
 proto_lang_toolchain rule.

This is needed for the proto_common.generate_code function.

This change doesn't modify any of proto_lang_toolchain public attributes.

PiperOrigin-RevId: 437713619
---
 .../proto/ProtoCompileActionBuilder.java      |  8 +--
 .../lib/rules/proto/ProtoLangToolchain.java   |  4 +-
 .../proto/ProtoLangToolchainProvider.java     | 54 ++++++++++++-------
 .../rules/proto/ProtoLangToolchainRule.java   | 18 +++++++
 .../proto/ProtoCompileActionBuilderTest.java  | 34 ++++++++----
 .../rules/proto/ProtoLangToolchainTest.java   |  6 +++
 6 files changed, 89 insertions(+), 35 deletions(-)

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 fe8301b4bd3454..b64a07a5d2d96e 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
@@ -331,7 +331,7 @@ public static void writeDescriptorSet(
             protoToolchain,
             ImmutableList.of(
                 createDescriptorSetToolchain(
-                    ruleContext.getFragment(ProtoConfiguration.class), output.getExecPathString())),
+                    ruleContext.getFragment(ProtoConfiguration.class), output.getExecPathString(), protoToolchain)),
             protoInfo,
             ruleContext.getLabel(),
             ImmutableList.of(output),
@@ -348,7 +348,7 @@ public static void writeDescriptorSet(
   }
 
   private static ToolchainInvocation createDescriptorSetToolchain(
-      ProtoConfiguration configuration, CharSequence outReplacement) {
+      ProtoConfiguration configuration, CharSequence outReplacement, ProtoToolchainInfo protoToolchain) {
     ImmutableList.Builder<String> protocOpts = ImmutableList.builder();
     if (configuration.experimentalProtoDescriptorSetsIncludeSourceInfo()) {
       protocOpts.add("--include_source_info");
@@ -365,7 +365,9 @@ private static ToolchainInvocation createDescriptorSetToolchain(
             /* pluginFormatFlag = */ null,
             /* pluginExecutable= */ null,
             /* runtime= */ null,
-            /* providedProtoSources= */ ImmutableList.of()),
+            /* providedProtoSources= */ ImmutableList.of(),
+            /* protoCompiler= */ protoToolchain.getCompiler(),
+            /* protoCopts=*/ configuration.protocOpts()),
         outReplacement,
         protocOpts.build());
   }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java
index 4ab6296920f70b..cc5fb1ebb32c8e 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java
@@ -62,7 +62,9 @@ public ConfiguredTarget create(RuleContext ruleContext)
                 // targets. While this has the potential to use more memory than using a NestedSet,
                 // there are only a few `proto_lang_toolchain` targets in every build, so the impact
                 // on memory consumption should be neglectable.
-                providedProtoSources.build().toList()))
+                providedProtoSources.build().toList(),
+                ruleContext.getPrerequisite(":proto_compiler", FilesToRunProvider.class),
+                ruleContext.getConfiguration().getFragment(ProtoConfiguration.class).protocOpts()))
         .setFilesToBuild(NestedSetBuilder.<Artifact>emptySet(STABLE_ORDER))
         .addProvider(RunfilesProvider.simple(Runfiles.EMPTY))
         .build();
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
index 3fb68a361bcf0d..bd07c214cba3be 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
@@ -22,8 +22,9 @@
 import com.google.devtools.build.lib.analysis.TransitiveInfoProvider;
 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 javax.annotation.Nullable;
+import net.starlark.java.annot.StarlarkMethod;
+import net.starlark.java.eval.StarlarkList;
 
 // Note: AutoValue v1.4-rc1 has AutoValue.CopyAnnotations which makes it work with Starlark. No need
 // to un-AutoValue this class to expose it to Starlark.
@@ -32,13 +33,26 @@
  * rules.
  */
 @AutoValue
-@AutoCodec
 public abstract class ProtoLangToolchainProvider implements TransitiveInfoProvider {
+  @StarlarkMethod(
+      name = "out_replacement_format_flag",
+      doc = "Format string used when passing output to the plugin used by proto compiler.",
+      structField = true)
   public abstract String outReplacementFormatFlag();
 
+  @StarlarkMethod(
+      name = "plugin_format_flag",
+      doc = "Format string used when passing plugin to proto compiler.",
+      structField = true,
+      allowReturnNones = true)
   @Nullable
   public abstract String pluginFormatFlag();
 
+  @StarlarkMethod(
+      name = "plugin",
+      doc = "Proto compiler plugin.",
+      structField = true,
+      allowReturnNones = true)
   @Nullable
   public abstract FilesToRunProvider pluginExecutable();
 
@@ -51,6 +65,19 @@ public abstract class ProtoLangToolchainProvider implements TransitiveInfoProvid
    */
   public abstract ImmutableList<ProtoSource> providedProtoSources();
 
+  @StarlarkMethod(name = "proto_compiler", doc = "Proto compiler.", structField = true)
+  public abstract FilesToRunProvider protoc();
+
+  @StarlarkMethod(
+      name = "protoc_opts",
+      doc = "Options to pass to proto compiler.",
+      structField = true)
+  public StarlarkList<String> protocOptsForStarlark() {
+    return StarlarkList.immutableCopyOf(protocOpts());
+  }
+
+  public abstract ImmutableList<String> protocOpts();
+
   /**
    * This makes the blacklisted_protos member available in the provider. It can be removed after
    * users are migrated and a sufficient time for Bazel rules to migrate has elapsed.
@@ -64,29 +91,14 @@ public NestedSet<Artifact> blacklistedProtos() {
   @Deprecated
   public abstract NestedSet<Artifact> forbiddenProtos();
 
-  @AutoCodec.Instantiator
-  public static ProtoLangToolchainProvider createForDeserialization(
-      String outReplacementFormatFlag,
-      String pluginFormatFlag,
-      FilesToRunProvider pluginExecutable,
-      TransitiveInfoCollection runtime,
-      ImmutableList<ProtoSource> providedProtoSources,
-      NestedSet<Artifact> blacklistedProtos) {
-    return new AutoValue_ProtoLangToolchainProvider(
-        outReplacementFormatFlag,
-        pluginFormatFlag,
-        pluginExecutable,
-        runtime,
-        providedProtoSources,
-        blacklistedProtos);
-  }
-
   public static ProtoLangToolchainProvider create(
       String outReplacementFormatFlag,
       String pluginFormatFlag,
       FilesToRunProvider pluginExecutable,
       TransitiveInfoCollection runtime,
-      ImmutableList<ProtoSource> providedProtoSources) {
+      ImmutableList<ProtoSource> providedProtoSources,
+      FilesToRunProvider protoc,
+      ImmutableList<String> protocOpts) {
     NestedSetBuilder<Artifact> blacklistedProtos = NestedSetBuilder.stableOrder();
     for (ProtoSource protoSource : providedProtoSources) {
       blacklistedProtos.add(protoSource.getOriginalSourceFile());
@@ -97,6 +109,8 @@ public static ProtoLangToolchainProvider create(
         pluginExecutable,
         runtime,
         providedProtoSources,
+        protoc,
+        protocOpts,
         blacklistedProtos.build());
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java
index 324d5e90ce2e09..f5256ef2814efe 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java
@@ -22,12 +22,25 @@
 import com.google.devtools.build.lib.analysis.RuleDefinition;
 import com.google.devtools.build.lib.analysis.RuleDefinitionEnvironment;
 import com.google.devtools.build.lib.analysis.config.ExecutionTransitionFactory;
+import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.packages.Attribute;
 import com.google.devtools.build.lib.packages.RuleClass;
 import com.google.devtools.build.lib.packages.StarlarkProviderIdentifier;
 import com.google.devtools.build.lib.packages.Type;
 
 /** Implements {code proto_lang_toolchain}. */
 public class ProtoLangToolchainRule implements RuleDefinition {
+  private static final Label DEFAULT_PROTO_COMPILER =
+      Label.parseAbsoluteUnchecked(ProtoConstants.DEFAULT_PROTOC_LABEL);
+  private static final Attribute.LabelLateBoundDefault<?> PROTO_COMPILER =
+      Attribute.LabelLateBoundDefault.fromTargetConfiguration(
+          ProtoConfiguration.class,
+          DEFAULT_PROTO_COMPILER,
+          (rule, attributes, protoConfig) ->
+              protoConfig.protoCompiler() != null
+                  ? protoConfig.protoCompiler()
+                  : DEFAULT_PROTO_COMPILER);
+
   @Override
   public RuleClass build(RuleClass.Builder builder, RuleDefinitionEnvironment environment) {
     return builder
@@ -78,6 +91,11 @@ public RuleClass build(RuleClass.Builder builder, RuleDefinitionEnvironment envi
             attr("blacklisted_protos", LABEL_LIST)
                 .allowedFileTypes()
                 .mandatoryProviders(StarlarkProviderIdentifier.forKey(ProtoInfo.PROVIDER.getKey())))
+        .add(
+            attr(":proto_compiler", LABEL)
+                .cfg(ExecutionTransitionFactory.create())
+                .exec()
+                .value(PROTO_COMPILER))
         .requiresConfigurationFragments(ProtoConfiguration.class)
         .advertiseProvider(ProtoLangToolchainProvider.class)
         .removeAttribute("data")
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 4dff20c5ae1f22..eb2aea83e6d565 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
@@ -108,15 +108,18 @@ public void commandLine_basic() throws Exception {
             /* pluginFormatFlag= */ null,
             /* pluginExecutable= */ null,
             /* runtime= */ mock(TransitiveInfoCollection.class),
-            /* providedProtoSources= */ ImmutableList.of());
-
+            /* providedProtoSources= */ ImmutableList.of(),
+            /* protoc= */ FilesToRunProvider.EMPTY,
+            /* protocOpts= */ ImmutableList.of());
     ProtoLangToolchainProvider toolchainWithPlugin =
         ProtoLangToolchainProvider.create(
             "--PLUGIN_pluginName_out=param3,param4:%s",
             /* pluginFormatFlag= */ "--plugin=protoc-gen-PLUGIN_pluginName=%s",
             plugin,
             /* runtime= */ mock(TransitiveInfoCollection.class),
-            /* providedProtoSources= */ ImmutableList.of());
+            /* providedProtoSources= */ ImmutableList.of(),
+            /* protoc= */ FilesToRunProvider.EMPTY,
+            /* protocOpts= */ ImmutableList.of());
 
     CustomCommandLine cmdLine =
         createCommandLineFromToolchains(
@@ -180,7 +183,9 @@ public void commandLine_strictDeps() throws Exception {
             /* pluginFormatFlag= */ null,
             /* pluginExecutable= */ null,
             /* runtime= */ mock(TransitiveInfoCollection.class),
-            /* providedProtoSources= */ ImmutableList.of());
+            /* providedProtoSources= */ ImmutableList.of(),
+            /* protoc= */ FilesToRunProvider.EMPTY,
+            /* protocOpts= */ ImmutableList.of());
 
     CustomCommandLine cmdLine =
         createCommandLineFromToolchains(
@@ -219,7 +224,9 @@ public void commandLine_exports() throws Exception {
             /* pluginFormatFlag= */ null,
             /* pluginExecutable= */ null,
             /* runtime= */ mock(TransitiveInfoCollection.class),
-            /* providedProtoSources= */ ImmutableList.of());
+            /* providedProtoSources= */ ImmutableList.of(),
+            /* protoc= */ FilesToRunProvider.EMPTY,
+            /* protocOpts= */ ImmutableList.of());
 
     CustomCommandLine cmdLine =
         createCommandLineFromToolchains(
@@ -290,7 +297,9 @@ public String toString() {
             /* pluginFormatFlag= */ null,
             /* pluginExecutable= */ null,
             /* runtime= */ mock(TransitiveInfoCollection.class),
-            /* providedProtoSources= */ ImmutableList.of());
+            /* providedProtoSources= */ ImmutableList.of(),
+            /* protoc= */ FilesToRunProvider.EMPTY,
+            /* protocOpts= */ ImmutableList.of());
 
     CustomCommandLine cmdLine =
         createCommandLineFromToolchains(
@@ -321,19 +330,22 @@ public String toString() {
   public void exceptionIfSameName() throws Exception {
     ProtoLangToolchainProvider toolchain1 =
         ProtoLangToolchainProvider.create(
-            "dontcare",
+            "dontcare=%s",
             /* pluginFormatFlag= */ null,
             /* pluginExecutable= */ null,
             /* runtime= */ mock(TransitiveInfoCollection.class),
-            /* providedProtoSources= */ ImmutableList.of());
-
+            /* providedProtoSources= */ ImmutableList.of(),
+            /* protoc= */ FilesToRunProvider.EMPTY,
+            /* protocOpts= */ ImmutableList.of());
     ProtoLangToolchainProvider toolchain2 =
         ProtoLangToolchainProvider.create(
-            "dontcare",
+            "dontcare=%s",
             /* pluginFormatFlag= */ null,
             /* pluginExecutable= */ null,
             /* runtime= */ mock(TransitiveInfoCollection.class),
-            /* providedProtoSources= */ ImmutableList.of());
+            /* providedProtoSources= */ ImmutableList.of(),
+            /* protoc= */ FilesToRunProvider.EMPTY,
+            /* protocOpts= */ ImmutableList.of());
 
     IllegalStateException e =
         assertThrows(
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java b/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
index d6eb3a8124caca..b03738c30f23d7 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
@@ -36,6 +36,7 @@ public class ProtoLangToolchainTest extends BuildViewTestCase {
   public void setUp() throws Exception {
     MockProtoSupport.setupWorkspace(scratch);
     MockProtoSupport.setup(mockToolsConfig);
+    useConfiguration("--protocopt=--myflag");
     invalidatePackages();
   }
 
@@ -54,6 +55,11 @@ private void validateProtoLangToolchain(ProtoLangToolchainProvider toolchain) th
             "third_party/x/metadata.proto",
             "third_party/x/descriptor.proto",
             "third_party/x/any.proto");
+
+    assertThat(toolchain.protocOpts()).containsExactly("--myflag");
+    Label protoc = Label.parseAbsoluteUnchecked(ProtoConstants.DEFAULT_PROTOC_LABEL);
+    assertThat(toolchain.protoc().getExecutable().prettyPrint())
+        .isEqualTo(protoc.toPathFragment().getPathString());
   }
 
   @Test

From d53efdd350e7eac68225a6c7b1c72d9cc38c98d1 Mon Sep 17 00:00:00 2001
From: ilist <ilist@google.com>
Date: Mon, 28 Mar 2022 03:38:18 -0700
Subject: [PATCH 05/24] Extend proto_lang_toolchain rule with progress_message
 and mnemonic attributes.

This is needed for proto_common.generate_sources function.

Two public attributes are added to proto_lang_toolchain rule. Both have defaults, so no immediate migration of targets is needed.

PiperOrigin-RevId: 437714094
---
 .../proto/ProtoCompileActionBuilder.java      |  5 +++-
 .../lib/rules/proto/ProtoLangToolchain.java   |  4 ++-
 .../proto/ProtoLangToolchainProvider.java     | 18 +++++++++++-
 .../rules/proto/ProtoLangToolchainRule.java   |  9 ++++++
 .../proto/ProtoCompileActionBuilderTest.java  | 28 ++++++++++++++-----
 .../rules/proto/ProtoLangToolchainTest.java   | 16 +++++++++--
 6 files changed, 67 insertions(+), 13 deletions(-)

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 b64a07a5d2d96e..8ebaa66cee452a 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
@@ -367,7 +367,10 @@ private static ToolchainInvocation createDescriptorSetToolchain(
             /* runtime= */ null,
             /* providedProtoSources= */ ImmutableList.of(),
             /* protoCompiler= */ protoToolchain.getCompiler(),
-            /* protoCopts=*/ configuration.protocOpts()),
+            /* protoCopts=*/ configuration.protocOpts(),
+            /* progressMessage=*/ "",
+            /* mnemonic=*/ "GenProtoDescriptorSet"),
+
         outReplacement,
         protocOpts.build());
   }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java
index cc5fb1ebb32c8e..8d409b8acc1418 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java
@@ -64,7 +64,9 @@ public ConfiguredTarget create(RuleContext ruleContext)
                 // on memory consumption should be neglectable.
                 providedProtoSources.build().toList(),
                 ruleContext.getPrerequisite(":proto_compiler", FilesToRunProvider.class),
-                ruleContext.getConfiguration().getFragment(ProtoConfiguration.class).protocOpts()))
+                ruleContext.getConfiguration().getFragment(ProtoConfiguration.class).protocOpts(),
+                ruleContext.attributes().get("progress_message", Type.STRING),
+                ruleContext.attributes().get("mnemonic", Type.STRING)))
         .setFilesToBuild(NestedSetBuilder.<Artifact>emptySet(STABLE_ORDER))
         .addProvider(RunfilesProvider.simple(Runfiles.EMPTY))
         .build();
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
index bd07c214cba3be..cf9d7971cf7a74 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
@@ -78,6 +78,18 @@ public StarlarkList<String> protocOptsForStarlark() {
 
   public abstract ImmutableList<String> protocOpts();
 
+  @StarlarkMethod(
+      name = "progress_message",
+      doc = "Progress message to set on the proto compiler action.",
+      structField = true)
+  public abstract String progressMessage();
+
+  @StarlarkMethod(
+      name = "mnemonic",
+      doc = "Mnemonic to set on the proto compiler action.",
+      structField = true)
+  public abstract String mnemonic();
+
   /**
    * This makes the blacklisted_protos member available in the provider. It can be removed after
    * users are migrated and a sufficient time for Bazel rules to migrate has elapsed.
@@ -98,7 +110,9 @@ public static ProtoLangToolchainProvider create(
       TransitiveInfoCollection runtime,
       ImmutableList<ProtoSource> providedProtoSources,
       FilesToRunProvider protoc,
-      ImmutableList<String> protocOpts) {
+      ImmutableList<String> protocOpts,
+      String progressMessage,
+      String mnemonic) {
     NestedSetBuilder<Artifact> blacklistedProtos = NestedSetBuilder.stableOrder();
     for (ProtoSource protoSource : providedProtoSources) {
       blacklistedProtos.add(protoSource.getOriginalSourceFile());
@@ -111,6 +125,8 @@ public static ProtoLangToolchainProvider create(
         providedProtoSources,
         protoc,
         protocOpts,
+        progressMessage,
+        mnemonic,
         blacklistedProtos.build());
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java
index f5256ef2814efe..70944fd613a9bb 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java
@@ -44,6 +44,15 @@ public class ProtoLangToolchainRule implements RuleDefinition {
   @Override
   public RuleClass build(RuleClass.Builder builder, RuleDefinitionEnvironment environment) {
     return builder
+        /* <!-- #BLAZE_RULE(proto_lang_toolchain).ATTRIBUTE(progress_message) -->
+        This value will be set as the progress message on protoc action.
+        <!-- #END_BLAZE_RULE.ATTRIBUTE --> */
+        .add(attr("progress_message", Type.STRING).value("Generating proto_library %{label}"))
+
+        /* <!-- #BLAZE_RULE(proto_lang_toolchain).ATTRIBUTE(mnemonic) -->
+        This value will be set as the mnemonic on protoc action.
+        <!-- #END_BLAZE_RULE.ATTRIBUTE --> */
+        .add(attr("mnemonic", Type.STRING).value("GenProto"))
 
         /* <!-- #BLAZE_RULE(proto_lang_toolchain).ATTRIBUTE(command_line) -->
         This value will be passed to proto-compiler to generate the code. Only include the parts
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 eb2aea83e6d565..5c180453a8debb 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
@@ -110,7 +110,9 @@ public void commandLine_basic() throws Exception {
             /* runtime= */ mock(TransitiveInfoCollection.class),
             /* providedProtoSources= */ ImmutableList.of(),
             /* protoc= */ FilesToRunProvider.EMPTY,
-            /* protocOpts= */ ImmutableList.of());
+            /* protocOpts= */ ImmutableList.of(),
+            /* progressMessage = */ "",
+            /* mnemonic= */ "");
     ProtoLangToolchainProvider toolchainWithPlugin =
         ProtoLangToolchainProvider.create(
             "--PLUGIN_pluginName_out=param3,param4:%s",
@@ -119,7 +121,9 @@ public void commandLine_basic() throws Exception {
             /* runtime= */ mock(TransitiveInfoCollection.class),
             /* providedProtoSources= */ ImmutableList.of(),
             /* protoc= */ FilesToRunProvider.EMPTY,
-            /* protocOpts= */ ImmutableList.of());
+            /* protocOpts= */ ImmutableList.of(),
+            /* progressMessage = */ "",
+            /* mnemonic= */ "");
 
     CustomCommandLine cmdLine =
         createCommandLineFromToolchains(
@@ -185,7 +189,9 @@ public void commandLine_strictDeps() throws Exception {
             /* runtime= */ mock(TransitiveInfoCollection.class),
             /* providedProtoSources= */ ImmutableList.of(),
             /* protoc= */ FilesToRunProvider.EMPTY,
-            /* protocOpts= */ ImmutableList.of());
+            /* protocOpts= */ ImmutableList.of(),
+            /* progressMessage = */ "",
+            /* mnemonic= */ "");
 
     CustomCommandLine cmdLine =
         createCommandLineFromToolchains(
@@ -226,7 +232,9 @@ public void commandLine_exports() throws Exception {
             /* runtime= */ mock(TransitiveInfoCollection.class),
             /* providedProtoSources= */ ImmutableList.of(),
             /* protoc= */ FilesToRunProvider.EMPTY,
-            /* protocOpts= */ ImmutableList.of());
+            /* protocOpts= */ ImmutableList.of(),
+            /* progressMessage = */ "",
+            /* mnemonic= */ "");
 
     CustomCommandLine cmdLine =
         createCommandLineFromToolchains(
@@ -299,7 +307,9 @@ public String toString() {
             /* runtime= */ mock(TransitiveInfoCollection.class),
             /* providedProtoSources= */ ImmutableList.of(),
             /* protoc= */ FilesToRunProvider.EMPTY,
-            /* protocOpts= */ ImmutableList.of());
+            /* protocOpts= */ ImmutableList.of(),
+            /* progressMessage = */ "",
+            /* mnemonic= */ "");
 
     CustomCommandLine cmdLine =
         createCommandLineFromToolchains(
@@ -336,7 +346,9 @@ public void exceptionIfSameName() throws Exception {
             /* runtime= */ mock(TransitiveInfoCollection.class),
             /* providedProtoSources= */ ImmutableList.of(),
             /* protoc= */ FilesToRunProvider.EMPTY,
-            /* protocOpts= */ ImmutableList.of());
+            /* protocOpts= */ ImmutableList.of(),
+            /* progressMessage = */ "",
+            /* mnemonic= */ "");
     ProtoLangToolchainProvider toolchain2 =
         ProtoLangToolchainProvider.create(
             "dontcare=%s",
@@ -345,7 +357,9 @@ public void exceptionIfSameName() throws Exception {
             /* runtime= */ mock(TransitiveInfoCollection.class),
             /* providedProtoSources= */ ImmutableList.of(),
             /* protoc= */ FilesToRunProvider.EMPTY,
-            /* protocOpts= */ ImmutableList.of());
+            /* protocOpts= */ ImmutableList.of(),
+            /* progressMessage = */ "",
+            /* mnemonic= */ "");
 
     IllegalStateException e =
         assertThrows(
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java b/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
index b03738c30f23d7..2d1130bd339600 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
@@ -60,6 +60,9 @@ private void validateProtoLangToolchain(ProtoLangToolchainProvider toolchain) th
     Label protoc = Label.parseAbsoluteUnchecked(ProtoConstants.DEFAULT_PROTOC_LABEL);
     assertThat(toolchain.protoc().getExecutable().prettyPrint())
         .isEqualTo(protoc.toPathFragment().getPathString());
+
+    assertThat(toolchain.progressMessage()).isEqualTo("Progress Message %{label}");
+    assertThat(toolchain.mnemonic()).isEqualTo("MyMnemonic");
   }
 
   @Test
@@ -83,7 +86,9 @@ public void protoToolchain() throws Exception {
         "    plugin_format_flag = '--plugin=%s',",
         "    plugin = '//third_party/x:plugin',",
         "    runtime = '//third_party/x:runtime',",
-        "    blacklisted_protos = ['//third_party/x:denied']",
+        "    blacklisted_protos = ['//third_party/x:denied'],",
+        "    progress_message = 'Progress Message %{label}',",
+        "    mnemonic = 'MyMnemonic',",
         ")");
 
     update(ImmutableList.of("//foo:toolchain"), false, 1, true, new EventBus());
@@ -112,7 +117,9 @@ public void protoToolchainBlacklistProtoLibraries() throws Exception {
         "    plugin_format_flag = '--plugin=%s',",
         "    plugin = '//third_party/x:plugin',",
         "    runtime = '//third_party/x:runtime',",
-        "    blacklisted_protos = ['//third_party/x:descriptors', '//third_party/x:any']",
+        "    blacklisted_protos = ['//third_party/x:descriptors', '//third_party/x:any'],",
+        "    progress_message = 'Progress Message %{label}',",
+        "    mnemonic = 'MyMnemonic',",
         ")");
 
     update(ImmutableList.of("//foo:toolchain"), false, 1, true, new EventBus());
@@ -141,7 +148,9 @@ public void protoToolchainBlacklistTransitiveProtos() throws Exception {
         "    plugin_format_flag = '--plugin=%s',",
         "    plugin = '//third_party/x:plugin',",
         "    runtime = '//third_party/x:runtime',",
-        "    blacklisted_protos = ['//third_party/x:any']",
+        "    blacklisted_protos = ['//third_party/x:any'],",
+        "    progress_message = 'Progress Message %{label}',",
+        "    mnemonic = 'MyMnemonic',",
         ")");
 
     update(ImmutableList.of("//foo:toolchain"), false, 1, true, new EventBus());
@@ -169,5 +178,6 @@ public void optionalFieldsAreEmpty() throws Exception {
     assertThat(toolchain.runtime()).isNull();
     assertThat(toolchain.blacklistedProtos().toList()).isEmpty();
     assertThat(toolchain.forbiddenProtos().toList()).isEmpty();
+    assertThat(toolchain.mnemonic()).isEqualTo("GenProto");
   }
 }

From b0181e75b8fda1713ffe979479b78eb2ade8a298 Mon Sep 17 00:00:00 2001
From: kotlaja <kotlaja@google.com>
Date: Mon, 25 Apr 2022 03:46:18 -0700
Subject: [PATCH 06/24] Starlarkify proto_lang_toolchain and
 ProtoLangToolchainInfo provider

Added StarlarkProtoLangToolchainTest which uses starlarkified rule for verification.

I've deleted blacklisted_protos and forbidden_protos since they are not needed anymore.

PiperOrigin-RevId: 444223497
---
 .../lib/rules/cpp/proto/CcProtoAspect.java    |   2 +-
 .../java/proto/JavaProtoAspectCommon.java     |   8 +-
 .../build/lib/rules/objc/J2ObjcAspect.java    |   2 +-
 .../proto/ProtoLangToolchainProvider.java     |  23 +-
 .../proto/ProtoSourceFileExcludeList.java     |   6 +-
 .../starlark/builtins_bzl/common/exports.bzl  |   2 +
 .../common/proto/proto_lang_toolchain.bzl     |  85 ++++++++
 .../builtins_bzl/common/proto/providers.bzl   |  30 +++
 .../devtools/build/lib/rules/proto/BUILD      |  20 +-
 .../rules/proto/ProtoLangToolchainTest.java   |  16 +-
 .../proto/StarlarkProtoLangToolchainTest.java | 197 ++++++++++++++++++
 11 files changed, 349 insertions(+), 42 deletions(-)
 create mode 100644 src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain.bzl
 create mode 100644 src/main/starlark/builtins_bzl/common/proto/providers.bzl
 create mode 100644 src/test/java/com/google/devtools/build/lib/rules/proto/StarlarkProtoLangToolchainTest.java

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 689d29f2623087..f66a8f1a75d1a7 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
@@ -284,7 +284,7 @@ private static void checkProtoLibrariesInDeps(
 
     private boolean areSrcsExcluded() {
       return !new ProtoSourceFileExcludeList(
-              ruleContext, getProtoToolchainProvider().forbiddenProtos())
+              ruleContext, getProtoToolchainProvider().providedProtoSources())
           .checkSrcs(protoInfo.getDirectSources(), "cc_proto_library");
     }
 
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/proto/JavaProtoAspectCommon.java b/src/main/java/com/google/devtools/build/lib/rules/java/proto/JavaProtoAspectCommon.java
index dedfb0308ba994..ffb21afa9b337a 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/proto/JavaProtoAspectCommon.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/proto/JavaProtoAspectCommon.java
@@ -205,14 +205,8 @@ boolean shouldGenerateCode(ProtoInfo protoInfo, String ruleName) {
       return false;
     }
 
-    NestedSetBuilder<Artifact> forbiddenProtos = NestedSetBuilder.stableOrder();
-    forbiddenProtos.addTransitive(getProtoToolchainProvider().forbiddenProtos());
-    if (rpcSupport != null) {
-      forbiddenProtos.addTransitive(rpcSupport.getForbiddenProtos(ruleContext));
-    }
-
     final ProtoSourceFileExcludeList protoExcludeList =
-        new ProtoSourceFileExcludeList(ruleContext, forbiddenProtos.build());
+        new ProtoSourceFileExcludeList(ruleContext, getProtoToolchainProvider().providedProtoSources());
 
     return protoExcludeList.checkSrcs(protoInfo.getDirectSources(), ruleName);
   }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/objc/J2ObjcAspect.java b/src/main/java/com/google/devtools/build/lib/rules/objc/J2ObjcAspect.java
index de89d98f883f81..dd2607cd9a2056 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/objc/J2ObjcAspect.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/objc/J2ObjcAspect.java
@@ -393,7 +393,7 @@ private ConfiguredAspect proto(ConfiguredTarget base, RuleContext ruleContext)
         ruleContext.getPrerequisite(J2OBJC_PROTO_TOOLCHAIN_ATTR, ProtoLangToolchainProvider.class);
     // Avoid pulling in any generated files from forbidden protos.
     ProtoSourceFileExcludeList protoExcludeList =
-        new ProtoSourceFileExcludeList(ruleContext, protoToolchain.forbiddenProtos());
+        new ProtoSourceFileExcludeList(ruleContext, protoToolchain.providedProtoSources());
 
     ImmutableList<Artifact> filteredProtoSources =
         ImmutableList.copyOf(protoExcludeList.filter(protoSources));
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
index cf9d7971cf7a74..edc750f76cbdec 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
@@ -16,12 +16,13 @@
 
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
-import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.analysis.FilesToRunProvider;
 import com.google.devtools.build.lib.analysis.TransitiveInfoCollection;
 import com.google.devtools.build.lib.analysis.TransitiveInfoProvider;
 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.BuiltinProvider;
+import com.google.devtools.build.lib.packages.NativeInfo;
 import javax.annotation.Nullable;
 import net.starlark.java.annot.StarlarkMethod;
 import net.starlark.java.eval.StarlarkList;
@@ -90,19 +91,6 @@ public StarlarkList<String> protocOptsForStarlark() {
       structField = true)
   public abstract String mnemonic();
 
-  /**
-   * This makes the blacklisted_protos member available in the provider. It can be removed after
-   * users are migrated and a sufficient time for Bazel rules to migrate has elapsed.
-   */
-  @Deprecated
-  public NestedSet<Artifact> blacklistedProtos() {
-    return forbiddenProtos();
-  }
-
-  // TODO(yannic): Remove after migrating all users to `providedProtoSources()`.
-  @Deprecated
-  public abstract NestedSet<Artifact> forbiddenProtos();
-
   public static ProtoLangToolchainProvider create(
       String outReplacementFormatFlag,
       String pluginFormatFlag,
@@ -113,10 +101,6 @@ public static ProtoLangToolchainProvider create(
       ImmutableList<String> protocOpts,
       String progressMessage,
       String mnemonic) {
-    NestedSetBuilder<Artifact> blacklistedProtos = NestedSetBuilder.stableOrder();
-    for (ProtoSource protoSource : providedProtoSources) {
-      blacklistedProtos.add(protoSource.getOriginalSourceFile());
-    }
     return new AutoValue_ProtoLangToolchainProvider(
         outReplacementFormatFlag,
         pluginFormatFlag,
@@ -126,7 +110,6 @@ public static ProtoLangToolchainProvider create(
         protoc,
         protocOpts,
         progressMessage,
-        mnemonic,
-        blacklistedProtos.build());
+        mnemonic);
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoSourceFileExcludeList.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoSourceFileExcludeList.java
index 98d006450fd0e9..f477825fff5a26 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoSourceFileExcludeList.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoSourceFileExcludeList.java
@@ -54,12 +54,12 @@ public class ProtoSourceFileExcludeList {
    * @param noGenerateProtoFiles a list of .proto files. The list will be iterated.
    */
   public ProtoSourceFileExcludeList(
-      RuleContext ruleContext, NestedSet<Artifact> noGenerateProtoFiles) {
+      RuleContext ruleContext, ImmutableList<ProtoSource> noGenerateProtoFiles) {
     this.ruleContext = ruleContext;
     ImmutableSet.Builder<PathFragment> noGenerateProtoFilePathsBuilder =
         new ImmutableSet.Builder<>();
-    for (Artifact noGenerateProtoFile : noGenerateProtoFiles.toList()) {
-      PathFragment execPath = noGenerateProtoFile.getExecPath();
+    for (ProtoSource noGenerateProtoFile : noGenerateProtoFiles) {
+      PathFragment execPath = noGenerateProtoFile.getOriginalSourceFile().getExecPath();
       // 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
       // user repositories will not have that prefix.
diff --git a/src/main/starlark/builtins_bzl/common/exports.bzl b/src/main/starlark/builtins_bzl/common/exports.bzl
index b2690c98ba3764..11985d20cdb6bc 100644
--- a/src/main/starlark/builtins_bzl/common/exports.bzl
+++ b/src/main/starlark/builtins_bzl/common/exports.bzl
@@ -22,6 +22,7 @@ load("@_builtins//:common/objc/objc_import.bzl", "objc_import")
 load("@_builtins//:common/objc/objc_library.bzl", "objc_library")
 load("@_builtins//:common/objc/apple_static_library.bzl", "apple_static_library")
 load("@_builtins//:common/objc/compilation_support.bzl", "compilation_support")
+load("@_builtins//:common/proto/proto_lang_toolchain.bzl", "proto_lang_toolchain")
 
 exported_toplevels = {
     # This dummy symbol is not part of the public API; it is only used to test
@@ -38,6 +39,7 @@ exported_rules = {
     "+apple_static_library": apple_static_library,
     "+cc_shared_library": cc_shared_library,
     "+cc_shared_library_permissions": cc_shared_library_permissions,
+    "-proto_lang_toolchain": proto_lang_toolchain,
 }
 exported_to_java = {
     "register_compile_and_archive_actions_for_j2objc": compilation_support.register_compile_and_archive_actions_for_j2objc,
diff --git a/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain.bzl b/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain.bzl
new file mode 100644
index 00000000000000..d8960b41b540f5
--- /dev/null
+++ b/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain.bzl
@@ -0,0 +1,85 @@
+# Copyright 2021 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A Starlark implementation of the proto_lang_toolchain rule."""
+
+load(":common/proto/providers.bzl", "ProtoLangToolchainInfo")
+load(":common/proto/proto_semantics.bzl", "semantics")
+
+ProtoInfo = _builtins.toplevel.ProtoInfo
+proto_common = _builtins.toplevel.proto_common
+
+def _rule_impl(ctx):
+    provided_proto_sources = []
+    transitive_files = depset(transitive = [bp[ProtoInfo].transitive_sources for bp in ctx.attr.blacklisted_protos])
+    for file in transitive_files.to_list():
+        source_root = file.root.path
+        provided_proto_sources.append(proto_common.ProtoSource(file, file, source_root))
+
+    flag = ctx.attr.command_line
+    if flag.find("$(PLUGIN_OUT)") > -1:
+        fail("in attribute 'command_line': Placeholder '$(PLUGIN_OUT)' is not supported.")
+    flag = flag.replace("$(OUT)", "%s")
+
+    plugin = None
+    if ctx.attr.plugin != None:
+        plugin = ctx.attr.plugin[DefaultInfo].files_to_run
+
+    return [
+        DefaultInfo(
+            files = depset(),
+            runfiles = ctx.runfiles(),
+        ),
+        ProtoLangToolchainInfo(
+            out_replacement_format_flag = flag,
+            plugin_format_flag = ctx.attr.plugin_format_flag,
+            plugin = plugin,
+            runtime = ctx.attr.runtime,
+            provided_proto_sources = provided_proto_sources,
+            proto_compiler = ctx.attr._proto_compiler.files_to_run,
+            protoc_opts = ctx.fragments.proto.experimental_protoc_opts,
+            progress_message = ctx.attr.progress_message,
+            mnemonic = ctx.attr.mnemonic,
+        ),
+    ]
+
+proto_lang_toolchain = rule(
+    implementation = _rule_impl,
+    attrs = {
+        "progress_message": attr.string(default = "Generating proto_library %{label}"),
+        "mnemonic": attr.string(default = "GenProto"),
+        "command_line": attr.string(mandatory = True),
+        "plugin_format_flag": attr.string(),
+        "plugin": attr.label(
+            executable = True,
+            cfg = "exec",
+            allow_files = True,
+        ),
+        "runtime": attr.label(
+            allow_files = True,
+        ),
+        "blacklisted_protos": attr.label_list(
+            allow_files = True,
+            providers = [ProtoInfo],
+        ),
+        "_proto_compiler": attr.label(
+            cfg = "exec",
+            executable = True,
+            allow_files = True,
+            default = configuration_field("proto", "proto_compiler"),
+        ),
+    },
+    provides = [ProtoLangToolchainInfo],
+    fragments = ["proto"] + semantics.EXTRA_FRAGMENTS,
+)
diff --git a/src/main/starlark/builtins_bzl/common/proto/providers.bzl b/src/main/starlark/builtins_bzl/common/proto/providers.bzl
new file mode 100644
index 00000000000000..aba5fe0577abe0
--- /dev/null
+++ b/src/main/starlark/builtins_bzl/common/proto/providers.bzl
@@ -0,0 +1,30 @@
+# Copyright 2021 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Bazel providers for proto rules."""
+
+ProtoLangToolchainInfo = provider(
+    doc = "Specifies how to generate language-specific code from .proto files. Used by LANG_proto_library rules.",
+    fields = dict(
+        out_replacement_format_flag = "(str) Format string used when passing output to the plugin used by proto compiler.",
+        plugin_format_flag = "(str) Format string used when passing plugin to proto compiler.",
+        plugin = "(FilesToRunProvider) Proto compiler plugin.",
+        runtime = "(Target) Runtime.",
+        provided_proto_sources = "(list[ProtoSource]) Proto sources provided by the toolchain.",
+        proto_compiler = "(FilesToRunProvider) Proto compiler.",
+        protoc_opts = "(list[str]) Options to pass to proto compiler.",
+        progress_message = "(str) Progress message to set on the proto compiler action.",
+        mnemonic = "(str) Mnemonic to set on the proto compiler action.",
+    ),
+)
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD b/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD
index a3adb78806ba3a..bfa57bb8f1ea29 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD
@@ -44,7 +44,6 @@ java_test(
         "//src/main/java/com/google/devtools/build/lib/analysis:transitive_info_collection",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/rules/proto",
-        "//src/test/java/com/google/devtools/build/lib/actions/util",
         "//src/test/java/com/google/devtools/build/lib/analysis/util",
         "//src/test/java/com/google/devtools/build/lib/packages:testutil",
         "//src/test/java/com/google/devtools/build/lib/testutil:TestConstants",
@@ -54,6 +53,25 @@ java_test(
     ],
 )
 
+java_test(
+    name = "StarlarkProtoLangToolchainTest",
+    srcs = ["StarlarkProtoLangToolchainTest.java"],
+    deps = [
+        ":ProtoLangToolchainTest",
+        "//src/main/java/com/google/devtools/build/lib/analysis:analysis_cluster",
+        "//src/main/java/com/google/devtools/build/lib/analysis:transitive_info_collection",
+        "//src/main/java/com/google/devtools/build/lib/cmdline",
+        "//src/main/java/com/google/devtools/build/lib/packages",
+        "//src/main/java/com/google/devtools/build/lib/rules/proto",
+        "//src/main/java/net/starlark/java/eval",
+        "//src/test/java/com/google/devtools/build/lib/analysis/util",
+        "//src/test/java/com/google/devtools/build/lib/testutil:TestConstants",
+        "//third_party:guava",
+        "//third_party:junit4",
+        "//third_party:truth",
+    ],
+)
+
 java_test(
     name = "BazelProtoLibraryTest",
     srcs = ["BazelProtoLibraryTest.java"],
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java b/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
index 2d1130bd339600..b43d753a034d1c 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
@@ -15,7 +15,6 @@
 package com.google.devtools.build.lib.rules.proto;
 
 import static com.google.common.truth.Truth.assertThat;
-import static com.google.devtools.build.lib.actions.util.ActionsTestUtil.prettyArtifactNames;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -30,8 +29,15 @@
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
+/** Unit tests for {@code proto_lang_toolchain}. */
 @RunWith(JUnit4.class)
 public class ProtoLangToolchainTest extends BuildViewTestCase {
+
+  @Before
+  public void setupStarlarkRule() throws Exception {
+    setBuildLanguageOptions("--experimental_builtins_injection_override=-proto_lang_toolchain");
+  }
+
   @Before
   public void setUp() throws Exception {
     MockProtoSupport.setupWorkspace(scratch);
@@ -50,12 +56,6 @@ private void validateProtoLangToolchain(ProtoLangToolchainProvider toolchain) th
     assertThat(runtimes.getLabel())
         .isEqualTo(Label.parseAbsolute("//third_party/x:runtime", ImmutableMap.of()));
 
-    assertThat(prettyArtifactNames(toolchain.forbiddenProtos()))
-        .containsExactly(
-            "third_party/x/metadata.proto",
-            "third_party/x/descriptor.proto",
-            "third_party/x/any.proto");
-
     assertThat(toolchain.protocOpts()).containsExactly("--myflag");
     Label protoc = Label.parseAbsoluteUnchecked(ProtoConstants.DEFAULT_PROTOC_LABEL);
     assertThat(toolchain.protoc().getExecutable().prettyPrint())
@@ -176,8 +176,6 @@ public void optionalFieldsAreEmpty() throws Exception {
 
     assertThat(toolchain.pluginExecutable()).isNull();
     assertThat(toolchain.runtime()).isNull();
-    assertThat(toolchain.blacklistedProtos().toList()).isEmpty();
-    assertThat(toolchain.forbiddenProtos().toList()).isEmpty();
     assertThat(toolchain.mnemonic()).isEqualTo("GenProto");
   }
 }
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/StarlarkProtoLangToolchainTest.java b/src/test/java/com/google/devtools/build/lib/rules/proto/StarlarkProtoLangToolchainTest.java
new file mode 100644
index 00000000000000..de33c3881fde17
--- /dev/null
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/StarlarkProtoLangToolchainTest.java
@@ -0,0 +1,197 @@
+// Copyright 2016 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.rules.proto;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.eventbus.EventBus;
+import com.google.devtools.build.lib.analysis.FilesToRunProvider;
+import com.google.devtools.build.lib.analysis.TransitiveInfoCollection;
+import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
+import com.google.devtools.build.lib.packages.Provider;
+import com.google.devtools.build.lib.packages.StarlarkInfo;
+import com.google.devtools.build.lib.packages.StarlarkProvider;
+import com.google.devtools.build.lib.testutil.TestConstants;
+import net.starlark.java.eval.Starlark;
+import net.starlark.java.eval.StarlarkList;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Unit tests for {@code proto_lang_toolchain}. */
+@RunWith(JUnit4.class)
+public class StarlarkProtoLangToolchainTest extends ProtoLangToolchainTest {
+
+  @Override
+  @Before
+  public void setupStarlarkRule() throws Exception {
+    setBuildLanguageOptions("--experimental_builtins_injection_override=+proto_lang_toolchain");
+  }
+
+  Provider.Key getStarlarkProtoLangToolchainInfoKey() throws LabelSyntaxException {
+    return new StarlarkProvider.Key(
+        Label.parseAbsolute("@_builtins//:common/proto/providers.bzl", ImmutableMap.of()),
+        "ProtoLangToolchainInfo");
+  }
+
+  @SuppressWarnings("unchecked")
+  private void validateStarlarkProtoLangToolchain(StarlarkInfo toolchain) throws Exception {
+    assertThat(toolchain.getValue("out_replacement_format_flag")).isEqualTo("cmd-line:%s");
+    assertThat(toolchain.getValue("plugin_format_flag")).isEqualTo("--plugin=%s");
+    assertThat(toolchain.getValue("progress_message")).isEqualTo("Progress Message %{label}");
+    assertThat(toolchain.getValue("mnemonic")).isEqualTo("MyMnemonic");
+    assertThat(ImmutableList.copyOf((StarlarkList<String>) toolchain.getValue("protoc_opts")))
+        .containsExactly("--myflag");
+    assertThat(
+            ((FilesToRunProvider) toolchain.getValue("plugin"))
+                .getExecutable()
+                .getRootRelativePathString())
+        .isEqualTo("third_party/x/plugin");
+
+    TransitiveInfoCollection runtimes = (TransitiveInfoCollection) toolchain.getValue("runtime");
+    assertThat(runtimes.getLabel())
+        .isEqualTo(Label.parseAbsolute("//third_party/x:runtime", ImmutableMap.of()));
+
+    Label protoc = Label.parseAbsoluteUnchecked(ProtoConstants.DEFAULT_PROTOC_LABEL);
+    assertThat(
+            ((FilesToRunProvider) toolchain.getValue("proto_compiler"))
+                .getExecutable()
+                .prettyPrint())
+        .isEqualTo(protoc.toPathFragment().getPathString());
+  }
+
+  @Override
+  @Test
+  public void protoToolchain() throws Exception {
+    scratch.file(
+        "third_party/x/BUILD",
+        "licenses(['unencumbered'])",
+        "cc_binary(name = 'plugin', srcs = ['plugin.cc'])",
+        "cc_library(name = 'runtime', srcs = ['runtime.cc'])",
+        "filegroup(name = 'descriptors', srcs = ['metadata.proto', 'descriptor.proto'])",
+        "filegroup(name = 'any', srcs = ['any.proto'])",
+        "proto_library(name = 'denied', srcs = [':descriptors', ':any'])");
+
+    scratch.file(
+        "foo/BUILD",
+        TestConstants.LOAD_PROTO_LANG_TOOLCHAIN,
+        "licenses(['unencumbered'])",
+        "proto_lang_toolchain(",
+        "    name = 'toolchain',",
+        "    command_line = 'cmd-line:$(OUT)',",
+        "    plugin_format_flag = '--plugin=%s',",
+        "    plugin = '//third_party/x:plugin',",
+        "    runtime = '//third_party/x:runtime',",
+        "    progress_message = 'Progress Message %{label}',",
+        "    mnemonic = 'MyMnemonic',",
+        ")");
+
+    update(ImmutableList.of("//foo:toolchain"), false, 1, true, new EventBus());
+
+    validateStarlarkProtoLangToolchain(
+        (StarlarkInfo)
+            getConfiguredTarget("//foo:toolchain").get(getStarlarkProtoLangToolchainInfoKey()));
+  }
+
+  @Override
+  @Test
+  public void protoToolchainBlacklistProtoLibraries() throws Exception {
+    scratch.file(
+        "third_party/x/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "licenses(['unencumbered'])",
+        "cc_binary(name = 'plugin', srcs = ['plugin.cc'])",
+        "cc_library(name = 'runtime', srcs = ['runtime.cc'])",
+        "proto_library(name = 'descriptors', srcs = ['metadata.proto', 'descriptor.proto'])",
+        "proto_library(name = 'any', srcs = ['any.proto'], strip_import_prefix = '/third_party')");
+
+    scratch.file(
+        "foo/BUILD",
+        TestConstants.LOAD_PROTO_LANG_TOOLCHAIN,
+        "proto_lang_toolchain(",
+        "    name = 'toolchain',",
+        "    command_line = 'cmd-line:$(OUT)',",
+        "    plugin_format_flag = '--plugin=%s',",
+        "    plugin = '//third_party/x:plugin',",
+        "    runtime = '//third_party/x:runtime',",
+        "    progress_message = 'Progress Message %{label}',",
+        "    mnemonic = 'MyMnemonic',",
+        ")");
+
+    update(ImmutableList.of("//foo:toolchain"), false, 1, true, new EventBus());
+
+    validateStarlarkProtoLangToolchain(
+        (StarlarkInfo)
+            getConfiguredTarget("//foo:toolchain").get(getStarlarkProtoLangToolchainInfoKey()));
+  }
+
+  @Override
+  @Test
+  public void protoToolchainBlacklistTransitiveProtos() throws Exception {
+    scratch.file(
+        "third_party/x/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "licenses(['unencumbered'])",
+        "cc_binary(name = 'plugin', srcs = ['plugin.cc'])",
+        "cc_library(name = 'runtime', srcs = ['runtime.cc'])",
+        "proto_library(name = 'descriptors', srcs = ['metadata.proto', 'descriptor.proto'])",
+        "proto_library(name = 'any', srcs = ['any.proto'], deps = [':descriptors'])");
+
+    scratch.file(
+        "foo/BUILD",
+        TestConstants.LOAD_PROTO_LANG_TOOLCHAIN,
+        "proto_lang_toolchain(",
+        "    name = 'toolchain',",
+        "    command_line = 'cmd-line:$(OUT)',",
+        "    plugin_format_flag = '--plugin=%s',",
+        "    plugin = '//third_party/x:plugin',",
+        "    runtime = '//third_party/x:runtime',",
+        "    progress_message = 'Progress Message %{label}',",
+        "    mnemonic = 'MyMnemonic',",
+        ")");
+
+    update(ImmutableList.of("//foo:toolchain"), false, 1, true, new EventBus());
+
+    validateStarlarkProtoLangToolchain(
+        (StarlarkInfo)
+            getConfiguredTarget("//foo:toolchain").get(getStarlarkProtoLangToolchainInfoKey()));
+  }
+
+  @Override
+  @Test
+  public void optionalFieldsAreEmpty() throws Exception {
+    scratch.file(
+        "foo/BUILD",
+        TestConstants.LOAD_PROTO_LANG_TOOLCHAIN,
+        "proto_lang_toolchain(",
+        "    name = 'toolchain',",
+        "    command_line = 'cmd-line:$(OUT)',",
+        ")");
+
+    update(ImmutableList.of("//foo:toolchain"), false, 1, true, new EventBus());
+
+    StarlarkInfo toolchain =
+        (StarlarkInfo)
+            getConfiguredTarget("//foo:toolchain").get(getStarlarkProtoLangToolchainInfoKey());
+
+    assertThat(toolchain.getValue("plugin")).isEqualTo(Starlark.NONE);
+    assertThat(toolchain.getValue("runtime")).isEqualTo(Starlark.NONE);
+    assertThat(toolchain.getValue("mnemonic")).isEqualTo("GenProto");
+  }
+}

From 92b95181a9d1f087f998b4d5a8842c17aeb9449f Mon Sep 17 00:00:00 2001
From: kotlaja <kotlaja@google.com>
Date: Wed, 4 May 2022 02:20:15 -0700
Subject: [PATCH 07/24] Roll forward of
 https://github.com/bazelbuild/bazel/commit/ae349e9368f90bc5cf27f51ecc773ada9e3a514c:
 Export ProtoLangToolchainInfo provider and flip proto_lang_toolchain rule
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

NEW: I've moved ProtoLangToolchainProvider from providers.bzl file to proto_common.bzl file since proto_common doesn't allow loading other components. I've also deleted providers.bzl file since we don't need it anymore.

Automated rollback of commit 5a6b1a85e82f4b034cc2e96fdc84c0c3c0b5571a.

*** Reason for rollback ***

Good to submit since the blocking error has been resolved.

*** Original change description ***

Automated rollback of commit ae349e9368f90bc5cf27f51ecc773ada9e3a514c.

*** Reason for rollback ***

This CL breaks proto_common.bzl file which seems that can't load providers.bzl file.

*** Original change description ***

Export ProtoLangToolchainInfo provider and flip proto_lang_toolchain rule

I’ve exported ProtoLangToolchainInfo provider from it’s native class by adding two new functions: one that’s creating starlark provider (create function), and the other that’s wrapping the starlark provider as a nat

***

PiperOrigin-RevId: 446401388
---
 .../proto/BazelJavaLiteProtoLibraryRule.java  |   4 +-
 .../java/proto/BazelJavaProtoAspect.java      |   9 +-
 .../lib/rules/android/DexArchiveAspect.java   |   4 +-
 .../lib/rules/cpp/proto/CcProtoAspect.java    |   6 +-
 .../devtools/build/lib/rules/java/BUILD       |   1 +
 .../rules/java/proto/JavaLiteProtoAspect.java |   4 +-
 .../java/proto/JavaProtoAspectCommon.java     |   6 +-
 .../java/proto/JavaProtoStarlarkCommon.java   |   2 +-
 .../lib/rules/java/proto/RpcSupport.java      |   4 +-
 .../build/lib/rules/objc/J2ObjcAspect.java    |   3 +-
 .../devtools/build/lib/rules/proto/BUILD      |   1 +
 .../build/lib/rules/proto/ProtoCommon.java    |   1 -
 .../proto/ProtoCompileActionBuilder.java      |   4 +-
 .../lib/rules/proto/ProtoLangToolchain.java   |   2 +-
 .../proto/ProtoLangToolchainProvider.java     | 179 +++++++++++++-----
 .../rules/proto/ProtoLangToolchainRule.java   |   1 -
 .../build/lib/rules/proto/ProtoSource.java    |   3 +-
 .../starlark/builtins_bzl/common/exports.bzl  |   2 +-
 .../proto/{providers.bzl => proto_common.bzl} |   6 +-
 .../common/proto/proto_lang_toolchain.bzl     |   2 +-
 .../common/proto/proto_semantics.bzl          |  21 ++
 .../devtools/build/lib/rules/proto/BUILD      |  20 +-
 .../proto/ProtoCompileActionBuilderTest.java  |  14 +-
 .../rules/proto/ProtoLangToolchainTest.java   |  85 +++++++--
 24 files changed, 264 insertions(+), 120 deletions(-)
 rename src/main/starlark/builtins_bzl/common/proto/{providers.bzl => proto_common.bzl} (89%)
 create mode 100644 src/main/starlark/builtins_bzl/common/proto/proto_semantics.bzl

diff --git a/src/main/java/com/google/devtools/build/lib/bazel/rules/java/proto/BazelJavaLiteProtoLibraryRule.java b/src/main/java/com/google/devtools/build/lib/bazel/rules/java/proto/BazelJavaLiteProtoLibraryRule.java
index b5fe5a7b1d4566..3504120aad1edb 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/rules/java/proto/BazelJavaLiteProtoLibraryRule.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/rules/java/proto/BazelJavaLiteProtoLibraryRule.java
@@ -58,9 +58,7 @@ public RuleClass build(RuleClass.Builder builder, RuleDefinitionEnvironment envi
                 .aspect(javaProtoAspect))
         .add(
             attr(JavaProtoAspectCommon.LITE_PROTO_TOOLCHAIN_ATTR, LABEL)
-                .mandatoryBuiltinProviders(
-                    ImmutableList.<Class<? extends TransitiveInfoProvider>>of(
-                        ProtoLangToolchainProvider.class))
+                .mandatoryProviders(ProtoLangToolchainProvider.PROVIDER_ID)
                 .value(getProtoToolchainLabel(DEFAULT_PROTO_TOOLCHAIN_LABEL)))
         .advertiseStarlarkProvider(StarlarkProviderIdentifier.forKey(JavaInfo.PROVIDER.getKey()))
         .build();
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/rules/java/proto/BazelJavaProtoAspect.java b/src/main/java/com/google/devtools/build/lib/bazel/rules/java/proto/BazelJavaProtoAspect.java
index d74eb5f097ae80..4d82d27032a3b2 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/rules/java/proto/BazelJavaProtoAspect.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/rules/java/proto/BazelJavaProtoAspect.java
@@ -16,6 +16,7 @@
 
 import static com.google.devtools.build.lib.collect.nestedset.Order.STABLE_ORDER;
 
+import com.google.common.base.Optional;
 import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.analysis.RuleContext;
@@ -26,9 +27,11 @@
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.packages.AspectDefinition;
 import com.google.devtools.build.lib.packages.AspectParameters;
+import com.google.devtools.build.lib.packages.StarlarkInfo;
 import com.google.devtools.build.lib.rules.java.proto.JavaProtoAspect;
 import com.google.devtools.build.lib.rules.java.proto.RpcSupport;
 import com.google.devtools.build.lib.rules.proto.ProtoCompileActionBuilder;
+import com.google.devtools.build.lib.rules.proto.ProtoCompileActionBuilder.ToolchainInvocation;
 import java.util.List;
 
 /** An Aspect which BazelJavaProtoLibrary injects to build Java SPEED protos. */
@@ -45,7 +48,7 @@ public BazelJavaProtoAspect(RuleDefinitionEnvironment env) {
   private static class NoopRpcSupport
       implements RpcSupport {
     @Override
-    public List<ProtoCompileActionBuilder.ToolchainInvocation> getToolchainInvocation(
+    public List<ToolchainInvocation> getToolchainInvocation(
         RuleContext ruleContext, Artifact sourceJar) {
       return ImmutableList.of();
     }
@@ -56,8 +59,8 @@ public boolean allowServices(RuleContext ruleContext) {
     }
 
     @Override
-    public NestedSet<Artifact> getForbiddenProtos(RuleContext ruleContext) {
-      return NestedSetBuilder.emptySet(STABLE_ORDER);
+    public Optional<StarlarkInfo> getToolchain(RuleContext ruleContext) {
+      return Optional.absent();
     }
 
     @Override
diff --git a/src/main/java/com/google/devtools/build/lib/rules/android/DexArchiveAspect.java b/src/main/java/com/google/devtools/build/lib/rules/android/DexArchiveAspect.java
index f81fbd534e3be1..ac9a96be504e6c 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/android/DexArchiveAspect.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/android/DexArchiveAspect.java
@@ -176,11 +176,11 @@ public AspectDefinition getDefinition(AspectParameters params) {
             // For android_sdk rules, where we just want to get at aidl runtime deps.
             .requireStarlarkProviders(forKey(AndroidSdkProvider.PROVIDER.getKey()))
             .requireStarlarkProviders(forKey(ProtoInfo.PROVIDER.getKey()))
-            .requireProviderSets(
+            .requireStarlarkProviderSets(
                 ImmutableList.of(
                     // For proto_lang_toolchain rules, where we just want to get at their runtime
                     // deps.
-                    ImmutableSet.of(ProtoLangToolchainProvider.class)))
+                    ImmutableSet.of(ProtoLangToolchainProvider.PROVIDER_ID)))
             .addRequiredToolchains(
                 Label.parseAbsoluteUnchecked(toolsRepository + sdkToolchainLabel))
             // Parse labels since we don't have RuleDefinitionEnvironment.getLabel like in a rule
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 f66a8f1a75d1a7..34333475528fd2 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
@@ -41,6 +41,7 @@
 import com.google.devtools.build.lib.packages.Attribute.LabelLateBoundDefault;
 import com.google.devtools.build.lib.packages.NativeAspectClass;
 import com.google.devtools.build.lib.packages.RuleClass.ConfiguredTargetFactory.RuleErrorException;
+import com.google.devtools.build.lib.packages.StarlarkInfo;
 import com.google.devtools.build.lib.packages.TargetUtils;
 import com.google.devtools.build.lib.rules.cpp.AspectLegalCppSemantics;
 import com.google.devtools.build.lib.rules.cpp.CcCommon;
@@ -127,7 +128,7 @@ public AspectDefinition getDefinition(AspectParameters aspectParameters) {
             .useToolchainTransition(true)
             .add(
                 attr(PROTO_TOOLCHAIN_ATTR, LABEL)
-                    .mandatoryBuiltinProviders(ImmutableList.of(ProtoLangToolchainProvider.class))
+                    .mandatoryProviders(ProtoLangToolchainProvider.PROVIDER_ID)
                     .value(PROTO_TOOLCHAIN_LABEL))
             .add(
                 attr(CcToolchain.CC_TOOLCHAIN_DEFAULT_ATTRIBUTE_NAME, LABEL)
@@ -449,7 +450,6 @@ private void createProtoCompileAction(Collection<Artifact> outputs) {
       } else {
         genfilesPath = genfilesFragment.getRelative(protoRootFragment).getPathString();
       }
-
       ImmutableList.Builder<ToolchainInvocation> invocations = ImmutableList.builder();
       invocations.add(
           new ToolchainInvocation("C++", checkNotNull(getProtoToolchainProvider()), genfilesPath));
@@ -465,7 +465,7 @@ private void createProtoCompileAction(Collection<Artifact> outputs) {
     }
 
     private ProtoLangToolchainProvider getProtoToolchainProvider() {
-      return ruleContext.getPrerequisite(PROTO_TOOLCHAIN_ATTR, ProtoLangToolchainProvider.class);
+      return ProtoLangToolchainProvider.get(ruleContext, PROTO_TOOLCHAIN_ATTR);
     }
 
     public void addProviders(ConfiguredAspect.Builder builder) {
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/BUILD b/src/main/java/com/google/devtools/build/lib/rules/java/BUILD
index fc4c128971aa2b..caae66d357a896 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/BUILD
@@ -105,6 +105,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/collect/nestedset",
         "//src/main/java/com/google/devtools/build/lib/packages",
         "//src/main/java/com/google/devtools/build/lib/rules/proto",
+        "//src/main/java/net/starlark/java/eval",
         "//third_party:guava",
     ],
 )
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/proto/JavaLiteProtoAspect.java b/src/main/java/com/google/devtools/build/lib/rules/java/proto/JavaLiteProtoAspect.java
index f504e8c5b99d52..80a5411a7f95ef 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/proto/JavaLiteProtoAspect.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/proto/JavaLiteProtoAspect.java
@@ -111,9 +111,7 @@ public AspectDefinition getDefinition(AspectParameters aspectParameters) {
                 ImmutableList.of(StarlarkProviderIdentifier.forKey(JavaInfo.PROVIDER.getKey())))
             .add(
                 attr(JavaProtoAspectCommon.LITE_PROTO_TOOLCHAIN_ATTR, LABEL)
-                    .mandatoryBuiltinProviders(
-                        ImmutableList.<Class<? extends TransitiveInfoProvider>>of(
-                            ProtoLangToolchainProvider.class))
+                    .mandatoryProviders(ProtoLangToolchainProvider.PROVIDER_ID)
                     .value(getProtoToolchainLabel(defaultProtoToolchainLabel)))
             .add(
                 attr(JavaRuleClasses.JAVA_TOOLCHAIN_ATTRIBUTE_NAME, LABEL)
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/proto/JavaProtoAspectCommon.java b/src/main/java/com/google/devtools/build/lib/rules/java/proto/JavaProtoAspectCommon.java
index ffb21afa9b337a..2f9021e499c59a 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/proto/JavaProtoAspectCommon.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/proto/JavaProtoAspectCommon.java
@@ -21,7 +21,6 @@
 import com.google.devtools.build.lib.analysis.RuleContext;
 import com.google.devtools.build.lib.analysis.TransitiveInfoCollection;
 import com.google.devtools.build.lib.analysis.config.CoreOptionConverters.StrictDepsMode;
-import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.rules.java.JavaCompilationArgsProvider;
 import com.google.devtools.build.lib.rules.java.JavaCompilationArtifacts;
 import com.google.devtools.build.lib.rules.java.JavaInfo;
@@ -158,15 +157,14 @@ public ImmutableList<TransitiveInfoCollection> getProtoRuntimeDeps() {
 
   /** Returns the toolchain that specifies how to generate code from {@code .proto} files. */
   public ProtoLangToolchainProvider getProtoToolchainProvider() {
-    return checkNotNull(
-        ruleContext.getPrerequisite(protoToolchainAttr, ProtoLangToolchainProvider.class));
+    return checkNotNull(ProtoLangToolchainProvider.get(ruleContext, protoToolchainAttr));
   }
 
   /**
    * Returns the toolchain that specifies how to generate Java-lite code from {@code .proto} files.
    */
   static ProtoLangToolchainProvider getLiteProtoToolchainProvider(RuleContext ruleContext) {
-    return ruleContext.getPrerequisite(LITE_PROTO_TOOLCHAIN_ATTR, ProtoLangToolchainProvider.class);
+    return ProtoLangToolchainProvider.get(ruleContext, LITE_PROTO_TOOLCHAIN_ATTR);
   }
 
   /**
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/proto/JavaProtoStarlarkCommon.java b/src/main/java/com/google/devtools/build/lib/rules/java/proto/JavaProtoStarlarkCommon.java
index a9dcc65b17975d..2cc5d71174ab21 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/proto/JavaProtoStarlarkCommon.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/proto/JavaProtoStarlarkCommon.java
@@ -81,6 +81,6 @@ private static ProtoLangToolchainProvider getProtoToolchainProvider(
       StarlarkRuleContext starlarkRuleContext, String protoToolchainAttr) throws EvalException {
     ConfiguredTarget javaliteToolchain =
         (ConfiguredTarget) checkNotNull(starlarkRuleContext.getAttr().getValue(protoToolchainAttr));
-    return checkNotNull(javaliteToolchain.getProvider(ProtoLangToolchainProvider.class));
+    return checkNotNull(ProtoLangToolchainProvider.get(javaliteToolchain));
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/proto/RpcSupport.java b/src/main/java/com/google/devtools/build/lib/rules/java/proto/RpcSupport.java
index 7f989464071010..b373e233694116 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/proto/RpcSupport.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/proto/RpcSupport.java
@@ -14,6 +14,7 @@
 
 package com.google.devtools.build.lib.rules.java.proto;
 
+import com.google.common.base.Optional;
 import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.analysis.RuleContext;
@@ -21,6 +22,7 @@
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.packages.AspectDefinition;
 import com.google.devtools.build.lib.packages.AspectParameters;
+import com.google.devtools.build.lib.packages.StarlarkInfo;
 import com.google.devtools.build.lib.rules.proto.ProtoCompileActionBuilder;
 import java.util.List;
 
@@ -31,7 +33,7 @@ List<ProtoCompileActionBuilder.ToolchainInvocation> getToolchainInvocation(
 
   boolean allowServices(RuleContext ruleContext);
 
-  NestedSet<Artifact> getForbiddenProtos(RuleContext ruleContext);
+  Optional<StarlarkInfo> getToolchain(RuleContext ruleContext);
 
   ImmutableList<TransitiveInfoCollection> getRuntimes(RuleContext ruleContext);
 
diff --git a/src/main/java/com/google/devtools/build/lib/rules/objc/J2ObjcAspect.java b/src/main/java/com/google/devtools/build/lib/rules/objc/J2ObjcAspect.java
index dd2607cd9a2056..22d9bb7c454e16 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/objc/J2ObjcAspect.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/objc/J2ObjcAspect.java
@@ -51,6 +51,7 @@
 import com.google.devtools.build.lib.packages.BuildType;
 import com.google.devtools.build.lib.packages.NativeAspectClass;
 import com.google.devtools.build.lib.packages.RuleClass.ConfiguredTargetFactory.RuleErrorException;
+import com.google.devtools.build.lib.packages.StarlarkInfo;
 import com.google.devtools.build.lib.packages.StarlarkProviderIdentifier;
 import com.google.devtools.build.lib.rules.apple.AppleConfiguration;
 import com.google.devtools.build.lib.rules.apple.AppleToolchain;
@@ -390,7 +391,7 @@ private ConfiguredAspect proto(ConfiguredTarget base, RuleContext ruleContext)
     ImmutableList<Artifact> protoSources = protoInfo.getDirectProtoSources();
 
     ProtoLangToolchainProvider protoToolchain =
-        ruleContext.getPrerequisite(J2OBJC_PROTO_TOOLCHAIN_ATTR, ProtoLangToolchainProvider.class);
+        ProtoLangToolchainProvider.get(ruleContext, J2OBJC_PROTO_TOOLCHAIN_ATTR);
     // Avoid pulling in any generated files from forbidden protos.
     ProtoSourceFileExcludeList protoExcludeList =
         new ProtoSourceFileExcludeList(ruleContext, protoToolchain.providedProtoSources());
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/BUILD b/src/main/java/com/google/devtools/build/lib/rules/proto/BUILD
index 25e76419a7a5de..b0f29c42a60f55 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/BUILD
@@ -54,6 +54,7 @@ java_library(
         "//src/main/java/com/google/devtools/common/options",
         "//src/main/java/net/starlark/java/annot",
         "//src/main/java/net/starlark/java/eval",
+        "//src/main/java/net/starlark/java/syntax",
         "//third_party:auto_value",
         "//third_party:guava",
         "//third_party:jsr305",
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 0d468dd69d4a75..8c1d348d8ffe4c 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
@@ -32,7 +32,6 @@
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.concurrent.BlazeInterners;
 import com.google.devtools.build.lib.packages.BuildType;
-import com.google.devtools.build.lib.util.Pair;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import javax.annotation.Nullable;
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 8ebaa66cee452a..6c2e8d1c45bee0 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
@@ -356,13 +356,13 @@ private static ToolchainInvocation createDescriptorSetToolchain(
 
     return new ToolchainInvocation(
         "dontcare",
-        ProtoLangToolchainProvider.create(
+        ProtoLangToolchainProvider.createNative(
             // Note: adding --include_imports here was requested multiple times, but it'll cause the
             // output size to become quadratic, so don't.
             // A rule that concatenates the artifacts from ctx.deps.proto.transitive_descriptor_sets
             // provides similar results.
             "--descriptor_set_out=%s",
-            /* pluginFormatFlag = */ null,
+            /* pluginFormatFlag = */ "",
             /* pluginExecutable= */ null,
             /* runtime= */ null,
             /* providedProtoSources= */ ImmutableList.of(),
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java
index 8d409b8acc1418..e4abfe9bc194a2 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java
@@ -49,7 +49,7 @@ public ConfiguredTarget create(RuleContext ruleContext)
     flag = flag.replace("$(OUT)", "%s");
 
     return new RuleConfiguredTargetBuilder(ruleContext)
-        .addProvider(
+        .addStarlarkDeclaredProvider(
             ProtoLangToolchainProvider.create(
                 flag,
                 ruleContext.attributes().get("plugin_format_flag", Type.STRING),
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
index edc750f76cbdec..8753baa0a8b9f3 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
@@ -17,15 +17,23 @@
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.analysis.FilesToRunProvider;
+import com.google.devtools.build.lib.analysis.RuleContext;
 import com.google.devtools.build.lib.analysis.TransitiveInfoCollection;
-import com.google.devtools.build.lib.analysis.TransitiveInfoProvider;
-import com.google.devtools.build.lib.collect.nestedset.NestedSet;
+import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.collect.nestedset.Depset;
+import com.google.devtools.build.lib.collect.nestedset.Depset.ElementType;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
-import com.google.devtools.build.lib.packages.BuiltinProvider;
-import com.google.devtools.build.lib.packages.NativeInfo;
+import com.google.devtools.build.lib.packages.StarlarkInfo;
+import com.google.devtools.build.lib.packages.StarlarkProvider;
+import com.google.devtools.build.lib.packages.StarlarkProviderIdentifier;
+import java.util.LinkedHashMap;
+import java.util.Map;
 import javax.annotation.Nullable;
-import net.starlark.java.annot.StarlarkMethod;
+import net.starlark.java.eval.EvalException;
+import net.starlark.java.eval.NoneType;
+import net.starlark.java.eval.Starlark;
 import net.starlark.java.eval.StarlarkList;
+import net.starlark.java.syntax.Location;
 
 // Note: AutoValue v1.4-rc1 has AutoValue.CopyAnnotations which makes it work with Starlark. No need
 // to un-AutoValue this class to expose it to Starlark.
@@ -34,26 +42,23 @@
  * rules.
  */
 @AutoValue
-public abstract class ProtoLangToolchainProvider implements TransitiveInfoProvider {
-  @StarlarkMethod(
-      name = "out_replacement_format_flag",
-      doc = "Format string used when passing output to the plugin used by proto compiler.",
-      structField = true)
+public abstract class ProtoLangToolchainProvider {
+  public static final String PROVIDER_NAME = "ProtoLangToolchainInfo";
+  public static final StarlarkProvider.Key starlarkProtoLangToolchainKey =
+      new StarlarkProvider.Key(
+          Label.parseAbsoluteUnchecked("@_builtins//:common/proto/proto_common.bzl"),
+          PROVIDER_NAME);
+  public static final StarlarkProviderIdentifier PROVIDER_ID =
+      StarlarkProviderIdentifier.forKey(starlarkProtoLangToolchainKey);
+
+  // Format string used when passing output to the plugin used by proto compiler.
   public abstract String outReplacementFormatFlag();
 
-  @StarlarkMethod(
-      name = "plugin_format_flag",
-      doc = "Format string used when passing plugin to proto compiler.",
-      structField = true,
-      allowReturnNones = true)
+  // Format string used when passing plugin to proto compiler.
   @Nullable
   public abstract String pluginFormatFlag();
 
-  @StarlarkMethod(
-      name = "plugin",
-      doc = "Proto compiler plugin.",
-      structField = true,
-      allowReturnNones = true)
+  // Proto compiler plugin.
   @Nullable
   public abstract FilesToRunProvider pluginExecutable();
 
@@ -64,34 +69,128 @@ public abstract class ProtoLangToolchainProvider implements TransitiveInfoProvid
    * Returns a list of {@link ProtoSource}s that are already provided by the protobuf runtime (i.e.
    * for which {@code <lang>_proto_library} should not generate bindings.
    */
+  // Proto sources provided by the toolchain.
   public abstract ImmutableList<ProtoSource> providedProtoSources();
 
-  @StarlarkMethod(name = "proto_compiler", doc = "Proto compiler.", structField = true)
+  // Proto compiler.
   public abstract FilesToRunProvider protoc();
 
-  @StarlarkMethod(
-      name = "protoc_opts",
-      doc = "Options to pass to proto compiler.",
-      structField = true)
+  // Options to pass to proto compiler.
   public StarlarkList<String> protocOptsForStarlark() {
     return StarlarkList.immutableCopyOf(protocOpts());
   }
 
   public abstract ImmutableList<String> protocOpts();
 
-  @StarlarkMethod(
-      name = "progress_message",
-      doc = "Progress message to set on the proto compiler action.",
-      structField = true)
+  // Progress message to set on the proto compiler action.
   public abstract String progressMessage();
 
-  @StarlarkMethod(
-      name = "mnemonic",
-      doc = "Mnemonic to set on the proto compiler action.",
-      structField = true)
+  // Mnemonic to set on the proto compiler action.
   public abstract String mnemonic();
 
-  public static ProtoLangToolchainProvider create(
+  public static StarlarkInfo create(
+      String outReplacementFormatFlag,
+      String pluginFormatFlag,
+      FilesToRunProvider pluginExecutable,
+      TransitiveInfoCollection runtime,
+      ImmutableList<ProtoSource> providedProtoSources,
+      FilesToRunProvider protoc,
+      ImmutableList<String> protocOpts,
+      String progressMessage,
+      String mnemonic) {
+
+    NestedSetBuilder<ProtoSource> providedProtoSourcesSet = NestedSetBuilder.stableOrder();
+    providedProtoSources.forEach(providedProtoSourcesSet::add);
+    NestedSetBuilder<String> protocOptsSet = NestedSetBuilder.stableOrder();
+    protocOpts.forEach(protocOptsSet::add);
+
+    Map<String, Object> m = new LinkedHashMap<>();
+    m.put("plugin", pluginExecutable == null ? Starlark.NONE : pluginExecutable);
+    m.put("plugin_format_flag", pluginFormatFlag == null ? Starlark.NONE : pluginFormatFlag);
+    m.put("proto_compiler", protoc == null ? Starlark.NONE : protoc);
+    m.put(
+        "provided_proto_sources",
+        StarlarkList.immutableCopyOf(providedProtoSources));
+    m.put("protoc_opts", StarlarkList.immutableCopyOf(protocOpts));
+    m.put("out_replacement_format_flag", outReplacementFormatFlag);
+    m.put("progress_message", progressMessage);
+    m.put("mnemonic", mnemonic);
+    m.put("plugin", pluginExecutable == null ? Starlark.NONE : pluginExecutable);
+    m.put("runtime", runtime == null ? Starlark.NONE : runtime);
+
+    StarlarkProvider provider =
+        StarlarkProvider.createExportedSchemaless(starlarkProtoLangToolchainKey,
+            Location.BUILTIN);
+
+    return StarlarkInfo.create(provider, m, Location.BUILTIN);
+  }
+
+  private static ImmutableList<ProtoLangToolchainProvider> getToolchains(
+      RuleContext ruleContext, String attributeName) {
+    ImmutableList.Builder<ProtoLangToolchainProvider> result = ImmutableList.builder();
+    for (TransitiveInfoCollection prerequisite : ruleContext.getPrerequisites(attributeName)) {
+      ProtoLangToolchainProvider toolchain = get(prerequisite);
+      if (toolchain != null) {
+        result.add(toolchain);
+      }
+    }
+    return result.build();
+  }
+
+  public static ProtoLangToolchainProvider get(RuleContext ruleContext, String attributeName) {
+    return getToolchains(ruleContext, attributeName).stream().findFirst().orElse(null);
+  }
+
+  public static ProtoLangToolchainProvider get(TransitiveInfoCollection prerequisite) {
+    StarlarkInfo provider = (StarlarkInfo) prerequisite.get(starlarkProtoLangToolchainKey);
+    return wrapStarlarkProviderWithNativeProvider(provider);
+  }
+
+  public static StarlarkInfo getStarlarkProvider(RuleContext ruleContext, String attributeName) {
+    for (TransitiveInfoCollection prerequisite : ruleContext.getPrerequisites(attributeName)) {
+      StarlarkInfo provider = (StarlarkInfo) prerequisite.get(starlarkProtoLangToolchainKey);
+      if (provider != null) {
+        return provider;
+      }
+    }
+    return null;
+  }
+
+  public static StarlarkInfo getStarlarkProvider(TransitiveInfoCollection prerequisite) {
+    return (StarlarkInfo) prerequisite.get(starlarkProtoLangToolchainKey);
+  }
+
+  @SuppressWarnings("unchecked")
+  private static ProtoLangToolchainProvider wrapStarlarkProviderWithNativeProvider(
+      StarlarkInfo provider) {
+    if (provider != null) {
+      try {
+        return new AutoValue_ProtoLangToolchainProvider(
+            provider.getValue("out_replacement_format_flag", String.class),
+            provider.getValue("plugin_format_flag") instanceof NoneType ?
+                null
+            : provider.getValue("plugin_format_flag", String.class),
+            provider.getValue("plugin") instanceof NoneType
+                ? null
+                : provider.getValue("plugin", FilesToRunProvider.class),
+            provider.getValue("runtime") instanceof NoneType
+                ? null
+                : provider.getValue("runtime", TransitiveInfoCollection.class),
+            ImmutableList.copyOf(
+                (StarlarkList<ProtoSource>) provider.getValue("provided_proto_sources")),
+            provider.getValue("proto_compiler", FilesToRunProvider.class),
+            ImmutableList.copyOf((StarlarkList<String>) provider.getValue("protoc_opts")),
+            provider.getValue("progress_message", String.class),
+            provider.getValue("mnemonic", String.class));
+      } catch (EvalException e) {
+        return null;
+      }
+    }
+    return null;
+  }
+
+
+  public static ProtoLangToolchainProvider createNative(
       String outReplacementFormatFlag,
       String pluginFormatFlag,
       FilesToRunProvider pluginExecutable,
@@ -101,15 +200,7 @@ public static ProtoLangToolchainProvider create(
       ImmutableList<String> protocOpts,
       String progressMessage,
       String mnemonic) {
-    return new AutoValue_ProtoLangToolchainProvider(
-        outReplacementFormatFlag,
-        pluginFormatFlag,
-        pluginExecutable,
-        runtime,
-        providedProtoSources,
-        protoc,
-        protocOpts,
-        progressMessage,
-        mnemonic);
+    return wrapStarlarkProviderWithNativeProvider(create(outReplacementFormatFlag,pluginFormatFlag,pluginExecutable,runtime,providedProtoSources, protoc, protocOpts, progressMessage, mnemonic));
   }
+
 }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java
index 70944fd613a9bb..7ac6344ea2b6ce 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java
@@ -106,7 +106,6 @@ public RuleClass build(RuleClass.Builder builder, RuleDefinitionEnvironment envi
                 .exec()
                 .value(PROTO_COMPILER))
         .requiresConfigurationFragments(ProtoConfiguration.class)
-        .advertiseProvider(ProtoLangToolchainProvider.class)
         .removeAttribute("data")
         .removeAttribute("deps")
         .build();
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoSource.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoSource.java
index 754c37521a6fc7..50d49455e0f8b5 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoSource.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoSource.java
@@ -18,11 +18,12 @@
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.build.lib.vfs.PathFragment;
+import net.starlark.java.eval.StarlarkValue;
 
 /** Represents a single {@code .proto} source file. */
 @Immutable
 @AutoCodec
-class ProtoSource {
+class ProtoSource implements StarlarkValue {
   private final Artifact sourceFile;
   private final Artifact originalSourceFile;
   private final PathFragment sourceRoot;
diff --git a/src/main/starlark/builtins_bzl/common/exports.bzl b/src/main/starlark/builtins_bzl/common/exports.bzl
index 11985d20cdb6bc..b58257ba5dda5e 100644
--- a/src/main/starlark/builtins_bzl/common/exports.bzl
+++ b/src/main/starlark/builtins_bzl/common/exports.bzl
@@ -39,7 +39,7 @@ exported_rules = {
     "+apple_static_library": apple_static_library,
     "+cc_shared_library": cc_shared_library,
     "+cc_shared_library_permissions": cc_shared_library_permissions,
-    "-proto_lang_toolchain": proto_lang_toolchain,
+    "+proto_lang_toolchain": proto_lang_toolchain,
 }
 exported_to_java = {
     "register_compile_and_archive_actions_for_j2objc": compilation_support.register_compile_and_archive_actions_for_j2objc,
diff --git a/src/main/starlark/builtins_bzl/common/proto/providers.bzl b/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
similarity index 89%
rename from src/main/starlark/builtins_bzl/common/proto/providers.bzl
rename to src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
index aba5fe0577abe0..bcdd8dd60c8669 100644
--- a/src/main/starlark/builtins_bzl/common/proto/providers.bzl
+++ b/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Bazel providers for proto rules."""
+"""Definition of proto_common module, together with bazel providers for proto rules."""
 
 ProtoLangToolchainInfo = provider(
     doc = "Specifies how to generate language-specific code from .proto files. Used by LANG_proto_library rules.",
@@ -28,3 +28,7 @@ ProtoLangToolchainInfo = provider(
         mnemonic = "(str) Mnemonic to set on the proto compiler action.",
     ),
 )
+
+proto_common_do_not_use = struct(
+    ProtoLangToolchainInfo = ProtoLangToolchainInfo,
+)
diff --git a/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain.bzl b/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain.bzl
index d8960b41b540f5..ef2d4aee87b291 100644
--- a/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain.bzl
+++ b/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain.bzl
@@ -14,7 +14,7 @@
 
 """A Starlark implementation of the proto_lang_toolchain rule."""
 
-load(":common/proto/providers.bzl", "ProtoLangToolchainInfo")
+load(":common/proto/proto_common.bzl", "ProtoLangToolchainInfo")
 load(":common/proto/proto_semantics.bzl", "semantics")
 
 ProtoInfo = _builtins.toplevel.ProtoInfo
diff --git a/src/main/starlark/builtins_bzl/common/proto/proto_semantics.bzl b/src/main/starlark/builtins_bzl/common/proto/proto_semantics.bzl
new file mode 100644
index 00000000000000..2d96efab970963
--- /dev/null
+++ b/src/main/starlark/builtins_bzl/common/proto/proto_semantics.bzl
@@ -0,0 +1,21 @@
+# Copyright 2021 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Proto Semantics
+"""
+
+semantics = struct(
+    EXTRA_FRAGMENTS = [],
+)
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD b/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD
index bfa57bb8f1ea29..74b050dc619e1c 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD
@@ -41,30 +41,12 @@ java_test(
     name = "ProtoLangToolchainTest",
     srcs = ["ProtoLangToolchainTest.java"],
     deps = [
-        "//src/main/java/com/google/devtools/build/lib/analysis:transitive_info_collection",
-        "//src/main/java/com/google/devtools/build/lib/cmdline",
-        "//src/main/java/com/google/devtools/build/lib/rules/proto",
-        "//src/test/java/com/google/devtools/build/lib/analysis/util",
-        "//src/test/java/com/google/devtools/build/lib/packages:testutil",
-        "//src/test/java/com/google/devtools/build/lib/testutil:TestConstants",
-        "//third_party:guava",
-        "//third_party:junit4",
-        "//third_party:truth",
-    ],
-)
-
-java_test(
-    name = "StarlarkProtoLangToolchainTest",
-    srcs = ["StarlarkProtoLangToolchainTest.java"],
-    deps = [
-        ":ProtoLangToolchainTest",
-        "//src/main/java/com/google/devtools/build/lib/analysis:analysis_cluster",
         "//src/main/java/com/google/devtools/build/lib/analysis:transitive_info_collection",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/packages",
         "//src/main/java/com/google/devtools/build/lib/rules/proto",
-        "//src/main/java/net/starlark/java/eval",
         "//src/test/java/com/google/devtools/build/lib/analysis/util",
+        "//src/test/java/com/google/devtools/build/lib/packages:testutil",
         "//src/test/java/com/google/devtools/build/lib/testutil:TestConstants",
         "//third_party:guava",
         "//third_party:junit4",
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 5c180453a8debb..a887b52bb60b63 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
@@ -103,7 +103,7 @@ public void commandLine_basic() throws Exception {
             artifact("//:dont-care", "protoc-gen-javalite.exe"));
 
     ProtoLangToolchainProvider toolchainNoPlugin =
-        ProtoLangToolchainProvider.create(
+        ProtoLangToolchainProvider.createNative(
             "--java_out=param1,param2:%s",
             /* pluginFormatFlag= */ null,
             /* pluginExecutable= */ null,
@@ -114,7 +114,7 @@ public void commandLine_basic() throws Exception {
             /* progressMessage = */ "",
             /* mnemonic= */ "");
     ProtoLangToolchainProvider toolchainWithPlugin =
-        ProtoLangToolchainProvider.create(
+        ProtoLangToolchainProvider.createNative(
             "--PLUGIN_pluginName_out=param3,param4:%s",
             /* pluginFormatFlag= */ "--plugin=protoc-gen-PLUGIN_pluginName=%s",
             plugin,
@@ -182,7 +182,7 @@ public void commandline_derivedArtifact() throws Exception {
   @Test
   public void commandLine_strictDeps() throws Exception {
     ProtoLangToolchainProvider toolchain =
-        ProtoLangToolchainProvider.create(
+        ProtoLangToolchainProvider.createNative(
             "--java_out=param1,param2:%s",
             /* pluginFormatFlag= */ null,
             /* pluginExecutable= */ null,
@@ -225,7 +225,7 @@ public void commandLine_strictDeps() throws Exception {
   @Test
   public void commandLine_exports() throws Exception {
     ProtoLangToolchainProvider toolchain =
-        ProtoLangToolchainProvider.create(
+        ProtoLangToolchainProvider.createNative(
             "--java_out=param1,param2:%s",
             /* pluginFormatFlag= */ null,
             /* pluginExecutable= */ null,
@@ -300,7 +300,7 @@ public String toString() {
         };
 
     ProtoLangToolchainProvider toolchain =
-        ProtoLangToolchainProvider.create(
+        ProtoLangToolchainProvider.createNative(
             "--java_out=param1,param2:%s",
             /* pluginFormatFlag= */ null,
             /* pluginExecutable= */ null,
@@ -339,7 +339,7 @@ public String toString() {
   @Test
   public void exceptionIfSameName() throws Exception {
     ProtoLangToolchainProvider toolchain1 =
-        ProtoLangToolchainProvider.create(
+        ProtoLangToolchainProvider.createNative(
             "dontcare=%s",
             /* pluginFormatFlag= */ null,
             /* pluginExecutable= */ null,
@@ -350,7 +350,7 @@ public void exceptionIfSameName() throws Exception {
             /* progressMessage = */ "",
             /* mnemonic= */ "");
     ProtoLangToolchainProvider toolchain2 =
-        ProtoLangToolchainProvider.create(
+        ProtoLangToolchainProvider.createNative(
             "dontcare=%s",
             /* pluginFormatFlag= */ null,
             /* pluginExecutable= */ null,
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java b/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
index b43d753a034d1c..a6860bdc2437d9 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
@@ -22,6 +22,9 @@
 import com.google.devtools.build.lib.analysis.TransitiveInfoCollection;
 import com.google.devtools.build.lib.analysis.util.BuildViewTestCase;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
+import com.google.devtools.build.lib.packages.Provider;
+import com.google.devtools.build.lib.packages.StarlarkProvider;
 import com.google.devtools.build.lib.packages.util.MockProtoSupport;
 import com.google.devtools.build.lib.testutil.TestConstants;
 import org.junit.Before;
@@ -32,12 +35,6 @@
 /** Unit tests for {@code proto_lang_toolchain}. */
 @RunWith(JUnit4.class)
 public class ProtoLangToolchainTest extends BuildViewTestCase {
-
-  @Before
-  public void setupStarlarkRule() throws Exception {
-    setBuildLanguageOptions("--experimental_builtins_injection_override=-proto_lang_toolchain");
-  }
-
   @Before
   public void setUp() throws Exception {
     MockProtoSupport.setupWorkspace(scratch);
@@ -46,6 +43,12 @@ public void setUp() throws Exception {
     invalidatePackages();
   }
 
+  Provider.Key getStarlarkProtoLangToolchainInfoKey() throws LabelSyntaxException {
+    return new StarlarkProvider.Key(
+        Label.parseAbsolute("@_builtins//:common/proto/proto_common.bzl", ImmutableMap.of()),
+        "ProtoLangToolchainInfo");
+  }
+
   private void validateProtoLangToolchain(ProtoLangToolchainProvider toolchain) throws Exception {
     assertThat(toolchain.outReplacementFormatFlag()).isEqualTo("cmd-line");
     assertThat(toolchain.pluginFormatFlag()).isEqualTo("--plugin=%s");
@@ -57,14 +60,16 @@ private void validateProtoLangToolchain(ProtoLangToolchainProvider toolchain) th
         .isEqualTo(Label.parseAbsolute("//third_party/x:runtime", ImmutableMap.of()));
 
     assertThat(toolchain.protocOpts()).containsExactly("--myflag");
-    Label protoc = Label.parseAbsoluteUnchecked(ProtoConstants.DEFAULT_PROTOC_LABEL);
-    assertThat(toolchain.protoc().getExecutable().prettyPrint())
-        .isEqualTo(protoc.toPathFragment().getPathString());
 
     assertThat(toolchain.progressMessage()).isEqualTo("Progress Message %{label}");
     assertThat(toolchain.mnemonic()).isEqualTo("MyMnemonic");
   }
 
+  private void validateProtoCompiler(ProtoLangToolchainProvider toolchain, Label protoCompiler) {
+    assertThat(toolchain.protoc().getExecutable().prettyPrint())
+        .isEqualTo(protoCompiler.toPathFragment().getPathString());
+  }
+
   @Test
   public void protoToolchain() throws Exception {
     scratch.file(
@@ -86,15 +91,52 @@ public void protoToolchain() throws Exception {
         "    plugin_format_flag = '--plugin=%s',",
         "    plugin = '//third_party/x:plugin',",
         "    runtime = '//third_party/x:runtime',",
-        "    blacklisted_protos = ['//third_party/x:denied'],",
         "    progress_message = 'Progress Message %{label}',",
         "    mnemonic = 'MyMnemonic',",
         ")");
 
     update(ImmutableList.of("//foo:toolchain"), false, 1, true, new EventBus());
+    ProtoLangToolchainProvider toolchain =
+        ProtoLangToolchainProvider.get(getConfiguredTarget("//foo:toolchain"));
+    Label protoc = Label.parseAbsoluteUnchecked(ProtoConstants.DEFAULT_PROTOC_LABEL);
+
+    validateProtoLangToolchain(toolchain);
+    validateProtoCompiler(toolchain, protoc);
+  }
+
+  @Test
+  public void protoToolchain_setProtoCompiler() throws Exception {
+    scratch.file(
+        "third_party/x/BUILD",
+        "licenses(['unencumbered'])",
+        "cc_binary(name = 'plugin', srcs = ['plugin.cc'])",
+        "cc_library(name = 'runtime', srcs = ['runtime.cc'])",
+        "filegroup(name = 'descriptors', srcs = ['metadata.proto', 'descriptor.proto'])",
+        "filegroup(name = 'any', srcs = ['any.proto'])",
+        "proto_library(name = 'denied', srcs = [':descriptors', ':any'])",
+        "cc_binary(name = 'compiler')");
+
+    scratch.file(
+        "foo/BUILD",
+        TestConstants.LOAD_PROTO_LANG_TOOLCHAIN,
+        "licenses(['unencumbered'])",
+        "proto_lang_toolchain(",
+        "    name = 'toolchain',",
+        "    command_line = 'cmd-line:$(OUT)',",
+        "    plugin_format_flag = '--plugin=%s',",
+        "    plugin = '//third_party/x:plugin',",
+        "    runtime = '//third_party/x:runtime',",
+        "    progress_message = 'Progress Message %{label}',",
+        "    mnemonic = 'MyMnemonic',",
+        "    proto_compiler = '//third_party/x:compiler',",
+        ")");
+
+    ProtoLangToolchainProvider toolchain =
+        ProtoLangToolchainProvider.get(getConfiguredTarget("//foo:toolchain"));
+    Label protoc = Label.parseAbsoluteUnchecked("//third_party/x:compiler");
 
-    validateProtoLangToolchain(
-        getConfiguredTarget("//foo:toolchain").getProvider(ProtoLangToolchainProvider.class));
+    validateProtoLangToolchain(toolchain);
+    validateProtoCompiler(toolchain, protoc);
   }
 
   @Test
@@ -117,15 +159,17 @@ public void protoToolchainBlacklistProtoLibraries() throws Exception {
         "    plugin_format_flag = '--plugin=%s',",
         "    plugin = '//third_party/x:plugin',",
         "    runtime = '//third_party/x:runtime',",
-        "    blacklisted_protos = ['//third_party/x:descriptors', '//third_party/x:any'],",
         "    progress_message = 'Progress Message %{label}',",
         "    mnemonic = 'MyMnemonic',",
         ")");
 
     update(ImmutableList.of("//foo:toolchain"), false, 1, true, new EventBus());
+    ProtoLangToolchainProvider toolchain =
+        ProtoLangToolchainProvider.get(getConfiguredTarget("//foo:toolchain"));
+    Label protoc = Label.parseAbsoluteUnchecked(ProtoConstants.DEFAULT_PROTOC_LABEL);
 
-    validateProtoLangToolchain(
-        getConfiguredTarget("//foo:toolchain").getProvider(ProtoLangToolchainProvider.class));
+    validateProtoLangToolchain(toolchain);
+    validateProtoCompiler(toolchain, protoc);
   }
 
   @Test
@@ -148,15 +192,17 @@ public void protoToolchainBlacklistTransitiveProtos() throws Exception {
         "    plugin_format_flag = '--plugin=%s',",
         "    plugin = '//third_party/x:plugin',",
         "    runtime = '//third_party/x:runtime',",
-        "    blacklisted_protos = ['//third_party/x:any'],",
         "    progress_message = 'Progress Message %{label}',",
         "    mnemonic = 'MyMnemonic',",
         ")");
 
     update(ImmutableList.of("//foo:toolchain"), false, 1, true, new EventBus());
+    ProtoLangToolchainProvider toolchain =
+        ProtoLangToolchainProvider.get(getConfiguredTarget("//foo:toolchain"));
+    Label protoc = Label.parseAbsoluteUnchecked(ProtoConstants.DEFAULT_PROTOC_LABEL);
 
-    validateProtoLangToolchain(
-        getConfiguredTarget("//foo:toolchain").getProvider(ProtoLangToolchainProvider.class));
+    validateProtoLangToolchain(toolchain);
+    validateProtoCompiler(toolchain, protoc);
   }
 
   @Test
@@ -170,9 +216,8 @@ public void optionalFieldsAreEmpty() throws Exception {
         ")");
 
     update(ImmutableList.of("//foo:toolchain"), false, 1, true, new EventBus());
-
     ProtoLangToolchainProvider toolchain =
-        getConfiguredTarget("//foo:toolchain").getProvider(ProtoLangToolchainProvider.class);
+        ProtoLangToolchainProvider.get(getConfiguredTarget("//foo:toolchain"));
 
     assertThat(toolchain.pluginExecutable()).isNull();
     assertThat(toolchain.runtime()).isNull();

From 1682183597f80fd475b4b87ef5f0990db4aad9ed Mon Sep 17 00:00:00 2001
From: ilist <ilist@google.com>
Date: Fri, 6 May 2022 01:16:32 -0700
Subject: [PATCH 08/24] Use data from transitive_proto_sources instead of
 transitive_sources in proto_lang_toolchain.

The problem with latter is, that transitive_source can contain "renamed" files (in _virtual_includes subdirectory), which doesn't work for the detection that needs original files (ProtoSource.original_source_file).

PiperOrigin-RevId: 446924924
---
 .../builtins_bzl/common/proto/proto_lang_toolchain.bzl      | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain.bzl b/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain.bzl
index ef2d4aee87b291..ad0c7221b86fb4 100644
--- a/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain.bzl
+++ b/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain.bzl
@@ -21,11 +21,7 @@ ProtoInfo = _builtins.toplevel.ProtoInfo
 proto_common = _builtins.toplevel.proto_common
 
 def _rule_impl(ctx):
-    provided_proto_sources = []
-    transitive_files = depset(transitive = [bp[ProtoInfo].transitive_sources for bp in ctx.attr.blacklisted_protos])
-    for file in transitive_files.to_list():
-        source_root = file.root.path
-        provided_proto_sources.append(proto_common.ProtoSource(file, file, source_root))
+    provided_proto_sources = depset(transitive = [bp[ProtoInfo].transitive_proto_sources() for bp in ctx.attr.blacklisted_protos]).to_list()
 
     flag = ctx.attr.command_line
     if flag.find("$(PLUGIN_OUT)") > -1:

From fa217b1d6b06555dc7b04ff09af86b636fbf02ed Mon Sep 17 00:00:00 2001
From: kotlaja <kotlaja@google.com>
Date: Fri, 6 May 2022 08:05:56 -0700
Subject: [PATCH 09/24] Remove allow_files from proto_lang_toolchain's
 attributes

It's the same behaviour in native code.

PiperOrigin-RevId: 446988764
---
 .../builtins_bzl/common/proto/proto_lang_toolchain.bzl     | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain.bzl b/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain.bzl
index ad0c7221b86fb4..7c51cb2433d314 100644
--- a/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain.bzl
+++ b/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain.bzl
@@ -60,19 +60,14 @@ proto_lang_toolchain = rule(
         "plugin": attr.label(
             executable = True,
             cfg = "exec",
-            allow_files = True,
-        ),
-        "runtime": attr.label(
-            allow_files = True,
         ),
+        "runtime": attr.label(),
         "blacklisted_protos": attr.label_list(
-            allow_files = True,
             providers = [ProtoInfo],
         ),
         "_proto_compiler": attr.label(
             cfg = "exec",
             executable = True,
-            allow_files = True,
             default = configuration_field("proto", "proto_compiler"),
         ),
     },

From 4c0640b003a699da505d06277d7548ec06377610 Mon Sep 17 00:00:00 2001
From: kotlaja <kotlaja@google.com>
Date: Fri, 6 May 2022 08:45:12 -0700
Subject: [PATCH 10/24] Remove native implementation of proto_lang_toolchain
 rule

PiperOrigin-RevId: 446995789
---
 .../lib/rules/proto/ProtoLangToolchain.java   | 74 -------------------
 .../rules/proto/ProtoLangToolchainRule.java   |  8 +-
 2 files changed, 6 insertions(+), 76 deletions(-)
 delete mode 100644 src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java

diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java
deleted file mode 100644
index e4abfe9bc194a2..00000000000000
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchain.java
+++ /dev/null
@@ -1,74 +0,0 @@
-// Copyright 2016 The Bazel Authors. All rights reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//    http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package com.google.devtools.build.lib.rules.proto;
-
-import static com.google.devtools.build.lib.collect.nestedset.Order.STABLE_ORDER;
-
-import com.google.devtools.build.lib.actions.Artifact;
-import com.google.devtools.build.lib.actions.MutableActionGraph.ActionConflictException;
-import com.google.devtools.build.lib.analysis.ConfiguredTarget;
-import com.google.devtools.build.lib.analysis.FilesToRunProvider;
-import com.google.devtools.build.lib.analysis.RuleConfiguredTargetBuilder;
-import com.google.devtools.build.lib.analysis.RuleConfiguredTargetFactory;
-import com.google.devtools.build.lib.analysis.RuleContext;
-import com.google.devtools.build.lib.analysis.Runfiles;
-import com.google.devtools.build.lib.analysis.RunfilesProvider;
-import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
-import com.google.devtools.build.lib.packages.Type;
-
-/** Implements {code proto_lang_toolchain}. */
-public class ProtoLangToolchain implements RuleConfiguredTargetFactory {
-  @Override
-  public ConfiguredTarget create(RuleContext ruleContext)
-      throws InterruptedException, RuleErrorException, ActionConflictException {
-    NestedSetBuilder<ProtoSource> providedProtoSources = NestedSetBuilder.stableOrder();
-    for (ProtoInfo protoInfo :
-        ruleContext.getPrerequisites("blacklisted_protos", ProtoInfo.PROVIDER)) {
-      providedProtoSources.addTransitive(protoInfo.getTransitiveSources());
-    }
-
-    String flag = ruleContext.attributes().get("command_line", Type.STRING);
-
-    if (flag.contains("$(PLUGIN_OUT)")) {
-      ruleContext.attributeError("command_line", "Placeholder '$(PLUGIN_OUT)' is not supported.");
-      return null;
-    }
-
-    flag = flag.replace("$(OUT)", "%s");
-
-    return new RuleConfiguredTargetBuilder(ruleContext)
-        .addStarlarkDeclaredProvider(
-            ProtoLangToolchainProvider.create(
-                flag,
-                ruleContext.attributes().get("plugin_format_flag", Type.STRING),
-                ruleContext.getPrerequisite("plugin", FilesToRunProvider.class),
-                ruleContext.getPrerequisite("runtime"),
-                // We intentionally flatten the NestedSet here.
-                //
-                // `providedProtoSources` are read during analysis, so flattening the set here once
-                // avoid flattening it in every `<lang>_proto_aspect` applied to `proto_library`
-                // targets. While this has the potential to use more memory than using a NestedSet,
-                // there are only a few `proto_lang_toolchain` targets in every build, so the impact
-                // on memory consumption should be neglectable.
-                providedProtoSources.build().toList(),
-                ruleContext.getPrerequisite(":proto_compiler", FilesToRunProvider.class),
-                ruleContext.getConfiguration().getFragment(ProtoConfiguration.class).protocOpts(),
-                ruleContext.attributes().get("progress_message", Type.STRING),
-                ruleContext.attributes().get("mnemonic", Type.STRING)))
-        .setFilesToBuild(NestedSetBuilder.<Artifact>emptySet(STABLE_ORDER))
-        .addProvider(RunfilesProvider.simple(Runfiles.EMPTY))
-        .build();
-  }
-}
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java
index 7ac6344ea2b6ce..20b91245c286bb 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainRule.java
@@ -28,7 +28,11 @@
 import com.google.devtools.build.lib.packages.StarlarkProviderIdentifier;
 import com.google.devtools.build.lib.packages.Type;
 
-/** Implements {code proto_lang_toolchain}. */
+/**
+ * Implements {code proto_lang_toolchain}.
+ *
+ * <p>This rule is implemented in Starlark. This class remains only for doc-gen purposes.
+ */
 public class ProtoLangToolchainRule implements RuleDefinition {
   private static final Label DEFAULT_PROTO_COMPILER =
       Label.parseAbsoluteUnchecked(ProtoConstants.DEFAULT_PROTOC_LABEL);
@@ -116,7 +120,7 @@ public Metadata getMetadata() {
     return RuleDefinition.Metadata.builder()
         .name("proto_lang_toolchain")
         .ancestors(BaseRuleClasses.NativeActionCreatingRule.class)
-        .factoryClass(ProtoLangToolchain.class)
+        .factoryClass(BaseRuleClasses.EmptyRuleConfiguredTargetFactory.class)
         .build();
   }
 }

From a14e613cff628c40e9d8a7301ad482bad7860bd3 Mon Sep 17 00:00:00 2001
From: kotlaja <kotlaja@google.com>
Date: Fri, 6 May 2022 10:23:13 -0700
Subject: [PATCH 11/24] Fix name in proto_lang_toolchain rule

Added a wrapper for proto_lang_toolchain in order to have two implementations - one with public proto_compiler attribute, and the other one with the private one.
Restructured proto_lang_toolchain rules' implementation - merged two rule's definitions into one.

PiperOrigin-RevId: 447016335
---
 .../starlark/builtins_bzl/common/exports.bzl  |  2 +-
 .../common/proto/proto_lang_toolchain.bzl     | 63 +++++++++++--------
 .../proto_lang_toolchain_custom_protoc.bzl    | 23 +++++++
 .../proto_lang_toolchain_default_protoc.bzl   | 23 +++++++
 .../proto/proto_lang_toolchain_wrapper.bzl    | 35 +++++++++++
 .../rules/proto/ProtoLangToolchainTest.java   |  8 +--
 6 files changed, 124 insertions(+), 30 deletions(-)
 create mode 100644 src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain_custom_protoc.bzl
 create mode 100644 src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain_default_protoc.bzl
 create mode 100644 src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain_wrapper.bzl

diff --git a/src/main/starlark/builtins_bzl/common/exports.bzl b/src/main/starlark/builtins_bzl/common/exports.bzl
index b58257ba5dda5e..a4e95bb1913728 100644
--- a/src/main/starlark/builtins_bzl/common/exports.bzl
+++ b/src/main/starlark/builtins_bzl/common/exports.bzl
@@ -22,7 +22,7 @@ load("@_builtins//:common/objc/objc_import.bzl", "objc_import")
 load("@_builtins//:common/objc/objc_library.bzl", "objc_library")
 load("@_builtins//:common/objc/apple_static_library.bzl", "apple_static_library")
 load("@_builtins//:common/objc/compilation_support.bzl", "compilation_support")
-load("@_builtins//:common/proto/proto_lang_toolchain.bzl", "proto_lang_toolchain")
+load("@_builtins//:common/proto/proto_lang_toolchain_wrapper.bzl", "proto_lang_toolchain")
 
 exported_toplevels = {
     # This dummy symbol is not part of the public API; it is only used to test
diff --git a/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain.bzl b/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain.bzl
index 7c51cb2433d314..1fd63f7632298e 100644
--- a/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain.bzl
+++ b/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain.bzl
@@ -18,7 +18,6 @@ load(":common/proto/proto_common.bzl", "ProtoLangToolchainInfo")
 load(":common/proto/proto_semantics.bzl", "semantics")
 
 ProtoInfo = _builtins.toplevel.ProtoInfo
-proto_common = _builtins.toplevel.proto_common
 
 def _rule_impl(ctx):
     provided_proto_sources = depset(transitive = [bp[ProtoInfo].transitive_proto_sources() for bp in ctx.attr.blacklisted_protos]).to_list()
@@ -32,6 +31,9 @@ def _rule_impl(ctx):
     if ctx.attr.plugin != None:
         plugin = ctx.attr.plugin[DefaultInfo].files_to_run
 
+    proto_compiler = getattr(ctx.attr, "proto_compiler", None)
+    proto_compiler = getattr(ctx.attr, "_proto_compiler", proto_compiler)
+
     return [
         DefaultInfo(
             files = depset(),
@@ -43,34 +45,45 @@ def _rule_impl(ctx):
             plugin = plugin,
             runtime = ctx.attr.runtime,
             provided_proto_sources = provided_proto_sources,
-            proto_compiler = ctx.attr._proto_compiler.files_to_run,
+            proto_compiler = proto_compiler.files_to_run,
             protoc_opts = ctx.fragments.proto.experimental_protoc_opts,
             progress_message = ctx.attr.progress_message,
             mnemonic = ctx.attr.mnemonic,
         ),
     ]
 
-proto_lang_toolchain = rule(
-    implementation = _rule_impl,
-    attrs = {
-        "progress_message": attr.string(default = "Generating proto_library %{label}"),
-        "mnemonic": attr.string(default = "GenProto"),
-        "command_line": attr.string(mandatory = True),
-        "plugin_format_flag": attr.string(),
-        "plugin": attr.label(
-            executable = True,
-            cfg = "exec",
-        ),
-        "runtime": attr.label(),
-        "blacklisted_protos": attr.label_list(
-            providers = [ProtoInfo],
-        ),
-        "_proto_compiler": attr.label(
-            cfg = "exec",
-            executable = True,
-            default = configuration_field("proto", "proto_compiler"),
+def make_proto_lang_toolchain(custom_proto_compiler):
+    return rule(
+        _rule_impl,
+        attrs = dict(
+            {
+                "progress_message": attr.string(default = "Generating proto_library %{label}"),
+                "mnemonic": attr.string(default = "GenProto"),
+                "command_line": attr.string(mandatory = True),
+                "plugin_format_flag": attr.string(),
+                "plugin": attr.label(
+                    executable = True,
+                    cfg = "exec",
+                ),
+                "runtime": attr.label(),
+                "blacklisted_protos": attr.label_list(
+                    providers = [ProtoInfo],
+                ),
+            },
+            **({
+                "proto_compiler": attr.label(
+                    cfg = "exec",
+                    executable = True,
+                ),
+            } if custom_proto_compiler else {
+                "_proto_compiler": attr.label(
+                    cfg = "exec",
+                    executable = True,
+                    allow_files = True,
+                    default = configuration_field("proto", "proto_compiler"),
+                ),
+            })
         ),
-    },
-    provides = [ProtoLangToolchainInfo],
-    fragments = ["proto"] + semantics.EXTRA_FRAGMENTS,
-)
+        provides = [ProtoLangToolchainInfo],
+        fragments = ["proto"] + semantics.EXTRA_FRAGMENTS,
+    )
diff --git a/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain_custom_protoc.bzl b/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain_custom_protoc.bzl
new file mode 100644
index 00000000000000..b4157f6aef66ad
--- /dev/null
+++ b/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain_custom_protoc.bzl
@@ -0,0 +1,23 @@
+# Copyright 2021 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Defines a proto_lang_toolchain rule class with custom proto compiler.
+
+There are two physical rule classes for proto_lang_toolchain and we want both of them
+to have a name string of "proto_lang_toolchain".
+"""
+
+load(":common/proto/proto_lang_toolchain.bzl", "make_proto_lang_toolchain")
+
+proto_lang_toolchain = make_proto_lang_toolchain(custom_proto_compiler = True)
diff --git a/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain_default_protoc.bzl b/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain_default_protoc.bzl
new file mode 100644
index 00000000000000..7345bc37fc0013
--- /dev/null
+++ b/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain_default_protoc.bzl
@@ -0,0 +1,23 @@
+# Copyright 2021 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Defines a proto_lang_toolchain rule class with default proto compiler.
+
+There are two physical rule classes for proto_lang_toolchain and we want both of them
+to have a name string of "proto_lang_toolchain".
+"""
+
+load(":common/proto/proto_lang_toolchain.bzl", "make_proto_lang_toolchain")
+
+proto_lang_toolchain = make_proto_lang_toolchain(custom_proto_compiler = False)
diff --git a/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain_wrapper.bzl b/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain_wrapper.bzl
new file mode 100644
index 00000000000000..075e6d73c1c2c0
--- /dev/null
+++ b/src/main/starlark/builtins_bzl/common/proto/proto_lang_toolchain_wrapper.bzl
@@ -0,0 +1,35 @@
+# Copyright 2021 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Macro encapsulating the proto_lang_toolchain implementation.
+
+This is needed since proto compiler can be defined, or used as a default one.
+There are two implementations of proto_lang_toolchain - one with public proto_compiler attribute, and the other one with private compiler.
+"""
+
+load(":common/proto/proto_lang_toolchain_default_protoc.bzl", toolchain_default_protoc = "proto_lang_toolchain")
+load(":common/proto/proto_lang_toolchain_custom_protoc.bzl", toolchain_custom_protoc = "proto_lang_toolchain")
+
+def proto_lang_toolchain(
+        proto_compiler = None,
+        **kwargs):
+    if proto_compiler != None:
+        toolchain_custom_protoc(
+            proto_compiler = proto_compiler,
+            **kwargs
+        )
+    else:
+        toolchain_default_protoc(
+            **kwargs
+        )
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java b/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
index a6860bdc2437d9..2a0d47fe52f792 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainTest.java
@@ -50,7 +50,7 @@ Provider.Key getStarlarkProtoLangToolchainInfoKey() throws LabelSyntaxException
   }
 
   private void validateProtoLangToolchain(ProtoLangToolchainProvider toolchain) throws Exception {
-    assertThat(toolchain.outReplacementFormatFlag()).isEqualTo("cmd-line");
+    assertThat(toolchain.outReplacementFormatFlag()).isEqualTo("cmd-line:%s");
     assertThat(toolchain.pluginFormatFlag()).isEqualTo("--plugin=%s");
     assertThat(toolchain.pluginExecutable().getExecutable().getRootRelativePathString())
         .isEqualTo("third_party/x/plugin");
@@ -87,7 +87,7 @@ public void protoToolchain() throws Exception {
         "licenses(['unencumbered'])",
         "proto_lang_toolchain(",
         "    name = 'toolchain',",
-        "    command_line = 'cmd-line',",
+        "    command_line = 'cmd-line:$(OUT)',",
         "    plugin_format_flag = '--plugin=%s',",
         "    plugin = '//third_party/x:plugin',",
         "    runtime = '//third_party/x:runtime',",
@@ -155,7 +155,7 @@ public void protoToolchainBlacklistProtoLibraries() throws Exception {
         TestConstants.LOAD_PROTO_LANG_TOOLCHAIN,
         "proto_lang_toolchain(",
         "    name = 'toolchain',",
-        "    command_line = 'cmd-line',",
+        "    command_line = 'cmd-line:$(OUT)',",
         "    plugin_format_flag = '--plugin=%s',",
         "    plugin = '//third_party/x:plugin',",
         "    runtime = '//third_party/x:runtime',",
@@ -188,7 +188,7 @@ public void protoToolchainBlacklistTransitiveProtos() throws Exception {
         TestConstants.LOAD_PROTO_LANG_TOOLCHAIN,
         "proto_lang_toolchain(",
         "    name = 'toolchain',",
-        "    command_line = 'cmd-line',",
+        "    command_line = 'cmd-line:$(OUT)',",
         "    plugin_format_flag = '--plugin=%s',",
         "    plugin = '//third_party/x:plugin',",
         "    runtime = '//third_party/x:runtime',",

From cc6c523eaec2d0c216982f857e41b5b644b16efc Mon Sep 17 00:00:00 2001
From: Ivo List <ilist@google.com>
Date: Wed, 13 Jul 2022 13:18:58 +0000
Subject: [PATCH 12/24] Cherry-pick missing things.

---
 .../build/lib/rules/proto/ProtoCommon.java    | 15 ++++++++++++++
 .../build/lib/rules/proto/ProtoInfo.java      | 14 +++++++++++++
 .../build/lib/rules/proto/ProtoSource.java    | 20 +++++++++++++++----
 .../lib/starlarkbuildapi/ProtoInfoApi.java    |  8 ++++++++
 .../starlark/builtins_bzl/common/exports.bzl  |  2 +-
 5 files changed, 54 insertions(+), 5 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 8c1d348d8ffe4c..733b41dad8d275 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
@@ -31,10 +31,15 @@
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.concurrent.BlazeInterners;
+import com.google.devtools.build.lib.packages.BazelModuleContext;
 import com.google.devtools.build.lib.packages.BuildType;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import javax.annotation.Nullable;
+import net.starlark.java.eval.EvalException;
+import net.starlark.java.eval.Module;
+import net.starlark.java.eval.Starlark;
+import net.starlark.java.eval.StarlarkThread;
 
 /** Utility functions for proto_library and proto aspect implementations. */
 public class ProtoCommon {
@@ -518,4 +523,14 @@ public static boolean areDepsStrict(RuleContext ruleContext) {
     StrictDepsMode getBool = ruleContext.getFragment(ProtoConfiguration.class).strictProtoDeps();
     return getBool != StrictDepsMode.OFF && getBool != StrictDepsMode.DEFAULT;
   }
+
+  public static void checkPrivateStarlarkificationAllowlist(StarlarkThread thread)
+      throws EvalException {
+    Label label =
+        ((BazelModuleContext) Module.ofInnermostEnclosingStarlarkFunction(thread).getClientData())
+            .label();
+    if (!label.getPackageIdentifier().getRepository().toString().equals("@_builtins")) {
+      throw Starlark.errorf("Rule in '%s' cannot use private API", label.getPackageName());
+    }
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoInfo.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoInfo.java
index 15a44346c8aca9..0593c5a0b3759f 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoInfo.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoInfo.java
@@ -25,6 +25,8 @@
 import com.google.devtools.build.lib.starlarkbuildapi.ProtoInfoApi;
 import com.google.devtools.build.lib.starlarkbuildapi.proto.ProtoBootstrap;
 import com.google.devtools.build.lib.vfs.PathFragment;
+import net.starlark.java.eval.EvalException;
+import net.starlark.java.eval.StarlarkThread;
 
 /**
  * Configured target classes that implement this class can contribute .proto files to the
@@ -133,6 +135,12 @@ public String getDirectProtoSourceRoot() {
     return Depset.of(Artifact.TYPE, getTransitiveProtoSources());
   }
 
+  @Override
+  public Depset getTransitiveSourcesForStarlark(StarlarkThread thread) throws EvalException {
+    ProtoCommon.checkPrivateStarlarkificationAllowlist(thread);
+    return Depset.of(ProtoSource.TYPE, transitiveSources);
+  }
+
   /**
    * The {@code .proto} source files in this {@code proto_library}'s {@code srcs} and all of its
    * transitive dependencies.
@@ -222,6 +230,12 @@ public NestedSet<ProtoSource> getStrictImportableSources() {
     return strictImportableSources;
   }
 
+  @Override
+  public Depset getStrictImportableSourcesForStarlark(StarlarkThread thread) throws EvalException {
+    ProtoCommon.checkPrivateStarlarkificationAllowlist(thread);
+    return Depset.of(ProtoSource.TYPE, strictImportableSources);
+  }
+
   /**
    * Returns a set of {@code .proto} sources that may be re-exported by this {@code proto_library}'s
    * direct sources.
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoSource.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoSource.java
index 50d49455e0f8b5..1833d3004362ca 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoSource.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoSource.java
@@ -15,15 +15,16 @@
 package com.google.devtools.build.lib.rules.proto;
 
 import com.google.devtools.build.lib.actions.Artifact;
+import com.google.devtools.build.lib.collect.nestedset.Depset;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
-import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import net.starlark.java.eval.StarlarkValue;
 
 /** Represents a single {@code .proto} source file. */
 @Immutable
-@AutoCodec
-class ProtoSource implements StarlarkValue {
+final class ProtoSource implements StarlarkValue {
+  public static final Depset.ElementType TYPE = Depset.ElementType.of(ProtoSource.class);
+
   private final Artifact sourceFile;
   private final Artifact originalSourceFile;
   private final PathFragment sourceRoot;
@@ -32,7 +33,6 @@ public ProtoSource(Artifact sourceFile, PathFragment sourceRoot) {
     this(sourceFile, sourceFile, sourceRoot);
   }
 
-  @AutoCodec.Instantiator
   ProtoSource(Artifact sourceFile, Artifact originalSourceFile, PathFragment sourceRoot) {
     this.sourceFile = sourceFile;
     this.originalSourceFile = originalSourceFile;
@@ -43,6 +43,12 @@ public Artifact getSourceFile() {
     return sourceFile;
   }
 
+  @StarlarkMethod(name = "source_file", documented = false, useStarlarkThread = true)
+  public Artifact getSourceFileForStarlark(StarlarkThread thread) throws EvalException {
+    ProtoCommon.checkPrivateStarlarkificationAllowlist(thread);
+    return sourceFile;
+  }
+
   /** Returns the original source file. Only for forbidding protos! */
   @Deprecated
   Artifact getOriginalSourceFile() {
@@ -53,6 +59,12 @@ public PathFragment getSourceRoot() {
     return sourceRoot;
   }
 
+  @StarlarkMethod(name = "import_path", documented = false, useStarlarkThread = true)
+  public String getImportPathForStarlark(StarlarkThread thread) throws EvalException {
+    ProtoCommon.checkPrivateStarlarkificationAllowlist(thread);
+    return getImportPath().getPathString();
+  }
+
   public PathFragment getImportPath() {
     return sourceFile.getExecPath().relativeTo(sourceRoot);
   }
diff --git a/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/ProtoInfoApi.java b/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/ProtoInfoApi.java
index ce1fad68e96a5d..ea7214f3251db3 100644
--- a/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/ProtoInfoApi.java
+++ b/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/ProtoInfoApi.java
@@ -21,6 +21,8 @@
 import com.google.devtools.build.lib.starlarkbuildapi.core.StructApi;
 import net.starlark.java.annot.StarlarkBuiltin;
 import net.starlark.java.annot.StarlarkMethod;
+import net.starlark.java.eval.EvalException;
+import net.starlark.java.eval.StarlarkThread;
 
 /** Info object propagating information about protocol buffer sources. */
 @StarlarkBuiltin(
@@ -105,4 +107,10 @@ interface ProtoInfoProviderApi extends ProviderApi {
               + " as a source, that source file would be imported as 'import c/d.proto'",
       structField = true)
   String getDirectProtoSourceRoot();
+
+  @StarlarkMethod(name = "strict_importable_sources", documented = false, useStarlarkThread = true)
+  Depset getStrictImportableSourcesForStarlark(StarlarkThread thread) throws EvalException;
+
+  @StarlarkMethod(name = "transitive_proto_sources", documented = false, useStarlarkThread = true)
+  Depset getTransitiveSourcesForStarlark(StarlarkThread thread) throws EvalException;
 }
diff --git a/src/main/starlark/builtins_bzl/common/exports.bzl b/src/main/starlark/builtins_bzl/common/exports.bzl
index a4e95bb1913728..7d0f209e6b2c70 100644
--- a/src/main/starlark/builtins_bzl/common/exports.bzl
+++ b/src/main/starlark/builtins_bzl/common/exports.bzl
@@ -39,7 +39,7 @@ exported_rules = {
     "+apple_static_library": apple_static_library,
     "+cc_shared_library": cc_shared_library,
     "+cc_shared_library_permissions": cc_shared_library_permissions,
-    "+proto_lang_toolchain": proto_lang_toolchain,
+    "proto_lang_toolchain": proto_lang_toolchain,
 }
 exported_to_java = {
     "register_compile_and_archive_actions_for_j2objc": compilation_support.register_compile_and_archive_actions_for_j2objc,

From d2ded6ccc6deeece9ecdb146514b721f04fc31f7 Mon Sep 17 00:00:00 2001
From: wilwell <wilwell@google.com>
Date: Tue, 30 Nov 2021 03:25:47 -0800
Subject: [PATCH 13/24] Separate ExecException.java from main actions target.

PiperOrigin-RevId: 413105213
---
 .../lib/actions/ActionExecutionException.java | 42 +++++++++++++++++
 .../google/devtools/build/lib/actions/BUILD   |  6 +++
 .../build/lib/actions/ExecException.java      | 46 -------------------
 .../lib/actions/LostInputsExecException.java  |  3 +-
 .../actions/AbstractFileWriteAction.java      |  6 +--
 .../lib/analysis/actions/SpawnAction.java     |  4 +-
 .../actions/TemplateExpansionAction.java      |  3 +-
 .../lib/analysis/test/TestRunnerAction.java   | 13 +++---
 .../coverage/CoverageReportActionBuilder.java | 34 +++++++-------
 .../build/lib/exec/SymlinkTreeStrategy.java   |  6 +--
 .../build/lib/rules/cpp/CppCompileAction.java | 11 +++--
 .../build/lib/rules/cpp/CppLinkAction.java    |  5 +-
 .../lib/rules/java/JavaCompileAction.java     | 18 ++++----
 .../lib/skyframe/SkyframeActionExecutor.java  |  4 +-
 .../StandaloneSpawnStrategyTest.java          |  2 +-
 15 files changed, 102 insertions(+), 101 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/actions/ActionExecutionException.java b/src/main/java/com/google/devtools/build/lib/actions/ActionExecutionException.java
index 0e32500cf9f657..613a15187a6d47 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/ActionExecutionException.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/ActionExecutionException.java
@@ -24,6 +24,7 @@
 import com.google.devtools.build.lib.skyframe.DetailedException;
 import com.google.devtools.build.lib.util.DetailedExitCode;
 import com.google.devtools.build.lib.util.ExitCode;
+import javax.annotation.Nullable;
 import net.starlark.java.syntax.Location;
 
 /**
@@ -115,6 +116,47 @@ private static NestedSet<Cause> rootCausesFromAction(
                 detailedExitCode));
   }
 
+  public static ActionExecutionException fromExecException(ExecException exception, Action action) {
+    return fromExecException(exception, null, action);
+  }
+
+  /**
+   * Returns a new ActionExecutionException given an optional action subtask describing which part
+   * of the action failed (should be null for standard action failures). When appropriate (we use
+   * some heuristics to decide), produces an abbreviated message incorporating just the termination
+   * status if available.
+   *
+   * @param exception initial ExecException
+   * @param actionSubtask additional information about the action
+   * @param action failed action
+   * @return ActionExecutionException object describing the action failure
+   */
+  public static ActionExecutionException fromExecException(
+      ExecException exception, @Nullable String actionSubtask, Action action) {
+    // Message from ActionExecutionException will be prepended with action.describe() where
+    // necessary: because not all ActionExecutionExceptions come from this codepath, it is safer
+    // for consumers to manually prepend. We still put action.describe() in the failure detail
+    // message argument.
+    String message =
+        (actionSubtask == null ? "" : actionSubtask + ": ")
+            + exception.getMessageForActionExecutionException();
+
+    DetailedExitCode code =
+        DetailedExitCode.of(exception.getFailureDetail(action.describe() + " failed: " + message));
+
+    if (exception instanceof LostInputsExecException) {
+      return ((LostInputsExecException) exception).fromExecException(message, action, code);
+    }
+
+    return fromExecException(exception, message, action, code);
+  }
+
+  public static ActionExecutionException fromExecException(
+      ExecException exception, String message, Action action, DetailedExitCode code) {
+    return new ActionExecutionException(
+        message, exception, action, exception.isCatastrophic(), code);
+  }
+
   /** Returns the action that failed. */
   public ActionAnalysisMetadata getAction() {
     return action;
diff --git a/src/main/java/com/google/devtools/build/lib/actions/BUILD b/src/main/java/com/google/devtools/build/lib/actions/BUILD
index 8ae7bb4d2d2779..e352f9b3f8f64e 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/actions/BUILD
@@ -344,3 +344,9 @@ java_library(
         "//third_party:jsr305",
     ],
 )
+
+java_library(
+    name = "exec_exception",
+    srcs = ["ExecException.java"],
+    deps = ["//src/main/protobuf:failure_details_java_proto"],
+)
diff --git a/src/main/java/com/google/devtools/build/lib/actions/ExecException.java b/src/main/java/com/google/devtools/build/lib/actions/ExecException.java
index 43ba000e3632f3..c3544877031246 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/ExecException.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/ExecException.java
@@ -15,9 +15,6 @@
 package com.google.devtools.build.lib.actions;
 
 import com.google.devtools.build.lib.server.FailureDetails.FailureDetail;
-import com.google.devtools.build.lib.util.DetailedExitCode;
-import com.google.errorprone.annotations.ForOverride;
-import javax.annotation.Nullable;
 
 /**
  * An exception indication that the execution of an action has failed OR could not be attempted OR
@@ -82,52 +79,9 @@ public boolean isCatastrophic() {
     return catastrophe;
   }
 
-  /**
-   * Returns a new ActionExecutionException without a message prefix.
-   *
-   * @param action failed action
-   * @return ActionExecutionException object describing the action failure
-   */
-  public final ActionExecutionException toActionExecutionException(Action action) {
-    return toActionExecutionException(null, action);
-  }
-
-  /**
-   * Returns a new ActionExecutionException given an optional action subtask describing which part
-   * of the action failed (should be null for standard action failures). When appropriate (we use
-   * some heuristics to decide), produces an abbreviated message incorporating just the termination
-   * status if available.
-   *
-   * @param actionSubtask additional information about the action
-   * @param action failed action
-   * @return ActionExecutionException object describing the action failure
-   */
-  public final ActionExecutionException toActionExecutionException(
-      @Nullable String actionSubtask, Action action) {
-    // Message from ActionExecutionException will be prepended with action.describe() where
-    // necessary: because not all ActionExecutionExceptions come from this codepath, it is safer
-    // for consumers to manually prepend. We still put action.describe() in the failure detail
-    // message argument.
-    String message =
-        (actionSubtask == null ? "" : actionSubtask + ": ")
-            + getMessageForActionExecutionException();
-    return toActionExecutionException(
-        message,
-        action,
-        DetailedExitCode.of(getFailureDetail(action.describe() + " failed: " + message)));
-  }
-
-  @ForOverride
-  protected ActionExecutionException toActionExecutionException(
-      String message, Action action, DetailedExitCode code) {
-    return new ActionExecutionException(message, this, action, isCatastrophic(), code);
-  }
-
-  @ForOverride
   protected String getMessageForActionExecutionException() {
     return getMessage();
   }
 
-  @ForOverride
   protected abstract FailureDetail getFailureDetail(String message);
 }
diff --git a/src/main/java/com/google/devtools/build/lib/actions/LostInputsExecException.java b/src/main/java/com/google/devtools/build/lib/actions/LostInputsExecException.java
index 0be336183d5bef..f992d5d3543e7f 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/LostInputsExecException.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/LostInputsExecException.java
@@ -64,8 +64,7 @@ public ActionInputDepOwners getOwners() {
     return owners;
   }
 
-  @Override
-  protected ActionExecutionException toActionExecutionException(
+  protected ActionExecutionException fromExecException(
       String message, Action action, DetailedExitCode code) {
     return new LostInputsActionExecutionException(
         message, lostInputs, owners, action, /*cause=*/ this, code);
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/AbstractFileWriteAction.java b/src/main/java/com/google/devtools/build/lib/analysis/actions/AbstractFileWriteAction.java
index 170f97cafd2f49..a0974b7813d162 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/actions/AbstractFileWriteAction.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/actions/AbstractFileWriteAction.java
@@ -89,16 +89,14 @@ public ActionContinuationOrResult execute()
               return this;
             }
           } catch (ExecException e) {
-            throw e.toActionExecutionException(
-                AbstractFileWriteAction.this);
+            throw ActionExecutionException.fromExecException(e, AbstractFileWriteAction.this);
           }
           afterWrite(actionExecutionContext);
           return ActionContinuationOrResult.of(ActionResult.create(nextContinuation.get()));
         }
       };
     } catch (ExecException e) {
-      throw e.toActionExecutionException(
-          this);
+      throw ActionExecutionException.fromExecException(e, this);
     }
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java b/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java
index ca3c553479c915..094e7e73c8235e 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java
@@ -318,7 +318,7 @@ public final ActionContinuationOrResult beginExecution(
       beforeExecute(actionExecutionContext);
       spawn = getSpawn(actionExecutionContext);
     } catch (ExecException e) {
-      throw e.toActionExecutionException(this);
+      throw ActionExecutionException.fromExecException(e, this);
     } catch (CommandLineExpansionException e) {
       throw createCommandLineException(e);
     }
@@ -1428,7 +1428,7 @@ public ActionContinuationOrResult execute()
         }
         return new SpawnActionContinuation(actionExecutionContext, nextContinuation);
       } catch (ExecException e) {
-        throw e.toActionExecutionException(SpawnAction.this);
+        throw ActionExecutionException.fromExecException(e, SpawnAction.this);
       }
     }
   }
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/TemplateExpansionAction.java b/src/main/java/com/google/devtools/build/lib/analysis/actions/TemplateExpansionAction.java
index b84902218fd17f..5ac3c526eb5b78 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/actions/TemplateExpansionAction.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/actions/TemplateExpansionAction.java
@@ -164,8 +164,7 @@ public ActionContinuationOrResult execute()
             return this;
           }
         } catch (ExecException e) {
-          throw e.toActionExecutionException(
-              TemplateExpansionAction.this);
+          throw ActionExecutionException.fromExecException(e, TemplateExpansionAction.this);
         }
         return ActionContinuationOrResult.of(ActionResult.create(nextContinuation.get()));
       }
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/test/TestRunnerAction.java b/src/main/java/com/google/devtools/build/lib/analysis/test/TestRunnerAction.java
index 01795eb8658a6b..2c037f6996f940 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/test/TestRunnerAction.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/test/TestRunnerAction.java
@@ -899,10 +899,10 @@ public ActionContinuationOrResult beginExecution(
           testActionContext.isTestKeepGoing(),
           cancelFuture);
     } catch (ExecException e) {
-      throw e.toActionExecutionException(this);
+      throw ActionExecutionException.fromExecException(e, this);
     } catch (IOException e) {
-      throw new EnvironmentalExecException(e, Code.TEST_RUNNER_IO_EXCEPTION)
-          .toActionExecutionException(this);
+      throw ActionExecutionException.fromExecException(
+          new EnvironmentalExecException(e, Code.TEST_RUNNER_IO_EXCEPTION), this);
     }
   }
 
@@ -1203,10 +1203,11 @@ public ActionContinuationOrResult execute()
         Preconditions.checkState(actualMaxAttempts != 0);
         return process(result, actualMaxAttempts);
       } catch (ExecException e) {
-        throw e.toActionExecutionException(this.testRunnerAction);
+        throw ActionExecutionException.fromExecException(e, this.testRunnerAction);
       } catch (IOException e) {
-        throw new EnvironmentalExecException(e, Code.TEST_RUNNER_IO_EXCEPTION)
-            .toActionExecutionException(this.testRunnerAction);
+        throw ActionExecutionException.fromExecException(
+            new EnvironmentalExecException(e, Code.TEST_RUNNER_IO_EXCEPTION),
+            this.testRunnerAction);
       }
     }
 
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/coverage/CoverageReportActionBuilder.java b/src/main/java/com/google/devtools/build/lib/bazel/coverage/CoverageReportActionBuilder.java
index 6e09f9743caeef..242e45a2eac57a 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/coverage/CoverageReportActionBuilder.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/coverage/CoverageReportActionBuilder.java
@@ -65,29 +65,29 @@
 
 /**
  * A class to create the coverage report generator action.
- * 
- * <p>The coverage report action is created after every test shard action is created, at the
- * very end of the analysis phase. There is only one coverage report action per coverage
- * command invocation. It can also be viewed as a single sink node of the action graph.
- * 
+ *
+ * <p>The coverage report action is created after every test shard action is created, at the very
+ * end of the analysis phase. There is only one coverage report action per coverage command
+ * invocation. It can also be viewed as a single sink node of the action graph.
+ *
  * <p>Its inputs are the individual coverage.dat files from the test outputs (each shard produces
- * one) and the baseline coverage artifacts. Note that each ConfiguredTarget among the
- * transitive dependencies of the top level test targets may provide baseline coverage artifacts.
- * 
+ * one) and the baseline coverage artifacts. Note that each ConfiguredTarget among the transitive
+ * dependencies of the top level test targets may provide baseline coverage artifacts.
+ *
  * <p>The coverage report generation can have two phases, though they both run in the same action.
  * The source code of the coverage report tool {@code lcov_merger} is in the {@code
- * testing/coverage/lcov_merger} directory. The deployed binaries used by Blaze are under
- * {@code tools/coverage}.
- * 
+ * testing/coverage/lcov_merger} directory. The deployed binaries used by Blaze are under {@code
+ * tools/coverage}.
+ *
  * <p>The first phase is merging the individual coverage files into a single report file. The
  * location of this file is reported by Blaze. This phase always happens if the {@code
  * --combined_report=lcov} or {@code --combined_report=html}.
- * 
+ *
  * <p>The second phase is generating an html report. It only happens if {@code
- * --combined_report=html}. The action generates an html output file potentially for every
- * tested source file into the report. Since this set of files is unknown in the analysis
- * phase (the tool figures it out from the contents of the merged coverage report file)
- * the action always runs locally when {@code --combined_report=html}.
+ * --combined_report=html}. The action generates an html output file potentially for every tested
+ * source file into the report. Since this set of files is unknown in the analysis phase (the tool
+ * figures it out from the contents of the merged coverage report file) the action always runs
+ * locally when {@code --combined_report=html}.
  */
 public final class CoverageReportActionBuilder {
 
@@ -148,7 +148,7 @@ public ActionResult execute(ActionExecutionContext actionExecutionContext)
         actionExecutionContext.getEventHandler().handle(Event.info(locationMessage));
         return ActionResult.create(spawnResults);
       } catch (ExecException e) {
-        throw e.toActionExecutionException(this);
+        throw ActionExecutionException.fromExecException(e, this);
       }
     }
 
diff --git a/src/main/java/com/google/devtools/build/lib/exec/SymlinkTreeStrategy.java b/src/main/java/com/google/devtools/build/lib/exec/SymlinkTreeStrategy.java
index de90cc6e754810..fc5dc61676466f 100644
--- a/src/main/java/com/google/devtools/build/lib/exec/SymlinkTreeStrategy.java
+++ b/src/main/java/com/google/devtools/build/lib/exec/SymlinkTreeStrategy.java
@@ -116,8 +116,8 @@ public void createSymlinks(
                 .createSymlinksDirectly(
                     action.getOutputManifest().getPath().getParentDirectory(), runfiles);
           } catch (IOException e) {
-            throw new EnvironmentalExecException(e, Code.SYMLINK_TREE_CREATION_IO_EXCEPTION)
-                .toActionExecutionException(action);
+            throw ActionExecutionException.fromExecException(
+                new EnvironmentalExecException(e, Code.SYMLINK_TREE_CREATION_IO_EXCEPTION), action);
           }
 
           Path inputManifest =
@@ -136,7 +136,7 @@ public void createSymlinks(
                   actionExecutionContext.getFileOutErr());
         }
       } catch (ExecException e) {
-        throw e.toActionExecutionException(action);
+        throw ActionExecutionException.fromExecException(e, action);
       }
     }
   }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileAction.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileAction.java
index ed39beb9d8c965..e713a5eb1e96b6 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileAction.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileAction.java
@@ -478,7 +478,7 @@ private NestedSet<Artifact> findUsedHeaders(
             createFailureDetail("Find used headers failure", Code.FIND_USED_HEADERS_IO_EXCEPTION));
       }
     } catch (ExecException e) {
-      throw e.toActionExecutionException("include scanning", this);
+      throw ActionExecutionException.fromExecException(e, "include scanning", this);
     }
   }
 
@@ -1853,7 +1853,7 @@ public ActionContinuationOrResult execute()
         dotDContents = getDotDContents(spawnResults.get(0));
       } catch (ExecException e) {
         copyTempOutErrToActionOutErr();
-        throw e.toActionExecutionException(CppCompileAction.this);
+        throw ActionExecutionException.fromExecException(e, CppCompileAction.this);
       } catch (InterruptedException e) {
         copyTempOutErrToActionOutErr();
         throw e;
@@ -1931,9 +1931,10 @@ private void copyTempOutErrToActionOutErr() throws ActionExecutionException {
             }
           }
         } catch (IOException e) {
-          throw new EnvironmentalExecException(
-                  e, createFailureDetail("OutErr copy failure", Code.COPY_OUT_ERR_FAILURE))
-              .toActionExecutionException(CppCompileAction.this);
+          throw ActionExecutionException.fromExecException(
+              new EnvironmentalExecException(
+                  e, createFailureDetail("OutErr copy failure", Code.COPY_OUT_ERR_FAILURE)),
+              CppCompileAction.this);
         }
       }
     }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppLinkAction.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppLinkAction.java
index e5a43dd72974bb..d8e92b262b8b6b 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppLinkAction.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppLinkAction.java
@@ -111,7 +111,7 @@ public Artifact create(
       };
 
   private static final String LINK_GUID = "58ec78bd-1176-4e36-8143-439f656b181d";
-  
+
   @Nullable private final String mnemonic;
   private final LibraryToLink outputLibrary;
   private final Artifact linkOutput;
@@ -453,8 +453,7 @@ public ActionContinuationOrResult execute()
         }
         return ActionContinuationOrResult.of(ActionResult.create(nextContinuation.get()));
       } catch (ExecException e) {
-        throw e.toActionExecutionException(
-            CppLinkAction.this);
+        throw ActionExecutionException.fromExecException(e, CppLinkAction.this);
       }
     }
   }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaCompileAction.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaCompileAction.java
index e34e51d01b738e..d8c80a99cd4a9f 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/JavaCompileAction.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/JavaCompileAction.java
@@ -677,9 +677,10 @@ private Deps.Dependencies readFullOutputDeps(
       return createFullOutputDeps(
           Iterables.getOnlyElement(results), outputDepsProto, getInputs(), actionExecutionContext);
     } catch (IOException e) {
-      throw new EnvironmentalExecException(
-              e, createFailureDetail(".jdeps read IOException", Code.JDEPS_READ_IO_EXCEPTION))
-          .toActionExecutionException(this);
+      throw ActionExecutionException.fromExecException(
+          new EnvironmentalExecException(
+              e, createFailureDetail(".jdeps read IOException", Code.JDEPS_READ_IO_EXCEPTION)),
+          this);
     }
   }
 
@@ -755,12 +756,13 @@ public ActionContinuationOrResult execute()
               // We don't create any tree artifacts anyway.
               /*cleanupArchivedArtifacts=*/ false);
         } catch (IOException e) {
-          throw new EnvironmentalExecException(
+          throw ActionExecutionException.fromExecException(
+              new EnvironmentalExecException(
                   e,
                   createFailureDetail(
                       "Failed to delete reduced action outputs",
-                      Code.REDUCED_CLASSPATH_FALLBACK_CLEANUP_FAILURE))
-              .toActionExecutionException(JavaCompileAction.this);
+                      Code.REDUCED_CLASSPATH_FALLBACK_CLEANUP_FAILURE)),
+              JavaCompileAction.this);
         }
         actionExecutionContext.getMetadataHandler().resetOutputs(getOutputs());
         Spawn spawn;
@@ -777,7 +779,7 @@ public ActionContinuationOrResult execute()
         return new JavaFallbackActionContinuation(
             actionExecutionContext, results, fallbackContinuation);
       } catch (ExecException e) {
-        throw e.toActionExecutionException(JavaCompileAction.this);
+        throw ActionExecutionException.fromExecException(e, JavaCompileAction.this);
       }
     }
   }
@@ -821,7 +823,7 @@ public ActionContinuationOrResult execute()
             ActionResult.create(
                 ImmutableList.copyOf(Iterables.concat(primaryResults, fallbackResults))));
       } catch (ExecException e) {
-        throw e.toActionExecutionException(JavaCompileAction.this);
+        throw ActionExecutionException.fromExecException(e, JavaCompileAction.this);
       }
     }
   }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java
index acdb1b0ec53fd5..06edeb33977fbb 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java
@@ -613,7 +613,7 @@ Token checkActionCache(
               remoteDefaultProperties,
               isRemoteCacheEnabled);
     } catch (UserExecException e) {
-      throw e.toActionExecutionException(action);
+      throw ActionExecutionException.fromExecException(e, action);
     }
     if (token == null) {
       boolean eventPosted = false;
@@ -693,7 +693,7 @@ void updateActionCache(
               ? remoteOptions.getRemoteDefaultExecProperties()
               : ImmutableSortedMap.of();
     } catch (UserExecException e) {
-      throw e.toActionExecutionException(action);
+      throw ActionExecutionException.fromExecException(e, action);
     }
 
     try {
diff --git a/src/test/java/com/google/devtools/build/lib/standalone/StandaloneSpawnStrategyTest.java b/src/test/java/com/google/devtools/build/lib/standalone/StandaloneSpawnStrategyTest.java
index a94299c3b43296..2bdfb5a52f978d 100644
--- a/src/test/java/com/google/devtools/build/lib/standalone/StandaloneSpawnStrategyTest.java
+++ b/src/test/java/com/google/devtools/build/lib/standalone/StandaloneSpawnStrategyTest.java
@@ -321,7 +321,7 @@ public void testStandardError() throws Exception {
   public void testVerboseFailures() {
     ExecException e = assertThrows(ExecException.class, () -> run(createSpawn(getFalseCommand())));
     ActionExecutionException actionExecutionException =
-        e.toActionExecutionException(new NullAction());
+        ActionExecutionException.fromExecException(e, new NullAction());
     assertWithMessage("got: " + actionExecutionException.getMessage())
         .that(actionExecutionException.getMessage().contains("failed: error executing command"))
         .isTrue();

From d887fe46e4c455b8bc3852620d7beec68ab8a4a1 Mon Sep 17 00:00:00 2001
From: wilwell <wilwell@google.com>
Date: Thu, 9 Dec 2021 03:49:57 -0800
Subject: [PATCH 14/24] ResourceSet in StarlarkAction API

Added optional `resource_set` parameter to `run` and `run_shell` in StarlarkActionApi. `resource_set` is `StarlarkCallable` object that returns dict with resource set (cpu, memory, local_test).

PiperOrigin-RevId: 415224490
---
 .../google/devtools/build/lib/actions/BUILD   |   8 +-
 .../devtools/build/lib/actions/BaseSpawn.java |  10 +-
 .../build/lib/actions/DelegateSpawn.java      |   2 +-
 .../build/lib/actions/ResourceSet.java        |   4 +-
 .../lib/actions/ResourceSetOrBuilder.java     |   4 +-
 .../devtools/build/lib/actions/Spawn.java     |   6 +-
 .../lib/analysis/actions/SpawnAction.java     |   3 +-
 .../starlark/StarlarkActionFactory.java       | 148 ++++++++++-
 .../semantics/BuildLanguageOptions.java       |  16 ++
 .../devtools/build/lib/rules/proto/BUILD      |   2 +
 .../proto/ProtoCompileActionBuilder.java      |   6 +-
 .../StarlarkActionFactoryApi.java             |  40 ++-
 .../packages/semantics/ConsistencyTest.java   |   2 +
 .../devtools/build/lib/rules/proto/BUILD      |   1 +
 .../proto/ProtoCompileActionBuilderTest.java  |  10 +-
 .../google/devtools/build/lib/starlark/BUILD  |   2 +
 .../lib/starlark/StarlarkRuleContextTest.java | 229 ++++++++++++++++++
 17 files changed, 462 insertions(+), 31 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/actions/BUILD b/src/main/java/com/google/devtools/build/lib/actions/BUILD
index e352f9b3f8f64e..5f2f6a06bc180e 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/actions/BUILD
@@ -297,8 +297,7 @@ java_library(
         "ResourceSetOrBuilder.java",
     ],
     deps = [
-        ":artifacts",
-        "//src/main/java/com/google/devtools/build/lib/collect/nestedset",
+        ":exec_exception",
         "//src/main/java/com/google/devtools/build/lib/concurrent",
         "//src/main/java/com/google/devtools/build/lib/jni",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec",
@@ -347,6 +346,9 @@ java_library(
 
 java_library(
     name = "exec_exception",
-    srcs = ["ExecException.java"],
+    srcs = [
+        "ExecException.java",
+        "UserExecException.java",
+    ],
     deps = ["//src/main/protobuf:failure_details_java_proto"],
 )
diff --git a/src/main/java/com/google/devtools/build/lib/actions/BaseSpawn.java b/src/main/java/com/google/devtools/build/lib/actions/BaseSpawn.java
index 36ab4aded7ad8b..89d55f1726c912 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/BaseSpawn.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/BaseSpawn.java
@@ -20,6 +20,7 @@
 import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.analysis.platform.PlatformInfo;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
+import com.google.devtools.build.lib.util.OS;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.util.List;
 import java.util.Map;
@@ -35,7 +36,7 @@ public class BaseSpawn implements Spawn {
   private final ImmutableMap<String, String> executionInfo;
   private final RunfilesSupplier runfilesSupplier;
   private final ActionExecutionMetadata action;
-  private final ResourceSet localResources;
+  private final ResourceSetOrBuilder localResources;
 
   public BaseSpawn(
       List<String> arguments,
@@ -43,7 +44,7 @@ public BaseSpawn(
       Map<String, String> executionInfo,
       RunfilesSupplier runfilesSupplier,
       ActionExecutionMetadata action,
-      ResourceSet localResources) {
+      ResourceSetOrBuilder localResources) {
     this.arguments = ImmutableList.copyOf(arguments);
     this.environment = ImmutableMap.copyOf(environment);
     this.executionInfo = ImmutableMap.copyOf(executionInfo);
@@ -123,8 +124,9 @@ public ActionExecutionMetadata getResourceOwner() {
   }
 
   @Override
-  public ResourceSet getLocalResources() {
-    return localResources;
+  public ResourceSet getLocalResources() throws ExecException {
+    return localResources.buildResourceSet(
+        OS.getCurrent(), action.getInputs().memoizedFlattenAndGetSize());
   }
 
   @Override
diff --git a/src/main/java/com/google/devtools/build/lib/actions/DelegateSpawn.java b/src/main/java/com/google/devtools/build/lib/actions/DelegateSpawn.java
index 112452344d3d26..1cc002e666bd78 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/DelegateSpawn.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/DelegateSpawn.java
@@ -84,7 +84,7 @@ public ActionExecutionMetadata getResourceOwner() {
   }
 
   @Override
-  public ResourceSet getLocalResources() {
+  public ResourceSet getLocalResources() throws ExecException {
     return spawn.getLocalResources();
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/actions/ResourceSet.java b/src/main/java/com/google/devtools/build/lib/actions/ResourceSet.java
index b62f74d416907a..9fe4870b8de063 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/ResourceSet.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/ResourceSet.java
@@ -16,8 +16,8 @@
 
 import com.google.common.base.Splitter;
 import com.google.common.primitives.Doubles;
-import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
+import com.google.devtools.build.lib.util.OS;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.common.options.Converter;
 import com.google.devtools.common.options.OptionsParsingException;
@@ -175,7 +175,7 @@ public String getTypeDescription() {
   }
 
   @Override
-  public ResourceSet buildResourceSet(NestedSet<Artifact> inputs) {
+  public ResourceSet buildResourceSet(OS os, int inputsSize) throws ExecException {
     return this;
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/actions/ResourceSetOrBuilder.java b/src/main/java/com/google/devtools/build/lib/actions/ResourceSetOrBuilder.java
index fee458896a1e6d..71242c8d8ee7fc 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/ResourceSetOrBuilder.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/ResourceSetOrBuilder.java
@@ -14,7 +14,7 @@
 
 package com.google.devtools.build.lib.actions;
 
-import com.google.devtools.build.lib.collect.nestedset.NestedSet;
+import com.google.devtools.build.lib.util.OS;
 
 /** Common interface for ResourceSet and builder. */
 @FunctionalInterface
@@ -24,5 +24,5 @@ public interface ResourceSetOrBuilder {
    * will flatten NestedSet. This action could create a lot of garbagge, so use it as close as
    * possible to the execution phase,
    */
-  public ResourceSet buildResourceSet(NestedSet<Artifact> inputs);
+  public ResourceSet buildResourceSet(OS os, int inputsSize) throws ExecException;
 }
diff --git a/src/main/java/com/google/devtools/build/lib/actions/Spawn.java b/src/main/java/com/google/devtools/build/lib/actions/Spawn.java
index cf47e4203f63ff..3b55e26e1a34a0 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/Spawn.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/Spawn.java
@@ -130,10 +130,8 @@ default boolean isMandatoryOutput(ActionInput output) {
   /** Returns the resource owner for local fallback. */
   ActionExecutionMetadata getResourceOwner();
 
-  /**
-   * Returns the amount of resources needed for local fallback.
-   */
-  ResourceSet getLocalResources();
+  /** Returns the amount of resources needed for local fallback. */
+  ResourceSet getLocalResources() throws ExecException;
 
   /**
    * Returns a mnemonic (string constant) for this kind of spawn.
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java b/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java
index 094e7e73c8235e..402b418a994266 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java
@@ -590,8 +590,7 @@ private ActionSpawn(
           executionInfo,
           SpawnAction.this.getRunfilesSupplier(),
           SpawnAction.this,
-          SpawnAction.this.resourceSetOrBuilder.buildResourceSet(inputs));
-
+          SpawnAction.this.resourceSetOrBuilder);
       NestedSetBuilder<ActionInput> inputsBuilder = NestedSetBuilder.stableOrder();
       ImmutableList<Artifact> manifests = getRunfilesSupplier().getManifests();
       for (Artifact input : inputs.toList()) {
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkActionFactory.java b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkActionFactory.java
index 5fb62d507dddde..de77b7e696a5d5 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkActionFactory.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkActionFactory.java
@@ -15,8 +15,10 @@
 
 import static com.google.devtools.build.lib.packages.semantics.BuildLanguageOptions.EXPERIMENTAL_SIBLING_REPOSITORY_LAYOUT;
 
+import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Sets;
 import com.google.devtools.build.lib.actions.Action;
 import com.google.devtools.build.lib.actions.ActionAnalysisMetadata;
 import com.google.devtools.build.lib.actions.ActionLookupKey;
@@ -24,8 +26,12 @@
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.ArtifactRoot;
 import com.google.devtools.build.lib.actions.CommandLine;
+import com.google.devtools.build.lib.actions.ExecException;
 import com.google.devtools.build.lib.actions.ParamFileInfo;
+import com.google.devtools.build.lib.actions.ResourceSet;
+import com.google.devtools.build.lib.actions.ResourceSetOrBuilder;
 import com.google.devtools.build.lib.actions.RunfilesSupplier;
+import com.google.devtools.build.lib.actions.UserExecException;
 import com.google.devtools.build.lib.actions.extra.ExtraActionInfo;
 import com.google.devtools.build.lib.actions.extra.SpawnInfo;
 import com.google.devtools.build.lib.analysis.BashCommandConstructor;
@@ -49,23 +55,35 @@
 import com.google.devtools.build.lib.collect.nestedset.Order;
 import com.google.devtools.build.lib.packages.TargetUtils;
 import com.google.devtools.build.lib.packages.semantics.BuildLanguageOptions;
+import com.google.devtools.build.lib.server.FailureDetails;
+import com.google.devtools.build.lib.server.FailureDetails.FailureDetail;
+import com.google.devtools.build.lib.server.FailureDetails.Interrupted;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.SerializationConstant;
 import com.google.devtools.build.lib.starlarkbuildapi.FileApi;
 import com.google.devtools.build.lib.starlarkbuildapi.StarlarkActionFactoryApi;
+import com.google.devtools.build.lib.util.OS;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.protobuf.GeneratedMessage;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.UUID;
 import net.starlark.java.eval.Dict;
 import net.starlark.java.eval.EvalException;
+import net.starlark.java.eval.Mutability;
 import net.starlark.java.eval.Printer;
 import net.starlark.java.eval.Sequence;
 import net.starlark.java.eval.Starlark;
+import net.starlark.java.eval.StarlarkCallable;
+import net.starlark.java.eval.StarlarkFloat;
+import net.starlark.java.eval.StarlarkFunction;
+import net.starlark.java.eval.StarlarkInt;
 import net.starlark.java.eval.StarlarkSemantics;
 import net.starlark.java.eval.StarlarkThread;
 
@@ -75,6 +93,10 @@ public class StarlarkActionFactory implements StarlarkActionFactoryApi {
   /** Counter for actions.run_shell helper scripts. Every script must have a unique name. */
   private int runShellOutputCounter = 0;
 
+  private static final ResourceSet DEFAULT_RESOURCE_SET = ResourceSet.createWithRamCpu(250, 1);
+  private static final Set<String> validResources =
+      new HashSet<>(Arrays.asList("cpu", "memory", "local_test"));
+
   public StarlarkActionFactory(StarlarkRuleContext context) {
     this.context = context;
   }
@@ -339,7 +361,8 @@ public void run(
       Object executionRequirementsUnchecked,
       Object inputManifestsUnchecked,
       Object execGroupUnchecked,
-      Object shadowedActionUnchecked)
+      Object shadowedActionUnchecked,
+      Object resourceSetUnchecked)
       throws EvalException {
     context.checkMutable("actions.run");
 
@@ -377,6 +400,7 @@ public void run(
         inputManifestsUnchecked,
         execGroupUnchecked,
         shadowedActionUnchecked,
+        resourceSetUnchecked,
         builder);
   }
 
@@ -430,7 +454,8 @@ public void runShell(
       Object executionRequirementsUnchecked,
       Object inputManifestsUnchecked,
       Object execGroupUnchecked,
-      Object shadowedActionUnchecked)
+      Object shadowedActionUnchecked,
+      Object resourceSetUnchecked)
       throws EvalException {
     context.checkMutable("actions.run_shell");
     RuleContext ruleContext = getRuleContext();
@@ -497,6 +522,7 @@ public void runShell(
         inputManifestsUnchecked,
         execGroupUnchecked,
         shadowedActionUnchecked,
+        resourceSetUnchecked,
         builder);
   }
 
@@ -543,6 +569,7 @@ private void registerStarlarkAction(
       Object inputManifestsUnchecked,
       Object execGroupUnchecked,
       Object shadowedActionUnchecked,
+      Object resourceSetUnchecked,
       StarlarkAction.Builder builder)
       throws EvalException {
     if (inputs instanceof Sequence) {
@@ -648,10 +675,127 @@ private void registerStarlarkAction(
       builder.setShadowedAction(Optional.of((Action) shadowedActionUnchecked));
     }
 
+    if (getSemantics().getBool(BuildLanguageOptions.EXPERIMENTAL_ACTION_RESOURCE_SET)
+        && resourceSetUnchecked != Starlark.NONE) {
+      validateResourceSetBuilder(resourceSetUnchecked);
+      builder.setResources(
+          new StarlarkActionResourceSetBuilder(
+              (StarlarkFunction) resourceSetUnchecked, mnemonic, getSemantics()));
+    }
+
     // Always register the action
     registerAction(builder.build(ruleContext));
   }
 
+  private static class StarlarkActionResourceSetBuilder implements ResourceSetOrBuilder {
+    private final StarlarkCallable fn;
+    private final String mnemonic;
+    private final StarlarkSemantics semantics;
+
+    public StarlarkActionResourceSetBuilder(
+        StarlarkCallable fn, String mnemonic, StarlarkSemantics semantics) {
+      this.fn = fn;
+      this.mnemonic = mnemonic;
+      this.semantics = semantics;
+    }
+
+    @Override
+    public ResourceSet buildResourceSet(OS os, int inputsSize) throws ExecException {
+      try (Mutability mu = Mutability.create("resource_set_builder_function")) {
+        StarlarkThread thread = new StarlarkThread(mu, semantics);
+        StarlarkInt inputInt = StarlarkInt.of(inputsSize);
+        Object response =
+            Starlark.call(
+                thread,
+                this.fn,
+                ImmutableList.of(os.getCanonicalName(), inputInt),
+                ImmutableMap.of());
+        Map<String, Object> resourceSetMapRaw =
+            Dict.cast(response, String.class, Object.class, "resource_set");
+
+        if (!validResources.containsAll(resourceSetMapRaw.keySet())) {
+          String message =
+              String.format(
+                  "Illegal resource keys: (%s)",
+                  Joiner.on(",").join(Sets.difference(resourceSetMapRaw.keySet(), validResources)));
+          throw new EvalException(message);
+        }
+
+        return ResourceSet.create(
+            getNumericOrDefault(resourceSetMapRaw, "memory", DEFAULT_RESOURCE_SET.getMemoryMb()),
+            getNumericOrDefault(resourceSetMapRaw, "cpu", DEFAULT_RESOURCE_SET.getCpuUsage()),
+            (int)
+                getNumericOrDefault(
+                    resourceSetMapRaw,
+                    "local_test",
+                    (double) DEFAULT_RESOURCE_SET.getLocalTestCount()));
+      } catch (EvalException e) {
+        throw new UserExecException(
+            FailureDetail.newBuilder()
+                .setMessage(
+                    String.format("Could not build resources for %s. %s", mnemonic, e.getMessage()))
+                .setStarlarkAction(
+                    FailureDetails.StarlarkAction.newBuilder()
+                        .setCode(FailureDetails.StarlarkAction.Code.STARLARK_ACTION_UNKNOWN)
+                        .build())
+                .build());
+      } catch (InterruptedException e) {
+        throw new UserExecException(
+            FailureDetail.newBuilder()
+                .setMessage(e.getMessage())
+                .setInterrupted(
+                    Interrupted.newBuilder().setCode(Interrupted.Code.INTERRUPTED).build())
+                .build());
+      }
+    }
+
+    private static double getNumericOrDefault(
+        Map<String, Object> resourceSetMap, String key, double defaultValue) throws EvalException {
+      if (!resourceSetMap.containsKey(key)) {
+        return defaultValue;
+      }
+
+      Object value = resourceSetMap.get(key);
+      if (value instanceof StarlarkInt) {
+        return ((StarlarkInt) value).toDouble();
+      }
+
+      if (value instanceof StarlarkFloat) {
+        return ((StarlarkFloat) value).toDouble();
+      }
+      throw new EvalException(
+          String.format(
+              "Illegal resource value type for key %s: got %s, want int or float",
+              key, Starlark.type(value)));
+    }
+  }
+
+  private static StarlarkFunction validateResourceSetBuilder(Object fn) throws EvalException {
+    if (!(fn instanceof StarlarkFunction)) {
+      throw Starlark.errorf(
+          "resource_set should be a Starlark-defined function, but got %s instead",
+          Starlark.type(fn));
+    }
+
+    StarlarkFunction sfn = (StarlarkFunction) fn;
+
+    // Reject non-global functions, because arbitrary closures may cause large
+    // analysis-phase data structures to remain live into the execution phase.
+    // We require that the function is "global" as opposed to "not a closure"
+    // because a global function may be closure if it refers to load bindings.
+    // This unfortunately disallows such trivially safe non-global
+    // functions as "lambda x: x".
+    // See https://github.com/bazelbuild/bazel/issues/12701.
+    if (sfn.getModule().getGlobal(sfn.getName()) != sfn) {
+      throw Starlark.errorf(
+          "to avoid unintended retention of analysis data structures, "
+              + "the resource_set function (declared at %s) must be declared "
+              + "by a top-level def statement",
+          sfn.getLocation());
+    }
+    return (StarlarkFunction) fn;
+  }
+
   private String getMnemonic(Object mnemonicUnchecked) {
     String mnemonic = mnemonicUnchecked == Starlark.NONE ? "Action" : (String) mnemonicUnchecked;
     if (getRuleContext().getConfiguration().getReservedActionMnemonics().contains(mnemonic)) {
diff --git a/src/main/java/com/google/devtools/build/lib/packages/semantics/BuildLanguageOptions.java b/src/main/java/com/google/devtools/build/lib/packages/semantics/BuildLanguageOptions.java
index 8373d4fe87fb20..dba3e3827b4fdc 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/semantics/BuildLanguageOptions.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/semantics/BuildLanguageOptions.java
@@ -283,6 +283,20 @@ public class BuildLanguageOptions extends OptionsBase implements Serializable {
               + " https://github.com/bazelbuild/bazel/issues/8830 for details.")
   public boolean experimentalAllowTagsPropagation;
 
+  @Option(
+      name = "experimental_action_resource_set",
+      defaultValue = "true",
+      documentationCategory = OptionDocumentationCategory.STARLARK_SEMANTICS,
+      effectTags = {OptionEffectTag.EXECUTION, OptionEffectTag.BUILD_FILE_SEMANTICS},
+      metadataTags = {
+        OptionMetadataTag.EXPERIMENTAL,
+      },
+      help =
+          "If set to true, ctx.actions.run() and ctx.actions.run_shell() accept a resource_set"
+              + " parameter for local execution. Otherwise it will default to 250 MB for memory"
+              + " and 1 cpu.")
+  public boolean experimentalActionResourceSet;
+
   @Option(
       name = "incompatible_struct_has_no_methods",
       defaultValue = "false",
@@ -571,6 +585,7 @@ public StarlarkSemantics toStarlarkSemantics() {
             .setBool(INCOMPATIBLE_ENABLE_EXPORTS_PROVIDER, incompatibleEnableExportsProvider)
             .setBool(
                 INCOMPATIBLE_EXISTING_RULES_IMMUTABLE_VIEW, incompatibleExistingRulesImmutableView)
+            .setBool(EXPERIMENTAL_ACTION_RESOURCE_SET, experimentalActionResourceSet)
             .setBool(EXPERIMENTAL_GOOGLE_LEGACY_API, experimentalGoogleLegacyApi)
             .setBool(EXPERIMENTAL_NINJA_ACTIONS, experimentalNinjaActions)
             .setBool(EXPERIMENTAL_PLATFORMS_API, experimentalPlatformsApi)
@@ -646,6 +661,7 @@ public StarlarkSemantics toStarlarkSemantics() {
   public static final String EXPERIMENTAL_REPO_REMOTE_EXEC = "-experimental_repo_remote_exec";
   public static final String EXPERIMENTAL_SIBLING_REPOSITORY_LAYOUT =
       "-experimental_sibling_repository_layout";
+  public static final String EXPERIMENTAL_ACTION_RESOURCE_SET = "+experimental_action_resource_set";
   public static final String INCOMPATIBLE_ALLOW_TAGS_PROPAGATION =
       "-incompatible_allow_tags_propagation";
   public static final String INCOMPATIBLE_ALWAYS_CHECK_DEPSET_ELEMENTS =
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/BUILD b/src/main/java/com/google/devtools/build/lib/rules/proto/BUILD
index b0f29c42a60f55..598d632f1844a0 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/BUILD
@@ -45,10 +45,12 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/concurrent",
         "//src/main/java/com/google/devtools/build/lib/packages",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec",
+        "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec:serialization-constant-annotation",
         "//src/main/java/com/google/devtools/build/lib/starlarkbuildapi",
         "//src/main/java/com/google/devtools/build/lib/starlarkbuildapi/proto",
         "//src/main/java/com/google/devtools/build/lib/util",
         "//src/main/java/com/google/devtools/build/lib/util:filetype",
+        "//src/main/java/com/google/devtools/build/lib/util:os",
         "//src/main/java/com/google/devtools/build/lib/vfs",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//src/main/java/com/google/devtools/common/options",
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 6c2e8d1c45bee0..4bf12845baf1b7 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
@@ -41,6 +41,8 @@
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
+import com.google.devtools.build.lib.skyframe.serialization.autocodec.SerializationConstant;
+import com.google.devtools.build.lib.util.OS;
 import com.google.devtools.build.lib.util.OnDemandString;
 import java.util.HashSet;
 import java.util.List;
@@ -189,9 +191,9 @@ public String lookupFunction(String name, String param)
   /** Builds a ResourceSet based on the number of inputs. */
   public static class ProtoCompileResourceSetBuilder implements ResourceSetOrBuilder {
     @Override
-    public ResourceSet buildResourceSet(NestedSet<Artifact> inputs) {
+    public ResourceSet buildResourceSet(OS os, int inputsSize) {
       return ResourceSet.createWithRamCpu(
-          /* memoryMb= */ 25 + 0.15 * inputs.memoizedFlattenAndGetSize(), /* cpuUsage= */ 1);
+          /* memoryMb= */ 25 + 0.15 * inputsSize, /* cpuUsage= */ 1);
     }
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/StarlarkActionFactoryApi.java b/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/StarlarkActionFactoryApi.java
index f2205cca4f47e0..27fee5e251c79a 100644
--- a/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/StarlarkActionFactoryApi.java
+++ b/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/StarlarkActionFactoryApi.java
@@ -24,6 +24,7 @@
 import net.starlark.java.eval.EvalException;
 import net.starlark.java.eval.NoneType;
 import net.starlark.java.eval.Sequence;
+import net.starlark.java.eval.StarlarkCallable;
 import net.starlark.java.eval.StarlarkThread;
 import net.starlark.java.eval.StarlarkValue;
 
@@ -435,6 +436,27 @@ void symlink(
                     + " environment added to the action's inputs list and environment. The action"
                     + " environment can overwrite any of the shadowed action's environment"
                     + " variables. If none, uses only the action's inputs and given environment."),
+        @Param(
+            name = "resource_set",
+            allowedTypes = {
+              @ParamType(type = StarlarkCallable.class),
+              @ParamType(type = NoneType.class),
+            },
+            defaultValue = "None",
+            named = true,
+            positional = false,
+            doc =
+                "A callback function that returns a resource set dictionary, used to estimate"
+                    + " resource usage at execution time if this action is run locally.<p>The"
+                    + " function accepts two positional arguments: a string representing an OS name"
+                    + " (e.g. \"osx\"), and an integer representing the number of inputs to the"
+                    + " action. The returned dictionary may contain the following entries, each of"
+                    + " which may be a float or an int:<ul><li>\"cpu\": number of CPUs; default"
+                    + " 1<li>\"memory\": in MB; default 250<li>\"local_test\": number of local"
+                    + " tests; default 1</ul><p>If this parameter is set to <code>None</code> or if"
+                    + " <code>--experimental_action_resource_set</code> is false, the default"
+                    + " values are used.<p>The callback must be top-level (lambda and nested"
+                    + " functions aren't allowed)."),
       })
   void run(
       Sequence<?> outputs,
@@ -450,7 +472,8 @@ void run(
       Object executionRequirementsUnchecked,
       Object inputManifestsUnchecked,
       Object execGroupUnchecked,
-      Object shadowedAction)
+      Object shadowedAction,
+      Object resourceSetUnchecked)
       throws EvalException;
 
   @StarlarkMethod(
@@ -639,6 +662,18 @@ void run(
                 "Runs the action using the given shadowed action's discovered inputs"
                     + " added to the action's inputs list. If none, uses only the action's"
                     + " inputs."),
+        @Param(
+            name = "resource_set",
+            allowedTypes = {
+              @ParamType(type = StarlarkCallable.class),
+              @ParamType(type = NoneType.class),
+            },
+            defaultValue = "None",
+            named = true,
+            positional = false,
+            doc =
+                "A callback function for estimating resource usage if run locally. See"
+                    + "<a href=\"#run.resource_set\"><code>ctx.actions.run()</code></a>."),
       })
   void runShell(
       Sequence<?> outputs,
@@ -653,7 +688,8 @@ void runShell(
       Object executionRequirementsUnchecked,
       Object inputManifestsUnchecked,
       Object execGroupUnchecked,
-      Object shadowedAction)
+      Object shadowedAction,
+      Object resourceSetUnchecked)
       throws EvalException;
 
   @StarlarkMethod(
diff --git a/src/test/java/com/google/devtools/build/lib/packages/semantics/ConsistencyTest.java b/src/test/java/com/google/devtools/build/lib/packages/semantics/ConsistencyTest.java
index 110667fe880406..10e62bf50d5994 100644
--- a/src/test/java/com/google/devtools/build/lib/packages/semantics/ConsistencyTest.java
+++ b/src/test/java/com/google/devtools/build/lib/packages/semantics/ConsistencyTest.java
@@ -130,6 +130,7 @@ private static BuildLanguageOptions buildRandomOptions(Random rand) throws Excep
         "--incompatible_allow_tags_propagation=" + rand.nextBoolean(), // flag, Java names differ
         "--experimental_cc_shared_library=" + rand.nextBoolean(),
         "--experimental_repo_remote_exec=" + rand.nextBoolean(),
+        "--experimental_action_resource_set=" + rand.nextBoolean(),
         "--incompatible_always_check_depset_elements=" + rand.nextBoolean(),
         "--incompatible_depset_for_libraries_to_link_getter=" + rand.nextBoolean(),
         "--incompatible_disable_target_provider_fields=" + rand.nextBoolean(),
@@ -171,6 +172,7 @@ private static StarlarkSemantics buildRandomSemantics(Random rand) {
         .setBool(BuildLanguageOptions.EXPERIMENTAL_ALLOW_TAGS_PROPAGATION, rand.nextBoolean())
         .setBool(BuildLanguageOptions.EXPERIMENTAL_CC_SHARED_LIBRARY, rand.nextBoolean())
         .setBool(BuildLanguageOptions.EXPERIMENTAL_REPO_REMOTE_EXEC, rand.nextBoolean())
+        .setBool(BuildLanguageOptions.EXPERIMENTAL_ACTION_RESOURCE_SET, rand.nextBoolean())
         .setBool(BuildLanguageOptions.INCOMPATIBLE_ALWAYS_CHECK_DEPSET_ELEMENTS, rand.nextBoolean())
         .setBool(
             BuildLanguageOptions.INCOMPATIBLE_DEPSET_FOR_LIBRARIES_TO_LINK_GETTER,
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD b/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD
index 74b050dc619e1c..ac97630531a214 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD
@@ -25,6 +25,7 @@ java_test(
         "//src/main/java/com/google/devtools/build/lib/collect/nestedset",
         "//src/main/java/com/google/devtools/build/lib/rules/proto",
         "//src/main/java/com/google/devtools/build/lib/util",
+        "//src/main/java/com/google/devtools/build/lib/util:os",
         "//src/main/java/com/google/devtools/build/lib/vfs",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//src/main/java/com/google/devtools/build/lib/vfs/inmemoryfs",
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 a887b52bb60b63..f9b0168c8cdc3f 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
@@ -39,6 +39,7 @@
 import com.google.devtools.build.lib.rules.proto.ProtoCompileActionBuilder.Exports;
 import com.google.devtools.build.lib.rules.proto.ProtoCompileActionBuilder.Services;
 import com.google.devtools.build.lib.rules.proto.ProtoCompileActionBuilder.ToolchainInvocation;
+import com.google.devtools.build.lib.util.OS;
 import com.google.devtools.build.lib.util.OnDemandString;
 import com.google.devtools.build.lib.vfs.DigestHashFunction;
 import com.google.devtools.build.lib.vfs.PathFragment;
@@ -493,17 +494,12 @@ public void testEstimateResourceConsumptionLocal() throws Exception {
 
     assertThat(
             new ProtoCompileActionBuilder.ProtoCompileResourceSetBuilder()
-                .buildResourceSet(NestedSetBuilder.emptySet(STABLE_ORDER)))
+                .buildResourceSet(OS.DARWIN, 0))
         .isEqualTo(ResourceSet.createWithRamCpu(25, 1));
 
     assertThat(
             new ProtoCompileActionBuilder.ProtoCompileResourceSetBuilder()
-                .buildResourceSet(
-                    NestedSetBuilder.wrap(
-                        STABLE_ORDER,
-                        ImmutableList.of(
-                            artifact("//:dont-care", "protoc-gen-javalite.exe"),
-                            artifact("//:dont-care-2", "protoc-gen-javalite-2.exe")))))
+                .buildResourceSet(OS.LINUX, 2))
         .isEqualTo(ResourceSet.createWithRamCpu(25.3, 1));
   }
 }
diff --git a/src/test/java/com/google/devtools/build/lib/starlark/BUILD b/src/test/java/com/google/devtools/build/lib/starlark/BUILD
index 0f8a78aecbf428..f86649cd7c21ff 100644
--- a/src/test/java/com/google/devtools/build/lib/starlark/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/starlark/BUILD
@@ -29,6 +29,7 @@ java_test(
         "//src/main/java/com/google/devtools/build/lib/actions:action_lookup_key",
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
         "//src/main/java/com/google/devtools/build/lib/actions:commandline_item",
+        "//src/main/java/com/google/devtools/build/lib/actions:localhost_capacity",
         "//src/main/java/com/google/devtools/build/lib/actions:thread_state_receiver",
         "//src/main/java/com/google/devtools/build/lib/analysis:actions/parameter_file_write_action",
         "//src/main/java/com/google/devtools/build/lib/analysis:actions/substitution",
@@ -67,6 +68,7 @@ java_test(
         "//src/main/java/com/google/devtools/build/lib/util",
         "//src/main/java/com/google/devtools/build/lib/util:classpath",
         "//src/main/java/com/google/devtools/build/lib/util:filetype",
+        "//src/main/java/com/google/devtools/build/lib/util:os",
         "//src/main/java/com/google/devtools/build/lib/vfs",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//src/main/java/com/google/devtools/common/options",
diff --git a/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleContextTest.java b/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleContextTest.java
index dda0eb9c83ccfc..94964936fa1ae1 100644
--- a/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleContextTest.java
+++ b/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleContextTest.java
@@ -29,6 +29,8 @@
 import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.actions.ActionAnalysisMetadata;
 import com.google.devtools.build.lib.actions.Artifact;
+import com.google.devtools.build.lib.actions.ExecException;
+import com.google.devtools.build.lib.actions.ResourceSet;
 import com.google.devtools.build.lib.analysis.ActionsProvider;
 import com.google.devtools.build.lib.analysis.ConfiguredRuleClassProvider;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
@@ -54,6 +56,7 @@
 import com.google.devtools.build.lib.starlark.util.BazelEvaluationTestCase;
 import com.google.devtools.build.lib.testutil.TestRuleClassProvider;
 import com.google.devtools.build.lib.util.FileTypeSet;
+import com.google.devtools.build.lib.util.OS;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.util.ArrayList;
@@ -652,6 +655,232 @@ public void testCreateStarlarkActionArgumentsWithUnusedInputsList() throws Excep
     assertThat(action.isShareable()).isFalse();
   }
 
+  @Test
+  public void testCreateStarlarkActionArgumentsWithResourceSet_success() throws Exception {
+    setBuildLanguageOptions("--experimental_action_resource_set");
+    StarlarkRuleContext ruleContext = createRuleContext("//foo:foo");
+    setRuleContext(ruleContext);
+
+    ev.exec(
+        "def get_resources(os, inputs_size):",
+        "  if os == \"osx\":",
+        "    return {\"cpu\": 2., \"memory\": 350. + inputs_size * 20, \"local_test\": 2.}",
+        "  return {\"cpu\": 1., \"memory\": 350. + inputs_size * 10, \"local_test\": 0.}",
+        "ruleContext.actions.run(",
+        "  inputs = ruleContext.files.srcs,",
+        "  outputs = ruleContext.files.srcs,",
+        "  resource_set = get_resources,",
+        "  executable = 'executable')");
+    StarlarkAction action =
+        (StarlarkAction)
+            Iterables.getOnlyElement(
+                ruleContext.getRuleContext().getAnalysisEnvironment().getRegisteredActions());
+
+    assertThat(action.getResourceSetOrBuilder().buildResourceSet(OS.LINUX, 2))
+        .isEqualTo(ResourceSet.create(370, 1, 0));
+    assertThat(action.getResourceSetOrBuilder().buildResourceSet(OS.DARWIN, 2))
+        .isEqualTo(ResourceSet.create(390, 2, 2));
+  }
+
+  @Test
+  public void testCreateStarlarkActionArgumentsWithResourceSet_flagDisabled() throws Exception {
+    setBuildLanguageOptions("--noexperimental_action_resource_set");
+    StarlarkRuleContext ruleContext = createRuleContext("//foo:foo");
+    setRuleContext(ruleContext);
+
+    ev.exec(
+        "def get_resources(os, inputs_size):",
+        "  if os == \"osx\":",
+        "    return {\"cpu\": 2., \"memory\": 350. + inputs_size * 20, \"local_test\": 2.}",
+        "  return {\"cpu\": 1., \"memory\": 350. + inputs_size * 10, \"local_test\": 0.}",
+        "ruleContext.actions.run(",
+        "  inputs = ruleContext.files.srcs,",
+        "  outputs = ruleContext.files.srcs,",
+        "  resource_set = get_resources,",
+        "  executable = 'executable')");
+    StarlarkAction action =
+        (StarlarkAction)
+            Iterables.getOnlyElement(
+                ruleContext.getRuleContext().getAnalysisEnvironment().getRegisteredActions());
+
+    assertThat(action.getResourceSetOrBuilder().buildResourceSet(OS.LINUX, 2))
+        .isEqualTo(ResourceSet.create(250, 1, 0));
+  }
+
+  @Test
+  public void testCreateStarlarkActionArgumentsWithResourceSet_lambdaForbidden() throws Exception {
+    setBuildLanguageOptions("--experimental_action_resource_set");
+    StarlarkRuleContext ruleContext = createRuleContext("//foo:foo");
+    setRuleContext(ruleContext);
+
+    Exception thrown =
+        assertThrows(
+            EvalException.class,
+            () ->
+                ev.exec(
+                    "ruleContext.actions.run(",
+                    "  inputs = ruleContext.files.srcs,",
+                    "  outputs = ruleContext.files.srcs,",
+                    "  resource_set = lambda os, inputs_size : {\"cpu\": 1., \"memory\": 1.,"
+                        + " \"local_test\": 1.} ,",
+                    "  executable = 'executable')"));
+
+    assertThat(thrown).hasMessageThat().contains("must be declared by a top-level def statement");
+  }
+
+  @Test
+  public void testCreateStarlarkActionArgumentsWithResourceSet_illegalResource() throws Exception {
+    setBuildLanguageOptions("--experimental_action_resource_set");
+    StarlarkRuleContext ruleContext = createRuleContext("//foo:foo");
+    setRuleContext(ruleContext);
+
+    ev.exec(
+        "def get_resources(os, inputs_size):",
+        "  return {\"cpu\": 2., \"memory\": 350., \"local_test\": 2., \"gpu\": 1.}",
+        "ruleContext.actions.run(",
+        "  inputs = ruleContext.files.srcs,",
+        "  outputs = ruleContext.files.srcs,",
+        "  resource_set = get_resources,",
+        "  executable = 'executable')");
+    StarlarkAction action =
+        (StarlarkAction)
+            Iterables.getOnlyElement(
+                ruleContext.getRuleContext().getAnalysisEnvironment().getRegisteredActions());
+
+    Exception thrown =
+        assertThrows(
+            ExecException.class,
+            () -> action.getResourceSetOrBuilder().buildResourceSet(OS.LINUX, 2));
+    assertThat(thrown).hasMessageThat().contains("Illegal resource keys: (gpu)");
+  }
+
+  @Test
+  public void testCreateStarlarkActionArgumentsWithResourceSet_defaultValue() throws Exception {
+    setBuildLanguageOptions("--experimental_action_resource_set");
+    StarlarkRuleContext ruleContext = createRuleContext("//foo:foo");
+    setRuleContext(ruleContext);
+
+    ev.exec(
+        "def get_resources(os, inputs_size):",
+        "  return {\"cpu\": 2., \"local_test\": 2.}",
+        "ruleContext.actions.run(",
+        "  inputs = ruleContext.files.srcs,",
+        "  outputs = ruleContext.files.srcs,",
+        "  resource_set = get_resources,",
+        "  executable = 'executable')");
+    StarlarkAction action =
+        (StarlarkAction)
+            Iterables.getOnlyElement(
+                ruleContext.getRuleContext().getAnalysisEnvironment().getRegisteredActions());
+
+    assertThat(action.getResourceSetOrBuilder().buildResourceSet(OS.LINUX, 2))
+        .isEqualTo(ResourceSet.create(250, 2, 2));
+  }
+
+  @Test
+  public void testCreateStarlarkActionArgumentsWithResourceSet_intDict() throws Exception {
+    setBuildLanguageOptions("--experimental_action_resource_set");
+    StarlarkRuleContext ruleContext = createRuleContext("//foo:foo");
+    setRuleContext(ruleContext);
+
+    ev.exec(
+        "def get_resources(os, inputs_size):",
+        "  return {\"cpu\": 1, \"memory\": 2, \"local_test\": 3}",
+        "ruleContext.actions.run(",
+        "  inputs = ruleContext.files.srcs,",
+        "  outputs = ruleContext.files.srcs,",
+        "  resource_set = get_resources,",
+        "  executable = 'executable')");
+    StarlarkAction action =
+        (StarlarkAction)
+            Iterables.getOnlyElement(
+                ruleContext.getRuleContext().getAnalysisEnvironment().getRegisteredActions());
+
+    assertThat(action.getResourceSetOrBuilder().buildResourceSet(OS.LINUX, 0))
+        .isEqualTo(ResourceSet.create(2, 1, 3));
+  }
+
+  @Test
+  public void testCreateStarlarkActionArgumentsWithResourceSet_notDict() throws Exception {
+    setBuildLanguageOptions("--experimental_action_resource_set");
+    StarlarkRuleContext ruleContext = createRuleContext("//foo:foo");
+    setRuleContext(ruleContext);
+
+    ev.exec(
+        "def get_resources(os, inputs_size):",
+        "  return \"keks\"",
+        "ruleContext.actions.run(",
+        "  inputs = ruleContext.files.srcs,",
+        "  outputs = ruleContext.files.srcs,",
+        "  resource_set = get_resources,",
+        "  executable = 'executable')");
+    StarlarkAction action =
+        (StarlarkAction)
+            Iterables.getOnlyElement(
+                ruleContext.getRuleContext().getAnalysisEnvironment().getRegisteredActions());
+
+    Exception thrown =
+        assertThrows(
+            ExecException.class,
+            () -> action.getResourceSetOrBuilder().buildResourceSet(OS.LINUX, 2));
+    assertThat(thrown).hasMessageThat().contains("got string for 'resource_set', want dict");
+  }
+
+  @Test
+  public void testCreateStarlarkActionArgumentsWithResourceSet_wrongDict() throws Exception {
+    setBuildLanguageOptions("--experimental_action_resource_set");
+    StarlarkRuleContext ruleContext = createRuleContext("//foo:foo");
+    setRuleContext(ruleContext);
+
+    ev.exec(
+        "def get_resources(os, inputs_size):",
+        "  return {\"cpu\": 1, \"memory\": 2, \"local_test\": \"hi\"}",
+        "ruleContext.actions.run(",
+        "  inputs = ruleContext.files.srcs,",
+        "  outputs = ruleContext.files.srcs,",
+        "  resource_set = get_resources,",
+        "  executable = 'executable')");
+    StarlarkAction action =
+        (StarlarkAction)
+            Iterables.getOnlyElement(
+                ruleContext.getRuleContext().getAnalysisEnvironment().getRegisteredActions());
+
+    Exception thrown =
+        assertThrows(
+            ExecException.class,
+            () -> action.getResourceSetOrBuilder().buildResourceSet(OS.LINUX, 2));
+    assertThat(thrown).hasMessageThat().contains("Illegal resource value type for key local_test");
+  }
+
+  @Test
+  public void testCreateStarlarkActionArgumentsWithResourceSet_incorrectSignature()
+      throws Exception {
+    setBuildLanguageOptions("--experimental_action_resource_set");
+    StarlarkRuleContext ruleContext = createRuleContext("//foo:foo");
+    setRuleContext(ruleContext);
+
+    ev.exec(
+        "def get_resources(os):",
+        "  return {\"cpu\": 1, \"memory\": 2, \"local_test\": \"hi\"}",
+        "ruleContext.actions.run(",
+        "  inputs = ruleContext.files.srcs,",
+        "  outputs = ruleContext.files.srcs,",
+        "  resource_set = get_resources,",
+        "  executable = 'executable')");
+    StarlarkAction action =
+        (StarlarkAction)
+            Iterables.getOnlyElement(
+                ruleContext.getRuleContext().getAnalysisEnvironment().getRegisteredActions());
+
+    Exception thrown =
+        assertThrows(
+            ExecException.class,
+            () -> action.getResourceSetOrBuilder().buildResourceSet(OS.LINUX, 2));
+    assertThat(thrown)
+        .hasMessageThat()
+        .contains("get_resources() accepts no more than 1 positional argument but got 2");
+  }
+
   @Test
   public void testCreateStarlarkActionArgumentsWithoutUnusedInputsList() throws Exception {
     StarlarkRuleContext ruleContext = createRuleContext("//foo:foo");

From 193399daa147a3f779d545b2db7630364117fd7a Mon Sep 17 00:00:00 2001
From: ilist <ilist@google.com>
Date: Thu, 7 Apr 2022 07:15:29 -0700
Subject: [PATCH 15/24] Implement and expose proto_common.compile call.

Design doc: https://docs.google.com/document/d/1dY_jfRvnH8SjRXGIfg8av-vquyWsvIZydXJOywvaR1A/edit

PiperOrigin-RevId: 440098122
---
 .../proto/ProtoBootstrap.java                 |   3 +
 .../starlark/builtins_bzl/common/exports.bzl  |   2 +
 .../common/proto/proto_common.bzl             |  85 ++++
 .../devtools/build/lib/rules/proto/BUILD      |  18 +
 .../lib/rules/proto/BazelProtoCommonTest.java | 464 ++++++++++++++++++
 5 files changed, 572 insertions(+)
 create mode 100644 src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java

diff --git a/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/proto/ProtoBootstrap.java b/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/proto/ProtoBootstrap.java
index fca7d5aad1b63c..774ff9d3c6c2ee 100644
--- a/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/proto/ProtoBootstrap.java
+++ b/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/proto/ProtoBootstrap.java
@@ -33,6 +33,8 @@ public class ProtoBootstrap implements Bootstrap {
   /** The name of the proto namespace in Starlark. */
   public static final String PROTO_COMMON_NAME = "proto_common";
 
+  public static final String PROTO_COMMON_SECOND_NAME = "proto_common_do_not_use";
+
   private final ProtoInfoProviderApi protoInfoApiProvider;
   private final ProtoToolchainInfoApi.Provider<? extends FilesToRunProviderApi<? extends FileApi>>
       protoToolchainInfoApi;
@@ -59,6 +61,7 @@ public void addBindingsToBuilder(ImmutableMap.Builder<String, Object> builder) {
     builder.put(PROTO_INFO_STARLARK_NAME, protoInfoApiProvider);
     builder.put(ProtoToolchainInfoApi.NAME, protoToolchainInfoApi);
     builder.put(PROTO_COMMON_NAME, protoCommon);
+    builder.put(PROTO_COMMON_SECOND_NAME, protoCommon);
     builder.put(
         "ProtoRegistryAspect",
         FlagGuardedValue.onlyWhenExperimentalFlagIsTrue(
diff --git a/src/main/starlark/builtins_bzl/common/exports.bzl b/src/main/starlark/builtins_bzl/common/exports.bzl
index 7d0f209e6b2c70..249e554ff2ac5f 100644
--- a/src/main/starlark/builtins_bzl/common/exports.bzl
+++ b/src/main/starlark/builtins_bzl/common/exports.bzl
@@ -23,12 +23,14 @@ load("@_builtins//:common/objc/objc_library.bzl", "objc_library")
 load("@_builtins//:common/objc/apple_static_library.bzl", "apple_static_library")
 load("@_builtins//:common/objc/compilation_support.bzl", "compilation_support")
 load("@_builtins//:common/proto/proto_lang_toolchain_wrapper.bzl", "proto_lang_toolchain")
+load("@_builtins//:common/proto/proto_common.bzl", "proto_common_do_not_use")
 
 exported_toplevels = {
     # This dummy symbol is not part of the public API; it is only used to test
     # that builtins injection is working properly. Its built-in value is
     # "original value".
     "_builtins_dummy": "overridden value",
+    "proto_common_do_not_use": proto_common_do_not_use,
 }
 exported_rules = {
     "+cc_import": cc_import,
diff --git a/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl b/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
index bcdd8dd60c8669..96075abbd89371 100644
--- a/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
+++ b/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
@@ -29,6 +29,91 @@ ProtoLangToolchainInfo = provider(
     ),
 )
 
+def _proto_path_flag(path):
+    if path == ".":
+        return None
+    return "--proto_path=%s" % path
+
+def _Iimport_path_equals_fullpath(proto_source):
+    return "-I%s=%s" % (proto_source.import_path(), proto_source.source_file().path)
+
+def _compile(
+        actions,
+        proto_library_target,
+        proto_lang_toolchain_info,
+        generated_files,
+        plugin_output = None,
+        additional_args = None,
+        additional_tools = [],
+        additional_inputs = depset(),
+        resource_set = None):
+    """Creates proto compile action for compiling *.proto files to language specific sources.
+
+    Args:
+      actions: (ActionFactory)  Obtained by ctx.actions, used to register the actions.
+      proto_library_target:  (Target) The proto_library to generate the sources for.
+        Obtained as the `target` parameter from an aspect's implementation.
+      proto_lang_toolchain_info: (ProtoLangToolchainInfo) The proto lang toolchain info.
+        Obtained from a `proto_lang_toolchain` target or constructed ad-hoc..
+      generated_files: (list[File]) The output files generated by the proto compiler.
+        Callee needs to declare files using `ctx.actions.declare_file`.
+        See also: `proto_common.declare_generated_files`.
+      plugin_output: (File|str) The file or directory passed to the plugin.
+        Formatted with `proto_lang_toolchain.out_replacement_format_flag`
+      additional_args: (Args) Additional arguments to add to the action.
+        Accepts an ctx.actions.args() object that is added at the beginning
+        of the command line.
+      additional_tools: (list[File]) Additional tools to add to the action.
+      additional_inputs: (Depset[File]) Additional input files to add to the action.
+      resource_set:
+        (func) A callback function that is passed to the created action.
+        See `ctx.actions.run`, `resource_set` parameter for full definition of
+        the callback.
+    """
+    proto_info = proto_library_target[_builtins.toplevel.ProtoInfo]
+
+    args = actions.args()
+    args.use_param_file(param_file_arg = "@%s")
+    args.set_param_file_format("multiline")
+    tools = list(additional_tools)
+
+    if plugin_output:
+        args.add(plugin_output, format = proto_lang_toolchain_info.out_replacement_format_flag)
+    if proto_lang_toolchain_info.plugin:
+        tools.append(proto_lang_toolchain_info.plugin)
+        args.add(proto_lang_toolchain_info.plugin.executable, format = proto_lang_toolchain_info.plugin_format_flag)
+
+    args.add_all(proto_info.transitive_proto_path, map_each = _proto_path_flag)
+    # Example: `--proto_path=--proto_path=bazel-bin/target/third_party/pkg/_virtual_imports/subpkg`
+
+    args.add_all(proto_lang_toolchain_info.protoc_opts)
+
+    # Include maps
+    # 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.
+    args.add_all(proto_info.transitive_proto_sources(), map_each = _Iimport_path_equals_fullpath)
+    # Example: `-Ia.proto=bazel-bin/target/third_party/pkg/_virtual_imports/subpkg/a.proto`
+
+    args.add_all(proto_info.direct_sources)
+
+    if additional_args:
+        additional_args.use_param_file(param_file_arg = "@%s")
+        additional_args.set_param_file_format("multiline")
+
+    actions.run(
+        mnemonic = proto_lang_toolchain_info.mnemonic,
+        progress_message = proto_lang_toolchain_info.progress_message,
+        executable = proto_lang_toolchain_info.proto_compiler,
+        arguments = [additional_args, args] if additional_args else [args],
+        inputs = depset(transitive = [proto_info.transitive_sources, additional_inputs]),
+        outputs = generated_files,
+        tools = tools,
+        use_default_shell_env = True,
+        resource_set = resource_set,
+    )
+
 proto_common_do_not_use = struct(
+    compile = _compile,
     ProtoLangToolchainInfo = ProtoLangToolchainInfo,
 )
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD b/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD
index ac97630531a214..38fce72efd4f44 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD
@@ -12,6 +12,24 @@ filegroup(
     visibility = ["//src:__subpackages__"],
 )
 
+java_test(
+    name = "BazelProtoCommonTest",
+    srcs = ["BazelProtoCommonTest.java"],
+    deps = [
+        "//src/main/java/com/google/devtools/build/lib/actions:localhost_capacity",
+        "//src/main/java/com/google/devtools/build/lib/analysis:analysis_cluster",
+        "//src/main/java/com/google/devtools/build/lib/analysis:configured_target",
+        "//src/main/java/com/google/devtools/build/lib/util:os",
+        "//src/test/java/com/google/devtools/build/lib/actions/util",
+        "//src/test/java/com/google/devtools/build/lib/analysis/util",
+        "//src/test/java/com/google/devtools/build/lib/packages:testutil",
+        "//src/test/java/com/google/devtools/build/lib/testutil:TestConstants",
+        "//third_party:junit4",
+        "//third_party:truth",
+        "@com_google_testparameterinjector//:testparameterinjector",
+    ],
+)
+
 java_test(
     name = "ProtoCompileActionBuilderTest",
     srcs = ["ProtoCompileActionBuilderTest.java"],
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java b/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
new file mode 100644
index 00000000000000..c1312f637cc0b9
--- /dev/null
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
@@ -0,0 +1,464 @@
+// Copyright 2022 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.rules.proto;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.devtools.build.lib.actions.util.ActionsTestUtil.prettyArtifactNames;
+
+import com.google.common.truth.Correspondence;
+import com.google.devtools.build.lib.actions.ResourceSet;
+import com.google.devtools.build.lib.analysis.ConfiguredTarget;
+import com.google.devtools.build.lib.analysis.actions.SpawnAction;
+import com.google.devtools.build.lib.analysis.util.BuildViewTestCase;
+import com.google.devtools.build.lib.packages.util.MockProtoSupport;
+import com.google.devtools.build.lib.testutil.TestConstants;
+import com.google.devtools.build.lib.util.OS;
+import com.google.testing.junit.testparameterinjector.TestParameterInjector;
+import com.google.testing.junit.testparameterinjector.TestParameters;
+import java.util.List;
+import java.util.regex.Pattern;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/** Unit test for proto_common module. */
+@RunWith(TestParameterInjector.class)
+public class BazelProtoCommonTest extends BuildViewTestCase {
+  private static final Correspondence<String, String> MATCHES_REGEX =
+      Correspondence.from((a, b) -> Pattern.matches(b, a), "matches");
+
+  @Before
+  public final void setup() throws Exception {
+    MockProtoSupport.setupWorkspace(scratch);
+    invalidatePackages();
+
+    MockProtoSupport.setup(mockToolsConfig);
+
+    scratch.file(
+        "third_party/x/BUILD",
+        "licenses(['unencumbered'])",
+        "cc_binary(name = 'plugin', srcs = ['plugin.cc'])",
+        "cc_library(name = 'runtime', srcs = ['runtime.cc'])",
+        "filegroup(name = 'descriptors', srcs = ['metadata.proto', 'descriptor.proto'])",
+        "filegroup(name = 'any', srcs = ['any.proto'])",
+        "proto_library(name = 'denied', srcs = [':descriptors', ':any'])");
+    scratch.file(
+        "foo/BUILD",
+        TestConstants.LOAD_PROTO_LANG_TOOLCHAIN,
+        "proto_lang_toolchain(",
+        "    name = 'toolchain',",
+        "    command_line = '--java_out=param1,param2:$(OUT)',",
+        "    plugin_format_flag = '--plugin=%s',",
+        "    plugin = '//third_party/x:plugin',",
+        "    runtime = '//third_party/x:runtime',",
+        "    blacklisted_protos = ['//third_party/x:denied'],",
+        "    progress_message = 'Progress Message %{label}',",
+        "    mnemonic = 'MyMnemonic',",
+        ")",
+        "proto_lang_toolchain(",
+        "    name = 'toolchain_noplugin',",
+        "    command_line = '--java_out=param1,param2:$(OUT)',",
+        "    runtime = '//third_party/x:runtime',",
+        "    blacklisted_protos = ['//third_party/x:denied'],",
+        "    progress_message = 'Progress Message %{label}',",
+        "    mnemonic = 'MyMnemonic',",
+        ")");
+
+    scratch.file(
+        "foo/generate.bzl",
+        "def _resource_set_callback(os, inputs_size):",
+        "   return {'memory': 25 + 0.15 * inputs_size, 'cpu': 1}",
+        "def _impl(ctx):",
+        "  outfile = ctx.actions.declare_file('out')",
+        "  kwargs = {}",
+        "  if ctx.attr.plugin_output:",
+        "    kwargs['plugin_output'] = ctx.attr.plugin_output",
+        "  if ctx.attr.additional_args:",
+        "    additional_args = ctx.actions.args()",
+        "    additional_args.add_all(ctx.attr.additional_args)",
+        "    kwargs['additional_args'] = additional_args",
+        "  if ctx.files.additional_tools:",
+        "    kwargs['additional_tools'] = ctx.files.additional_tools",
+        "  if ctx.files.additional_inputs:",
+        "    kwargs['additional_inputs'] = depset(ctx.files.additional_inputs)",
+        "  if ctx.attr.use_resource_set:",
+        "    kwargs['resource_set'] = _resource_set_callback",
+        "  proto_common_do_not_use.compile(",
+        "    ctx.actions,",
+        "    ctx.attr.proto_dep,",
+        "    ctx.attr.toolchain[proto_common_do_not_use.ProtoLangToolchainInfo],",
+        "    [outfile],",
+        "    **kwargs)",
+        "  return [DefaultInfo(files = depset([outfile]))]",
+        "generate_rule = rule(_impl,",
+        "  attrs = {",
+        "     'proto_dep': attr.label(),",
+        "     'plugin_output': attr.string(),",
+        "     'toolchain': attr.label(default = '//foo:toolchain'),",
+        "     'additional_args': attr.string_list(),",
+        "     'additional_tools': attr.label_list(cfg = 'exec'),",
+        "     'additional_inputs': attr.label_list(allow_files = True),",
+        "     'use_resource_set': attr.bool(),",
+        "  })");
+  }
+
+  /** Verifies basic usage of <code>proto_common.generate_code</code>. */
+  @Test
+  public void generateCode_basic() throws Exception {
+    scratch.file(
+        "bar/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "load('//foo:generate.bzl', 'generate_rule')",
+        "proto_library(name = 'proto', srcs = ['A.proto'])",
+        "generate_rule(name = 'simple', proto_dep = ':proto')");
+
+    ConfiguredTarget target = getConfiguredTarget("//bar:simple");
+
+    SpawnAction spawnAction = getGeneratingSpawnAction(getBinArtifact("out", target));
+    List<String> cmdLine = spawnAction.getRemainingArguments();
+    assertThat(cmdLine)
+        .comparingElementsUsing(MATCHES_REGEX)
+        .containsExactly(
+            "--plugin=bl?azel?-out/[^/]*-exec-[^/]*/bin/third_party/x/plugin",
+            "-Ibar/A.proto=bar/A.proto",
+            "bar/A.proto")
+        .inOrder();
+    assertThat(spawnAction.getMnemonic()).isEqualTo("MyMnemonic");
+    assertThat(spawnAction.getProgressMessage()).isEqualTo("Progress Message //bar:simple");
+  }
+
+  /** Verifies usage of proto_common.generate_code with no plugin specified by toolchain. */
+  @Test
+  public void generateCode_noPlugin() throws Exception {
+    scratch.file(
+        "bar/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "load('//foo:generate.bzl', 'generate_rule')",
+        "proto_library(name = 'proto', srcs = ['A.proto'])",
+        "generate_rule(name = 'simple', proto_dep = ':proto',",
+        "  toolchain = '//foo:toolchain_noplugin')");
+
+    ConfiguredTarget target = getConfiguredTarget("//bar:simple");
+
+    List<String> cmdLine =
+        getGeneratingSpawnAction(getBinArtifact("out", target)).getRemainingArguments();
+    assertThat(cmdLine)
+        .comparingElementsUsing(MATCHES_REGEX)
+        .containsExactly("-Ibar/A.proto=bar/A.proto", "bar/A.proto")
+        .inOrder();
+  }
+
+  /**
+   * Verifies usage of <code>proto_common.generate_code</code> with <code>plugin_output</code>
+   * parameter.
+   */
+  @Test
+  public void generateCode_withPluginOutput() throws Exception {
+    scratch.file(
+        "bar/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "load('//foo:generate.bzl', 'generate_rule')",
+        "proto_library(name = 'proto', srcs = ['A.proto'])",
+        "generate_rule(name = 'simple', proto_dep = ':proto', plugin_output = 'foo.srcjar')");
+
+    ConfiguredTarget target = getConfiguredTarget("//bar:simple");
+
+    List<String> cmdLine =
+        getGeneratingSpawnAction(getBinArtifact("out", target)).getRemainingArguments();
+    assertThat(cmdLine)
+        .comparingElementsUsing(MATCHES_REGEX)
+        .containsExactly(
+            "--java_out=param1,param2:foo.srcjar",
+            "--plugin=bl?azel?-out/[^/]*-exec-[^/]*/bin/third_party/x/plugin",
+            "-Ibar/A.proto=bar/A.proto",
+            "bar/A.proto")
+        .inOrder();
+  }
+
+  /**
+   * Verifies usage of <code>proto_common.generate_code</code> with <code>additional_args</code>
+   * parameter.
+   */
+  @Test
+  public void generateCode_additionalArgs() throws Exception {
+    scratch.file(
+        "bar/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "load('//foo:generate.bzl', 'generate_rule')",
+        "proto_library(name = 'proto', srcs = ['A.proto'])",
+        "generate_rule(name = 'simple', proto_dep = ':proto', additional_args = ['--a', '--b'])");
+
+    ConfiguredTarget target = getConfiguredTarget("//bar:simple");
+
+    List<String> cmdLine =
+        getGeneratingSpawnAction(getBinArtifact("out", target)).getRemainingArguments();
+    assertThat(cmdLine)
+        .comparingElementsUsing(MATCHES_REGEX)
+        .containsExactly(
+            "--a",
+            "--b",
+            "--plugin=bl?azel?-out/[^/]*-exec-[^/]*/bin/third_party/x/plugin",
+            "-Ibar/A.proto=bar/A.proto",
+            "bar/A.proto")
+        .inOrder();
+  }
+
+  /**
+   * Verifies usage of <code>proto_common.generate_code</code> with <code>additional_tools</code>
+   * parameter.
+   */
+  @Test
+  public void generateCode_additionalTools() throws Exception {
+    scratch.file(
+        "bar/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "load('//foo:generate.bzl', 'generate_rule')",
+        "proto_library(name = 'proto', srcs = ['A.proto'])",
+        "cc_binary(name = 'tool1', srcs = ['tool1.cc'])",
+        "cc_binary(name = 'tool2', srcs = ['tool2.cc'])",
+        "generate_rule(name = 'simple', proto_dep = ':proto',",
+        "  additional_tools = [':tool1', ':tool2'])");
+
+    ConfiguredTarget target = getConfiguredTarget("//bar:simple");
+
+    SpawnAction spawnAction = getGeneratingSpawnAction(getBinArtifact("out", target));
+    assertThat(prettyArtifactNames(spawnAction.getTools()))
+        .containsAtLeast("bar/tool1", "bar/tool2", "third_party/x/plugin");
+  }
+
+  /**
+   * Verifies usage of <code>proto_common.generate_code</code> with <code>additional_tools</code>
+   * parameter and no plugin on the toolchain.
+   */
+  @Test
+  public void generateCode_additionalToolsNoPlugin() throws Exception {
+    scratch.file(
+        "bar/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "load('//foo:generate.bzl', 'generate_rule')",
+        "proto_library(name = 'proto', srcs = ['A.proto'])",
+        "cc_binary(name = 'tool1', srcs = ['tool1.cc'])",
+        "cc_binary(name = 'tool2', srcs = ['tool2.cc'])",
+        "generate_rule(name = 'simple',",
+        "  proto_dep = ':proto',",
+        "  additional_tools = [':tool1', ':tool2'],",
+        "  toolchain = '//foo:toolchain_noplugin',",
+        ")");
+
+    ConfiguredTarget target = getConfiguredTarget("//bar:simple");
+
+    SpawnAction spawnAction = getGeneratingSpawnAction(getBinArtifact("out", target));
+    assertThat(prettyArtifactNames(spawnAction.getTools()))
+        .containsAtLeast("bar/tool1", "bar/tool2");
+  }
+
+  /**
+   * Verifies usage of <code>proto_common.generate_code</code> with <code>additional_inputs</code>
+   * parameter.
+   */
+  @Test
+  public void generateCode_additionalInputs() throws Exception {
+    scratch.file(
+        "bar/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "load('//foo:generate.bzl', 'generate_rule')",
+        "proto_library(name = 'proto', srcs = ['A.proto'])",
+        "generate_rule(name = 'simple', proto_dep = ':proto',",
+        "  additional_inputs = [':input1.txt', ':input2.txt'])");
+
+    ConfiguredTarget target = getConfiguredTarget("//bar:simple");
+
+    SpawnAction spawnAction = getGeneratingSpawnAction(getBinArtifact("out", target));
+    assertThat(prettyArtifactNames(spawnAction.getInputs()))
+        .containsAtLeast("bar/input1.txt", "bar/input2.txt");
+  }
+
+  /**
+   * Verifies usage of <code>proto_common.generate_code</code> with <code>resource_set</code>
+   * parameter.
+   */
+  @Test
+  public void generateCode_resourceSet() throws Exception {
+    scratch.file(
+        "bar/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "load('//foo:generate.bzl', 'generate_rule')",
+        "proto_library(name = 'proto', srcs = ['A.proto'])",
+        "generate_rule(name = 'simple', proto_dep = ':proto', use_resource_set = True)");
+
+    ConfiguredTarget target = getConfiguredTarget("//bar:simple");
+
+    SpawnAction spawnAction = getGeneratingSpawnAction(getBinArtifact("out", target));
+    assertThat(spawnAction.getResourceSetOrBuilder().buildResourceSet(OS.DARWIN, 0))
+        .isEqualTo(ResourceSet.createWithRamCpu(25, 1));
+    assertThat(spawnAction.getResourceSetOrBuilder().buildResourceSet(OS.LINUX, 2))
+        .isEqualTo(ResourceSet.createWithRamCpu(25.3, 1));
+  }
+
+  /** Verifies <code>--protocopts</code> are passed to command line. */
+  @Test
+  public void generateCode_protocOpts() throws Exception {
+    useConfiguration("--protocopt=--foo", "--protocopt=--bar");
+    scratch.file(
+        "bar/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "load('//foo:generate.bzl', 'generate_rule')",
+        "proto_library(name = 'proto', srcs = ['A.proto'])",
+        "generate_rule(name = 'simple', proto_dep = ':proto')");
+
+    ConfiguredTarget target = getConfiguredTarget("//bar:simple");
+
+    SpawnAction spawnAction = getGeneratingSpawnAction(getBinArtifact("out", target));
+    List<String> cmdLine = spawnAction.getRemainingArguments();
+    assertThat(cmdLine)
+        .comparingElementsUsing(MATCHES_REGEX)
+        .containsExactly(
+            "--plugin=bl?azel?-out/[^/]*-exec-[^/]*/bin/third_party/x/plugin",
+            "--foo",
+            "--bar",
+            "-Ibar/A.proto=bar/A.proto",
+            "bar/A.proto")
+        .inOrder();
+  }
+
+  /**
+   * Verifies <code>proto_common.generate_code</code> correctly handles direct generated <code>
+   * .proto</code> files.
+   */
+  @Test
+  public void generateCode_directGeneratedProtos() throws Exception {
+    useConfiguration("--noincompatible_generated_protos_in_virtual_imports");
+    scratch.file(
+        "bar/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "load('//foo:generate.bzl', 'generate_rule')",
+        "genrule(name = 'generate', srcs = ['A.txt'], cmd = '', outs = ['G.proto'])",
+        "proto_library(name = 'proto', srcs = ['A.proto', 'G.proto'])",
+        "generate_rule(name = 'simple', proto_dep = ':proto')");
+
+    ConfiguredTarget target = getConfiguredTarget("//bar:simple");
+
+    SpawnAction spawnAction = getGeneratingSpawnAction(getBinArtifact("out", target));
+    List<String> cmdLine = spawnAction.getRemainingArguments();
+    assertThat(cmdLine)
+        .comparingElementsUsing(MATCHES_REGEX)
+        .containsExactly(
+            "--plugin=bl?azel?-out/[^/]*-exec-[^/]*/bin/third_party/x/plugin",
+            "-Ibar/A.proto=bar/A.proto",
+            "-Ibar/G.proto=bl?azel?-out/k8-fastbuild/bin/bar/G.proto",
+            "bar/A.proto",
+            "bl?azel?-out/k8-fastbuild/bin/bar/G.proto")
+        .inOrder();
+  }
+
+  /**
+   * Verifies <code>proto_common.generate_code</code> correctly handles in-direct generated <code>
+   * .proto</code> files.
+   */
+  @Test
+  public void generateCode_inDirectGeneratedProtos() throws Exception {
+    useConfiguration("--noincompatible_generated_protos_in_virtual_imports");
+    scratch.file(
+        "bar/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "load('//foo:generate.bzl', 'generate_rule')",
+        "genrule(name = 'generate', srcs = ['A.txt'], cmd = '', outs = ['G.proto'])",
+        "proto_library(name = 'generated', srcs = ['G.proto'])",
+        "proto_library(name = 'proto', srcs = ['A.proto'], deps = [':generated'])",
+        "generate_rule(name = 'simple', proto_dep = ':proto')");
+
+    ConfiguredTarget target = getConfiguredTarget("//bar:simple");
+
+    SpawnAction spawnAction = getGeneratingSpawnAction(getBinArtifact("out", target));
+    List<String> cmdLine = spawnAction.getRemainingArguments();
+    assertThat(cmdLine)
+        .comparingElementsUsing(MATCHES_REGEX)
+        .containsExactly(
+            "--plugin=bl?azel?-out/[^/]*-exec-[^/]*/bin/third_party/x/plugin",
+            "-Ibar/A.proto=bar/A.proto",
+            "-Ibar/G.proto=bl?azel?-out/k8-fastbuild/bin/bar/G.proto",
+            "bar/A.proto")
+        .inOrder();
+  }
+
+  /**
+   * Verifies <code>proto_common.generate_code</code> correctly handles external <code>proto_library
+   * </code>-es.
+   */
+  @Test
+  @TestParameters({
+    "{virtual: false, sibling: false, generated: false, expectedFlags:"
+        + " ['--proto_path=external/foo','-Ie/E.proto=external/foo/e/E.proto']}",
+    "{virtual: false, sibling: false, generated: true, expectedFlags:"
+        + " ['--proto_path=external/foo',"
+        + " '-Ie/E.proto=bl?azel?-out/k8-fastbuild/bin/external/foo/e/E.proto']}",
+    "{virtual: true, sibling: false, generated: false,expectedFlags:"
+        + " ['--proto_path=external/foo','-Ie/E.proto=external/foo/e/E.proto']}",
+    "{virtual: true, sibling: false, generated: true, expectedFlags:"
+        + " ['--proto_path=bl?azel?-out/k8-fastbuild/bin/external/foo/e/_virtual_imports/e',"
+        + " '-Ie/E.proto=bl?azel?-out/k8-fastbuild/bin/external/foo/e/_virtual_imports/e/e/E.proto']}",
+    "{virtual: true, sibling: true, generated: false,expectedFlags:"
+        + " ['--proto_path=../foo','-I../foo/e/E.proto=../foo/e/E.proto']}",
+    "{virtual: true, sibling: true, generated: true, expectedFlags:"
+        + " ['--proto_path=bl?azel?-out/foo/k8-fastbuild/bin/e/_virtual_imports/e',"
+        + " '-Ie/E.proto=bl?azel?-out/foo/k8-fastbuild/bin/e/_virtual_imports/e/e/E.proto']}",
+    "{virtual: false, sibling: true, generated: false,expectedFlags:"
+        + " ['--proto_path=../foo','-I../foo/e/E.proto=../foo/e/E.proto']}",
+    "{virtual: false, sibling: true, generated: true, expectedFlags:"
+        + " ['--proto_path=../foo','-Ie/E.proto=bl?azel?-out/foo/k8-fastbuild/bin/e/E.proto']}",
+  })
+  public void generateCode_externalProtoLibrary(
+      boolean virtual, boolean sibling, boolean generated, List<String> expectedFlags)
+      throws Exception {
+    if (virtual) {
+      useConfiguration("--incompatible_generated_protos_in_virtual_imports");
+    } else {
+      useConfiguration("--noincompatible_generated_protos_in_virtual_imports");
+    }
+    if (sibling) {
+      setBuildLanguageOptions("--experimental_sibling_repository_layout");
+    }
+    scratch.appendFile("WORKSPACE", "local_repository(name = 'foo', path = '/foo')");
+    invalidatePackages();
+    scratch.file("/foo/WORKSPACE");
+    scratch.file(
+        "/foo/e/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "proto_library(name='e', srcs=['E.proto'])",
+        generated
+            ? "genrule(name = 'generate', srcs = ['A.txt'], cmd = '', outs = ['E.proto'])"
+            : "");
+    scratch.file(
+        "bar/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "load('//foo:generate.bzl', 'generate_rule')",
+        "proto_library(name = 'proto', srcs = ['A.proto'], deps = ['@foo//e:e'])",
+        "generate_rule(name = 'simple', proto_dep = ':proto')");
+
+    ConfiguredTarget target = getConfiguredTarget("//bar:simple");
+
+    SpawnAction spawnAction = getGeneratingSpawnAction(getBinArtifact("out", target));
+    List<String> cmdLine = spawnAction.getRemainingArguments();
+    assertThat(cmdLine)
+        .comparingElementsUsing(MATCHES_REGEX)
+        .containsExactly(
+            "--plugin=bl?azel?-out/[^/]*-exec-[^/]*/bin/third_party/x/plugin",
+            expectedFlags.get(0),
+            "-Ibar/A.proto=bar/A.proto",
+            expectedFlags.get(1),
+            "bar/A.proto")
+        .inOrder();
+  }
+}

From 00497015b138e4a4c10e125fb4c5ec2d8b6b4091 Mon Sep 17 00:00:00 2001
From: Ivo List <ilist@google.com>
Date: Thu, 14 Jul 2022 13:19:47 +0000
Subject: [PATCH 16/24] Fix ProtoCommon tests.

---
 .../build/lib/rules/proto/BazelProtoCommonTest.java   | 11 +----------
 1 file changed, 1 insertion(+), 10 deletions(-)

diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java b/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
index c1312f637cc0b9..162557f71a13f0 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
@@ -90,7 +90,7 @@ public final void setup() throws Exception {
         "    additional_args.add_all(ctx.attr.additional_args)",
         "    kwargs['additional_args'] = additional_args",
         "  if ctx.files.additional_tools:",
-        "    kwargs['additional_tools'] = ctx.files.additional_tools",
+        "    kwargs['additional_tools'] = [f[DefaultInfo].files_to_run for f in ctx.attr.additional_tools]",
         "  if ctx.files.additional_inputs:",
         "    kwargs['additional_inputs'] = depset(ctx.files.additional_inputs)",
         "  if ctx.attr.use_resource_set:",
@@ -401,23 +401,14 @@ public void generateCode_inDirectGeneratedProtos() throws Exception {
   @TestParameters({
     "{virtual: false, sibling: false, generated: false, expectedFlags:"
         + " ['--proto_path=external/foo','-Ie/E.proto=external/foo/e/E.proto']}",
-    "{virtual: false, sibling: false, generated: true, expectedFlags:"
-        + " ['--proto_path=external/foo',"
-        + " '-Ie/E.proto=bl?azel?-out/k8-fastbuild/bin/external/foo/e/E.proto']}",
     "{virtual: true, sibling: false, generated: false,expectedFlags:"
         + " ['--proto_path=external/foo','-Ie/E.proto=external/foo/e/E.proto']}",
     "{virtual: true, sibling: false, generated: true, expectedFlags:"
         + " ['--proto_path=bl?azel?-out/k8-fastbuild/bin/external/foo/e/_virtual_imports/e',"
         + " '-Ie/E.proto=bl?azel?-out/k8-fastbuild/bin/external/foo/e/_virtual_imports/e/e/E.proto']}",
-    "{virtual: true, sibling: true, generated: false,expectedFlags:"
-        + " ['--proto_path=../foo','-I../foo/e/E.proto=../foo/e/E.proto']}",
     "{virtual: true, sibling: true, generated: true, expectedFlags:"
         + " ['--proto_path=bl?azel?-out/foo/k8-fastbuild/bin/e/_virtual_imports/e',"
         + " '-Ie/E.proto=bl?azel?-out/foo/k8-fastbuild/bin/e/_virtual_imports/e/e/E.proto']}",
-    "{virtual: false, sibling: true, generated: false,expectedFlags:"
-        + " ['--proto_path=../foo','-I../foo/e/E.proto=../foo/e/E.proto']}",
-    "{virtual: false, sibling: true, generated: true, expectedFlags:"
-        + " ['--proto_path=../foo','-Ie/E.proto=bl?azel?-out/foo/k8-fastbuild/bin/e/E.proto']}",
   })
   public void generateCode_externalProtoLibrary(
       boolean virtual, boolean sibling, boolean generated, List<String> expectedFlags)

From 46f06580651251450727ba8ecebf9a8c78f2193b Mon Sep 17 00:00:00 2001
From: ilist <ilist@google.com>
Date: Thu, 7 Apr 2022 07:18:14 -0700
Subject: [PATCH 17/24] Add experimental_progress_message parameter to
 proto_common.compile.

This will make migration to the new call easier (because we don't have progress_message set yet on the proto_lang_toolchain rules).

Design doc: https://docs.google.com/document/d/1dY_jfRvnH8SjRXGIfg8av-vquyWsvIZydXJOywvaR1A/edit

PiperOrigin-RevId: 440098602
---
 .../common/proto/proto_common.bzl             |  7 +++--
 .../lib/rules/proto/BazelProtoCommonTest.java | 28 +++++++++++++++++++
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl b/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
index 96075abbd89371..644a1f414cb31b 100644
--- a/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
+++ b/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
@@ -46,7 +46,8 @@ def _compile(
         additional_args = None,
         additional_tools = [],
         additional_inputs = depset(),
-        resource_set = None):
+        resource_set = None,
+        experimental_progress_message = None):
     """Creates proto compile action for compiling *.proto files to language specific sources.
 
     Args:
@@ -69,6 +70,8 @@ def _compile(
         (func) A callback function that is passed to the created action.
         See `ctx.actions.run`, `resource_set` parameter for full definition of
         the callback.
+      experimental_progress_message: Overrides progres_message from the toolchain.
+        Don't use this parameter. It's only intended for the transition.
     """
     proto_info = proto_library_target[_builtins.toplevel.ProtoInfo]
 
@@ -103,7 +106,7 @@ def _compile(
 
     actions.run(
         mnemonic = proto_lang_toolchain_info.mnemonic,
-        progress_message = proto_lang_toolchain_info.progress_message,
+        progress_message = experimental_progress_message if experimental_progress_message else proto_lang_toolchain_info.progress_message,
         executable = proto_lang_toolchain_info.proto_compiler,
         arguments = [additional_args, args] if additional_args else [args],
         inputs = depset(transitive = [proto_info.transitive_sources, additional_inputs]),
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java b/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
index 162557f71a13f0..ad56212a46afad 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
@@ -95,6 +95,8 @@ public final void setup() throws Exception {
         "    kwargs['additional_inputs'] = depset(ctx.files.additional_inputs)",
         "  if ctx.attr.use_resource_set:",
         "    kwargs['resource_set'] = _resource_set_callback",
+        "  if ctx.attr.progress_message:",
+        "    kwargs['experimental_progress_message'] = ctx.attr.progress_message",
         "  proto_common_do_not_use.compile(",
         "    ctx.actions,",
         "    ctx.attr.proto_dep,",
@@ -111,6 +113,7 @@ public final void setup() throws Exception {
         "     'additional_tools': attr.label_list(cfg = 'exec'),",
         "     'additional_inputs': attr.label_list(allow_files = True),",
         "     'use_resource_set': attr.bool(),",
+        "     'progress_message': attr.string(),",
         "  })");
   }
 
@@ -452,4 +455,29 @@ public void generateCode_externalProtoLibrary(
             "bar/A.proto")
         .inOrder();
   }
+
+  /** Verifies <code>experimental_progress_message</code> parameters. */
+  @Test
+  public void generateCode_overrideProgressMessage() throws Exception {
+    scratch.file(
+        "bar/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "load('//foo:generate.bzl', 'generate_rule')",
+        "proto_library(name = 'proto', srcs = ['A.proto'])",
+        "generate_rule(name = 'simple', proto_dep = ':proto', progress_message = 'My %{label}')");
+
+    ConfiguredTarget target = getConfiguredTarget("//bar:simple");
+
+    SpawnAction spawnAction = getGeneratingSpawnAction(getBinArtifact("out", target));
+    List<String> cmdLine = spawnAction.getRemainingArguments();
+    assertThat(cmdLine)
+        .comparingElementsUsing(MATCHES_REGEX)
+        .containsExactly(
+            "--plugin=bl?azel?-out/[^/]*-exec-[^/]*/bin/third_party/x/plugin",
+            "-Ibar/A.proto=bar/A.proto",
+            "bar/A.proto")
+        .inOrder();
+    assertThat(spawnAction.getMnemonic()).isEqualTo("MyMnemonic");
+    assertThat(spawnAction.getProgressMessage()).isEqualTo("My //bar:simple");
+  }
 }

From c4b8b95515ca4cdeb102600e084d5871bd8a276d Mon Sep 17 00:00:00 2001
From: ilist <ilist@google.com>
Date: Thu, 7 Apr 2022 08:10:29 -0700
Subject: [PATCH 18/24] Implement
 proto_common.experimental_should_generate_code.

Design doc: https://docs.google.com/document/d/1dY_jfRvnH8SjRXGIfg8av-vquyWsvIZydXJOywvaR1A/edit

PiperOrigin-RevId: 440109298
---
 .../build/lib/rules/proto/ProtoCommon.java    |  2 +
 .../build/lib/rules/proto/ProtoInfo.java      |  7 ++
 .../proto/ProtoLangToolchainProvider.java     |  6 +-
 .../build/lib/rules/proto/ProtoSource.java    | 10 ++-
 .../lib/starlarkbuildapi/ProtoInfoApi.java    |  3 +
 .../common/proto/proto_common.bzl             | 73 ++++++++++++++++++
 .../devtools/build/lib/rules/proto/BUILD      |  2 +
 .../lib/rules/proto/BazelProtoCommonTest.java | 77 +++++++++++++++++++
 8 files changed, 178 insertions(+), 2 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 733b41dad8d275..49a985f1ccba9b 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
@@ -38,8 +38,10 @@
 import javax.annotation.Nullable;
 import net.starlark.java.eval.EvalException;
 import net.starlark.java.eval.Module;
+import net.starlark.java.eval.Sequence;
 import net.starlark.java.eval.Starlark;
 import net.starlark.java.eval.StarlarkThread;
+import net.starlark.java.eval.Tuple;
 
 /** Utility functions for proto_library and proto aspect implementations. */
 public class ProtoCommon {
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoInfo.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoInfo.java
index 0593c5a0b3759f..376e28f64ee713 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoInfo.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoInfo.java
@@ -118,6 +118,13 @@ public ImmutableList<Artifact> getDirectProtoSources() {
     return directProtoSources;
   }
 
+  @Override
+  public ImmutableList<ProtoSource> getDirectProtoSourcesForStarlark(StarlarkThread thread)
+      throws EvalException {
+    ProtoCommon.checkPrivateStarlarkificationAllowlist(thread);
+    return directSources;
+  }
+
   /**
    * The source root of the current library.
    *
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
index 8753baa0a8b9f3..4ce277d5af7e81 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoLangToolchainProvider.java
@@ -29,6 +29,7 @@
 import java.util.LinkedHashMap;
 import java.util.Map;
 import javax.annotation.Nullable;
+import net.starlark.java.annot.StarlarkMethod;
 import net.starlark.java.eval.EvalException;
 import net.starlark.java.eval.NoneType;
 import net.starlark.java.eval.Starlark;
@@ -69,7 +70,10 @@ public abstract class ProtoLangToolchainProvider {
    * Returns a list of {@link ProtoSource}s that are already provided by the protobuf runtime (i.e.
    * for which {@code <lang>_proto_library} should not generate bindings.
    */
-  // Proto sources provided by the toolchain.
+  @StarlarkMethod(
+      name = "provided_proto_sources",
+      doc = "Proto sources provided by the toolchain.",
+      structField = true)
   public abstract ImmutableList<ProtoSource> providedProtoSources();
 
   // Proto compiler.
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoSource.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoSource.java
index 1833d3004362ca..01bdee4cd0972a 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoSource.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoSource.java
@@ -18,6 +18,9 @@
 import com.google.devtools.build.lib.collect.nestedset.Depset;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.vfs.PathFragment;
+import net.starlark.java.annot.StarlarkMethod;
+import net.starlark.java.eval.EvalException;
+import net.starlark.java.eval.StarlarkThread;
 import net.starlark.java.eval.StarlarkValue;
 
 /** Represents a single {@code .proto} source file. */
@@ -50,7 +53,12 @@ public Artifact getSourceFileForStarlark(StarlarkThread thread) throws EvalExcep
   }
 
   /** Returns the original source file. Only for forbidding protos! */
-  @Deprecated
+  @StarlarkMethod(name = "original_source_file", documented = false, useStarlarkThread = true)
+  public Artifact getOriginalSourceFileForStarlark(StarlarkThread thread) throws EvalException {
+    ProtoCommon.checkPrivateStarlarkificationAllowlist(thread);
+    return originalSourceFile;
+  }
+
   Artifact getOriginalSourceFile() {
     return originalSourceFile;
   }
diff --git a/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/ProtoInfoApi.java b/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/ProtoInfoApi.java
index ea7214f3251db3..6289e8f688025d 100644
--- a/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/ProtoInfoApi.java
+++ b/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/ProtoInfoApi.java
@@ -64,6 +64,9 @@ interface ProtoInfoProviderApi extends ProviderApi {
       structField = true)
   ImmutableList<FileT> getDirectProtoSources();
 
+  @StarlarkMethod(name = "direct_proto_sources", documented = false, useStarlarkThread = true)
+  ImmutableList<?> getDirectProtoSourcesForStarlark(StarlarkThread thread) throws EvalException;
+
   @StarlarkMethod(
       name = "check_deps_sources",
       doc =
diff --git a/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl b/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
index 644a1f414cb31b..6ce825e75713af 100644
--- a/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
+++ b/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
@@ -116,7 +116,80 @@ def _compile(
         resource_set = resource_set,
     )
 
+_BAZEL_TOOLS_PREFIX = "external/bazel_tools/"
+
+def _experimental_filter_sources(proto_library_target, proto_lang_toolchain_info):
+    proto_info = proto_library_target[_builtins.toplevel.ProtoInfo]
+    if not proto_info.direct_sources:
+        return [], []
+
+    # Collect a set of provided protos
+    provided_proto_sources = proto_lang_toolchain_info.provided_proto_sources
+    provided_paths = {}
+    for src in provided_proto_sources:
+        path = src.original_source_file().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
+        # user repositories will not have that prefix.
+        if path.startswith(_BAZEL_TOOLS_PREFIX):
+            provided_paths[path[len(_BAZEL_TOOLS_PREFIX):]] = None
+        else:
+            provided_paths[path] = None
+
+    # Filter proto files
+    proto_files = [src.original_source_file() for src in proto_info.direct_proto_sources()]
+    excluded = []
+    included = []
+    for proto_file in proto_files:
+        if proto_file.path in provided_paths:
+            excluded.append(proto_file)
+        else:
+            included.append(proto_file)
+    return included, excluded
+
+def _experimental_should_generate_code(
+        proto_library_target,
+        proto_lang_toolchain_info,
+        rule_name):
+    """Checks if the code should be generated for the given proto_library.
+
+    The code shouldn't be generated only when the toolchain already provides it
+    to the language through its runtime dependency.
+
+    It fails when the proto_library contains mixed proto files, that should and
+    shouldn't generate code.
+
+    Args:
+      proto_library_target:
+        (Target) The proto_library to generate the sources for.
+        Obtained as the `target` parameter from an aspect's implementation.
+      proto_lang_toolchain_info:
+        (ProtoLangToolchainInfo) The proto lang toolchain info.
+        Obtained from a `proto_lang_toolchain` target or constructed ad-hoc.
+      rule_name: (str) Name of the rule used in the failure message.
+
+    Returns:
+      (bool) True when the code should be generated.
+    """
+    included, excluded = _experimental_filter_sources(proto_library_target, proto_lang_toolchain_info)
+
+    if included and excluded:
+        fail(("The 'srcs' attribute of '%s' contains protos for which '%s' " +
+              "shouldn't generate code (%s), in addition to protos for which it should (%s).\n" +
+              "Separate '%s' into 2 proto_library rules.") % (
+            proto_library_target.label,
+            rule_name,
+            ", ".join([f.short_path for f in excluded]),
+            ", ".join([f.short_path for f in included]),
+            proto_library_target.label,
+        ))
+
+    return bool(included)
+
 proto_common_do_not_use = struct(
     compile = _compile,
+    experimental_should_generate_code = _experimental_should_generate_code,
+    experimental_filter_sources = _experimental_filter_sources,
     ProtoLangToolchainInfo = ProtoLangToolchainInfo,
 )
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD b/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD
index 38fce72efd4f44..4b9cce0b8820cf 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD
@@ -19,6 +19,8 @@ java_test(
         "//src/main/java/com/google/devtools/build/lib/actions:localhost_capacity",
         "//src/main/java/com/google/devtools/build/lib/analysis:analysis_cluster",
         "//src/main/java/com/google/devtools/build/lib/analysis:configured_target",
+        "//src/main/java/com/google/devtools/build/lib/cmdline",
+        "//src/main/java/com/google/devtools/build/lib/packages",
         "//src/main/java/com/google/devtools/build/lib/util:os",
         "//src/test/java/com/google/devtools/build/lib/actions/util",
         "//src/test/java/com/google/devtools/build/lib/analysis/util",
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java b/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
index ad56212a46afad..6978a65c91aec1 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
@@ -22,6 +22,10 @@
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.actions.SpawnAction;
 import com.google.devtools.build.lib.analysis.util.BuildViewTestCase;
+import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.packages.StarlarkInfo;
+import com.google.devtools.build.lib.packages.StarlarkProvider;
+import com.google.devtools.build.lib.packages.StarlarkProviderIdentifier;
 import com.google.devtools.build.lib.packages.util.MockProtoSupport;
 import com.google.devtools.build.lib.testutil.TestConstants;
 import com.google.devtools.build.lib.util.OS;
@@ -39,6 +43,11 @@ public class BazelProtoCommonTest extends BuildViewTestCase {
   private static final Correspondence<String, String> MATCHES_REGEX =
       Correspondence.from((a, b) -> Pattern.matches(b, a), "matches");
 
+  private static final StarlarkProviderIdentifier boolProviderId =
+      StarlarkProviderIdentifier.forKey(
+          new StarlarkProvider.Key(
+              Label.parseAbsoluteUnchecked("//foo:should_generate.bzl"), "BoolProvider"));
+
   @Before
   public final void setup() throws Exception {
     MockProtoSupport.setupWorkspace(scratch);
@@ -53,6 +62,8 @@ public final void setup() throws Exception {
         "cc_library(name = 'runtime', srcs = ['runtime.cc'])",
         "filegroup(name = 'descriptors', srcs = ['metadata.proto', 'descriptor.proto'])",
         "filegroup(name = 'any', srcs = ['any.proto'])",
+        "filegroup(name = 'something', srcs = ['something.proto'])",
+        "proto_library(name = 'mixed', srcs = [':descriptors', ':something'])",
         "proto_library(name = 'denied', srcs = [':descriptors', ':any'])");
     scratch.file(
         "foo/BUILD",
@@ -115,6 +126,21 @@ public final void setup() throws Exception {
         "     'use_resource_set': attr.bool(),",
         "     'progress_message': attr.string(),",
         "  })");
+
+    scratch.file(
+        "foo/should_generate.bzl",
+        "BoolProvider = provider()",
+        "def _impl(ctx):",
+        "  result = proto_common_do_not_use.experimental_should_generate_code(",
+        "    ctx.attr.proto_dep,",
+        "    ctx.attr.toolchain[proto_common_do_not_use.ProtoLangToolchainInfo],",
+        "    'MyRule')",
+        "  return [BoolProvider(value = result)]",
+        "should_generate_rule = rule(_impl,",
+        "  attrs = {",
+        "     'proto_dep': attr.label(),",
+        "     'toolchain': attr.label(default = '//foo:toolchain'),",
+        "  })");
   }
 
   /** Verifies basic usage of <code>proto_common.generate_code</code>. */
@@ -480,4 +506,55 @@ public void generateCode_overrideProgressMessage() throws Exception {
     assertThat(spawnAction.getMnemonic()).isEqualTo("MyMnemonic");
     assertThat(spawnAction.getProgressMessage()).isEqualTo("My //bar:simple");
   }
+
+  /** Verifies <code>proto_common.should_generate_code</code> call. */
+  @Test
+  public void shouldGenerateCode_basic() throws Exception {
+    scratch.file(
+        "bar/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "load('//foo:should_generate.bzl', 'should_generate_rule')",
+        "proto_library(name = 'proto', srcs = ['A.proto'])",
+        "should_generate_rule(name = 'simple', proto_dep = ':proto')");
+
+    ConfiguredTarget target = getConfiguredTarget("//bar:simple");
+
+    StarlarkInfo boolProvider = (StarlarkInfo) target.get(boolProviderId);
+    assertThat(boolProvider.getValue("value", Boolean.class)).isTrue();
+  }
+
+  /** Verifies <code>proto_common.should_generate_code</code> call. */
+  @Test
+  public void shouldGenerateCode_dontGenerate() throws Exception {
+    scratch.file(
+        "bar/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "load('//foo:should_generate.bzl', 'should_generate_rule')",
+        "should_generate_rule(name = 'simple', proto_dep = '//third_party/x:denied')");
+
+    ConfiguredTarget target = getConfiguredTarget("//bar:simple");
+
+    StarlarkInfo boolProvider = (StarlarkInfo) target.get(boolProviderId);
+    assertThat(boolProvider.getValue("value", Boolean.class)).isFalse();
+  }
+
+  /** Verifies <code>proto_common.should_generate_code</code> call. */
+  @Test
+  public void shouldGenerateCode_mixed() throws Exception {
+    scratch.file(
+        "bar/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "load('//foo:should_generate.bzl', 'should_generate_rule')",
+        "should_generate_rule(name = 'simple', proto_dep = '//third_party/x:mixed')");
+
+    reporter.removeHandler(failFastHandler);
+    getConfiguredTarget("//bar:simple");
+
+    assertContainsEvent(
+        "The 'srcs' attribute of '//third_party/x:mixed' contains protos for which 'MyRule'"
+            + " shouldn't generate code (third_party/x/metadata.proto,"
+            + " third_party/x/descriptor.proto), in addition to protos for which it should"
+            + " (third_party/x/something.proto).\n"
+            + "Separate '//third_party/x:mixed' into 2 proto_library rules.");
+  }
 }

From 600bd5e1720ac5f6fd00bd0f1ca7d553cbe3a63d Mon Sep 17 00:00:00 2001
From: ilist <ilist@google.com>
Date: Mon, 11 Apr 2022 23:44:34 -0700
Subject: [PATCH 19/24] Implement proto_common.declare_generated_files.

Design doc: https://docs.google.com/document/d/1dY_jfRvnH8SjRXGIfg8av-vquyWsvIZydXJOywvaR1A/edit

PiperOrigin-RevId: 441097041
---
 .../common/proto/proto_common.bzl             | 45 ++++++++++++++++
 .../devtools/build/lib/rules/proto/BUILD      |  1 +
 .../lib/rules/proto/BazelProtoCommonTest.java | 52 +++++++++++++++++++
 3 files changed, 98 insertions(+)

diff --git a/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl b/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
index 6ce825e75713af..885fa5d5c3360e 100644
--- a/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
+++ b/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
@@ -187,8 +187,53 @@ def _experimental_should_generate_code(
 
     return bool(included)
 
+def _declare_generated_files(
+        actions,
+        proto_library_target,
+        extension,
+        name_mapper = None):
+    """Declares generated files with a specific extension.
+
+    Use this in lang_proto_library-es when protocol compiler generates files
+    that correspond to .proto file names.
+
+    The function removes ".proto" extension with given one (e.g. ".pb.cc") and
+    declares new output files.
+
+    Args:
+      actions:
+        (ActionFactory) Obtained by ctx.actions, used to declare the files.
+      proto_library_target:
+        (Target) The proto_library to generate the files for.
+        Obtained as the `target` parameter from an aspect's implementation.
+      extension: (str) The extension to use for generated files.
+      name_mapper: (str->str) A function mapped over the base filename without
+        the extension. Used it to replace characters in the name that
+        cause problems in a specific programming language.
+
+    Returns:
+      (list[File]) The list of declared files.
+    """
+    proto_info = proto_library_target[_builtins.toplevel.ProtoInfo]
+    proto_sources = proto_info.direct_sources
+    outputs = []
+
+    for src in proto_sources:
+        basename_no_ext = src.basename[:-(len(src.extension) + 1)]
+
+        if name_mapper:
+            basename_no_ext = name_mapper(basename_no_ext)
+
+        # Note that two proto_library rules can have the same source file, so this is actually a
+        # shared action. NB: This can probably result in action conflicts if the proto_library rules
+        # are not the same.
+        outputs.append(actions.declare_file(basename_no_ext + extension, sibling = src))
+
+    return outputs
+
 proto_common_do_not_use = struct(
     compile = _compile,
+    declare_generated_files = _declare_generated_files,
     experimental_should_generate_code = _experimental_should_generate_code,
     experimental_filter_sources = _experimental_filter_sources,
     ProtoLangToolchainInfo = ProtoLangToolchainInfo,
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD b/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD
index 4b9cce0b8820cf..52cab96a1ff8bf 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/BUILD
@@ -19,6 +19,7 @@ java_test(
         "//src/main/java/com/google/devtools/build/lib/actions:localhost_capacity",
         "//src/main/java/com/google/devtools/build/lib/analysis:analysis_cluster",
         "//src/main/java/com/google/devtools/build/lib/analysis:configured_target",
+        "//src/main/java/com/google/devtools/build/lib/analysis:file_provider",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/packages",
         "//src/main/java/com/google/devtools/build/lib/util:os",
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java b/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
index 6978a65c91aec1..753fb8fd4c4b0f 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
@@ -20,6 +20,7 @@
 import com.google.common.truth.Correspondence;
 import com.google.devtools.build.lib.actions.ResourceSet;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
+import com.google.devtools.build.lib.analysis.FileProvider;
 import com.google.devtools.build.lib.analysis.actions.SpawnAction;
 import com.google.devtools.build.lib.analysis.util.BuildViewTestCase;
 import com.google.devtools.build.lib.cmdline.Label;
@@ -141,6 +142,24 @@ public final void setup() throws Exception {
         "     'proto_dep': attr.label(),",
         "     'toolchain': attr.label(default = '//foo:toolchain'),",
         "  })");
+
+    scratch.file(
+        "foo/declare_generated_files.bzl",
+        "def _impl(ctx):",
+        "  files = proto_common_do_not_use.declare_generated_files(",
+        "    ctx.actions,",
+        "    ctx.attr.proto_dep,",
+        "    ctx.attr.extension,",
+        "    (lambda s: s.replace('-','_').replace('.','/')) if ctx.attr.python_names else None)",
+        "  for f in files:",
+        "    ctx.actions.write(f, '')",
+        "  return [DefaultInfo(files = depset(files))]",
+        "declare_generated_files = rule(_impl,",
+        "  attrs = {",
+        "     'proto_dep': attr.label(),",
+        "     'extension': attr.string(),",
+        "     'python_names': attr.bool(default = False),",
+        "  })");
   }
 
   /** Verifies basic usage of <code>proto_common.generate_code</code>. */
@@ -557,4 +576,37 @@ public void shouldGenerateCode_mixed() throws Exception {
             + " (third_party/x/something.proto).\n"
             + "Separate '//third_party/x:mixed' into 2 proto_library rules.");
   }
+
+  /** Verifies <code>proto_common.declare_generated_files</code> call. */
+  @Test
+  public void declareGenerateFiles_basic() throws Exception {
+    scratch.file(
+        "bar/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "load('//foo:declare_generated_files.bzl', 'declare_generated_files')",
+        "proto_library(name = 'proto', srcs = ['A.proto', 'b/B.proto'])",
+        "declare_generated_files(name = 'simple', proto_dep = ':proto', extension = '.cc')");
+
+    ConfiguredTarget target = getConfiguredTarget("//bar:simple");
+
+    assertThat(prettyArtifactNames(target.getProvider(FileProvider.class).getFilesToBuild()))
+        .containsExactly("bar/A.cc", "bar/b/B.cc");
+  }
+
+  /** Verifies <code>proto_common.declare_generated_files</code> call for Python. */
+  @Test
+  public void declareGenerateFiles_pythonc() throws Exception {
+    scratch.file(
+        "bar/BUILD",
+        TestConstants.LOAD_PROTO_LIBRARY,
+        "load('//foo:declare_generated_files.bzl', 'declare_generated_files')",
+        "proto_library(name = 'proto', srcs = ['my-proto.gen.proto'])",
+        "declare_generated_files(name = 'simple', proto_dep = ':proto', extension = '_pb2.py',",
+        "  python_names = True)");
+
+    ConfiguredTarget target = getConfiguredTarget("//bar:simple");
+
+    assertThat(prettyArtifactNames(target.getProvider(FileProvider.class).getFilesToBuild()))
+        .containsExactly("bar/my_proto/gen_pb2.py");
+  }
 }

From 8b8f4e8454b1bdb5c0f16426b464a8ab90df24c8 Mon Sep 17 00:00:00 2001
From: ilist <ilist@google.com>
Date: Thu, 28 Apr 2022 07:45:45 -0700
Subject: [PATCH 20/24] Fix docstring in proto_common

PiperOrigin-RevId: 445149402
---
 .../builtins_bzl/common/proto/proto_common.bzl    | 15 +++++----------
 1 file changed, 5 insertions(+), 10 deletions(-)

diff --git a/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl b/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
index 885fa5d5c3360e..e78d4fa0d232fb 100644
--- a/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
+++ b/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
@@ -66,8 +66,7 @@ def _compile(
         of the command line.
       additional_tools: (list[File]) Additional tools to add to the action.
       additional_inputs: (Depset[File]) Additional input files to add to the action.
-      resource_set:
-        (func) A callback function that is passed to the created action.
+      resource_set: (func) A callback function that is passed to the created action.
         See `ctx.actions.run`, `resource_set` parameter for full definition of
         the callback.
       experimental_progress_message: Overrides progres_message from the toolchain.
@@ -161,11 +160,9 @@ def _experimental_should_generate_code(
     shouldn't generate code.
 
     Args:
-      proto_library_target:
-        (Target) The proto_library to generate the sources for.
+      proto_library_target: (Target) The proto_library to generate the sources for.
         Obtained as the `target` parameter from an aspect's implementation.
-      proto_lang_toolchain_info:
-        (ProtoLangToolchainInfo) The proto lang toolchain info.
+      proto_lang_toolchain_info: (ProtoLangToolchainInfo) The proto lang toolchain info.
         Obtained from a `proto_lang_toolchain` target or constructed ad-hoc.
       rule_name: (str) Name of the rule used in the failure message.
 
@@ -201,10 +198,8 @@ def _declare_generated_files(
     declares new output files.
 
     Args:
-      actions:
-        (ActionFactory) Obtained by ctx.actions, used to declare the files.
-      proto_library_target:
-        (Target) The proto_library to generate the files for.
+      actions: (ActionFactory) Obtained by ctx.actions, used to declare the files.
+      proto_library_target: (Target) The proto_library to generate the files for.
         Obtained as the `target` parameter from an aspect's implementation.
       extension: (str) The extension to use for generated files.
       name_mapper: (str->str) A function mapped over the base filename without

From 80d9ae284fdbbc56d8aa267392661f2f13478fc5 Mon Sep 17 00:00:00 2001
From: Googler <ilist@google.com>
Date: Fri, 17 Jun 2022 07:45:36 -0700
Subject: [PATCH 21/24] Add proto_info parameter to proto_common.compile

Adding it will make it possible to migrate the uses from proto_library_target to proto_info. When the migration is done proto_library_target will be removed.

The cost of changing this now is still low, because it's not yet released/used in Bazel.

PiperOrigin-RevId: 455616355
Change-Id: Ieb0f03b0600e1f90b72a61f90420675075c79a9e
---
 .../builtins_bzl/common/proto/proto_common.bzl      | 13 ++++++++-----
 .../build/lib/rules/proto/BazelProtoCommonTest.java |  2 +-
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl b/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
index e78d4fa0d232fb..cda3f8b2f25b60 100644
--- a/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
+++ b/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
@@ -39,7 +39,6 @@ def _Iimport_path_equals_fullpath(proto_source):
 
 def _compile(
         actions,
-        proto_library_target,
         proto_lang_toolchain_info,
         generated_files,
         plugin_output = None,
@@ -47,13 +46,14 @@ def _compile(
         additional_tools = [],
         additional_inputs = depset(),
         resource_set = None,
-        experimental_progress_message = None):
+        experimental_progress_message = None,
+        proto_info = None,
+        proto_library_target = None):
     """Creates proto compile action for compiling *.proto files to language specific sources.
 
     Args:
       actions: (ActionFactory)  Obtained by ctx.actions, used to register the actions.
-      proto_library_target:  (Target) The proto_library to generate the sources for.
-        Obtained as the `target` parameter from an aspect's implementation.
+      proto_info: (ProtoInfo) The ProtoInfo from proto_library to generate the sources for.
       proto_lang_toolchain_info: (ProtoLangToolchainInfo) The proto lang toolchain info.
         Obtained from a `proto_lang_toolchain` target or constructed ad-hoc..
       generated_files: (list[File]) The output files generated by the proto compiler.
@@ -71,8 +71,11 @@ def _compile(
         the callback.
       experimental_progress_message: Overrides progres_message from the toolchain.
         Don't use this parameter. It's only intended for the transition.
+      proto_library_target:  (Target) Deprecated: use proto_info instead.
+        Exactly one need to be set: `proto_info` or `proto_library_target`.
     """
-    proto_info = proto_library_target[_builtins.toplevel.ProtoInfo]
+    if not proto_info:
+        proto_info = proto_library_target[_builtins.toplevel.ProtoInfo]
 
     args = actions.args()
     args.use_param_file(param_file_arg = "@%s")
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java b/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
index 753fb8fd4c4b0f..0cabf2ac3749ca 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
@@ -111,9 +111,9 @@ public final void setup() throws Exception {
         "    kwargs['experimental_progress_message'] = ctx.attr.progress_message",
         "  proto_common_do_not_use.compile(",
         "    ctx.actions,",
-        "    ctx.attr.proto_dep,",
         "    ctx.attr.toolchain[proto_common_do_not_use.ProtoLangToolchainInfo],",
         "    [outfile],",
+        "    proto_info = ctx.attr.proto_dep[ProtoInfo],",
         "    **kwargs)",
         "  return [DefaultInfo(files = depset([outfile]))]",
         "generate_rule = rule(_impl,",

From bc452c1ebe2248ef0d83ac44bcbc9a4ddc974ee1 Mon Sep 17 00:00:00 2001
From: Googler <ilist@google.com>
Date: Wed, 22 Jun 2022 03:37:47 -0700
Subject: [PATCH 22/24] Migrate proto_library_target to proto_info in
 proto_common.declare_generated_files.

PiperOrigin-RevId: 456475071
Change-Id: I2882d80cd4f7fdd9b8dcba3347930eaf1f194d0a
---
 .../builtins_bzl/common/proto/proto_common.bzl  | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl b/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
index cda3f8b2f25b60..501ec7588306fc 100644
--- a/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
+++ b/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
@@ -189,9 +189,10 @@ def _experimental_should_generate_code(
 
 def _declare_generated_files(
         actions,
-        proto_library_target,
-        extension,
-        name_mapper = None):
+        proto_library_target = None,
+        extension = None,
+        name_mapper = None,
+        proto_info = None):
     """Declares generated files with a specific extension.
 
     Use this in lang_proto_library-es when protocol compiler generates files
@@ -202,17 +203,21 @@ def _declare_generated_files(
 
     Args:
       actions: (ActionFactory) Obtained by ctx.actions, used to declare the files.
-      proto_library_target: (Target) The proto_library to generate the files for.
-        Obtained as the `target` parameter from an aspect's implementation.
+      proto_info: (ProtoInfo) The ProtoInfo to declare the files for.
       extension: (str) The extension to use for generated files.
       name_mapper: (str->str) A function mapped over the base filename without
         the extension. Used it to replace characters in the name that
         cause problems in a specific programming language.
+      proto_library_target: (Target) Deprecated: use proto_info instead.
 
     Returns:
       (list[File]) The list of declared files.
     """
-    proto_info = proto_library_target[_builtins.toplevel.ProtoInfo]
+    if not proto_info:
+        proto_info = proto_library_target[_builtins.toplevel.ProtoInfo]
+    if not extension:
+        fail("missing parameter extension")
+
     proto_sources = proto_info.direct_sources
     outputs = []
 

From 0b381ace41110a17a78607d165c2fb8b89b220cb Mon Sep 17 00:00:00 2001
From: Googler <ilist@google.com>
Date: Mon, 11 Jul 2022 02:36:18 -0700
Subject: [PATCH 23/24] Migrate proto_library_target to proto_info in
 proto_common.experimental_should_generate_code

PiperOrigin-RevId: 460162832
Change-Id: I57a6fa4c6e6c9618cf9edb8518e17b46fc90be9f
---
 .../builtins_bzl/common/proto/proto_common.bzl | 18 +++++++++---------
 .../lib/rules/proto/BazelProtoCommonTest.java  |  5 +++--
 2 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl b/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
index 501ec7588306fc..c5e4451a649ebf 100644
--- a/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
+++ b/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
@@ -120,8 +120,7 @@ def _compile(
 
 _BAZEL_TOOLS_PREFIX = "external/bazel_tools/"
 
-def _experimental_filter_sources(proto_library_target, proto_lang_toolchain_info):
-    proto_info = proto_library_target[_builtins.toplevel.ProtoInfo]
+def _experimental_filter_sources(proto_info, proto_lang_toolchain_info):
     if not proto_info.direct_sources:
         return [], []
 
@@ -151,9 +150,10 @@ def _experimental_filter_sources(proto_library_target, proto_lang_toolchain_info
     return included, excluded
 
 def _experimental_should_generate_code(
-        proto_library_target,
+        proto_info,
         proto_lang_toolchain_info,
-        rule_name):
+        rule_name,
+        target_label):
     """Checks if the code should be generated for the given proto_library.
 
     The code shouldn't be generated only when the toolchain already provides it
@@ -163,26 +163,26 @@ def _experimental_should_generate_code(
     shouldn't generate code.
 
     Args:
-      proto_library_target: (Target) The proto_library to generate the sources for.
-        Obtained as the `target` parameter from an aspect's implementation.
+      proto_info: (ProtoInfo) The ProtoInfo from proto_library to check the generation for.
       proto_lang_toolchain_info: (ProtoLangToolchainInfo) The proto lang toolchain info.
         Obtained from a `proto_lang_toolchain` target or constructed ad-hoc.
       rule_name: (str) Name of the rule used in the failure message.
+      target_label: (Label) The label of the target used in the failure message.
 
     Returns:
       (bool) True when the code should be generated.
     """
-    included, excluded = _experimental_filter_sources(proto_library_target, proto_lang_toolchain_info)
+    included, excluded = _experimental_filter_sources(proto_info, proto_lang_toolchain_info)
 
     if included and excluded:
         fail(("The 'srcs' attribute of '%s' contains protos for which '%s' " +
               "shouldn't generate code (%s), in addition to protos for which it should (%s).\n" +
               "Separate '%s' into 2 proto_library rules.") % (
-            proto_library_target.label,
+            target_label,
             rule_name,
             ", ".join([f.short_path for f in excluded]),
             ", ".join([f.short_path for f in included]),
-            proto_library_target.label,
+            target_label,
         ))
 
     return bool(included)
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java b/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
index 0cabf2ac3749ca..5f4e46bc8ed42e 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
@@ -133,9 +133,10 @@ public final void setup() throws Exception {
         "BoolProvider = provider()",
         "def _impl(ctx):",
         "  result = proto_common_do_not_use.experimental_should_generate_code(",
-        "    ctx.attr.proto_dep,",
+        "    ctx.attr.proto_dep[ProtoInfo],",
         "    ctx.attr.toolchain[proto_common_do_not_use.ProtoLangToolchainInfo],",
-        "    'MyRule')",
+        "    'MyRule',",
+        "    ctx.attr.proto_dep.label)",
         "  return [BoolProvider(value = result)]",
         "should_generate_rule = rule(_impl,",
         "  attrs = {",

From 3b78ca9aab641c8b3b8c07e852e55366b2a4da6b Mon Sep 17 00:00:00 2001
From: Googler <ilist@google.com>
Date: Tue, 12 Jul 2022 02:06:37 -0700
Subject: [PATCH 24/24] Remove proto_library_target from proto_common

PiperOrigin-RevId: 460406536
Change-Id: I10021f32fb40e163ded02ebab8297902b63760fa
---
 .../common/proto/proto_common.bzl             | 23 ++++---------------
 .../lib/rules/proto/BazelProtoCommonTest.java |  4 ++--
 2 files changed, 7 insertions(+), 20 deletions(-)

diff --git a/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl b/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
index c5e4451a649ebf..a3fa72e010fafc 100644
--- a/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
+++ b/src/main/starlark/builtins_bzl/common/proto/proto_common.bzl
@@ -39,6 +39,7 @@ def _Iimport_path_equals_fullpath(proto_source):
 
 def _compile(
         actions,
+        proto_info,
         proto_lang_toolchain_info,
         generated_files,
         plugin_output = None,
@@ -46,9 +47,7 @@ def _compile(
         additional_tools = [],
         additional_inputs = depset(),
         resource_set = None,
-        experimental_progress_message = None,
-        proto_info = None,
-        proto_library_target = None):
+        experimental_progress_message = None):
     """Creates proto compile action for compiling *.proto files to language specific sources.
 
     Args:
@@ -71,12 +70,7 @@ def _compile(
         the callback.
       experimental_progress_message: Overrides progres_message from the toolchain.
         Don't use this parameter. It's only intended for the transition.
-      proto_library_target:  (Target) Deprecated: use proto_info instead.
-        Exactly one need to be set: `proto_info` or `proto_library_target`.
     """
-    if not proto_info:
-        proto_info = proto_library_target[_builtins.toplevel.ProtoInfo]
-
     args = actions.args()
     args.use_param_file(param_file_arg = "@%s")
     args.set_param_file_format("multiline")
@@ -189,10 +183,9 @@ def _experimental_should_generate_code(
 
 def _declare_generated_files(
         actions,
-        proto_library_target = None,
-        extension = None,
-        name_mapper = None,
-        proto_info = None):
+        proto_info,
+        extension,
+        name_mapper = None):
     """Declares generated files with a specific extension.
 
     Use this in lang_proto_library-es when protocol compiler generates files
@@ -208,16 +201,10 @@ def _declare_generated_files(
       name_mapper: (str->str) A function mapped over the base filename without
         the extension. Used it to replace characters in the name that
         cause problems in a specific programming language.
-      proto_library_target: (Target) Deprecated: use proto_info instead.
 
     Returns:
       (list[File]) The list of declared files.
     """
-    if not proto_info:
-        proto_info = proto_library_target[_builtins.toplevel.ProtoInfo]
-    if not extension:
-        fail("missing parameter extension")
-
     proto_sources = proto_info.direct_sources
     outputs = []
 
diff --git a/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java b/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
index 5f4e46bc8ed42e..161d5e27b5fa7e 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
+++ b/src/test/java/com/google/devtools/build/lib/rules/proto/BazelProtoCommonTest.java
@@ -111,9 +111,9 @@ public final void setup() throws Exception {
         "    kwargs['experimental_progress_message'] = ctx.attr.progress_message",
         "  proto_common_do_not_use.compile(",
         "    ctx.actions,",
+        "    ctx.attr.proto_dep[ProtoInfo],",
         "    ctx.attr.toolchain[proto_common_do_not_use.ProtoLangToolchainInfo],",
         "    [outfile],",
-        "    proto_info = ctx.attr.proto_dep[ProtoInfo],",
         "    **kwargs)",
         "  return [DefaultInfo(files = depset([outfile]))]",
         "generate_rule = rule(_impl,",
@@ -149,7 +149,7 @@ public final void setup() throws Exception {
         "def _impl(ctx):",
         "  files = proto_common_do_not_use.declare_generated_files(",
         "    ctx.actions,",
-        "    ctx.attr.proto_dep,",
+        "    ctx.attr.proto_dep[ProtoInfo],",
         "    ctx.attr.extension,",
         "    (lambda s: s.replace('-','_').replace('.','/')) if ctx.attr.python_names else None)",
         "  for f in files:",