Allow precomputed annotations

This allows performance improvement in the incremental case. This
way annotations for the library dependencies can be reused until
they change and only project classes need recomputation.

Bug: N/A
Test: MultipreviewTest
Change-Id: I24913c7e75336b8926e46a573248fae6954e457e
diff --git a/multipreview-asm/src/com/android/tools/preview/multipreview/Graph.kt b/multipreview-asm/src/com/android/tools/preview/multipreview/Graph.kt
index dceb5cd..4f38eae 100644
--- a/multipreview-asm/src/com/android/tools/preview/multipreview/Graph.kt
+++ b/multipreview-asm/src/com/android/tools/preview/multipreview/Graph.kt
@@ -20,8 +20,15 @@
     return references
   }
 
+  fun addAnnotations(annotations: Map<DerivedAnnotationRepresentation, AnnotationReferences>) {
+    annotationNodes.putAll(annotations)
+  }
+
   override val methods: Set<MethodRepresentation> = methodNodes.keys
 
+  override val annotations: Map<DerivedAnnotationRepresentation, AnnotationReferences>
+        get() = annotationNodes
+
   /**
    * Cleans the graph off the redundant annotations (that are neither base nor derived annotations)
    * and methods (that are annotated by neither base nor derived annotations). This should be run
diff --git a/multipreview-asm/src/com/android/tools/preview/multipreview/Multipreview.kt b/multipreview-asm/src/com/android/tools/preview/multipreview/Multipreview.kt
index d3dc96a..3acee29 100644
--- a/multipreview-asm/src/com/android/tools/preview/multipreview/Multipreview.kt
+++ b/multipreview-asm/src/com/android/tools/preview/multipreview/Multipreview.kt
@@ -21,8 +21,14 @@
  * resolved instances of the base annotations.
  */
 interface Multipreview {
-
+  /** All the methods that are annotated by either base or derived annotations. */
   val methods: Set<MethodRepresentation>
 
+  /** All the derived annotations. */
+  val annotations: Map<DerivedAnnotationRepresentation, AnnotationReferences>
+
+  /**
+   * Get all the resolved annotations for the [method] as a set of parameterized base annotations.
+   */
   fun getAnnotations(method: MethodRepresentation): Set<BaseAnnotationRepresentation>
 }
diff --git a/multipreview-asm/src/com/android/tools/preview/multipreview/MultipreviewBuilder.kt b/multipreview-asm/src/com/android/tools/preview/multipreview/MultipreviewBuilder.kt
index 919ec16..17e7f9b 100644
--- a/multipreview-asm/src/com/android/tools/preview/multipreview/MultipreviewBuilder.kt
+++ b/multipreview-asm/src/com/android/tools/preview/multipreview/MultipreviewBuilder.kt
@@ -28,14 +28,17 @@
 
 /**
  * Builds multipreview structure based on the [settings] and bytecode data from [provider] for
- * methods allowed by [methodsFilter].
+ * methods allowed by [methodsFilter]. Allows specifying additional precomputed [annotations] to
+ * speed up incremental processing.
  */
 fun buildMultipreview(
   settings: MultipreviewSettings,
+  annotations: Map<DerivedAnnotationRepresentation, AnnotationReferences> = emptyMap(),
   methodsFilter: MethodsFilter = MethodsFilter { true },
   provider: ClassBytecodeProvider,
 ): Multipreview {
   val multipreviewGraph = Graph()
+  multipreviewGraph.addAnnotations(annotations)
   provider.forEachClass {
     val cr = ClassReader(it)
     val classVisitor = if ((cr.access and Opcodes.ACC_ANNOTATION) != 0) {
@@ -71,9 +74,10 @@
 fun buildMultipreview(
   settings: MultipreviewSettings,
   paths: Collection<String>,
+  annotations: Map<DerivedAnnotationRepresentation, AnnotationReferences> = emptyMap(),
   methodsFilter: MethodsFilter = MethodsFilter { true }
 ): Multipreview {
-  return buildMultipreview(settings, methodsFilter) { processor ->
+  return buildMultipreview(settings, annotations, methodsFilter) { processor ->
     forEachClass(paths, processor::onClassBytecode)
   }
 }
diff --git a/multipreview-asm/testSrc/com/android/tools/preview/multipreview/MultipreviewTest.kt b/multipreview-asm/testSrc/com/android/tools/preview/multipreview/MultipreviewTest.kt
index c41630f..5bd3035 100644
--- a/multipreview-asm/testSrc/com/android/tools/preview/multipreview/MultipreviewTest.kt
+++ b/multipreview-asm/testSrc/com/android/tools/preview/multipreview/MultipreviewTest.kt
@@ -157,7 +157,11 @@
     }
 
     run {
-      val multipreview = buildMultipreview(settings, { it.startsWith("$PKG.filtered") }) { processor ->
+      val multipreview = buildMultipreview(
+          settings,
+          emptyMap(),
+          { it.startsWith("$PKG.filtered") },
+      ) { processor ->
         classStreams.forEach { processor.onClassBytecode(it) }
       }
 
@@ -261,6 +265,50 @@
     }
   }
 
+  @Test
+  fun testPrecomputedAnnotations() {
+    val classStreams = listOf(
+      DerivedAnnotation1Lvl1::class.java,
+      Class.forName("$PKG.AnnotatedMethods1Kt"),
+    ).map { loadClassBytecode(it) }
+
+    val annotations = mapOf(
+      DerivedAnnotationRepresentation(DerivedAnnotation1Lvl1::class.java.name) to object : AnnotationReferences {
+          override val baseAnnotations: MutableList<BaseAnnotationRepresentation>
+              get() = mutableListOf(
+                  BaseAnnotationRepresentation(emptyMap()),
+                  BaseAnnotationRepresentation(mapOf("paramName1" to "foo"))
+              )
+          override val derivedAnnotations: MutableSet<DerivedAnnotationRepresentation>
+              get() = mutableSetOf()
+      }
+    )
+
+    run {
+      val multipreview = buildMultipreview(settings, annotations = annotations) { processor ->
+        classStreams.forEach { processor.onClassBytecode(it) }
+      }
+
+      val sortedMethods = multipreview.methods.sortedWith(MethodComparator).toTypedArray()
+
+      assertArrayEquals(
+        arrayOf(
+          method("$PKG.AnnotatedMethods1Kt.method1"),
+          method("$PKG.AnnotatedMethods1Kt.method2"),
+        ),
+        sortedMethods
+      )
+
+      assertArrayEquals(
+        arrayOf(
+          BaseAnnotationRepresentation(emptyMap()),
+          BaseAnnotationRepresentation(mapOf("paramName1" to "foo")),
+        ),
+        multipreview.getAnnotations(sortedMethods[1]).sortedWith(AnnotationComparator).toTypedArray()
+      )
+    }
+  }
+
   companion object {
 
     private val settings = MultipreviewSettings(BASE_ANNOTATION, PARAMETER_ANNOTATION)