twitter-ktlint-nested_lambda.kt

import com.pinterest.ktlint.core.Rule
import org.jetbrains.kotlin.KtNodeTypes.BLOCK
import org.jetbrains.kotlin.KtNodeTypes.CALL_EXPRESSION
import org.jetbrains.kotlin.KtNodeTypes.DOT_QUALIFIED_EXPRESSION
import org.jetbrains.kotlin.KtNodeTypes.FUNCTION_LITERAL
import org.jetbrains.kotlin.KtNodeTypes.LAMBDA_ARGUMENT
import org.jetbrains.kotlin.KtNodeTypes.LAMBDA_EXPRESSION
import org.jetbrains.kotlin.KtNodeTypes.PROPERTY_DELEGATE
import org.jetbrains.kotlin.KtNodeTypes.REFERENCE_EXPRESSION
import org.jetbrains.kotlin.KtNodeTypes.VALUE_PARAMETER_LIST
import org.jetbrains.kotlin.com.intellij.lang.ASTNode
import org.jetbrains.kotlin.com.intellij.psi.impl.source.tree.CompositeElement
import org.jetbrains.kotlin.com.intellij.psi.impl.source.tree.TreeUtil
import org.jetbrains.kotlin.com.intellij.psi.tree.TokenSet
import package.Utils

class NestedLambdaNamingRule(private val shouldError: Boolean = true) : Rule(RULE_ID) {

    companion object {
        const val RULE_ID = "nested-lambda-naming"
        private val specialIgnoreCases = listOf(
            "main",
            "diskIO",
            "background",
            "setOnClickListener",
            "apply",
            "withContext",
            "launch",
            "runBlocking",
            "stub",
            "then",
            "runBlocking",
            "runBlockingTest"
        )
    }

    override fun visit(
        node: ASTNode,
        autoCorrect: Boolean,
        emit: (offset: Int, errorMessage: String, canBeAutoCorrected: Boolean) -> Unit
    ) {
        if (autoCorrect) return
        if (node is CompositeElement && node.elementType == LAMBDA_EXPRESSION &&
            TreeUtil.findParent(
                node, TokenSet.create(PROPERTY_DELEGATE), TokenSet.create(LAMBDA_EXPRESSION)
            ) == null
        ) {

            val functionLiteral = node.findChildByType(FUNCTION_LITERAL)!!
            val argumentList = functionLiteral
                .getChildren(TokenSet.create(VALUE_PARAMETER_LIST))

            val block = functionLiteral.findChildByType(BLOCK)!!
            if (argumentList.isEmpty() &&
                doesBlockHaveLambda(block) &&
                !isLambdaInExceptList(node)
            ) {
                val errorMessage = "Lambda naming: ${node.text}"
                if (shouldError) {
                    emit(
                        node.startOffset,
                        "Error - $errorMessage",
                        false
                    )
                } else {
                    val clazz = Utils.getOuterClass(node)
                    if (clazz == null) {
                        System.out.print("WARNING - $errorMessage\n")
                    } else {
                        System.out.print("${clazz.name}: WARNING - $errorMessage\n")
                    }
                }
            }
        }
    }

    private fun doesBlockHaveLambda(block: ASTNode): Boolean {
        if (block.elementType != BLOCK)
            throw IllegalArgumentException("Must pass block")

        val dotQualifiedExpression = block.findChildByType(DOT_QUALIFIED_EXPRESSION)
        val callExpression = dotQualifiedExpression?.findChildByType(CALL_EXPRESSION)
        val lambdaArgument = callExpression?.findChildByType(LAMBDA_ARGUMENT)
        val lambdaExpression = lambdaArgument?.findChildByType(LAMBDA_EXPRESSION)

        return lambdaExpression != null
    }

    private fun isLambdaInExceptList(lambda: ASTNode): Boolean {
        if (lambda.elementType != LAMBDA_EXPRESSION)
            throw IllegalArgumentException("Must pass lambda")

        val callExpression = TreeUtil.findParent(lambda, TokenSet.create(CALL_EXPRESSION))
        val referenceExpression = callExpression?.findChildByType(REFERENCE_EXPRESSION)
        return specialIgnoreCases.contains(referenceExpression?.text)
    }
}