[kotlin] KTIJ-36991 refactored K2CompletionContributor.addElement to use context parameters rather than context receivers

This is required because both the KaSession and the K2CompletionSectionContext might be needed as parameters

Space-RevId: 744e2869be4e08e2a352fbe756b18e9d871841d4

GitOrigin-RevId: 03ccd41b7803d3de0f2fefca4d862900cb68c6f4

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/K2CompletionContributor.kt

@@ -254,13 +254,15 @@ internal abstract class K2CompletionContributor<P : KotlinRawPositionContext>(
     context(_: KaSession, context: K2CompletionSectionContext<P>)
     open fun shouldExecute(): Boolean = true
 
-    protected fun K2CompletionSectionContext<P>.addElement(element: LookupElement) {
-        sink.addElement(decorateLookupElement(element))
+    context(_: KaSession, context: K2CompletionSectionContext<P>)
+    protected fun addElement(element: LookupElement) {
+        context.sink.addElement(decorateLookupElement(element))
     }
 
-    protected fun K2CompletionSectionContext<P>.addElements(elements: Iterable<LookupElement>) {
+    context(_: KaSession, context: K2CompletionSectionContext<P>)
+    protected fun addElements(elements: Iterable<LookupElement>) {
         val decoratedElements = elements.map { decorateLookupElement(it) }
-        sink.addElements(decoratedElements)
+        context.sink.addElements(decoratedElements)
     }
 
     /**
@@ -272,24 +274,25 @@ internal abstract class K2CompletionContributor<P : KotlinRawPositionContext>(
      */
     protected open fun K2CompletionSectionContext<P>.getGroupPriority(): Int = 0
 
-    private fun K2CompletionSectionContext<P>.decorateLookupElement(
+    context(_: KaSession, context: K2CompletionSectionContext<P>)
+    private fun decorateLookupElement(
         element: LookupElement,
     ): LookupElement {
-        element.groupPriority = getGroupPriority()
+        element.groupPriority = context.getGroupPriority()
         element.contributorClass = this::class.java
 
-        if (isAtFunctionLiteralStart(parameters.position)) {
+        if (isAtFunctionLiteralStart(context.parameters.position)) {
             element.suppressItemSelectionByCharsOnTyping = true
         }
 
-        val bracesInsertHandler = when (parameters.type) {
+        val bracesInsertHandler = when (context.parameters.type) {
             KotlinFirCompletionParameters.CorrectionType.BRACES_FOR_STRING_TEMPLATE -> InsertStringTemplateBracesInsertHandler
             else -> WrapSingleStringTemplateEntryWithBracesInsertHandler
         }
 
         return LookupElementDecorator.withDelegateInsertHandler(
             LookupElementDecorator.withDelegateInsertHandler(element, bracesInsertHandler),
-            CompletionCharInsertHandler(parameters.delegate.isAutoPopup),
+            CompletionCharInsertHandler(context.parameters.delegate.isAutoPopup),
         )
     }
 }

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/K2ActualDeclarationContributor.kt

@@ -38,7 +38,7 @@ internal class K2ActualDeclarationContributor : K2SimpleCompletionContributor<Ko
             importStrategyDetector = context.importStrategyDetector,
             declaration = declaration,
         ).createActualLookups(context.parameters, context.project)
-        context.addElements(elements)
+        addElements(elements)
     }
 
     private fun KtTypeReference.getDeclaration(): KtCallableDeclaration? {

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/K2CallableCompletionContributor.kt

@@ -179,7 +179,7 @@ internal abstract class K2AbstractCallableCompletionContributor<P : KotlinNameRe
         if (receiver == null) return
         val elements = collectDotCompletionFromLocalScope(receiver)
         elements.createFilteredLookupElements(shadowedCallablesFilter)
-            .forEach { context.addElement(it) }
+            .forEach { addElement(it) }
     }
 
     context(_: KaSession, context: K2CompletionSectionContext<P>)
@@ -258,7 +258,7 @@ internal abstract class K2AbstractCallableCompletionContributor<P : KotlinNameRe
             collectDotCompletionFromIndex(receiver)
         }
         elements.createFilteredLookupElements(shadowedCallablesFilter)
-            .forEach { context.addElement(it) }
+            .forEach { addElement(it) }
     }
 
     context(_: KaSession, context: K2CompletionSectionContext<P>)
@@ -753,7 +753,7 @@ internal abstract class K2AbstractCallableCompletionContributor<P : KotlinNameRe
 
         createAndFilterMetadataForMemberCallables(enumEntries)
             .createFilteredLookupElements(shadowedCallablesFilter)
-            .forEach { context.addElement(it) }
+            .forEach { addElement(it) }
     }
 
     context(_: KaSession, context: K2CompletionSectionContext<P>)
