main.GenericUsages.kt Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of meta Show documentation
"Annotation Processing and Code Generation Utils"
There is a newer version: 0.76.0
package de.peekandpoke.ultra.meta

import com.squareup.kotlinpoet.ClassName
import com.squareup.kotlinpoet.DelicateKotlinPoetApi
import com.squareup.kotlinpoet.ParameterizedTypeName
import com.squareup.kotlinpoet.asTypeName
import de.peekandpoke.ultra.meta.model.MType
import de.peekandpoke.ultra.meta.model.Model
import javax.lang.model.element.TypeElement
import javax.lang.model.element.VariableElement
import javax.lang.model.type.DeclaredType
import javax.lang.model.type.TypeMirror
import javax.lang.model.type.TypeVariable

class GenericUsages(val model: Model, types: List) : ProcessorUtils {

    override val ctx = model.ctx

    private val registry = mutableMapOf>()

    private val visitedDeclaredTypes = mutableSetOf()

    init {
        types.forEach { visit(it) }
    }

    fun get(cls: ClassName): Set = registry[cls] ?: setOf()

    private fun DeclaredType.isReified() = typeArguments.all { it is DeclaredType }

    @OptIn(DelicateKotlinPoetApi::class)
    private fun put(declaredType: DeclaredType) = put(
        MType(
            model = model,
            type = declaredType.asElement() as TypeElement,
            typeName = declaredType.asTypeName() as ParameterizedTypeName
        )
    )

    private fun put(mType: MType) {
        registry.getOrPut(mType.className) { mutableSetOf() }.add(mType)
    }

    private fun visit(type: TypeElement) {
        type.variables.forEach { visit(it) }
    }

    private fun visit(variable: VariableElement) {
        visit(variable.asType())
    }

    private fun visit(typeMirror: TypeMirror) {
        if (typeMirror is DeclaredType) {
            visit(typeMirror)
        }
    }

    private fun visit(declaredType: DeclaredType) {

        // avoid endless loops and blacklisted types
        if (visitedDeclaredTypes.contains(declaredType) ||
            declaredType.typeArguments.isEmpty() ||
            !declaredType.isReified()
        ) {
            return
        }

        // put this type into the visited types
        visitedDeclaredTypes.add(declaredType)

        // put this type into the registry
        put(declaredType)

        // visit all type arguments of the type
        declaredType.typeArguments.forEach { visit(it) }

        // visit the type
        val type = declaredType.asElement() as TypeElement

        // visit all variable that have reified type parameters already
        type.variables
            .filter { v -> v.asType().let { it is DeclaredType && it.isReified() } }
            .forEach { visit(it) }

        // We have to reify type parameters for some variables
        type.variables
            .map { it.asType() }
            .filterIsInstance()
            .filter { !it.isReified() }
            .forEach {
                visit(
                    it.reify(declaredType)
                )
            }
    }

    /**
     * Reifies the type variables of the receiver with type information from the [provider]
     */
    private fun DeclaredType.reify(provider: DeclaredType): TypeMirror {

        // visit the type
        val providerElement = provider.asElement() as TypeElement

        // Create a new declared type by mapping type variables to real types
        return typeUtils.getDeclaredType(
            // the variable element as a type element
            asElement() as TypeElement,
            // we try to map all type arguments to real types
            *typeArguments.map { typeArg ->
                when (typeArg) {
                    is TypeVariable -> {
                        // Get the index of the type arg
                        val idx = providerElement.typeParameters.indexOfFirst {
                            it == typeArg.asElement()
                        }
                        // Get the real type from the provider by the idx we found
                        when (idx) {
                            -1 -> typeArg
                            else -> provider.typeArguments[idx]
                        }
                    }
                    else -> typeArg
                }
            }.toTypedArray()
        )
    }
}