Add source jar support to aar_import via the srcjar attribute

Fixes https://github.com/bazelbuild/bazel/issues/8671

Closes #8692.

Change-Id: I09b0266cf46775c448baadc09787e20a87004d55
PiperOrigin-RevId: 258840861
diff --git a/src/main/java/com/google/devtools/build/lib/rules/android/AarImport.java b/src/main/java/com/google/devtools/build/lib/rules/android/AarImport.java
index 624286b..8d2ff78 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/android/AarImport.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/android/AarImport.java
@@ -44,6 +44,8 @@
 import com.google.devtools.build.lib.rules.java.JavaRuntimeInfo;
 import com.google.devtools.build.lib.rules.java.JavaSemantics;
 import com.google.devtools.build.lib.rules.java.JavaSkylarkApiProvider;
+import com.google.devtools.build.lib.rules.java.JavaSourceInfoProvider;
+import com.google.devtools.build.lib.rules.java.JavaSourceJarsProvider;
 import com.google.devtools.build.lib.rules.java.JavaToolchainProvider;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import javax.annotation.Nullable;
@@ -192,12 +194,36 @@
             /* isNeverLink = */ JavaCommon.isNeverLink(ruleContext),
             /* srcLessDepsExport = */ false);
 
+    // Wire up the source jar for the current target and transitive source jars from dependencies.
+    ImmutableList<Artifact> srcJars = ImmutableList.of();
+    Artifact srcJar = ruleContext.getPrerequisiteArtifact("srcjar", Mode.TARGET);
+    NestedSetBuilder<Artifact> transitiveJavaSourceJarBuilder = NestedSetBuilder.stableOrder();
+    if (srcJar != null) {
+      srcJars = ImmutableList.of(srcJar);
+      transitiveJavaSourceJarBuilder.add(srcJar);
+    }
+    for (JavaSourceJarsProvider other :
+        JavaInfo.getProvidersFromListOfTargets(
+            JavaSourceJarsProvider.class, ruleContext.getPrerequisites("exports", Mode.TARGET))) {
+      transitiveJavaSourceJarBuilder.addTransitive(other.getTransitiveSourceJars());
+    }
+    NestedSet<Artifact> transitiveJavaSourceJars = transitiveJavaSourceJarBuilder.build();
+    JavaSourceJarsProvider javaSourceJarsProvider =
+        JavaSourceJarsProvider.create(transitiveJavaSourceJars, srcJars);
+    JavaSourceInfoProvider javaSourceInfoProvider =
+        new JavaSourceInfoProvider.Builder()
+            .setJarFiles(ImmutableList.of(mergedJar))
+            .setSourceJarsForJarFiles(srcJars)
+            .build();
+
     JavaInfo.Builder javaInfoBuilder =
         JavaInfo.Builder.create()
             .setRuntimeJars(ImmutableList.of(mergedJar))
             .setJavaConstraints(ImmutableList.of("android"))
             .setNeverlink(JavaCommon.isNeverLink(ruleContext))
             .addProvider(JavaCompilationArgsProvider.class, javaCompilationArgsProvider)
+            .addProvider(JavaSourceJarsProvider.class, javaSourceJarsProvider)
+            .addProvider(JavaSourceInfoProvider.class, javaSourceInfoProvider)
             .addProvider(JavaRuleOutputJarsProvider.class, jarProviderBuilder.build());
 
     common.addTransitiveInfoProviders(
diff --git a/src/main/java/com/google/devtools/build/lib/rules/android/AarImportBaseRule.java b/src/main/java/com/google/devtools/build/lib/rules/android/AarImportBaseRule.java
index 90eb552..291b293 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/android/AarImportBaseRule.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/android/AarImportBaseRule.java
@@ -27,6 +27,7 @@
 import com.google.devtools.build.lib.rules.android.AndroidRuleClasses.AndroidBaseRule;
 import com.google.devtools.build.lib.rules.java.JavaConfiguration;
 import com.google.devtools.build.lib.rules.java.JavaInfo;
+import com.google.devtools.build.lib.rules.java.JavaSemantics;
 import com.google.devtools.build.lib.util.FileType;
 
 /** Rule definition for the aar_import rule. */
@@ -53,6 +54,13 @@
                 .allowedRuleClasses("aar_import", "java_import")
                 .allowedFileTypes()
                 .validityPredicate(ANY_EDGE))
