Handle containerized annotations in Compose Preview

MultiPreviews imported from a library will put repeated annotations
of a given class (e.g. Previews) inside a container.

These change makes the Preview finding logic aware of such cases.

Bug: 279162835
Test: AnnotationFileComposePreviewElementFinderGradleTest
Change-Id: Ifb68a25cf93ec48c630cefae2095ad6e29381edc
(cherry picked from commit fc25a4345412d54cb8d8d46dc2ac89d693dacdac)

5 files changed, 60 insertions, 16 deletions

diff --git a/compose-designer/src/com/android/tools/idea/compose/preview/PreviewElementFinderUtils.kt b/compose-designer/src/com/android/tools/idea/compose/preview/PreviewElementFinderUtils.kt
index b6d4f4fc2c5..c226aeae09c 100644
--- a/compose-designer/src/com/android/tools/idea/compose/preview/PreviewElementFinderUtils.kt
+++ b/compose-designer/src/com/android/tools/idea/compose/preview/PreviewElementFinderUtils.kt
@@ -20,6 +20,7 @@ import com.android.tools.compose.COMPOSE_PREVIEW_ANNOTATION_FQN
 import com.android.tools.compose.COMPOSE_PREVIEW_ANNOTATION_NAME
 import com.android.tools.compose.COMPOSE_PREVIEW_PARAMETER_ANNOTATION_FQN
 import com.android.tools.idea.annotations.getContainingUMethodAnnotatedWith
+import com.android.tools.idea.annotations.getUAnnotations
 import com.android.tools.idea.annotations.isAnnotatedWith
 import com.android.tools.idea.compose.preview.analytics.MultiPreviewNode
 import com.android.tools.idea.compose.preview.analytics.MultiPreviewNodeImpl
@@ -34,7 +35,6 @@ import com.intellij.openapi.application.ReadAction
 import com.intellij.openapi.application.runReadAction
 import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiLiteralExpression
-import com.intellij.psi.PsiModifierListOwner
 import com.intellij.util.containers.sequenceOfNotNull
 import com.intellij.util.text.nullize
 import org.jetbrains.kotlin.asJava.classes.KtLightClass
@@ -46,7 +46,6 @@ import org.jetbrains.uast.UClassLiteralExpression
 import org.jetbrains.uast.UExpression
 import org.jetbrains.uast.UMethod
 import org.jetbrains.uast.UParameter
