Fix litho model compilation (#342)

epoxy-processor/src/main/java/com/airbnb/epoxy/KotlinUtils.kt

@@ -1,18 +1,47 @@
 package com.airbnb.epoxy
 
+import com.airbnb.epoxy.Utils.*
 import com.squareup.javapoet.*
 import javax.lang.model.*
 import javax.lang.model.element.*
 import javax.lang.model.type.*
 import javax.lang.model.util.*
 import kotlin.reflect.*
 
+fun getTypeMirror(
+        className: ClassName,
+        elements: Elements,
+        types: Types
+): TypeMirror {
+    val classElement = getElementByName(className, elements, types)
+            ?: throw IllegalArgumentException("Unknown class: " + className)
+
+    return classElement.asType()
+}
+
+fun getTypeMirror(
+        clazz: Class<*>,
+        elements: Elements
+): TypeMirror? = getTypeMirror(clazz.canonicalName, elements)
+
+fun getTypeMirror(
+        canonicalName: String,
+        elements: Elements
+): TypeMirror? {
+    return try {
+        elements.getTypeElement(canonicalName)?.asType()
+    } catch (mte: MirroredTypeException) {
+        mte.typeMirror
+    }
+
+}
+
 fun Class<*>.asTypeElement(
         elements: Elements,
         types: Types
-): TypeElement {
-    val typeMirror = Utils.getTypeMirror(this, elements)
-    return types.asElement(typeMirror) as TypeElement
+): TypeElement? {
+    val typeMirror = getTypeMirror(this, elements) ?: return null
+    return types.asElement(typeMirror) as? TypeElement
 }
 
 fun KClass<*>.asTypeElement(
@@ -60,9 +89,13 @@ fun AnnotatedConstruct.hasAnyAnnotation(annotationClasses: List<Class<out Annota
     }
 }
 
+
 fun AnnotatedConstruct.hasAnnotation(annotationClass: KClass<out Annotation>)
         = hasAnyAnnotation(listOf(annotationClass.java))
 
+fun AnnotatedConstruct.hasAnnotation(annotationClass: Class<out Annotation>)
+        = hasAnyAnnotation(listOf(annotationClass))
+
 inline fun <reified T : Annotation> AnnotatedConstruct.annotation(): T?
         = getAnnotation(T::class.java)
 

epoxy-processor/src/main/java/com/airbnb/epoxy/LithoSpecProcessor.kt

@@ -0,0 +1,108 @@
+package com.airbnb.epoxy
+
+import com.airbnb.epoxy.ClassNames.*
+import com.airbnb.epoxy.GeneratedModelWriter.*
+import com.airbnb.epoxy.Utils.*
+import com.squareup.javapoet.*
+import com.squareup.javapoet.TypeSpec.*
+import java.util.*
+import javax.annotation.processing.*
+import javax.lang.model.element.*
+import javax.lang.model.util.*
+
+internal class LithoSpecProcessor(
+        private val elementUtils: Elements,
+        private val typeUtils: Types,
+        private val configManager: ConfigManager,
+        private val errorLogger: ErrorLogger,
+        private val modelWriter: GeneratedModelWriter
+) {
+    private var layoutSpecAnnotationClass: Class<out Annotation>? = null
+
+    fun processSpecs(roundEnv: RoundEnvironment): Collection<LithoModelInfo> {
+
+        if (!hasLithoEpoxyDependency()) {
+            // If the epoxy-litho module has not been included then we don't have access to the Epoxy
+            // litho model and can't build a model that extends it
+            return emptyList()
+        }
+
+        layoutSpecAnnotationClass = getAnnotationClass(ClassNames.LITHO_ANNOTATION_LAYOUT_SPEC)
+        if (layoutSpecAnnotationClass == null) {
+            // There is no dependency on Litho so there aren't any litho components to check for
+            return emptyList()
+        }
+
+        val modelInfoMap = LinkedHashMap<TypeElement, LithoModelInfo>()
+        roundEnv.getElementsAnnotatedWith(layoutSpecAnnotationClass)
+                .filterIsInstance<TypeElement>()
+                .forEach { modelInfoMap.put(it, LithoModelInfo(typeUtils, elementUtils, it)) }
+
+        val propClass = getAnnotationClass(ClassNames.LITHO_ANNOTATION_PROP)
+        val hashCodeValidator = HashCodeValidator(typeUtils, elementUtils)
+        for (propElement in roundEnv.getElementsAnnotatedWith(propClass)) {
+            val lithoModelInfo = getModelInfoForProp(modelInfoMap, propElement)
+            lithoModelInfo?.addProp(propElement, hashCodeValidator)
+        }
+
+        for ((_, modelInfo) in modelInfoMap) {
+            try {
+                modelWriter.generateClassForModel(modelInfo, object : BuilderHooks() {
+                    override fun beforeFinalBuild(builder: Builder) {
+                        updateGeneratedClassForLithoComponent(modelInfo, builder)
+                    }
+                })
+            } catch (e: Exception) {
+                errorLogger.logError(e, "Error generating model classes")
+            }
+
+        }
+
+        return modelInfoMap.values
+    }
+
+    // Only true if the epoxy-litho module is included in dependencies
+    private fun hasLithoEpoxyDependency(): Boolean =
+            getTypeMirror(EPOXY_LITHO_MODEL.reflectionName(), elementUtils) != null
+
+    private fun updateGeneratedClassForLithoComponent(
+            modelInfo: LithoModelInfo,
+            classBuilder: Builder
+    ) {
+        // Adding the "buildComponent" method
+        val methodBuilder = MethodSpec.methodBuilder("buildComponent")
+                .addAnnotation(Override::class.java)
+                .addModifiers(Modifier.PROTECTED)
+                .returns(
+                        ParameterizedTypeName.get(ClassNames.LITHO_COMPONENT,
+                                                  modelInfo.lithoComponentName)
+                )
+                .addParameter(ClassNames.LITHO_COMPONENT_CONTEXT, "context")
+                .addCode("return \$T.create(context)", modelInfo.lithoComponentName)
+
+        for (attributeInfo in modelInfo.attributeInfo) {
+            methodBuilder.addCode(".\$L(\$L)", attributeInfo.getFieldName(),
+                                  attributeInfo.getFieldName())
+        }
+
+        methodBuilder.addStatement(".build()")
+
+        classBuilder.addMethod(methodBuilder.build())
+    }
+
+    private fun getModelInfoForProp(
+            modelInfoMap: Map<TypeElement, LithoModelInfo>,
+            propElement: Element
+    ): LithoModelInfo? {
+        val methodElement = propElement.enclosingElement ?: return null
+
+        val classElement = methodElement.enclosingElement
+
+        return if (classElement.getAnnotation(layoutSpecAnnotationClass) == null) {
+            null
+        } else {
+            modelInfoMap[classElement]
+        }
+
+    }
+}

epoxy-processor/src/main/java/com/airbnb/epoxy/ModelViewInfo.kt

@@ -136,7 +136,7 @@ internal class ModelViewInfo(
         }
 
         val typeMirror = bounds[0]
-        val viewType = Utils.getTypeMirror(ClassNames.ANDROID_VIEW, elements, typeUtils)
+        val viewType = getTypeMirror(ClassNames.ANDROID_VIEW, elements, typeUtils)
         return typeUtils.isAssignable(viewType, typeMirror) || typeUtils.isSubtype(typeMirror,
                                                                                    viewType)
     }

