name.remal.gradle_plugins.dsl.utils.ClassDependenciesCollector.kt Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of gradle-plugins-kotlin-dsl Show documentation
Remal Gradle plugins: gradle-plugins-kotlin-dsl
There is a newer version: 1.9.2
Show newest version
package name.remal.gradle_plugins.dsl.utils

import name.remal.accept
import name.remal.buildSet
import name.remal.concurrentMapOf
import org.objectweb.asm.*
import org.objectweb.asm.commons.ClassRemapper
import org.objectweb.asm.commons.FieldRemapper
import org.objectweb.asm.commons.MethodRemapper
import org.objectweb.asm.commons.Remapper
import org.objectweb.asm.tree.ClassNode
import java.io.Closeable
import java.util.concurrent.ForkJoinPool
import java.util.concurrent.RecursiveAction
import javax.annotation.concurrent.ThreadSafe

@ThreadSafe
class ClassDependenciesCollector(
    private val dependenciesFilter: ((dependencyClassName: String) -> Boolean)? = null,
    private val bytecodeRetriever: (className: String) -> ByteArray?
) : Closeable {

    companion object {
        private fun String.fromInternalName() = replace('/', '.')
        private fun String.toInternalName() = replace('.', '/')
        private val logger = getGradleLogger(ClassDependenciesCollector::class.java)
    }


    constructor(bytecodeRetriever: (className: String) -> ByteArray?) : this(null, bytecodeRetriever)


    private val cache = concurrentMapOf>()

    fun getDependencies(bytecode: ByteArray): Set {
        val classReader = ClassReader(bytecode);
        val classInternalName = classReader.className
        cache[classInternalName]?.let { return it }

        val dependencies = mutableSetOf()
        val remapper = object : Remapper() {
            override fun map(internalName: String): String {
                if (internalName != classInternalName) {
                    if (dependenciesFilter == null || dependenciesFilter.invoke(internalName.fromInternalName())) {
                        dependencies.add(internalName)
                    }
                }
                return internalName
            }
        }
        classReader.accept(CustomClassRemapper(remapper, ClassNode(), false))

        dependencies.mapTo(mutableSetOf(), { it.fromInternalName() }).let {
            logger.debug("{} depends on {}", classInternalName, it)
            cache[classInternalName] = it
            return it
        }
    }

    private class CustomClassRemapper(remapper: Remapper, classVisitor: ClassVisitor?, private val includeAnnotations: Boolean) : ClassRemapper(classVisitor, remapper) {

        override fun visitAnnotation(descriptor: String, visible: Boolean) = if (includeAnnotations) super.visitAnnotation(descriptor, visible) else null
        override fun visitTypeAnnotation(typeRef: Int, typePath: TypePath?, descriptor: String?, visible: Boolean) = if (includeAnnotations) super.visitTypeAnnotation(typeRef, typePath, descriptor, visible) else null

        override fun createFieldRemapper(fieldVisitor: FieldVisitor?): FieldVisitor? {
            return object : FieldRemapper(fieldVisitor, remapper) {
                override fun visitAnnotation(descriptor: String, visible: Boolean) = if (includeAnnotations) super.visitAnnotation(descriptor, visible) else null
                override fun visitTypeAnnotation(typeRef: Int, typePath: TypePath?, descriptor: String?, visible: Boolean) = if (includeAnnotations) super.visitTypeAnnotation(typeRef, typePath, descriptor, visible) else null
            }
        }

        override fun createMethodRemapper(methodVisitor: MethodVisitor?): MethodVisitor? {
            return object : MethodRemapper(methodVisitor, remapper) {
                override fun visitAnnotation(descriptor: String, visible: Boolean) = if (includeAnnotations) super.visitAnnotation(descriptor, visible) else null
                override fun visitTypeAnnotation(typeRef: Int, typePath: TypePath?, descriptor: String?, visible: Boolean) = if (includeAnnotations) super.visitTypeAnnotation(typeRef, typePath, descriptor, visible) else null
                override fun visitTryCatchAnnotation(typeRef: Int, typePath: TypePath?, descriptor: String?, visible: Boolean) = if (includeAnnotations) super.visitTryCatchAnnotation(typeRef, typePath, descriptor, visible) else null
                override fun visitInsnAnnotation(typeRef: Int, typePath: TypePath?, descriptor: String?, visible: Boolean) = if (includeAnnotations) super.visitInsnAnnotation(typeRef, typePath, descriptor, visible) else null
                override fun visitParameterAnnotation(parameter: Int, descriptor: String?, visible: Boolean) = if (includeAnnotations) super.visitParameterAnnotation(parameter, descriptor, visible) else null
                override fun visitLocalVariableAnnotation(typeRef: Int, typePath: TypePath?, start: Array?, end: Array?, index: IntArray?, descriptor: String?, visible: Boolean) = if (includeAnnotations) {
                    super.visitLocalVariableAnnotation(typeRef, typePath, start, end, index, descriptor, visible)
                } else {
                    null
                }
            }
        }

    }


    fun getDependencies(className: String): Set? {
        cache[className.toInternalName()]?.let { return it }
        val bytecode = bytecodeRetriever(className) ?: return null
        return getDependencies(bytecode)
    }


    private val forkJoinPool = ForkJoinPool()

    fun getAllDependenciesMap(className: String): Map> {
        val allDependenciesMap = concurrentMapOf>()

        class GetAllDependenciesAction(private val className: String) : RecursiveAction() {
            override fun compute() {
                val deps = getDependencies(this.className) ?: return
                val newDeps = deps.filter {
                    allDependenciesMap.computeIfAbsent(this.className, { mutableSetOf() }).add(it)
                }
                val actions = newDeps.map { GetAllDependenciesAction(it) }
                actions.forEach { it.fork() }
                actions.forEach { it.join() }
            }
        }
        forkJoinPool.submit(GetAllDependenciesAction(className)).get()

        return allDependenciesMap.toMap()
    }

    fun getAllDependencies(className: String) = buildSet {
        getAllDependenciesMap(className).values.forEach { addAll(it) }
    }


    override fun close() {
        forkJoinPool.shutdownNow()
    }

}