Refactoring: separate input / output models classes.

Previously transformation phase consumed and produced same
data classes, now inputs and outputs are described separately.
This refactoring enables support of parent classes declared in
jars.

bug:63474615
Test: refactoring.
Change-Id: Ibb91f4b2ed41afc86eb8752ca0d09192cd6835d2
diff --git a/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/input_collector.kt b/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/input_collector.kt
index 0c6f341..5183df5 100644
--- a/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/input_collector.kt
+++ b/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/input_collector.kt
@@ -16,8 +16,8 @@
 
 package android.arch.lifecycle
 
+import android.arch.lifecycle.model.EventMethod
 import android.arch.lifecycle.model.LifecycleObserverInfo
-import android.arch.lifecycle.model.StateMethod
 import com.google.auto.common.MoreElements
 import com.google.auto.common.MoreTypes
 import javax.annotation.processing.ProcessingEnvironment
@@ -44,7 +44,7 @@
             val method = MoreElements.asExecutable(elem)
             if (validator.validateClass(enclosingElement)
                     && validator.validateMethod(method, onState.value)) {
-                StateMethod(method, onState)
+                EventMethod(method, onState, MoreElements.asType(enclosingElement))
             } else {
                 null
             }
diff --git a/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/model/StateMethod.kt b/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/model/AdapterClass.kt
similarity index 78%
rename from lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/model/StateMethod.kt
rename to lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/model/AdapterClass.kt
index 77195b5..1e76fa8 100644
--- a/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/model/StateMethod.kt
+++ b/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/model/AdapterClass.kt
@@ -16,9 +16,9 @@
 
 package android.arch.lifecycle.model
 
-import android.arch.lifecycle.OnLifecycleEvent
 import javax.lang.model.element.ExecutableElement
 import javax.lang.model.element.TypeElement
 
-data class StateMethod(val method: ExecutableElement, val onLifecycleEvent: OnLifecycleEvent,
-                       val syntheticAccess: TypeElement? = null)
\ No newline at end of file
+data class AdapterClass(val type: TypeElement,
+                        val calls: List<EventMethodCall>,
+                        val syntheticMethods: Set<ExecutableElement>)
\ No newline at end of file
diff --git a/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/model/StateMethod.kt b/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/model/EventMethod.kt
similarity index 69%
copy from lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/model/StateMethod.kt
copy to lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/model/EventMethod.kt
index 77195b5..a3d4712 100644
--- a/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/model/StateMethod.kt
+++ b/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/model/EventMethod.kt
@@ -17,8 +17,15 @@
 package android.arch.lifecycle.model
 
 import android.arch.lifecycle.OnLifecycleEvent
+import android.arch.lifecycle.getPackageQName
 import javax.lang.model.element.ExecutableElement
 import javax.lang.model.element.TypeElement
 
-data class StateMethod(val method: ExecutableElement, val onLifecycleEvent: OnLifecycleEvent,
-                       val syntheticAccess: TypeElement? = null)
\ No newline at end of file
+data class EventMethod(val method: ExecutableElement,
+                       val onLifecycleEvent: OnLifecycleEvent,
+                       val type: TypeElement) {
+
+    fun packageName() = type.getPackageQName()
+}
+
+data class EventMethodCall(val method: EventMethod, val syntheticAccess: TypeElement? = null)
\ No newline at end of file
diff --git a/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/model/LifecycleObserverInfo.kt b/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/model/LifecycleObserverInfo.kt
index d7428f2..d8bc364 100644
--- a/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/model/LifecycleObserverInfo.kt
+++ b/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/model/LifecycleObserverInfo.kt
@@ -16,10 +16,8 @@
 
 package android.arch.lifecycle.model
 
-import javax.lang.model.element.ExecutableElement
 import javax.lang.model.element.TypeElement
 
 data class LifecycleObserverInfo(
         val type: TypeElement,
-        val methods: List<StateMethod>,
-        var syntheticMethods: MutableSet<ExecutableElement> = mutableSetOf())
\ No newline at end of file
+        val methods: List<EventMethod>)
\ No newline at end of file
diff --git a/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/transformation.kt b/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/transformation.kt
index 47013a4..66fabf7 100644
--- a/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/transformation.kt
+++ b/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/transformation.kt
@@ -16,12 +16,14 @@
 
 package android.arch.lifecycle
 
