Make java_toolchain attribute Environment aware.

RELNOTES:none:
PiperOrigin-RevId: 200988244
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/rules/BazelRuleClassProvider.java b/src/main/java/com/google/devtools/build/lib/bazel/rules/BazelRuleClassProvider.java
index 8be4821..227920d 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/rules/BazelRuleClassProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/rules/BazelRuleClassProvider.java
@@ -80,6 +80,7 @@
 import com.google.devtools.build.lib.rules.core.CoreRules;
 import com.google.devtools.build.lib.rules.cpp.proto.CcProtoAspect;
 import com.google.devtools.build.lib.rules.cpp.proto.CcProtoLibraryRule;
+import com.google.devtools.build.lib.rules.java.JavaConfiguration;
 import com.google.devtools.build.lib.rules.java.JavaSemantics;
 import com.google.devtools.build.lib.rules.platform.PlatformRules;
 import com.google.devtools.build.lib.rules.proto.BazelProtoLibraryRule;
@@ -257,10 +258,14 @@
       new RuleSet() {
         @Override
         public void init(ConfiguredRuleClassProvider.Builder builder) {
-          LabelLateBoundDefault<?> hostJdkAttribute = JavaSemantics.hostJdkAttribute(builder);
-          BazelJavaProtoAspect bazelJavaProtoAspect = new BazelJavaProtoAspect(hostJdkAttribute);
+          LabelLateBoundDefault<JavaConfiguration> hostJdkAttribute =
+              JavaSemantics.hostJdkAttribute(builder);
+          LabelLateBoundDefault<JavaConfiguration> javaToolchainAttribute =
+              JavaSemantics.javaToolchainAttribute(builder);
+          BazelJavaProtoAspect bazelJavaProtoAspect =
+              new BazelJavaProtoAspect(hostJdkAttribute, javaToolchainAttribute);
           BazelJavaLiteProtoAspect bazelJavaLiteProtoAspect =
-              new BazelJavaLiteProtoAspect(hostJdkAttribute);
+              new BazelJavaLiteProtoAspect(hostJdkAttribute, javaToolchainAttribute);
           builder.addNativeAspectClass(bazelJavaProtoAspect);
           builder.addNativeAspectClass(bazelJavaLiteProtoAspect);
           builder.addRuleDefinition(new BazelJavaProtoLibraryRule(bazelJavaProtoAspect));
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/rules/java/proto/BazelJavaLiteProtoAspect.java b/src/main/java/com/google/devtools/build/lib/bazel/rules/java/proto/BazelJavaLiteProtoAspect.java
index feec6ce..ec3df47 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/rules/java/proto/BazelJavaLiteProtoAspect.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/rules/java/proto/BazelJavaLiteProtoAspect.java
@@ -16,6 +16,7 @@
 
 import com.google.devtools.build.lib.bazel.rules.java.BazelJavaSemantics;
 import com.google.devtools.build.lib.packages.Attribute.LabelLateBoundDefault;
+import com.google.devtools.build.lib.rules.java.JavaConfiguration;
 import com.google.devtools.build.lib.rules.java.proto.JavaLiteProtoAspect;
 
 /** An Aspect which BazelJavaLiteProtoLibrary injects to build Java Lite protos. */
@@ -24,11 +25,14 @@
   public static final String DEFAULT_PROTO_TOOLCHAIN_LABEL =
       "@com_google_protobuf_javalite//:javalite_toolchain";
 
-  public BazelJavaLiteProtoAspect(LabelLateBoundDefault<?> hostJdkAttribute) {
+  public BazelJavaLiteProtoAspect(
+      LabelLateBoundDefault<JavaConfiguration> hostJdkAttribute,
+      LabelLateBoundDefault<JavaConfiguration> javaToolchainAttribute) {
     super(
         BazelJavaSemantics.INSTANCE,
         null /* jacocoLabel */,
         DEFAULT_PROTO_TOOLCHAIN_LABEL,
-        hostJdkAttribute);
+        hostJdkAttribute,
+        javaToolchainAttribute);
   }
 }
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 eeec9d3..51ad027 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
@@ -26,6 +26,7 @@
 import com.google.devtools.build.lib.packages.AspectDefinition;
 import com.google.devtools.build.lib.packages.AspectParameters;
 import com.google.devtools.build.lib.packages.Attribute.LabelLateBoundDefault;
+import com.google.devtools.build.lib.rules.java.JavaConfiguration;
 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;
@@ -34,13 +35,16 @@
 /** An Aspect which BazelJavaProtoLibrary injects to build Java SPEED protos. */
 public class BazelJavaProtoAspect extends JavaProtoAspect {
 
-  public BazelJavaProtoAspect(LabelLateBoundDefault<?> hostJdkAttribute) {
+  public BazelJavaProtoAspect(
+      LabelLateBoundDefault<JavaConfiguration> hostJdkAttribute,
+      LabelLateBoundDefault<JavaConfiguration> javaToolchainAttribute) {
     super(
         BazelJavaSemantics.INSTANCE,
         null, /* jacocoAttr */
         new NoopRpcSupport(),
         "@com_google_protobuf//:java_toolchain",
-        hostJdkAttribute);
+        hostJdkAttribute,
+        javaToolchainAttribute);
   }
 
   private static class NoopRpcSupport
diff --git a/src/main/java/com/google/devtools/build/lib/rules/android/AndroidRuleClasses.java b/src/main/java/com/google/devtools/build/lib/rules/android/AndroidRuleClasses.java
index d06c9b2..ae481f9 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/android/AndroidRuleClasses.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/android/AndroidRuleClasses.java
@@ -471,7 +471,7 @@
               attr(":java_toolchain", LABEL)
                   .useOutputLicenses()
                   .allowedRuleClasses("java_toolchain")
-                  .value(JavaSemantics.JAVA_TOOLCHAIN))
+                  .value(JavaSemantics.javaToolchainAttribute(environment)))
           .advertiseProvider(AndroidSdkProvider.class)
           .build();
     }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaRuleClasses.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaRuleClasses.java
