Refactor keyword completion handlers

- Separate completion keyword handlers stuff into files
- Introduce CompletionKeywordHandlerProvider as a common way to provide handlers
- Move CompletionKeywordHandler factory out of companion object

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/BasicCompletionSession.kt

@@ -24,6 +24,8 @@ import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.idea.analysis.analyzeInContext
 import org.jetbrains.kotlin.idea.caches.resolve.util.resolveToDescriptor
 import org.jetbrains.kotlin.idea.codeInsight.ReferenceVariantsHelper
+import org.jetbrains.kotlin.idea.completion.keywords.DefaultCompletionKeywordHandlerProvider
+import org.jetbrains.kotlin.idea.completion.keywords.createLookups
 import org.jetbrains.kotlin.idea.completion.smart.ExpectedInfoMatch
 import org.jetbrains.kotlin.idea.completion.smart.SMART_COMPLETION_ITEM_PRIORITY_KEY
 import org.jetbrains.kotlin.idea.completion.smart.SmartCompletion
@@ -583,7 +585,7 @@ class BasicCompletionSession(
                 val keyword = lookupElement.lookupString
                 if (keyword in keywordsToSkip) return@complete
 
-                val completionKeywordHandler = DefaultCompletionKeywordHandlers.defaultHandlers.getHandlerForKeyword(keyword)
+                val completionKeywordHandler = DefaultCompletionKeywordHandlerProvider.getHandlerForKeyword(keyword)
                 if (completionKeywordHandler != null) {
                     val lookups = completionKeywordHandler.createLookups(parameters, expression, lookupElement, project)
                     collector.addElements(lookups)

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/contributors/FirKeywordCompletionContributor.kt

@@ -16,6 +16,11 @@ import org.jetbrains.kotlin.idea.completion.context.FirPositionCompletionContext
 import org.jetbrains.kotlin.idea.completion.context.FirUnknownPositionContext
 import org.jetbrains.kotlin.idea.completion.contributors.keywords.OverrideKeywordHandler
 import org.jetbrains.kotlin.idea.completion.contributors.keywords.ReturnKeywordHandler
+import org.jetbrains.kotlin.idea.completion.contributors.keywords.ReturnKeywordHandler.createLookups
+import org.jetbrains.kotlin.idea.completion.keywords.CompletionKeywordHandlerProvider
+import org.jetbrains.kotlin.idea.completion.keywords.CompletionKeywordHandlers
+import org.jetbrains.kotlin.idea.completion.keywords.DefaultCompletionKeywordHandlerProvider
+import org.jetbrains.kotlin.idea.completion.keywords.createLookups
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.platform.jvm.isJvm
 import org.jetbrains.kotlin.psi.*
@@ -26,7 +31,7 @@ internal class FirKeywordCompletionContributor(basicContext: FirBasicCompletionC
         override fun getLanguageVersionSetting(module: Module) = LanguageVersionSettingsImpl.DEFAULT // TODO
     })
 
-    private val resolveDependentCompletionKeywordHandlers = ResolveDependentCompletionKeywordHandlers(basicContext)
+    private val resolveDependentCompletionKeywordHandlers = ResolveDependentCompletionKeywordHandlerProvider(basicContext)
 
     fun KtAnalysisSession.completeKeywords(
         positionContext: FirPositionCompletionContext
@@ -46,9 +51,9 @@ internal class FirKeywordCompletionContributor(basicContext: FirBasicCompletionC
 
     fun KtAnalysisSession.completeWithResolve(position: PsiElement, expression: KtExpression?) {
         complete(position) { lookupElement, keyword ->
-            val lookups = DefaultCompletionKeywordHandlers.defaultHandlers.getHandlerForKeyword(keyword)
+            val lookups = DefaultCompletionKeywordHandlerProvider.getHandlerForKeyword(keyword)
                 ?.createLookups(parameters, expression, lookupElement, project)
-                ?: resolveDependentCompletionKeywordHandlers.handlers.getHandlerForKeyword(keyword)?.run {
+                ?: resolveDependentCompletionKeywordHandlers.getHandlerForKeyword(keyword)?.run {
                     createLookups(parameters, expression, lookupElement, project)
                 }
                 ?: listOf(lookupElement)
@@ -64,8 +69,10 @@ internal class FirKeywordCompletionContributor(basicContext: FirBasicCompletionC
     }
 }
 
-private class ResolveDependentCompletionKeywordHandlers(basicContext: FirBasicCompletionContext) {
-    val handlers = CompletionKeywordHandlers(
+private class ResolveDependentCompletionKeywordHandlerProvider(
+    basicContext: FirBasicCompletionContext
+) : CompletionKeywordHandlerProvider<KtAnalysisSession>() {
+    override val handlers = CompletionKeywordHandlers(
         ReturnKeywordHandler,
         OverrideKeywordHandler(basicContext),
     )

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/contributors/keywords/OverrideKeywordHandler.kt

@@ -13,6 +13,7 @@ import com.intellij.ui.RowIcon
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.idea.completion.*
 import org.jetbrains.kotlin.idea.completion.context.FirBasicCompletionContext
+import org.jetbrains.kotlin.idea.completion.keywords.CompletionKeywordHandler
 import org.jetbrains.kotlin.idea.core.overrideImplement.*
 import org.jetbrains.kotlin.idea.core.overrideImplement.KtClassMember
 import org.jetbrains.kotlin.idea.core.overrideImplement.KtGenerateMembersHandler

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/contributors/keywords/ReturnKeywordHandler.kt

@@ -8,9 +8,9 @@ package org.jetbrains.kotlin.idea.completion.contributors.keywords
 import com.intellij.codeInsight.completion.CompletionParameters
 import com.intellij.codeInsight.lookup.LookupElement
 import com.intellij.openapi.project.Project
-import org.jetbrains.kotlin.idea.completion.CompletionKeywordHandler
 import org.jetbrains.kotlin.idea.completion.createKeywordElement
 import org.jetbrains.kotlin.idea.completion.createKeywordElementWithSpace
+import org.jetbrains.kotlin.idea.completion.keywords.CompletionKeywordHandler
 import org.jetbrains.kotlin.idea.completion.labelNameToTail
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtFunctionSymbol

idea/idea-frontend-independent/src/org/jetbrains/kotlin/idea/completion/keywords/CompletionKeywordHandler.kt

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.completion.keywords
+
+import com.intellij.codeInsight.completion.CompletionParameters
+import com.intellij.codeInsight.lookup.LookupElement
+import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.lexer.KtKeywordToken
+import org.jetbrains.kotlin.psi.KtExpression
+
+abstract class CompletionKeywordHandler<CONTEXT>(
+    val keyword: KtKeywordToken
+) {
+    object NO_CONTEXT
+
+    abstract fun CONTEXT.createLookups(
+        parameters: CompletionParameters,
+        expression: KtExpression?,
+        lookup: LookupElement,
+        project: Project
+    ): Collection<LookupElement>
+}
+
+inline fun <CONTEXT> completionKeywordHandler(
+    keyword: KtKeywordToken,
+    crossinline create: CONTEXT.(
+        parameters: CompletionParameters,
+        expression: KtExpression?,
+        lookup: LookupElement,
+        project: Project
+    ) -> Collection<LookupElement>
+) = object : CompletionKeywordHandler<CONTEXT>(keyword) {
+    override fun CONTEXT.createLookups(
+        parameters: CompletionParameters,
+        expression: KtExpression?,
+        lookup: LookupElement,
+        project: Project
+    ): Collection<LookupElement> = create(parameters, expression, lookup, project)
+}
+
+/**
+ * Create a list of [LookupElement] for [CompletionKeywordHandler] which has no context
+ *
+ * This function is needed to avoid writing `with(CompletionKeywordHandler.NO_CONTEXT) { ... }` to create such lookups
+ */
+fun CompletionKeywordHandler<CompletionKeywordHandler.NO_CONTEXT>.createLookups(
+    parameters: CompletionParameters,
+    expression: KtExpression?,
+    lookup: LookupElement,
+    project: Project
+): Collection<LookupElement> = with(CompletionKeywordHandler.NO_CONTEXT) { createLookups(parameters, expression, lookup, project) }

idea/idea-frontend-independent/src/org/jetbrains/kotlin/idea/completion/keywords/CompletionKeywordHandlerProvider.kt

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.completion.keywords
+
+abstract class CompletionKeywordHandlerProvider<CONTEXT> {
+    protected abstract val handlers: CompletionKeywordHandlers<CONTEXT>
+
+    fun getHandlerForKeyword(keyword: String): CompletionKeywordHandler<CONTEXT>? =
+        handlers.getHandlerForKeyword(keyword)
+}
+
+
+class CompletionKeywordHandlers<CONTEXT>(vararg handlers: CompletionKeywordHandler<CONTEXT>) {
+    private val handlerByKeyword = handlers.associateBy { it.keyword.value }
+
+    internal fun getHandlerForKeyword(keyword: String): CompletionKeywordHandler<CONTEXT>? =
+        handlerByKeyword[keyword]
+}

idea/idea-frontend-independent/src/org/jetbrains/kotlin/idea/completion/keywords/DefaultCompletionKeywordHandlerProvider.kt

@@ -3,48 +3,36 @@
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
-package org.jetbrains.kotlin.idea.completion
+package org.jetbrains.kotlin.idea.completion.keywords
 
 import com.intellij.codeInsight.completion.CompletionParameters
-import com.intellij.codeInsight.lookup.LookupElement
-import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.idea.completion.breakOrContinueExpressionItems
 import org.jetbrains.kotlin.idea.completion.handlers.createKeywordConstructLookupElement
-import org.jetbrains.kotlin.lexer.KtKeywordToken
 import org.jetbrains.kotlin.lexer.KtTokens
-import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.psiUtil.prevLeaf
 
-class CompletionKeywordHandlers<CONTEXT>(handlers: List<CompletionKeywordHandler<CONTEXT>>) {
-    constructor(vararg handlers: CompletionKeywordHandler<CONTEXT>) : this(handlers.toList())
-
-    private val handlerByKeyword = handlers.associateBy { it.keyword.value }
-
-    fun getHandlerForKeyword(keyword: String): CompletionKeywordHandler<CONTEXT>? =
-        handlerByKeyword[keyword]
-}
-
 @OptIn(ExperimentalStdlibApi::class)
-object DefaultCompletionKeywordHandlers {
+object DefaultCompletionKeywordHandlerProvider : CompletionKeywordHandlerProvider<CompletionKeywordHandler.NO_CONTEXT>() {
     private val BREAK_HANDLER =
-        CompletionKeywordHandler<CompletionKeywordHandler.NO_CONTEXT>(KtTokens.BREAK_KEYWORD) { _, expression, _, _ ->
+        completionKeywordHandler<CompletionKeywordHandler.NO_CONTEXT>(KtTokens.BREAK_KEYWORD) { _, expression, _, _ ->
             if (expression != null) {
                 breakOrContinueExpressionItems(expression, KtTokens.BREAK_KEYWORD.value)
             } else emptyList()
         }
 
     private val CONTINUE_HANDLER =
-        CompletionKeywordHandler<CompletionKeywordHandler.NO_CONTEXT>(KtTokens.CONTINUE_KEYWORD) { _, expression, _, _ ->
+        completionKeywordHandler<CompletionKeywordHandler.NO_CONTEXT>(KtTokens.CONTINUE_KEYWORD) { _, expression, _, _ ->
             if (expression != null) {
                 breakOrContinueExpressionItems(expression, KtTokens.CONTINUE_KEYWORD.value)
             } else emptyList()
         }
 
-    private val CONTRACT_HANDLER = CompletionKeywordHandler<CompletionKeywordHandler.NO_CONTEXT>(KtTokens.CONTRACT_KEYWORD) { _, _, _, _ ->
+    private val CONTRACT_HANDLER = completionKeywordHandler<CompletionKeywordHandler.NO_CONTEXT>(KtTokens.CONTRACT_KEYWORD) { _, _, _, _ ->
         emptyList()
     }
 
     private val GETTER_HANDLER =
-        CompletionKeywordHandler<CompletionKeywordHandler.NO_CONTEXT>(KtTokens.GET_KEYWORD) { parameters, _, lookupElement, project ->
+        completionKeywordHandler<CompletionKeywordHandler.NO_CONTEXT>(KtTokens.GET_KEYWORD) { parameters, _, lookupElement, project ->
             buildList {
                 add(lookupElement)
                 if (!parameters.isUseSiteAnnotationTarget) {
@@ -62,7 +50,7 @@ object DefaultCompletionKeywordHandlers {
         }
 
     private val SETTER_HANDLER =
-        CompletionKeywordHandler<CompletionKeywordHandler.NO_CONTEXT>(KtTokens.SET_KEYWORD) { parameters, _, lookupElement, project ->
+        completionKeywordHandler<CompletionKeywordHandler.NO_CONTEXT>(KtTokens.SET_KEYWORD) { parameters, _, lookupElement, project ->
             buildList {
                 add(lookupElement)
                 if (!parameters.isUseSiteAnnotationTarget) {
@@ -79,52 +67,13 @@ object DefaultCompletionKeywordHandlers {
             }
         }
 
-    val defaultHandlers = CompletionKeywordHandlers(
+    override val handlers = CompletionKeywordHandlers(
         BREAK_HANDLER, CONTINUE_HANDLER,
         GETTER_HANDLER, SETTER_HANDLER,
         CONTRACT_HANDLER,
     )
 }
 
-fun CompletionKeywordHandler<CompletionKeywordHandler.NO_CONTEXT>.createLookups(
-    parameters: CompletionParameters,
-    expression: KtExpression?,
-    lookup: LookupElement,
-    project: Project
-): Collection<LookupElement> = with(CompletionKeywordHandler.NO_CONTEXT) { createLookups(parameters, expression, lookup, project) }
 
 private val CompletionParameters.isUseSiteAnnotationTarget
-    get() = position.prevLeaf()?.node?.elementType == KtTokens.AT
-
-
-abstract class CompletionKeywordHandler<CONTEXT>(
-    val keyword: KtKeywordToken
-) {
-    object NO_CONTEXT
-
-    abstract fun CONTEXT.createLookups(
-        parameters: CompletionParameters,
-        expression: KtExpression?,
-        lookup: LookupElement,
-        project: Project
-    ): Collection<LookupElement>
-
-    companion object {
-        inline operator fun <CONTEXT> invoke(
-            keyword: KtKeywordToken,
-            crossinline create: CONTEXT.(
-                parameters: CompletionParameters,
-                expression: KtExpression?,
-                lookup: LookupElement,
-                project: Project
-            ) -> Collection<LookupElement>
-        ) = object : CompletionKeywordHandler<CONTEXT>(keyword) {
-            override fun CONTEXT.createLookups(
-                parameters: CompletionParameters,
-                expression: KtExpression?,
-                lookup: LookupElement,
-                project: Project
-            ): Collection<LookupElement> = create(parameters, expression, lookup, project)
-        }
-    }
-}
+    get() = position.prevLeaf()?.node?.elementType == KtTokens.AT