+import android.arch.lifecycle.model.AdapterClass
+import android.arch.lifecycle.model.EventMethod
+import android.arch.lifecycle.model.EventMethodCall
 import android.arch.lifecycle.model.LifecycleObserverInfo
-import android.arch.lifecycle.model.StateMethod
 import com.google.auto.common.MoreTypes
-import java.util.*
+import com.google.common.collect.HashMultimap
+import java.util.LinkedList
 import javax.annotation.processing.ProcessingEnvironment
-import javax.lang.model.element.ExecutableElement
 import javax.lang.model.element.TypeElement
 import javax.lang.model.type.NoType
 import javax.lang.model.type.TypeMirror
@@ -52,52 +54,34 @@
 }
 
 private fun mergeAndVerifyMethods(processingEnv: ProcessingEnvironment,
-                                  classMethods: List<StateMethod>,
-                                  parentMethods: List<StateMethod>): List<StateMethod> {
-    return parentMethods + classMethods.filter { (method, onLifecycleEvent) ->
-        val baseMethod = parentMethods.find { m ->
-            method.simpleName == m.method.simpleName
-                    && method.parameters.size == m.method.parameters.size
+                                  type: TypeElement,
+                                  classMethods: List<EventMethod>,
+                                  parentMethods: List<EventMethod>): List<EventMethod> {
+    // need to update parent methods like that because:
+    // 1. visibility can be expanded
+    // 2. we want to preserve order
+    val updatedParentMethods = parentMethods.map { parentMethod ->
+        val overrideMethod = classMethods.find { (method) ->
+            processingEnv.elementUtils.overrides(method, parentMethod.method, type)
         }
-        if (baseMethod != null
-                && baseMethod.onLifecycleEvent != onLifecycleEvent) {
-            processingEnv.messager.printMessage(Diagnostic.Kind.ERROR,
-                    ErrorMessages.INVALID_STATE_OVERRIDE_METHOD, method)
+        if (overrideMethod != null) {
+            if (overrideMethod.onLifecycleEvent != parentMethod.onLifecycleEvent) {
+                processingEnv.messager.printMessage(Diagnostic.Kind.ERROR,
+                        ErrorMessages.INVALID_STATE_OVERRIDE_METHOD, overrideMethod.method)
+            }
+            overrideMethod
+        } else {
+            parentMethod
         }
-        baseMethod == null
     }
-
+    return updatedParentMethods + classMethods.filterNot { updatedParentMethods.contains(it) }
 }
 
-fun transformToOutput(processingEnv: ProcessingEnvironment,
-                      world: Map<TypeElement, LifecycleObserverInfo>): List<LifecycleObserverInfo> {
-    val superObservers = world.mapValues { superObservers(world, it.value) }
-    val packagePrivateMethods = world.mapValues { observer ->
-        if (observer.value.type.kind.isInterface) {
-            emptyList()
-        } else {
-            observer.value.methods.filter {
-                it.method.isPackagePrivate() || it.method.isProtected()
-            }.map { it.method }
-        }
-    }
-
-    val ppMethodsToType = packagePrivateMethods.entries.fold(
-            mapOf<ExecutableElement, TypeElement>(), { map, entry ->
-        map + entry.value.associate { it to entry.key }
-    })
-
-    world.values.forEach {
-        val observers = superObservers[it.type]!!
-        val currentPackage = it.type.getPackageQName()
-        observers.filter { superObserver ->
-            superObserver.type.getPackageQName() != currentPackage
-                    && packagePrivateMethods[superObserver.type]!!.isNotEmpty()
-        }.forEach { it.syntheticMethods.addAll(packagePrivateMethods[it.type]!!) }
-    }
-
-
+fun flattenObservers(processingEnv: ProcessingEnvironment,
+                     world: Map<TypeElement, LifecycleObserverInfo>): List<LifecycleObserverInfo> {
     val flattened: MutableMap<LifecycleObserverInfo, LifecycleObserverInfo> = mutableMapOf()
+    val superObservers = world.mapValues { superObservers(world, it.value) }
+
     fun traverse(observer: LifecycleObserverInfo) {
         if (observer in flattened) {
             return
@@ -108,25 +92,43 @@
             return
         }
         observers.filter { it !in flattened }.forEach(::traverse)
-        val currentPackage = observer.type.getPackageQName()
         val methods = observers
-                .fold(emptyList<StateMethod>(), { list, parentObserver ->
-                    mergeAndVerifyMethods(processingEnv, parentObserver.methods, list)
-                })
-                .map {
-                    val packageName = ppMethodsToType[it.method]?.getPackageQName()
-                    if (packageName == null || packageName == currentPackage) {
-                        it
-                    } else {
-                        StateMethod(it.method, it.onLifecycleEvent, ppMethodsToType[it.method])
-                    }
+                .map(flattened::get)
+                .fold(emptyList<EventMethod>()) { list, parentObserver ->
+                    mergeAndVerifyMethods(processingEnv, observer.type, parentObserver!!.methods, list)
                 }
 
         flattened[observer] = LifecycleObserverInfo(observer.type,
-                mergeAndVerifyMethods(processingEnv, observer.methods, methods),
-                observer.syntheticMethods)
+                mergeAndVerifyMethods(processingEnv, observer.type, observer.methods, methods))
     }
 
     world.values.forEach(::traverse)
     return flattened.values.toList()
 }
+
+fun transformToOutput(processingEnv: ProcessingEnvironment,
+                      world: Map<TypeElement, LifecycleObserverInfo>): List<AdapterClass> {
+    val flatObservers = flattenObservers(processingEnv, world)
+    val syntheticMethods = HashMultimap.create<TypeElement, EventMethodCall>()
+    val adapterCalls = flatObservers.map { (type, methods) ->
+        val calls = methods.map { eventMethod ->
+            val executable = eventMethod.method
+            if (type.getPackageQName() != eventMethod.packageName()
+                    && (executable.isPackagePrivate() || executable.isProtected())) {
+                EventMethodCall(eventMethod, eventMethod.type)
+            } else {
+                EventMethodCall(eventMethod)
+            }
+        }
+        calls.filter { it.syntheticAccess != null }.forEach { eventMethod ->
+            syntheticMethods.put(eventMethod.method.type, eventMethod)
+        }
+        type to calls
+    }.toMap()
+
+    return adapterCalls.map { (type, calls) ->
+        val methods = syntheticMethods.get(type) ?: setOf()
+        val synthetic = methods.map { eventMethod -> eventMethod!!.method.method }.toSet()
+        AdapterClass(type, calls, synthetic)
+    }
+}
diff --git a/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/writer.kt b/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/writer.kt
index f971119..1cd20ed 100644
--- a/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/writer.kt
+++ b/lifecycle/compiler/src/main/kotlin/android/arch/lifecycle/writer.kt
@@ -16,16 +16,22 @@
 
 package android.arch.lifecycle
 
-import android.arch.lifecycle.model.LifecycleObserverInfo
-import android.arch.lifecycle.model.StateMethod
-import com.squareup.javapoet.*
+import android.arch.lifecycle.model.AdapterClass
+import android.arch.lifecycle.model.EventMethodCall
+import com.squareup.javapoet.ClassName
+import com.squareup.javapoet.FieldSpec
+import com.squareup.javapoet.JavaFile
+import com.squareup.javapoet.MethodSpec
+import com.squareup.javapoet.ParameterSpec
+import com.squareup.javapoet.TypeName
+import com.squareup.javapoet.TypeSpec
 import javax.annotation.processing.Filer
 import javax.lang.model.element.ExecutableElement
 import javax.lang.model.element.Modifier
 import javax.lang.model.element.TypeElement
 
-fun writeModels(infos: List<LifecycleObserverInfo>, filer: Filer) {
-    infos.forEach({ info -> writeAdapter(info, filer) })
+fun writeModels(infos: List<AdapterClass>, filer: Filer) {
+    infos.forEach({ adapter -> writeAdapter(adapter, filer) })
 }
 
 private val LIFECYCLE_OWNER = ClassName.get(LifecycleOwner::class.java)
@@ -35,11 +41,11 @@
 private val N = "\$N"
 private val L = "\$L"
 
-private fun writeAdapter(observer: LifecycleObserverInfo, filer: Filer) {
+private fun writeAdapter(adapter: AdapterClass, filer: Filer) {
     val ownerParam = ParameterSpec.builder(LIFECYCLE_OWNER, "owner").build()
     val eventParam = ParameterSpec.builder(ClassName.get(LIFECYCLE_EVENT), "event").build()
     val receiverName = "mReceiver"
-    val receiverField = FieldSpec.builder(ClassName.get(observer.type), receiverName,
+    val receiverField = FieldSpec.builder(ClassName.get(adapter.type), receiverName,
             Modifier.FINAL).build()
 
     val dispatchMethodBuilder = MethodSpec.methodBuilder("onStateChanged")
@@ -49,16 +55,16 @@
             .addModifiers(Modifier.PUBLIC)
             .addAnnotation(Override::class.java)
     val dispatchMethod = dispatchMethodBuilder.apply {
-        observer.methods
-                .groupBy { stateMethod -> stateMethod.onLifecycleEvent.value }
+        adapter.calls
+                .groupBy { (eventMethod) -> eventMethod.onLifecycleEvent.value }
                 .forEach { entry ->
                     val event = entry.key
-                    val methods = entry.value
+                    val calls = entry.value
                     if (event == Lifecycle.Event.ON_ANY) {
-                        writeMethodCalls(eventParam, methods, ownerParam, receiverField)
+                        writeMethodCalls(eventParam, calls, ownerParam, receiverField)
                     } else {
                         beginControlFlow("if ($N == $T.$L)", eventParam, LIFECYCLE_EVENT, event)
-                                .writeMethodCalls(eventParam, methods, ownerParam, receiverField)
+                                .writeMethodCalls(eventParam, calls, ownerParam, receiverField)
                         endControlFlow()
                     }
                 }
@@ -71,9 +77,10 @@
             .addStatement("return $N", receiverField)
             .build()
 
-    val receiverParam = ParameterSpec.builder(ClassName.get(observer.type), "receiver").build()
+    val receiverParam = ParameterSpec.builder(
+            ClassName.get(adapter.type), "receiver").build()
 
-    val syntheticMethods = observer.syntheticMethods.map {
+    val syntheticMethods = adapter.syntheticMethods.map {
         val method = MethodSpec.methodBuilder(syntheticName(it))
                 .returns(TypeName.VOID)
                 .addModifiers(Modifier.PUBLIC)
@@ -98,8 +105,8 @@
             .addStatement("this.$N = $N", receiverField, receiverParam)
             .build()
 
-    val adapterName = getAdapterName(observer.type)
-    val adapter = TypeSpec.classBuilder(adapterName)
+    val adapterName = getAdapterName(adapter.type)
+    val adapterTypeSpec = TypeSpec.classBuilder(adapterName)
             .addModifiers(Modifier.PUBLIC)
             .addSuperinterface(ClassName.get(GenericLifecycleObserver::class.java))
             .addField(receiverField)
@@ -108,24 +115,24 @@
             .addMethod(getWrappedMethod)
             .addMethods(syntheticMethods)
             .build()
-    JavaFile.builder(observer.type.getPackageQName(), adapter)
+    JavaFile.builder(adapter.type.getPackageQName(), adapterTypeSpec)
             .build().writeTo(filer)
 }
 
 private fun MethodSpec.Builder.writeMethodCalls(eventParam: ParameterSpec,
-                                                methods: List<StateMethod>,
+                                                calls: List<EventMethodCall>,
                                                 ownerParam: ParameterSpec,
                                                 receiverField: FieldSpec) {
-    methods.forEach { method ->
+    calls.forEach { (method, syntheticAccess) ->
         val count = method.method.parameters.size
-        if (method.syntheticAccess == null) {
+        if (syntheticAccess == null) {
             val paramString = generateParamString(count)
             addStatement("$N.$L($paramString)", receiverField,
                     method.method.name(),
                     *takeParams(count, ownerParam, eventParam))
 
         } else {
-            val originalType = method.syntheticAccess
+            val originalType = syntheticAccess
             val paramString = generateParamString(count + 1)
             val className = ClassName.get(originalType.getPackageQName(),
                     getAdapterName(originalType))