@@ -774,7 +774,7 @@ internal abstract class K2AbstractCallableCompletionContributor<P : KotlinNameRe
             }
 
         availableLocalAndMemberNonExtensions.createFilteredLookupElements(shadowedCallablesFilter)
-            .forEach { context.addElement(it) }
+            .forEach { addElement(it) }
     }
 
     context(_: KaSession, context: K2CompletionSectionContext<P>)
@@ -803,7 +803,7 @@ internal abstract class K2AbstractCallableCompletionContributor<P : KotlinNameRe
 
         (extensionsWhichCanBeCalled + availableStaticAndTopLevelNonExtensions)
             .createFilteredLookupElements(shadowedCallablesFilter)
-            .forEach { context.addElement(it) }
+            .forEach { addElement(it) }
     }
 
     context(_: KaSession, context: K2CompletionSectionContext<P>)
@@ -856,7 +856,7 @@ internal abstract class K2AbstractCallableCompletionContributor<P : KotlinNameRe
         }
         createAndFilterMetadataForMemberCallables(enumEntries + enumConstants)
             .createFilteredLookupElements(shadowedCallablesFilter)
-            .forEach { context.addElement(it) }
+            .forEach { addElement(it) }
     }
 
     context(_: KaSession, context: K2CompletionSectionContext<P>)

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/K2ClassReferenceCompletionContributor.kt

@@ -14,9 +14,9 @@ internal class K2ClassReferenceCompletionContributor : K2SimpleCompletionContrib
     context(_: KaSession, context: K2CompletionSectionContext<KotlinCallableReferencePositionContext>)
     override fun complete() {
         if (context.positionContext.explicitReceiver == null) return
-        context.addElement(createKeywordElement("class"))
+        addElement(createKeywordElement("class"))
         if (context.completionContext.targetPlatform.isJvm()) {
-            context.addElement(createKeywordElement("class", tail = ".java"))
+            addElement(createKeywordElement("class", tail = ".java"))
         }
     }
 }

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/K2ClassifierCompletionContributor.kt

@@ -182,7 +182,7 @@ internal open class K2ClassifierCompletionContributor : K2CompletionContributor<
                 }
             }
 
-        scopeClassifiers.forEach { sectionContext.addElement(it) }
+        scopeClassifiers.forEach { addElement(it) }
     }
 
     context(_: KaSession, sectionContext: K2CompletionSectionContext<KotlinNameReferencePositionContext>)
@@ -222,7 +222,7 @@ internal open class K2ClassifierCompletionContributor : K2CompletionContributor<
             emptySequence()
         }
 
-        indexClassifiers.forEach { sectionContext.addElement(it) }
+        indexClassifiers.forEach { addElement(it) }
     }
 
     context(_: KaSession, sectionContext: K2CompletionSectionContext<KotlinNameReferencePositionContext>)
@@ -244,7 +244,7 @@ internal open class K2ClassifierCompletionContributor : K2CompletionContributor<
                         positionContext = sectionContext.positionContext,
                         visibilityChecker = sectionContext.visibilityChecker,
                     ).map { it.applyWeighs(symbolWithOrigin) }
-                }.forEach { sectionContext.addElement(it) }
+                }.forEach { addElement(it) }
         } else {
             sectionContext.sink.registerChainContributor(this)
         }

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/K2DeclarationFromOverridableMembersContributor.kt

@@ -48,7 +48,7 @@ internal class K2DeclarationFromOverridableMembersContributor : K2SimpleCompleti
         if (declaration.hasModifier(KtTokens.OVERRIDE_KEYWORD)) {
             val elements = OverrideKeywordHandler(context.importStrategyDetector)
                 .createOverrideMemberLookups(context.parameters, declaration, context.project)
-            context.addElements(elements)
+            addElements(elements)
         }
     }
 

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/K2DeclarationFromUnresolvedNameContributor.kt

@@ -68,7 +68,7 @@ internal class K2DeclarationFromUnresolvedNameContributor : K2SimpleCompletionCo
         if (unresolvedRef.reference?.resolve() == null) {
             val lookupElement = LookupElementBuilder.create(name).suppressAutoInsertion()
                 .also { it.priority = ItemPriority.FROM_UNRESOLVED_NAME_SUGGESTION }
-            context.addElement(lookupElement)
+            addElement(lookupElement)
         }
     }
 

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/K2ImportDirectivePackageMembersCompletionContributor.kt

@@ -44,6 +44,6 @@ internal class K2ImportDirectivePackageMembersCompletionContributor : K2SimpleCo
                             scopeKind = scopeWithKind.kind,
                         )
                     }
-            }.forEach { context.addElement(it) }
+            }.forEach { addElement(it) }
     }
 }

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/K2KDocParameterNameContributor.kt