-import org.jetbrains.uast.toUElementOfType
 import org.jetbrains.uast.tryResolve
 
 /**
@@ -200,11 +199,7 @@ private fun UAnnotation.getPreviewNodes(
   visitedAnnotationClasses[annotationClassFqcn] =
     null // The MultiPreviewNodeInfo will be set later if needed
   val curAnnotationName = (this.tryResolve() as PsiClass).name
-  val annotations =
-    (this.tryResolve() as? PsiModifierListOwner)?.annotations?.mapNotNull {
-      it.toUElementOfType() as? UAnnotation
-    }
-      ?: return@runReadAction emptySequence()
+  val annotations = this.getUAnnotations()
 
   val nDirectPreviews = annotations.count { it.isPreviewAnnotation() }
   var nxtDirectPreviewId = 1
diff --git a/compose-designer/testData/classloader/multiPreviewTestLibrary/MultiPreviewAnnotations.kt b/compose-designer/testData/classloader/multiPreviewTestLibrary/MultiPreviewAnnotations.kt
index 089fe54b583..cc6dd83d912 100644
--- a/compose-designer/testData/classloader/multiPreviewTestLibrary/MultiPreviewAnnotations.kt
+++ b/compose-designer/testData/classloader/multiPreviewTestLibrary/MultiPreviewAnnotations.kt
@@ -9,4 +9,21 @@ import androidx.compose.ui.tooling.preview.Preview
   backgroundColor = 0xFF00FF00,
   device = "id:pixel_5",
 )
+@MultiPreviewWith2Previews
 annotation class MyMultiPreviewFromMyLibrary
+
+@Preview(
+  name = "preview 1",
+  group = "group 1-2",
+  showBackground = true,
+  backgroundColor = 0xFF00FF01,
+  device = "id:pixel_5",
+)
+@Preview(
+  name = "preview 2",
+  group = "group 1-2",
+  showBackground = true,
+  backgroundColor = 0xFF00FF02,
+  device = "id:pixel_5",
+)
+annotation class MultiPreviewWith2Previews
diff --git a/compose-designer/testData/classloader/multiPreviewTestLibrary/multiPreviewTestLibrary.aar b/compose-designer/testData/classloader/multiPreviewTestLibrary/multiPreviewTestLibrary.aar
index 7dd93c1f339..c8fd62fa28c 100644
--- a/compose-designer/testData/classloader/multiPreviewTestLibrary/multiPreviewTestLibrary.aar
+++ b/compose-designer/testData/classloader/multiPreviewTestLibrary/multiPreviewTestLibrary.aar
diff --git a/compose-designer/testSrc/com/android/tools/idea/compose/gradle/preview/AnnotationFileComposePreviewElementFinderGradleTest.kt b/compose-designer/testSrc/com/android/tools/idea/compose/gradle/preview/AnnotationFileComposePreviewElementFinderGradleTest.kt
index 9f8e55d756e..371695103ab 100644
--- a/compose-designer/testSrc/com/android/tools/idea/compose/gradle/preview/AnnotationFileComposePreviewElementFinderGradleTest.kt
+++ b/compose-designer/testSrc/com/android/tools/idea/compose/gradle/preview/AnnotationFileComposePreviewElementFinderGradleTest.kt
@@ -103,13 +103,29 @@ class AnnotationFileComposePreviewElementFinderGradleTest {
     projectRule.requestSyncAndWait()
     projectRule.buildAndAssertIsSuccessful()
 
-    val previewElement = runBlocking {
-      AnnotationFilePreviewElementFinder.findPreviewMethods(project, mainFile.virtualFile).first()
-    }
-    assertEquals("MyNewTestFun - test name", previewElement.displaySettings.name)
-    assertEquals("test group", previewElement.displaySettings.group)
-    assertTrue(previewElement.displaySettings.showBackground)
-    assertEquals("#FF00FF00", previewElement.displaySettings.backgroundColor!!.uppercase())
-    assertEquals("id:pixel_5", previewElement.configuration.deviceSpec)
+    val previewElements =
+      ArrayList(
+        runBlocking {
+          AnnotationFilePreviewElementFinder.findPreviewMethods(project, mainFile.virtualFile)
+        }
+      )
+
+    assertEquals("MyNewTestFun - preview 1", previewElements[0].displaySettings.name)
+    assertEquals("group 1-2", previewElements[0].displaySettings.group)
+    assertTrue(previewElements[0].displaySettings.showBackground)
+    assertEquals("#FF00FF01", previewElements[0].displaySettings.backgroundColor!!.uppercase())
+    assertEquals("id:pixel_5", previewElements[0].configuration.deviceSpec)
+
+    assertEquals("MyNewTestFun - preview 2", previewElements[1].displaySettings.name)
+    assertEquals("group 1-2", previewElements[1].displaySettings.group)
+    assertTrue(previewElements[1].displaySettings.showBackground)
+    assertEquals("#FF00FF02", previewElements[1].displaySettings.backgroundColor!!.uppercase())
+    assertEquals("id:pixel_5", previewElements[1].configuration.deviceSpec)
+
+    assertEquals("MyNewTestFun - test name", previewElements[2].displaySettings.name)
+    assertEquals("test group", previewElements[2].displaySettings.group)
+    assertTrue(previewElements[2].displaySettings.showBackground)
+    assertEquals("#FF00FF00", previewElements[2].displaySettings.backgroundColor!!.uppercase())
+    assertEquals("id:pixel_5", previewElements[2].configuration.deviceSpec)
   }
 }
diff --git a/designer/src/com/android/tools/idea/annotations/AnnotationsGraph.kt b/designer/src/com/android/tools/idea/annotations/AnnotationsGraph.kt
index 3bd7bdbffc0..ca4c4b9092d 100644
--- a/designer/src/com/android/tools/idea/annotations/AnnotationsGraph.kt
+++ b/designer/src/com/android/tools/idea/annotations/AnnotationsGraph.kt
@@ -18,9 +18,11 @@ package com.android.tools.idea.annotations
 import com.intellij.openapi.application.runReadAction
 import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiModifierListOwner
+import org.jetbrains.kotlin.utils.ifEmpty
 import org.jetbrains.uast.UAnnotation
 import org.jetbrains.uast.UElement
 import org.jetbrains.uast.UMethod
+import org.jetbrains.uast.toUElement
 import org.jetbrains.uast.toUElementOfType
 import org.jetbrains.uast.tryResolve
 
@@ -176,7 +178,21 @@ class AnnotationsGraph<S, T>(private val nodeInfoFactory: NodeInfoFactory<S>,
  * not guaranteed that it will work with other types of elements.
  */
 fun UElement.getUAnnotations() = runReadAction {
-  (this as? UMethod)?.uAnnotations
+  val annotations = (this as? UMethod)?.uAnnotations
   ?: (this.tryResolve() as? PsiModifierListOwner)?.annotations?.mapNotNull { it.toUElementOfType() as? UAnnotation }
   ?: emptyList()
+  annotations.flatMap { annotation ->
+    annotation.extractFromContainer().ifEmpty { listOf(annotation) }
+  }
+}
+
+/**
+ * MultiPreviews imported from a library will put repeated annotations of a given class (e.g. Previews)
+ * inside a container (more info at https://kotlinlang.org/docs/annotations.html#repeatable-annotations).
+ * This method extracts all annotations of a given container to have a list of individual annotations.
+ *
+ * When the annotation is not a container it returns an empty list.
+ */
+private fun UAnnotation.extractFromContainer() = runReadAction {
+  findDeclaredAttributeValue(null)?.sourcePsi?.children?.mapNotNull { it.toUElement() as? UAnnotation } ?: emptyList()
 }