index c12f946..d9a0d37 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/JavaRuleClasses.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/JavaRuleClasses.java
@@ -40,7 +40,7 @@
               attr(":java_toolchain", LABEL)
                   .useOutputLicenses()
                   .mandatoryProviders(ToolchainInfo.PROVIDER.id())
-                  .value(JavaSemantics.JAVA_TOOLCHAIN))
+                  .value(JavaSemantics.javaToolchainAttribute(env)))
           .setPreferredDependencyPredicate(JavaSemantics.JAVA_SOURCE)
           .build();
     }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaSemantics.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaSemantics.java
index 85cd979..d351102 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/JavaSemantics.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/JavaSemantics.java
@@ -114,11 +114,15 @@
   /** The java_toolchain.compatible_javacopts key for testonly compilations. */
   public static final String TESTONLY_JAVACOPTS_KEY = "testonly";
 
-  LabelLateBoundDefault<JavaConfiguration> JAVA_TOOLCHAIN =
-      LabelLateBoundDefault.fromTargetConfiguration(
-          JavaConfiguration.class,
-          Label.parseAbsoluteUnchecked(JAVA_TOOLCHAIN_LABEL),
-          (rule, attributes, javaConfig) -> javaConfig.getToolchainLabel());
+  static LabelLateBoundDefault<JavaConfiguration> javaToolchainAttribute(
+      RuleDefinitionEnvironment environment) {
+    return LabelLateBoundDefault.fromTargetConfiguration(
+        JavaConfiguration.class,
+        // TODO(b/79239052): replace by //environment.getToolsLabel(JAVA_TOOLCHAIN_LABEL)
+        // @bazel_tools//tools/defaults can not be resolved while DefaultPackage exists.
+        Label.parseAbsoluteUnchecked(JAVA_TOOLCHAIN_LABEL),
+        (rule, attributes, javaConfig) -> javaConfig.getToolchainLabel());
+  }
 
   /**
    * Name of the output group used for source jars.
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaToolchainAliasRule.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaToolchainAliasRule.java
index feca01c..41b6981 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/JavaToolchainAliasRule.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/JavaToolchainAliasRule.java
@@ -21,6 +21,6 @@
  */
 public class JavaToolchainAliasRule extends CommonAliasRule{
   public JavaToolchainAliasRule() {
-    super("java_toolchain_alias", env -> JavaSemantics.JAVA_TOOLCHAIN, JavaConfiguration.class);
+    super("java_toolchain_alias", JavaSemantics::javaToolchainAttribute, JavaConfiguration.class);
   }
 }
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 f531e0f7..6359334 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
@@ -72,17 +72,20 @@
 
   @Nullable private final String jacocoLabel;
   private final String defaultProtoToolchainLabel;