@@ -43,7 +43,7 @@ internal class K2KDocParameterNameContributor : K2SimpleCompletionContributor<KD
         getParametersForKDoc(ownerDeclaration.symbol)
             .filter { (it as KaNamedSymbol).name.asString() !in alreadyDocumentedParameters }
             .flatMap { createLookupElements(it) }
-            .forEach { context.addElement(it) }
+            .forEach { addElement(it) }
     }
 
     context(_: KaSession, context: K2CompletionSectionContext<KDocNameReferencePositionContext>)

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/K2KeywordCompletionContributor.kt

@@ -81,7 +81,7 @@ internal class K2KeywordCompletionContributor : K2SimpleCompletionContributor<Ko
                 ?: listOf(lookupElement)
 
             lookups.map { it.applyWeighs() }
-                .forEach { context.addElement(it) }
+                .forEach { addElement(it) }
         }
     }
 

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/K2NamedArgumentCompletionContributor.kt

@@ -95,7 +95,7 @@ internal class K2NamedArgumentCompletionContributor : K2SimpleCompletionContribu
                 }
             }
         }.map { it.applyWeighs() }
-            .forEach { context.addElement(it) }
+            .forEach { addElement(it) }
     }
 
     /**

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/K2OperatorNameCompletionContributor.kt

@@ -27,6 +27,6 @@ internal class K2OperatorNameCompletionContributor : K2SimpleCompletionContribut
         OperatorNameCompletion.getApplicableOperators {
             context.prefixMatcher.prefixMatches(it)
         }.map(OperatorNameLookupElementFactory::createLookup)
-            .forEach { context.addElement(it) }
+            .forEach { addElement(it) }
     }
 }

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/K2PackageCompletionContributor.kt

@@ -52,7 +52,7 @@ internal class K2PackageCompletionContributor : K2SimpleCompletionContributor<Ko
                             scopeKind = KtOutsideTowerScopeKinds.PackageMemberScope,
                         ),
                     )
-            }.forEach { context.addElement(it) }
+            }.forEach { addElement(it) }
     }
 
     override fun K2CompletionSetupScope<KotlinRawPositionContext>.isAppropriatePosition(): Boolean = when (position) {

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/K2SameAsFileClassifierNameCompletionContributor.kt

@@ -16,14 +16,15 @@ internal class K2SameAsFileClassifierNameCompletionContributor : K2SimpleComplet
 
     context(_: KaSession, context: K2CompletionSectionContext<KotlinClassifierNamePositionContext>)
     override fun complete() {
-        (context.positionContext.classLikeDeclaration as? KtClassOrObject)?.let { context.completeTopLevelClassName(it) }
+        (context.positionContext.classLikeDeclaration as? KtClassOrObject)?.let { completeTopLevelClassName(it) }
     }
 
-    private fun K2CompletionSectionContext<KotlinClassifierNamePositionContext>.completeTopLevelClassName(classOrObject: KtClassOrObject) {
+    context(_: KaSession, context: K2CompletionSectionContext<KotlinClassifierNamePositionContext>)
+    private fun completeTopLevelClassName(classOrObject: KtClassOrObject) {
         if (!classOrObject.isTopLevel()) return
-        val name = completionContext.originalFile.virtualFile.nameWithoutExtension
+        val name = context.completionContext.originalFile.virtualFile.nameWithoutExtension
         if (!isValidUpperCapitalizedClassName(name)) return
-        if (completionContext.originalFile.declarations.any { it is KtClassOrObject && it.name == name }) return
+        if (context.completionContext.originalFile.declarations.any { it is KtClassOrObject && it.name == name }) return
         addElement(LookupElementBuilder.create(name))
     }
 

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/K2SuperEntryContributor.kt

@@ -26,7 +26,7 @@ internal class K2SuperEntryContributor : K2SimpleCompletionContributor<KotlinSup
             LookupElementBuilder.create(SuperLookupObject(superType.name, superType.classId), superType.name.asString())
                 .withTailText(tailText)
                 .withInsertHandler(SuperCallInsertionHandler)
-                .let { context.addElement(it) }
+                .let { addElement(it) }
         }
     }
 }

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/K2SuperMemberCompletionContributor.kt

@@ -75,13 +75,13 @@ internal class K2SuperMemberCompletionContributor : K2SimpleCompletionContributo
                 namesNeedDisambiguation = namesNeedDisambiguation,
                 importStrategyDetector = context.importStrategyDetector,
             )
-        }.forEach { context.addElement(it) }
+        }.forEach { addElement(it) }
 
         collectDelegateCallToSuperMember(
             superReceiver = superReceiver,
             nonExtensionMembers = nonExtensionMembers,
             namesNeedDisambiguation = namesNeedDisambiguation,
-        ).forEach { context.addElement(it) }
+        ).forEach { addElement(it) }
     }
 
     context(_: KaSession, context: K2CompletionSectionContext<KotlinSuperReceiverNameReferencePositionContext>)

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/K2TrailingFunctionParameterNameCompletionContributorBase.kt

@@ -135,7 +135,7 @@ internal sealed class K2TrailingFunctionParameterNameCompletionContributorBase<P
                 elementBuilder
                     .apply { isTrailingLambdaParameter = true }
                     .applyWeighs()
-            }.forEach { context.addElement(it) }
+            }.forEach { addElement(it) }
     }
 
     context(_: KaSession, context: K2CompletionSectionContext<P>)

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/K2TypeParameterConstraintNameInWhereClauseCompletionContributor.kt

@@ -38,7 +38,7 @@ internal class K2TypeParameterConstraintNameInWhereClauseCompletionContributor :
                 .withInsertHandler(TypeParameterInWhenClauseInsertionHandler)
                 .withPsiElement(typeParameter.psi)
                 .withIcon(icon)
-                .let { context.addElement(it) }
+                .let { addElement(it) }
         }
     }
 }

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/K2VariableOrParameterNameWithTypeCompletionContributor.kt

@@ -135,7 +135,7 @@ internal class K2VariableOrParameterNameWithTypeCompletionContributor : K2Simple
             if (!shouldInsertType(contextElement) && lookupNamesAdded.contains(name)) continue
 
             lookupElement.nameWithTypePriority = -count // suggestions that appear more often than others get higher priority
-            context.addElement(lookupElement)
+            addElement(lookupElement)
             lookupNamesAdded.add(name)
         }
     }
@@ -225,7 +225,7 @@ internal class K2VariableOrParameterNameWithTypeCompletionContributor : K2Simple
             lookupElement.nameWithTypePriority = userPrefix.length // suggestions with longer user prefix get lower priority
             lookupElement.applyWeighs(KtSymbolWithOrigin(symbol, scopeKind))
 
-            context.addElement(lookupElement)
+            addElement(lookupElement)
         }
     }
 

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/K2WhenWithSubjectConditionContributor.kt

@@ -92,9 +92,9 @@ internal class K2WhenWithSubjectConditionContributor : K2SimpleCompletionContrib
         val isSingleCondition = whenCondition.isSingleConditionInEntry()
 
         createNullBranchLookupElement(subjectType)
-            ?.let { context.addElement(it) }
+            ?.let { addElement(it) }
         createElseBranchLookupElement(whenCondition)
-            ?.let { context.addElement(it) }
+            ?.let { addElement(it) }
 
         when {
             classSymbol?.classKind == KaClassKind.ENUM_CLASS -> {
@@ -163,7 +163,7 @@ internal class K2WhenWithSubjectConditionContributor : K2SimpleCompletionContrib
                     fqName = (classifier as? KaNamedClassSymbol)?.classId?.asSingleFqName(),
                     isSingleCondition = isSingleCondition,
                 )
-            }.forEach { context.addElement(it) }
+            }.forEach { addElement(it) }
 
         if (prefixMatcher.prefix.isNotEmpty()) {
             context.completeLaterInSameSession("Index", priority = K2ContributorSectionPriority.FROM_INDEX) {
@@ -183,7 +183,11 @@ internal class K2WhenWithSubjectConditionContributor : K2SimpleCompletionContrib
                             fqName = (classifier as? KaNamedClassSymbol)?.classId?.asSingleFqName(),
                             isSingleCondition = isSingleCondition,
                         )
-                    }.forEach { innerContext.addElement(it) }
+                    }.forEach {
+                        context(innerContext) {
+                            addElement(it)
+                        }
+                    }
             }
         }
     }
@@ -227,7 +231,7 @@ internal class K2WhenWithSubjectConditionContributor : K2SimpleCompletionContrib
                     fqName = classId.asSingleFqName(),
                     isSingleCondition = isSingleCondition,
                 )
-            }.forEach { context.addElement(it) }
+            }.forEach { addElement(it) }
 
         if (getAllSealedInheritors(classSymbol).any { it.modality == KaSymbolModality.ABSTRACT }) {
             completeAllTypes(
@@ -307,7 +311,7 @@ internal class K2WhenWithSubjectConditionContributor : K2SimpleCompletionContrib
                     fqName = entry.callableId?.asSingleFqName(),
                     isSingleCondition = isSingleCondition,
                 )
-            }.forEach { context.addElement(it) }
+            }.forEach { addElement(it) }
     }
 
     private fun KtWhenCondition.isSingleConditionInEntry(): Boolean {