StrictDepanalyzer.kt

package io.bazel.kotlin.compiler.plugins.strictdeps

import com.google.protobuf.TextFormat
import com.intellij.openapi.project.Project
import com.intellij.openapi.vfs.VfsUtilCore
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.psi.PsiElement
import io.bazel.kotlin.compiler.BazelCompiler
import io.bazel.kotlin.model.JvmCompilationTask
import io.bazel.kotlin.model.StrictDeps
import org.jetbrains.kotlin.analyzer.AnalysisResult
import org.jetbrains.kotlin.config.CompilerConfiguration
import org.jetbrains.kotlin.container.StorageComponentContainer
import org.jetbrains.kotlin.container.useInstance
import org.jetbrains.kotlin.descriptors.ClassifierDescriptor
import org.jetbrains.kotlin.descriptors.FunctionDescriptor
import org.jetbrains.kotlin.descriptors.ModuleDescriptor
import org.jetbrains.kotlin.descriptors.SourceElement
import org.jetbrains.kotlin.diagnostics.reportFromPlugin
import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
import org.jetbrains.kotlin.load.java.sources.JavaSourceElement
import org.jetbrains.kotlin.load.java.structure.impl.classFiles.BinaryJavaClass
import org.jetbrains.kotlin.load.kotlin.*
import org.jetbrains.kotlin.psi.KtFile
import org.jetbrains.kotlin.resolve.BindingTrace
import org.jetbrains.kotlin.resolve.TargetPlatform
import org.jetbrains.kotlin.resolve.calls.checkers.CallChecker
import org.jetbrains.kotlin.resolve.calls.checkers.CallCheckerContext
import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
import org.jetbrains.kotlin.resolve.checkers.ClassifierUsageChecker
import org.jetbrains.kotlin.resolve.checkers.ClassifierUsageCheckerContext
import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisHandlerExtension
import org.jetbrains.kotlin.resolve.source.KotlinSourceElement
import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedTypeAliasDescriptor
import java.nio.file.Paths

class StrictDepAnalyzer(
    private val tracker: StrictDepTracker,
    private val task: JvmCompilationTask
) : AnalysisHandlerExtension {
    private lateinit var container: StorageComponentContainer
    private val classifierUsageChecker = object : ClassifierUsageChecker {
        override fun check(
            targetDescriptor: ClassifierDescriptor,
            element: PsiElement,
            context: ClassifierUsageCheckerContext
        ) {
            when (targetDescriptor) {
                is DeserializedTypeAliasDescriptor -> {
                    // Add an entry for a foreign Alias.
                    val aliasJarPath = targetDescriptor.containerSource
                        ?.let { it as JvmPackagePartSource }
                        ?.knownJvmBinaryClass?.containingJarFile?.path
                        ?.also { tracker.recordDependency(element, it, context.trace) }

                    // If the target of the alias is in a different jar add a dependency to it.
                    targetDescriptor.underlyingType.constructor.declarationDescriptor?.apply {
                        resolveContainingJarFile(element, context.trace)?.path
                            ?.takeIf { it != aliasJarPath }
                            ?.also { tracker.recordDependency(element, it, context.trace) }
                    }
                }
                else -> {
                    tracker.recordDependency(
                        element,
                        targetDescriptor.resolveContainingJarFile(element, context.trace)?.path ?: return,
                        context.trace)
                }
            }
        }
    }

    private val callChecker = object : CallChecker {
        override fun check(resolvedCall: ResolvedCall<*>, reportOn: PsiElement, context: CallCheckerContext) {
            val descriptor = resolvedCall.candidateDescriptor
            when (descriptor) {
                is FunctionDescriptor -> descriptor.getContainingKotlinJvmBinaryClass()?.also {
                    tracker.recordDependency(reportOn, it.containingJarFile.path, context.trace)
                }
            }
        }
    }

    private val analysisHandlerExtension = object : AnalysisHandlerExtension {
        override fun analysisCompleted(
            project: Project,
            module: ModuleDescriptor,
            bindingTrace: BindingTrace,
            files: Collection<KtFile>
        ): AnalysisResult? {
            val jdeps = files.map { it.packageFqName.asString() }.distinct().let {
                tracker.buildDeps(it)
            }
            System.err.println(TextFormat.printToString(jdeps))
            // TODO figure out how to throw fatal exception.
            // TODO assert file did not exist
            Paths.get(task.outputs.jdeps).toFile().outputStream().use {
                jdeps.writeTo(it)
            }
            return null
        }
    }

    private fun ClassifierDescriptor.resolveContainingJarFile(element: PsiElement, trace: BindingTrace): VirtualFile? {
        if (source == SourceElement.NO_SOURCE) {
            return null
        } else {
            val resolveContainingBinaryFile = source.resolveContainingBinaryFile(element, trace) ?: return null

            return VfsUtilCore.getVirtualFileForJar(resolveContainingBinaryFile).also {
                if (it == null) {
                    trace.reportFromPlugin(StrictDepsDiagnostics.UNRESOLVABLE_JAR.on(element), StrictDepsDiagnostics)
                }
            }
        }
    }

    private fun SourceElement.resolveContainingBinaryFile(element: PsiElement, trace: BindingTrace): VirtualFile? =
        when (this) {
            is KotlinJvmBinarySourceElement -> (binaryClass as VirtualFileKotlinClass).file
            is JavaSourceElement -> {
                // the upper bound of JavaSourceElement is VirtualFileBoundJavaClass this can be a local java source
                // file.
                javaElement.let { it as? BinaryJavaClass }?.virtualFile
            }
            is KotlinSourceElement -> null
            else -> {
                trace.reportFromPlugin(StrictDepsDiagnostics.UNRESOLVABLE_FILE.on(element), StrictDepsDiagnostics)
                null
            }
        }

    private val KotlinJvmBinaryClass.containingJarFile: VirtualFile
        get() = checkNotNull(VfsUtilCore.getVirtualFileForJar((this as VirtualFileKotlinClass).file))

    companion object {
        fun register(project: Project, configuration: CompilerConfiguration) {
            val task = configuration[BazelCompiler.Keys.JVM_TASK]?.let {
                when (it.info.strictDeps) {
                    StrictDeps.STRICT_DEPS_OFF -> null
                    StrictDeps.UNRECOGNIZED -> throw IllegalStateException("unrecognised ${it.info.strictDeps}")
                    else -> it
                }
            } ?: return

            val sd = StrictDepTracker(task)
            val analyzer = StrictDepAnalyzer(sd, task)
//            configuration.put(JVMConfigurationKeys.JAVA_CLASSES_TRACKER, object: org.jetbrains.kotlin.load.java.JavaClassesTracker {
//                override fun onCompletedAnalysis(module: ModuleDescriptor) {
//                    System.err.println("done")
//                }
//
//                override fun reportClass(classDescriptor: JavaClassDescriptor) {
//                    System.err.println("java descriptor $classDescriptor")
//                }
//            })
            StorageComponentContainerContributor.registerExtension(project,
                object : StorageComponentContainerContributor {
                    override fun registerModuleComponents(
                        container: StorageComponentContainer,
                        platform: TargetPlatform,
                        moduleDescriptor: ModuleDescriptor
                    ) {
                        analyzer.container = container
                        container.useInstance(analyzer.classifierUsageChecker)
                        container.useInstance(analyzer.callChecker)
                    }
                })
            AnalysisHandlerExtension.registerExtension(project, analyzer.analysisHandlerExtension)
        }
    }
}