+        /* <!-- #BLAZE_RULE(aar_import).ATTRIBUTE(srcjar) -->
+        A JAR file that contains source code for the compiled JAR files in the AAR.
+        <!-- #END_BLAZE_RULE.ATTRIBUTE --> */
+        .add(
+            attr("srcjar", LABEL)
+                .allowedFileTypes(JavaSemantics.SOURCE_JAR, JavaSemantics.JAR)
+                .direct_compile_time_input())
         .add(
             attr(AAR_EMBEDDED_JARS_EXTACTOR, LABEL)
                 .cfg(HostTransition.createFactory())
diff --git a/src/test/java/com/google/devtools/build/lib/rules/android/AarImportTest.java b/src/test/java/com/google/devtools/build/lib/rules/android/AarImportTest.java
index 6b2b459..a8a9191 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/android/AarImportTest.java
+++ b/src/test/java/com/google/devtools/build/lib/rules/android/AarImportTest.java
@@ -37,6 +37,8 @@
 import com.google.devtools.build.lib.rules.java.JavaInfo;
 import com.google.devtools.build.lib.rules.java.JavaRuleOutputJarsProvider;
 import com.google.devtools.build.lib.rules.java.JavaRuleOutputJarsProvider.OutputJar;
+import com.google.devtools.build.lib.rules.java.JavaSourceInfoProvider;
+import com.google.devtools.build.lib.rules.java.JavaSourceJarsProvider;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Set;
@@ -74,17 +76,27 @@
         ")");
     scratch.file(
         "a/BUILD",
+        "java_import(",
+        "    name = 'foo_src',",
+        "    jars = ['foo-src.jar'],",
+        ")",
         "aar_import(",
         "    name = 'foo',",
         "    aar = 'foo.aar',",
+        "    srcjar = ':foo_src',",
         ")",
         "aar_import(",
         "    name = 'baz',",
         "    aar = 'baz.aar',",
         ")",
+        "java_import(",
+        "    name = 'bar_src',",
+        "    jars = ['bar-src.jar'],",
+        ")",
         "aar_import(",
         "    name = 'bar',",
         "    aar = 'bar.aar',",
+        "    srcjar = ':bar_src',",
         "    deps = [':baz'],",
         "    exports = [':foo', '//java:baz'],",
         ")",
@@ -185,6 +197,43 @@
   }
 
   @Test
+  public void testSourceJarsProvided() throws Exception {
+    ConfiguredTarget aarImportTarget = getConfiguredTarget("//a:foo");
+
+    Iterable<Artifact> srcJars =
+        JavaInfo.getProvider(JavaSourceJarsProvider.class, aarImportTarget).getSourceJars();
+    assertThat(srcJars).hasSize(1);
+    Artifact srcJar = Iterables.getOnlyElement(srcJars);
+    assertThat(srcJar.getExecPathString()).endsWith("foo-src.jar");
+
+    Iterable<Artifact> srcInfoJars =
+        JavaInfo.getProvider(JavaSourceInfoProvider.class, aarImportTarget)
+            .getSourceJarsForJarFiles();
+    assertThat(srcInfoJars).hasSize(1);
+    Artifact srcInfoJar = Iterables.getOnlyElement(srcInfoJars);
+    assertThat(srcInfoJar.getExecPathString()).endsWith("foo-src.jar");
+  }
+
+  @Test
+  public void testSourceJarsCollectedTransitively() throws Exception {
+    ConfiguredTarget aarImportTarget = getConfiguredTarget("//a:bar");
+
+    Iterable<Artifact> srcJars =
+        JavaInfo.getProvider(JavaSourceJarsProvider.class, aarImportTarget)
+            .getTransitiveSourceJars();
+    assertThat(srcJars).hasSize(2);
+    assertThat(ActionsTestUtil.baseArtifactNames(srcJars))
+        .containsExactly("foo-src.jar", "bar-src.jar");
+
+    Iterable<Artifact> srcInfoJars =
+        JavaInfo.getProvider(JavaSourceInfoProvider.class, aarImportTarget)
+            .getSourceJarsForJarFiles();
+    assertThat(srcInfoJars).hasSize(1);
+    Artifact srcInfoJar = Iterables.getOnlyElement(srcInfoJars);
+    assertThat(srcInfoJar.getExecPathString()).endsWith("bar-src.jar");
+  }
+
+  @Test
   public void testResourcesExtractor() throws Exception {
     ValidatedAndroidResources resourceContainer =
         getConfiguredTarget("//a:foo")