[kotlin] KTIJ-30784 explicit LookupElementContributor properties

GitOrigin-RevId: 66b1db5e9b8a744f659b21c5e2d545e3b8bbfb8f
2026-01-06 11:50:54 +07:00 · 2024-07-26 15:42:57 +02:00
parent 1795c26525
commit c78c933f54
6 changed files with 129 additions and 85 deletions
--- a/platform/lang-impl/src/com/intellij/codeInsight/completion/addingPolicy/ElementsAddingPolicy.kt
+++ b/platform/lang-impl/src/com/intellij/codeInsight/completion/addingPolicy/ElementsAddingPolicy.kt
@@ -41,14 +41,21 @@ interface ElementsAddingPolicy {
  fun onDeactivate(result: CompletionResultSet)

  interface Default : ElementsAddingPolicy {
+
    override fun onActivate(result: CompletionResultSet) {}

+    override fun onDeactivate(result: CompletionResultSet) {}
+
    override fun onResultStop(result: CompletionResultSet) {}

-    override fun addElement(result: CompletionResultSet, element: LookupElement) = addAllElements(result, listOf(element))
+    override fun addElement(
+      result: CompletionResultSet,
+      element: LookupElement,
+    ): Unit = result.addElement(element)

-    override fun addAllElements(result: CompletionResultSet, elements: Iterable<LookupElement>) = result.addAllElements(elements)
-
-    override fun onDeactivate(result: CompletionResultSet) {}
+    override fun addAllElements(
+      result: CompletionResultSet,
+      elements: Iterable<LookupElement>,
+    ): Unit = result.addAllElements(elements)
  }
 }
--- a/plugins/kotlin/base/code-insight/src/org/jetbrains/kotlin/idea/base/codeInsight/LookupElementContributor.kt
+++ b/plugins/kotlin/base/code-insight/src/org/jetbrains/kotlin/idea/base/codeInsight/LookupElementContributor.kt
@@ -1,6 +1,12 @@
 // Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package org.jetbrains.kotlin.idea.base.codeInsight

+import com.intellij.codeInsight.lookup.LookupElement
 import com.intellij.openapi.util.Key
+import org.jetbrains.kotlin.psi.UserDataProperty

-val LOOKUP_ELEMENT_CONTRIBUTOR = Key<Class<*>>("LookupElement.LOOKUP_ELEMENT_CONTRIBUTOR")
+var LookupElement.contributorClass: Class<*>? by UserDataProperty(Key.create<Class<*>>("LookupElement.LOOKUP_ELEMENT_CONTRIBUTOR"))
+
+fun LookupElement.withContributorClass(contributorClass: Class<*>): LookupElement = apply {
+    this.contributorClass = contributorClass
+}
--- a/plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/FirCompletionContributorFactory.kt
+++ b/plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/FirCompletionContributorFactory.kt
@@ -4,72 +4,93 @@ package org.jetbrains.kotlin.idea.completion.contributors

 import com.intellij.codeInsight.completion.addingPolicy.PolicyController
 import org.jetbrains.kotlin.idea.completion.context.FirBasicCompletionContext
