Inferred type detector

InferredTypeDetector.kt

package com.linecorp.lint.detectors

import com.android.tools.lint.client.api.UElementHandler
import com.android.tools.lint.detector.api.Category
import com.android.tools.lint.detector.api.Detector
import com.android.tools.lint.detector.api.Implementation
import com.android.tools.lint.detector.api.Issue
import com.android.tools.lint.detector.api.JavaContext
import com.android.tools.lint.detector.api.LintFix
import com.android.tools.lint.detector.api.Scope
import com.android.tools.lint.detector.api.Severity
import com.android.tools.lint.detector.api.SourceCodeScanner
import com.android.tools.lint.detector.api.isKotlin
import com.intellij.psi.GenericsUtil
import com.intellij.psi.PsiType
import com.intellij.psi.impl.source.PsiImmediateClassType
import com.intellij.psi.util.TypeConversionUtil
import org.jetbrains.kotlin.psi.KtNamedFunction
import org.jetbrains.kotlin.psi.KtProperty
import org.jetbrains.uast.UElement
import org.jetbrains.uast.UMethod

/**
 * A Detector to check if the compiler inferred the type of a function or a property in a given
 * file.
 */
class InferredTypeDetector : Detector(), SourceCodeScanner {

    companion object {
        @JvmField
        val ISSUE: Issue = Issue.create(
            "InferredType",
            "Inferred type",
            "This check highlights when you let the Kotlin compiler infer the variable type." +
                    "It is required to explicitly write the return type of a method by our code " +
                    "conventions.",
            Category.CORRECTNESS,
            6,
            Severity.WARNING,
            Implementation(InferredTypeDetector::class.java, Scope.JAVA_FILE_SCOPE)
        )
    }

    override fun getApplicableUastTypes(): MutableList<Class<out UElement>> =
        mutableListOf(UMethod::class.java)

    override fun createUastHandler(context: JavaContext): UElementHandler =
        InferredTypeHandler(context)

    /**
     * A [UElementHandler] to check whether the compiler inferred the function or property's type
     * in a given file. Visits every method of the file checking it's return type and reports to
     * the IDE accordingly. Provides a quick fix to resolve the problem.
     */
    class InferredTypeHandler(private val context: JavaContext) : UElementHandler() {

        /**
         * Reports to the IDE, provides a quick fix to resolve the problem.
         */
        private fun JavaContext.reportInferredType(
            node: UMethod,
            presentableText: String,
            canonicalText: String,
            inferredKind: InferredKind
        ) {
            // For primitive types we have to capitalize the type.
            val presentableTextCapitalized = presentableText.capitalize().replace("Integer", "Int")
            val fix: LintFix? = getFix(node, presentableTextCapitalized, inferredKind)
            report(
                ISSUE,
                node,
                getNameLocation(node),
                getReportMessage(canonicalText),
                fix
            )
        }

        /**
         * Creates an appropriate fix depending on the node's inferred kind (method or property)
         */
        private fun JavaContext.getFix(
            node: UMethod,
            inferredType: String,
            inferredKind: InferredKind
        ): LintFix? {
            val fixString = when (inferredKind) {
                InferredKind.METHOD -> {
                    val regex = "\\)\\s*=".toRegex()
                    regex.replaceFirst(node.text, "):$inferredType = ")
                }
                InferredKind.PROPERTY -> {
                    node.text.replaceFirst("=", ":$inferredType = ")
                }
                else -> null
            } ?: return null
            return LintFix.create().replace()
                .name("Add inferred type")
                .range(getLocation(node))
                .with(fixString)
                .reformat(true)
                .build()
        }

        private fun getReportMessage(inferredType: String) =
            "Explicit type declaration($inferredType) missing!"

        override fun visitMethod(node: UMethod) {
            if (!isKotlin(node)) return
            val returnType = if (node.isSuspendFunction) {
                node.suspendReturnType
            } else {
                node.returnType
            } ?: return
            val presentableText = returnType.presentableText
            val canonicalText = returnType.canonicalText
            val inferredKind = getInferredKind(node, returnType)
            if (inferredKind != InferredKind.NONE) {
                context.reportInferredType(node, presentableText, canonicalText, inferredKind)
            }
        }

        /**
         * Returns an enum representing whether the node's type has been inferred and what kind
         * of node it is.
         */
        private fun getInferredKind(
            node: UMethod,
            returnType: PsiType
        ): InferredKind = when {
            isInferredProperty(node) -> InferredKind.PROPERTY
            isInferredMethod(node, returnType) -> InferredKind.METHOD
            // The node doesn't have its type inferred.
            else -> InferredKind.NONE
        }

        /**
         * Returns true if the scanned [node] is a property and the compiler inferred its
         * return type.
         */
        private fun isInferredProperty(node: UMethod): Boolean {
            val property = node.sourcePsi as? KtProperty ?: return false
            return property.colon == null
        }

        /**
         * Returns true if the scanned [node] is a function and the compiler inferred its
         * return type.
         */
        private fun isInferredMethod(
            node: UMethod,
            returnType: PsiType
        ): Boolean {
            val function = node.sourcePsi as? KtNamedFunction ?: return false
            val isFunction = function.funKeyword != null
            val hasColon = function.colon != null
            val isVoid =
                returnType == PsiType.VOID || returnType.presentableText == "Unit"
            return isFunction && !hasColon && !isVoid
        }

        /**
         * Returns true if this node is a suspend function.
         */
        private val UMethod.isSuspendFunction: Boolean
            get() = sourcePsi is KtNamedFunction && text.contains("suspend")

        /**
         * Returns the real return type of a suspend function. When using suspend functions the
         * real signature of the function is modified in the byte code: a [Continuation] parameter
         * is added and the return type of the function becomes [Object].
         *
         * In order to calculate the real return type of the function, we have to annalise the
         * generated byte code and look for the [Continuation] parameter in the function's
         * parameter list. The type argument of the continuation will be the return type of the
         * suspend function.
         */
        private val UMethod.suspendReturnType: PsiType?
            get() {
                val continuationType = parameterList.parameters.filter {
                    TypeConversionUtil.erasure(it.type).presentableText == "Continuation"
                }.map { it.type }.first()
                val continuationTypeNoWildCards =
                    GenericsUtil.eliminateWildcards(continuationType) as? PsiImmediateClassType
                return continuationTypeNoWildCards?.typeArguments()?.first() as? PsiType
            }

        /**
         * An enum representing what kind of node was inferred. The compiler can basically infer two
         * kinds of nodes: methods and properties.
         */
        private enum class InferredKind {
            METHOD, PROPERTY, NONE
        }
    }
}