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.ClassReader
import org.objectweb.asm.ClassVisitor
import org.objectweb.asm.FieldVisitor
import org.objectweb.asm.Label
import org.objectweb.asm.MethodVisitor
import org.objectweb.asm.TypePath
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()
    }

}