-import org.jetbrains.kotlin.idea.completion.impl.k2.contributors.helpers.recording
+import org.jetbrains.kotlin.idea.completion.impl.k2.contributors.helpers.withPolicyController

 internal class FirCompletionContributorFactory(
    private val basicContext: FirBasicCompletionContext,
-    private val resultController: PolicyController,
+    private val policyController: PolicyController,
 ) {
    fun keywordContributor(priority: Int = 0) =
-        FirKeywordCompletionContributor(basicContext, priority).recording(resultController)
+        FirKeywordCompletionContributor(basicContext, priority)
+            .withPolicyController(policyController)

    fun classReferenceContributor(priority: Int = 0) =
-        FirClassReferenceCompletionContributor(basicContext, priority).recording(resultController)
+        FirClassReferenceCompletionContributor(basicContext, priority)
+            .withPolicyController(policyController)

    fun callableContributor(priority: Int = 0) =
-        FirCallableCompletionContributor(basicContext, priority).recording(resultController)
+        FirCallableCompletionContributor(basicContext, priority)
+            .withPolicyController(policyController)

    fun superMemberContributor(priority: Int) =
-        FirSuperMemberCompletionContributor(basicContext, priority).recording(resultController)
+        FirSuperMemberCompletionContributor(basicContext, priority)
+            .withPolicyController(policyController)

    fun infixCallableContributor(priority: Int = 0) =
-        FirInfixCallableCompletionContributor(basicContext, priority).recording(resultController)
+        FirInfixCallableCompletionContributor(basicContext, priority)
+            .withPolicyController(policyController)

    fun callableReferenceContributor(priority: Int = 0) =
-        FirCallableReferenceCompletionContributor(basicContext, priority).recording(resultController)
+        FirCallableReferenceCompletionContributor(basicContext, priority)
+            .withPolicyController(policyController)

    fun classifierContributor(priority: Int = 0) =
-        FirClassifierCompletionContributor(basicContext, priority).recording(resultController)
+        FirClassifierCompletionContributor(basicContext, priority)
+            .withPolicyController(policyController)

    fun classifierReferenceContributor(priority: Int = 0) =
-        FirClassifierReferenceCompletionContributor(basicContext, priority).recording(resultController)
+        FirClassifierReferenceCompletionContributor(basicContext, priority)
+            .withPolicyController(policyController)

    fun annotationsContributor(priority: Int = 0) =
-        FirAnnotationCompletionContributor(basicContext, priority).recording(resultController)
+        FirAnnotationCompletionContributor(basicContext, priority)
+            .withPolicyController(policyController)

    fun packageCompletionContributor(priority: Int = 0) =
-        FirPackageCompletionContributor(basicContext, priority).recording(resultController)
+        FirPackageCompletionContributor(basicContext, priority)
+            .withPolicyController(policyController)

    fun importDirectivePackageMembersContributor(priority: Int = 0) =
-        FirImportDirectivePackageMembersCompletionContributor(basicContext, priority).recording(resultController)
+        FirImportDirectivePackageMembersCompletionContributor(basicContext, priority)
+            .withPolicyController(policyController)

    fun typeParameterConstraintNameInWhereClauseContributor(priority: Int = 0) =
-        FirTypeParameterConstraintNameInWhereClauseCompletionContributor(basicContext, priority).recording(resultController)
+        FirTypeParameterConstraintNameInWhereClauseCompletionContributor(basicContext, priority)
+            .withPolicyController(policyController)

    fun classifierNameContributor(priority: Int = 0) =
-        FirSameAsFileClassifierNameCompletionContributor(basicContext, priority).recording(resultController)
+        FirSameAsFileClassifierNameCompletionContributor(basicContext, priority)
+            .withPolicyController(policyController)

    fun whenWithSubjectConditionContributor(priority: Int = 0) =
-        FirWhenWithSubjectConditionContributor(basicContext, priority).recording(resultController)
+        FirWhenWithSubjectConditionContributor(basicContext, priority)
+            .withPolicyController(policyController)

    fun superEntryContributor(priority: Int) =
-        FirSuperEntryContributor(basicContext, priority).recording(resultController)
+        FirSuperEntryContributor(basicContext, priority)
+            .withPolicyController(policyController)

    fun declarationFromUnresolvedNameContributor(priority: Int) =
-        FirDeclarationFromUnresolvedNameContributor(basicContext, priority).recording(resultController)
+        FirDeclarationFromUnresolvedNameContributor(basicContext, priority)
+            .withPolicyController(policyController)

    fun declarationFromOverridableMembersContributor(priority: Int) =
-        FirDeclarationFromOverridableMembersContributor(basicContext, priority).recording(resultController)
+        FirDeclarationFromOverridableMembersContributor(basicContext, priority)
+            .withPolicyController(policyController)

    fun namedArgumentContributor(priority: Int = 0) =
-        FirNamedArgumentCompletionContributor(basicContext, priority).recording(resultController)
+        FirNamedArgumentCompletionContributor(basicContext, priority)
+            .withPolicyController(policyController)

    fun variableOrParameterNameWithTypeContributor(priority: Int) =
-        FirVariableOrParameterNameWithTypeCompletionContributor(basicContext, priority).recording(resultController)
+        FirVariableOrParameterNameWithTypeCompletionContributor(basicContext, priority)
+            .withPolicyController(policyController)

    fun kDocParameterNameContributor(priority: Int) =
-        FirKDocParameterNameContributor(basicContext, priority).recording(resultController)
+        FirKDocParameterNameContributor(basicContext, priority)
+            .withPolicyController(policyController)

    fun kDocCallableContributor(priority: Int) =
-        FirKDocCallableCompletionContributor(basicContext, priority).recording(resultController)
+        FirKDocCallableCompletionContributor(basicContext, priority)
+            .withPolicyController(policyController)
 }
