ComposeUnpatrolledViewModelFunctionCallDetector.kt

/**
 * Should warn when making unattended a viewmodel function in a composable without considering
 * potential recompositions
 *
 *   Should warn when:
 *
 *   @Composable
 *   fun MyComposable(viewModel : MyViewModel) {
 *
 *       viewModel.startOperation()
 *
 *       Column {
 *       Text(....)
 *       Text(...)
 *       }
 *
 *   }
 *
 *   Should not warn when we have controlled function call
 *    @Composable
 *   fun MyComposable(viewModel : MyViewModel) {
 *       LaunchEffect(Unit) {
 *         viewModel.startOperation()
 *       }
 *       Column {
 *       Text(....)
 *       Text(...)
 *       }
 *
 *   }
 */
class ComposeUnpatrolledViewModelFunctionCallDetector : ComposeFunctionDetector() {

    companion object {

        val ISSUE = Issue.create(
            id = "compose-uncontrolled-viewmodel-function-call",
            briefDescription = "Unvetted view model function call in composable",
            explanation = "This called should be controlled how many times it is called in case of recomposition",
            category = Category.CUSTOM_LINT_CHECKS,
            priority = 5,
            severity = Severity.WARNING,
            androidSpecific = true,
            implementation = Implementation(
                ComposeUnpatrolledViewModelFunctionCallDetector::class.java,
                EnumSet.of(Scope.ALL_JAVA_FILES)
            )
        )

    }

    private lateinit var viewModelType: PsiType
    private lateinit var visitor: ViewModelFunctionCallVisitor
    override fun visitComposable(context: JavaContext, node: UMethod) {
        viewModelType = PsiType.getTypeByName(VIEW_MODEL, node.project, node.resolveScope)
        if (node.hasParametersOfSuperType(viewModelType)) {
            return
        }
        visitor = ViewModelFunctionCallVisitor(node)
        val body = node.uastBody as UBlockExpression
        body.expressions
            .onEach {
                val expression = it.getUCallExpression() ?: return@onEach
                if (visitor.visitCallExpression(expression)) {
                    context.report(
                        ISSUE,
                        node,
                        context.getLocation(it),
                        "Wrap this in a LaunchEffect or DisposableEffect"
                        //TODO offer fix
                    )
                } else {
                    println("Evaluating expression that is an ordinary method call")
                    //ordinary function call, skip if composable
                }
            }

    }


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

    inner class ViewModelFunctionCallVisitor(private val composable: UMethod) :
        AbstractUastVisitor() {


        override fun visitCallExpression(node: UCallExpression): Boolean {
            //TODO check if call must be direct method in the composable block
            return node.isReceiverOfSuperType(node.sourcePsi!!, VIEW_MODEL)
        }
    }

}

ExpressionKtx.kt

import com.android.tools.lint.detector.api.UastLintUtils.Companion.tryResolveUDeclaration
import com.intellij.psi.PsiAssignmentExpression
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiParameter
import com.intellij.psi.PsiType
import org.jetbrains.uast.UCallExpression
import org.jetbrains.uast.UMethod
import org.jetbrains.uast.toUElement

/**
 * @param type string of the type, later resolved as [PsiType]
 * @param node node currently being evaluated
 * Checks if receiver of the [UCallExpression] is of given type
 */
fun UCallExpression.isReceiverOfSuperType(node: PsiElement, type: String): Boolean {
    val receiver = receiver ?: return false
    val superTypes = receiverType?.superTypes.orEmpty()
    if (superTypes.isEmpty()) {
        return false
    }
    val resolvedType = PsiType.getTypeByName(type, node.project, node.resolveScope)
    return resolvedType in superTypes

}

private val UMethod.parametersSequences: Sequence<PsiParameter>
    get() = (0..parameterList.parametersCount)
        .asSequence().mapNotNull { this.parameterList.getParameter(it) }
const val COMPOSEABLE = "androidx.compose.runtime.Composable"
const val PREVIEW_COMPOSABLE = "androidx.compose.ui.tooling.preview.Preview"

data class ComposableFunctionEvaluationResult(
    val isComposable: Boolean,
    val isPreviewComposable: Boolean
)

/**
 * Checks if [UMethod] is composable or preview composeable
 * @return [ComposableFunctionEvaluationResult]
 *
 */
fun UMethod.isAnyKindOfComposable(): ComposableFunctionEvaluationResult {
    findAnnotation(COMPOSEABLE) ?: return ComposableFunctionEvaluationResult(false, false)
    return ComposableFunctionEvaluationResult(true, findAnnotation(PREVIEW_COMPOSABLE) != null)
}

/**
 * Checks if given [UCallExpression] is a call to a composable function
 */
fun UCallExpression.isComposableInvocation(): Boolean {
    return tryResolveUDeclaration()?.findAnnotation(COMPOSEABLE) != null
}

fun UMethod.getDefaultValueParameters(): Sequence<PsiParameter> {
    return parametersSequences.filter { it.initializer != null }
}

/**
 * @return if [UMethod] is composable
 */
fun UMethod.isComposable(): Boolean {
    return findAnnotation(COMPOSEABLE) != null
}

/**
 * @return if [UMethod] is of given type
 */
fun UMethod.isReturnTypeOfType(type: String): Boolean {
    return returnType?.let {
        val resolvedType = PsiType.getTypeByName(type, project, resolveScope)
        isReturnTypeOfType(resolvedType)
    } ?: false
}

/**
 * @return if [UMethod] is of given type
 */
fun UMethod.isReturnTypeOfType(type: PsiType): Boolean {
    return returnType?.let {
        it == type
    } ?: false
}

/**
 * @return Sequence of [PsiParameter] of given type
 */
fun UMethod.parametersOfSuperType(type: String): Sequence<PsiParameter> {
    if (!hasTypeParameters()) {
        return emptySequence()
    }
    val resolvedType = PsiType.getTypeByName(type, project, resolveScope)
    return parametersSequences
        .mapNotNull {
            val param = parameterList.getParameter(it)
            param?.takeIf { resolvedType in param.type.superTypes }
        }

}

/**
 * Example
 *
 *  fun myFunction(viewMolde : MyViewModel = sharedViewModel()) {
 *
 *
 *  }
 * @return if parameter initializer is a function call
 */
fun PsiParameter.isInitializerFunctionCall(): Boolean {
    val init = initializer
    return init is PsiAssignmentExpression && init.rExpression.toUElement() is UCallExpression
}

/**
 *
 * @return if parameter is lambda
 */
fun PsiParameter.isLambda(): Boolean {
    TODO("Not impl")
}

/**
 * @return true if method has parameter of given type
 */
fun UMethod.hasParametersOfSuperType(type: String): Boolean {
    if (!hasTypeParameters()) {
        return false
    }
    val resolvedType = PsiType.getTypeByName(type, project, resolveScope)
    return hasParametersOfSuperType(resolvedType)

}

/**
 * @return true if method has parameter of given type [PsiType]
 */
fun UMethod.hasParametersOfSuperType(type: PsiType): Boolean {
    if (!hasTypeParameters()) {
        return false
    }
    return (0..parameterList.parametersCount).any {
        val param = parameterList.getParameter(it)
        param?.takeIf { type in param.type.superTypes } != null
    }

}