room/compiler/src/main/kotlin/androidx/room/ext/processing_env_ext.kt

/*
 * Copyright 2020 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package androidx.room.ext

import com.google.auto.common.GeneratedAnnotations
import com.squareup.javapoet.TypeName
import javax.annotation.processing.ProcessingEnvironment
import javax.lang.model.element.TypeElement
import javax.lang.model.type.TypeMirror
import kotlin.reflect.KClass

/**
 * Query a type element by KClass and return null if it does not exist
 */
fun ProcessingEnvironment.findTypeElement(
    klass: KClass<*>
): TypeElement? = findTypeElement(klass.java.canonicalName!!)

/**
 * Query a type element by TypeName and return null if it does not exist
 */
fun ProcessingEnvironment.findTypeElement(
    typeName: TypeName
): TypeElement? = findTypeElement(typeName.toString())

/**
 * Query a type element by qualified name and return null if it does not exist
 */
fun ProcessingEnvironment.findTypeElement(
    qName: String
): TypeElement? = elementUtils.getTypeElement(qName)

/**
 * Query a type element by KClass and throw if it does not exist
 */
fun ProcessingEnvironment.requireTypeElement(
    klass: KClass<*>
): TypeElement = requireTypeElement(klass.java.canonicalName!!)

/**
 * Query a type element by TypeName and throw if it does not exist
 */
fun ProcessingEnvironment.requireTypeElement(
    typeName: TypeName
): TypeElement = requireTypeElement(typeName.toString())

/**
 * Query a type element by qualified name and throw if it does not exist
 */
fun ProcessingEnvironment.requireTypeElement(
    qName: String
): TypeElement = checkNotNull(elementUtils.getTypeElement(qName)) {
    // we do not throw MissingTypeException here as this should be called only if the type should
    // be there
    "Couldn't find required type $qName"
}

/**
 * Query a TypeMirror by KClass and throw if it does not exist
 */
fun ProcessingEnvironment.requireTypeMirror(
    klass: KClass<*>
): TypeMirror = requireTypeMirror(klass.java.canonicalName!!)

/**
 * Query a TypeMirror by TypeName and throw if it does not exist
 */
fun ProcessingEnvironment.requireTypeMirror(
    typeName: TypeName
): TypeMirror = requireTypeMirror(typeName.toString())

/**
 * Query a TypeMirror by qualified name and throw if it does not exist
 */
fun ProcessingEnvironment.requireTypeMirror(
    qName: String
): TypeMirror = requireTypeElement(qName).asType()

/**
 * Query a type mirror by KClass and return null if it does not exist
 */
fun ProcessingEnvironment.findTypeMirror(
    klass: KClass<*>
): TypeMirror? = findTypeMirror(klass.java.canonicalName!!)

/**
 * Query a type mirror by TypeName and return null if it does not exist
 */
fun ProcessingEnvironment.findTypeMirror(
    typeName: TypeName
): TypeMirror? = findTypeMirror(typeName.toString())

/**
 * Query a type mirror by qualified name and return null if it does not exist
 */
fun ProcessingEnvironment.findTypeMirror(
    qName: String
): TypeMirror? = findTypeElement(qName)?.asType()

fun ProcessingEnvironment.getGeneratedAnnotation(): TypeElement? {
    val element = GeneratedAnnotations.generatedAnnotation(elementUtils, sourceVersion)
    return if (element.isPresent) {
        element.get()
    } else {
        null
    }
}