--- a/plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/helpers/AttributedElementsAddingPolicy.kt
+++ b/plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/helpers/AttributedElementsAddingPolicy.kt
@@ -0,0 +1,57 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.kotlin.idea.completion.impl.k2.contributors.helpers
+
+import com.intellij.codeInsight.completion.CompletionResultSet
+import com.intellij.codeInsight.completion.addingPolicy.ElementsAddingPolicy
+import com.intellij.codeInsight.completion.addingPolicy.PolicyController
+import com.intellij.codeInsight.lookup.LookupElement
+import org.jetbrains.kotlin.analysis.api.KaSession
+import org.jetbrains.kotlin.idea.base.codeInsight.withContributorClass
+import org.jetbrains.kotlin.idea.completion.FirCompletionSessionParameters
+import org.jetbrains.kotlin.idea.completion.impl.k2.contributors.FirCompletionContributor
+import org.jetbrains.kotlin.idea.completion.weighers.WeighingContext
+import org.jetbrains.kotlin.idea.util.positionContext.KotlinRawPositionContext
+
+private class AttributedElementsAddingPolicy private constructor(
+    private val contributorClass: Class<FirCompletionContributor<*>>,
+) : ElementsAddingPolicy.Default {
+
+    constructor(contributor: FirCompletionContributor<*>) : this(contributor.javaClass)
+
+    override fun addElement(
+        result: CompletionResultSet,
+        element: LookupElement,
+    ) = super.addElement(
+        result = result,
+        element = element.withContributorClass(contributorClass),
+    )
+
+    override fun addAllElements(
+        result: CompletionResultSet,
+        elements: Iterable<LookupElement>,
+    ) = super.addAllElements(
+        result = result,
+        elements = elements.map { it.withContributorClass(contributorClass) },
+    )
+
+}
+
+internal fun <C : KotlinRawPositionContext> FirCompletionContributor<C>.withPolicyController(
+    policyController: PolicyController,
+): FirCompletionContributor<C> = object : FirCompletionContributor<C> {
+
+    context(KaSession)
+    override fun complete(
+        positionContext: C,
+        weighingContext: WeighingContext,
+        sessionParameters: FirCompletionSessionParameters,
+    ) {
+        policyController.invokeWithPolicy(AttributedElementsAddingPolicy(this@withPolicyController)) {
+            this@withPolicyController.complete(
+                positionContext,
+                weighingContext,
+                sessionParameters,
+            )
+        }
+    }
+}
--- a/plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/helpers/ToUserDataRecordingPolicy.kt
+++ b/plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/helpers/ToUserDataRecordingPolicy.kt
@@ -1,49 +0,0 @@
-// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
-package org.jetbrains.kotlin.idea.completion.impl.k2.contributors.helpers
-
-import com.intellij.codeInsight.completion.CompletionResultSet
-import com.intellij.codeInsight.completion.addingPolicy.ElementsAddingPolicy
-import com.intellij.codeInsight.completion.addingPolicy.PolicyController
-import com.intellij.codeInsight.lookup.LookupElement
-import org.jetbrains.kotlin.analysis.api.KaSession
-import org.jetbrains.kotlin.idea.base.codeInsight.LOOKUP_ELEMENT_CONTRIBUTOR
-import org.jetbrains.kotlin.idea.completion.FirCompletionSessionParameters
-import org.jetbrains.kotlin.idea.completion.impl.k2.contributors.FirCompletionContributor
-import org.jetbrains.kotlin.idea.completion.weighers.WeighingContext
-import org.jetbrains.kotlin.idea.util.positionContext.KotlinRawPositionContext
-
-internal class ToUserDataRecordingPolicy(private val contributor: FirCompletionContributor<*>) : ElementsAddingPolicy.Default {
-    override fun addElement(result: CompletionResultSet, element: LookupElement) {
-        result.addElement(element.decorate())
-    }
-
-    override fun addAllElements(result: CompletionResultSet, elements: Iterable<LookupElement>) {
-        elements.forEach { it.decorate() }
-        result.addAllElements(elements)
-    }
-
-    private fun LookupElement.decorate(): LookupElement {
-        this.putUserData(LOOKUP_ELEMENT_CONTRIBUTOR, contributor.javaClass)
-        return this
-    }
-}
-
-internal class RecordingFirCompletionContributorDelegate<C : KotlinRawPositionContext>(
-    private val originalContributor: FirCompletionContributor<C>,
-    private val resultController: PolicyController,
-) : FirCompletionContributor<C> {
-    context(KaSession)
-    override fun complete(
-        positionContext: C,
-        weighingContext: WeighingContext,
-        sessionParameters: FirCompletionSessionParameters
-    ) {
-        resultController.invokeWithPolicy(ToUserDataRecordingPolicy(originalContributor)) {
-            originalContributor.complete(positionContext, weighingContext, sessionParameters)
-        }
-    }
-}
-
-internal fun <C : KotlinRawPositionContext> FirCompletionContributor<C>.recording(resultController: PolicyController): FirCompletionContributor<C> {
-    return RecordingFirCompletionContributorDelegate(this, resultController)
-}
--- a/plugins/kotlin/ml-completion/src/org/jetbrains/kotlin/idea/mlCompletion/KotlinElementFeatureProvider.kt
+++ b/plugins/kotlin/ml-completion/src/org/jetbrains/kotlin/idea/mlCompletion/KotlinElementFeatureProvider.kt
@@ -6,20 +6,22 @@ import com.intellij.codeInsight.completion.ml.ContextFeatures
 import com.intellij.codeInsight.completion.ml.ElementFeatureProvider
 import com.intellij.codeInsight.completion.ml.MLFeatureValue
 import com.intellij.codeInsight.lookup.LookupElement
-import org.jetbrains.kotlin.idea.base.codeInsight.LOOKUP_ELEMENT_CONTRIBUTOR
+import org.jetbrains.kotlin.idea.base.codeInsight.contributorClass
+
+internal class KotlinElementFeatureProvider : ElementFeatureProvider {

-class KotlinElementFeatureProvider : ElementFeatureProvider {
    override fun getName(): String = "kotlin"

    override fun calculateFeatures(
        element: LookupElement,
        location: CompletionLocation,
-        contextFeatures: ContextFeatures
+        contextFeatures: ContextFeatures,
    ): Map<String, MLFeatureValue> {
-        val result: MutableMap<String, MLFeatureValue> = mutableMapOf()
-        element.getUserData(LOOKUP_ELEMENT_CONTRIBUTOR)?.let {
-            result["k2_contributor"] = MLFeatureValue.className(it)
-        }
-        return result
+        val contributor = element.contributorClass
+            ?: return emptyMap()
+
+        return mapOf(
+            "k2_contributor" to MLFeatureValue.className(contributor),
+        )
    }
 }