epoxy-processor/src/main/java/com/airbnb/epoxy/ParisStyleAttributeInfo.kt

@@ -32,7 +32,7 @@ internal class ParisStyleAttributeInfo(
         fieldName = PARIS_STYLE_ATTR_NAME
         modelName = modelInfo.generatedName.simpleName()
         modelPackageName = packageName
-        typeMirror = Utils.getTypeMirror(ClassNames.PARIS_STYLE, elements, types)
+        typeMirror = getTypeMirror(ClassNames.PARIS_STYLE, elements, types)
         styleBuilderClass = styleBuilderClassName
         ignoreRequireHashCode = true
         isGenerated = true

epoxy-processor/src/main/java/com/airbnb/epoxy/StyleWriter.kt

@@ -69,7 +69,7 @@ internal fun tryAddStyleBuilderAttribute(
                                               "StyleBuilder")
 
     val styleBuilderElement = try {
-        Utils.getTypeMirror(styleBuilderClassName, elements, types)
+        getTypeMirror(styleBuilderClassName, elements, types)
     } catch (e: IllegalArgumentException) {
         return false
     }

epoxy-processor/src/main/java/com/airbnb/epoxy/Utils.java

@@ -85,27 +85,6 @@ static Class<? extends Annotation> getAnnotationClass(ClassName name) {
     }
   }
 
-  static TypeMirror getTypeMirror(ClassName className, Elements elements, Types types) {
-    Element classElement = getElementByName(className, elements, types);
-    if (classElement == null) {
-      throw new IllegalArgumentException("Unknown class: " + className);
-    }
-
-    return classElement.asType();
-  }
-
-  static TypeMirror getTypeMirror(Class<?> clazz, Elements elements) {
-    return getTypeMirror(clazz.getCanonicalName(), elements);
-  }
-
-  static TypeMirror getTypeMirror(String canonicalName, Elements elements) {
-    try {
-      return elements.getTypeElement(canonicalName).asType();
-    } catch (MirroredTypeException mte) {
-      return mte.getTypeMirror();
-    }
-  }
-
   static Element getElementByName(ClassName name, Elements elements, Types types) {
     String canonicalName = name.reflectionName().replace("$", ".");
     return getElementByName(canonicalName, elements, types);
@@ -288,7 +267,8 @@ private static boolean areParamsTheSame(ExecutableElement method1, MethodSpec me
       ParameterSpec param2 = params2.get(i);
 
       TypeMirror param1Type = types.erasure(param1.asType());
-      TypeMirror param2Type = types.erasure(getTypeMirror(param2.type.toString(), elements));
+      TypeMirror param2Type =
+          types.erasure(KotlinUtilsKt.getTypeMirror(param2.type.toString(), elements));
 
       // If a param is a type variable then we don't need an exact type match, it just needs to
       // be assignable
@@ -432,7 +412,7 @@ static boolean isType(TypeMirror typeMirror, String otherType) {
   }
 
   static boolean isType(Elements elements, Types types, TypeMirror typeMirror, Class<?> clazz) {
-    TypeMirror classType = getTypeMirror(clazz, elements);
+    TypeMirror classType = KotlinUtilsKt.getTypeMirror(clazz, elements);
     return types.isSameType(typeMirror, classType);
   }