-  private final LabelLateBoundDefault<?> hostJdkAttribute;
+  private final LabelLateBoundDefault<JavaConfiguration> hostJdkAttribute;
+  private final LabelLateBoundDefault<JavaConfiguration> javaToolchainAttribute;
 
   public JavaLiteProtoAspect(
       JavaSemantics javaSemantics,
       @Nullable String jacocoLabel,
       String defaultProtoToolchainLabel,
-      LabelLateBoundDefault<?> hostJdkAttribute) {
+      LabelLateBoundDefault<JavaConfiguration> hostJdkAttribute,
+      LabelLateBoundDefault<JavaConfiguration> javaToolchainAttribute) {
     this.javaSemantics = javaSemantics;
     this.jacocoLabel = jacocoLabel;
     this.defaultProtoToolchainLabel = defaultProtoToolchainLabel;
     this.hostJdkAttribute = hostJdkAttribute;
+    this.javaToolchainAttribute = javaToolchainAttribute;
   }
 
   @Override
@@ -119,15 +122,16 @@
                         ImmutableList.<Class<? extends TransitiveInfoProvider>>of(
                             ProtoLangToolchainProvider.class))
                     .value(getProtoToolchainLabel(defaultProtoToolchainLabel)))
-            .add(attr(":host_jdk", LABEL)
-                .cfg(HostTransition.INSTANCE)
-                .value(hostJdkAttribute)
-                .mandatoryProviders(JavaRuntimeInfo.PROVIDER.id()))
+            .add(
+                attr(":host_jdk", LABEL)
+                    .cfg(HostTransition.INSTANCE)
+                    .value(hostJdkAttribute)
+                    .mandatoryProviders(JavaRuntimeInfo.PROVIDER.id()))
             .add(
                 attr(":java_toolchain", LABEL)
                     .useOutputLicenses()
                     .allowedRuleClasses("java_toolchain")
-                    .value(JavaSemantics.JAVA_TOOLCHAIN));
+                    .value(javaToolchainAttribute));
 
     Attribute.Builder<Label> jacocoAttr =
         attr("$jacoco_instrumentation", LABEL).cfg(HostTransition.INSTANCE);
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/proto/JavaProtoAspect.java b/src/main/java/com/google/devtools/build/lib/rules/java/proto/JavaProtoAspect.java
index 0e59ce8..60128d3 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/proto/JavaProtoAspect.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/proto/JavaProtoAspect.java
@@ -22,6 +22,7 @@
 import static com.google.devtools.build.lib.rules.java.proto.JplCcLinkParams.createCcLinkParamsStore;
 import static com.google.devtools.build.lib.rules.java.proto.StrictDepsUtils.createNonStrictCompilationArgsProvider;
 
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.analysis.ConfiguredAspect;
@@ -60,7 +61,8 @@
 /** An Aspect which JavaProtoLibrary injects to build Java SPEED protos. */
 public class JavaProtoAspect extends NativeAspectClass implements ConfiguredAspectFactory {
 
-  private final LabelLateBoundDefault<?> hostJdkAttribute;
+  private final LabelLateBoundDefault<JavaConfiguration> hostJdkAttribute;
+  private final LabelLateBoundDefault<JavaConfiguration> javaToolchainAttribute;
 
   private static LabelLateBoundDefault<?> getSpeedProtoToolchainLabel(String defaultValue) {
     return LabelLateBoundDefault.fromTargetConfiguration(
@@ -80,12 +82,15 @@
       @Nullable String jacocoLabel,
       RpcSupport rpcSupport,
       String defaultSpeedProtoToolchainLabel,
-      LabelLateBoundDefault<?> hostJdkAttribute) {
-    this.javaSemantics = javaSemantics;
+      LabelLateBoundDefault<JavaConfiguration> hostJdkAttribute,
+      LabelLateBoundDefault<JavaConfiguration> javaToolchainAttribute) {
+    this.javaSemantics = Preconditions.checkNotNull(javaSemantics);
     this.jacocoLabel = jacocoLabel;
-    this.rpcSupport = rpcSupport;
-    this.defaultSpeedProtoToolchainLabel = defaultSpeedProtoToolchainLabel;
-    this.hostJdkAttribute = hostJdkAttribute;
+    this.rpcSupport = Preconditions.checkNotNull(rpcSupport);
+    this.defaultSpeedProtoToolchainLabel =
+        Preconditions.checkNotNull(defaultSpeedProtoToolchainLabel);
+    this.hostJdkAttribute = Preconditions.checkNotNull(hostJdkAttribute);
+    this.javaToolchainAttribute = Preconditions.checkNotNull(javaToolchainAttribute);
   }
 
   @Override
@@ -130,7 +135,7 @@
                 attr(":java_toolchain", LABEL)
                     .useOutputLicenses()
                     .allowedRuleClasses("java_toolchain")
-                    .value(JavaSemantics.JAVA_TOOLCHAIN));
+                    .value(javaToolchainAttribute));
 
     rpcSupport.mutateAspectDefinition(result, aspectParameters);