K2: Get rid of AugmentedAssignmentCallOption resolution mode

It's been used only for regular function call nodes resolution,
so I made it more explicit.

The main idea behind this change is ResolutionMode hierarchy
simplification and the kind of statement
that "AugmentedAssignmentCallOption" is only relevant to calls.

Instead, CallResolutionMode is introduced that is only used for
`transformFunctionCallInternal` call sites.

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirCallResolver.kt

@@ -616,7 +616,7 @@ class FirCallResolver(
 
                 annotation.transformArgumentList { argumentsToParameters[it]?.returnTypeRef }
             } else {
-                annotation.replaceArgumentList(annotation.argumentList.transform(transformer, ResolutionMode.ContextDependent.Default))
+                annotation.replaceArgumentList(annotation.argumentList.transform(transformer, ResolutionMode.ContextDependent))
             }
 
             val callInfo = toCallInfo(annotation, reference)
@@ -633,7 +633,7 @@ class FirCallResolver(
                 explicitReceiver = null
             )
         } else {
-            annotation.replaceArgumentList(annotation.argumentList.transform(transformer, ResolutionMode.ContextDependent.Default))
+            annotation.replaceArgumentList(annotation.argumentList.transform(transformer, ResolutionMode.ContextDependent))
 
             val callInfo = toCallInfo(annotation, reference)
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolutionMode.kt

@@ -13,21 +13,9 @@ import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 
 sealed class ResolutionMode(
     val forceFullCompletion: Boolean,
-    // Only true AugmentedAssignmentCallOption, don't run an even slightest form of completion
-    val skipEvenPartialCompletion: Boolean = false,
 ) {
-    sealed class ContextDependent(
-        skipCompletion: Boolean = false,
-    ) : ResolutionMode(forceFullCompletion = false, skipCompletion) {
-        companion object Default : ContextDependent() {
-            override fun toString(): String = "ContextDependent"
-        }
-
-        data object AugmentedAssignmentCallOption : ContextDependent(skipCompletion = true)
-    }
-
+    data object ContextDependent : ResolutionMode(forceFullCompletion = false)
     data object Delegate : ResolutionMode(forceFullCompletion = false)
-
     data object ContextIndependent : ResolutionMode(forceFullCompletion = true)
 
     sealed class ReceiverResolution(val forCallableReference: Boolean) : ResolutionMode(forceFullCompletion = true) {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirCallCompleter.kt

@@ -55,7 +55,12 @@ class FirCallCompleter(
 
     val completer = ConstraintSystemCompleter(components, transformer.context)
 
-    fun <T> completeCall(call: T, resolutionMode: ResolutionMode): T where T : FirResolvable, T : FirStatement {
+    fun <T> completeCall(
+        call: T,
+        resolutionMode: ResolutionMode,
+        // Only expected to be true for resolving different versions of augmented assignments
+        skipEvenPartialCompletion: Boolean = false,
+    ): T where T : FirResolvable, T : FirStatement {
         val typeRef = components.typeFromCallee(call)
 
         val reference = call.calleeReference as? FirNamedReferenceWithCandidate ?: return call
@@ -74,7 +79,7 @@ class FirCallCompleter(
             resolutionMode,
         )
 
-        if (resolutionMode.skipEvenPartialCompletion) return call
+        if (skipEvenPartialCompletion) return call
 
         val completionMode = candidate.computeCompletionMode(
             session.inferenceComponents, resolutionMode, initialType

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt

@@ -403,7 +403,7 @@ open class FirDeclarationsResolveTransformer(
         // TODO: this generates some nodes in the control flow graph which we don't want if we
         //  end up not selecting this option, KT-59684
         transformer.expressionsTransformer?.transformFunctionCallInternal(
-            provideDelegateCall, ResolutionMode.ReceiverResolution, provideDelegate = true
+            provideDelegateCall, ResolutionMode.ReceiverResolution, FirExpressionsResolveTransformer.CallResolutionMode.PROVIDE_DELEGATE,
         )
 
         // If we got successful candidate for provideDelegate, let's select it

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -392,12 +392,29 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
     }
 
     override fun transformFunctionCall(functionCall: FirFunctionCall, data: ResolutionMode): FirStatement =
-        transformFunctionCallInternal(functionCall, data, provideDelegate = false)
+        transformFunctionCallInternal(functionCall, data, CallResolutionMode.REGULAR)
+
+    internal enum class CallResolutionMode {
+        REGULAR,
+
+        /**
+         * For PROVIDE_DELEGATE we skip transforming explicit receiver of the call since it's already been resolved
+         * at [org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirDeclarationsResolveTransformer.transformPropertyAccessorsWithDelegate]
+         */
+        PROVIDE_DELEGATE,
+
+        /**
+         * When we're resolving an operator like `a += b` we try to resolve it with different options of desugaring like
+         * `a = a.plus(b)` and `a.plusAssign(b)` until find something that looks successful.
+         * But at this stage, we skip transformation of receiver, arguments and skip completion in any form.
+         */
+        OPTION_FOR_AUGMENTED_ASSIGNMENT,
+    }
 
     internal fun transformFunctionCallInternal(
         functionCall: FirFunctionCall,
         data: ResolutionMode,
-        provideDelegate: Boolean,
+        callResolutionMode: CallResolutionMode,
     ): FirStatement =
         whileAnalysing(session, functionCall) {
             val calleeReference = functionCall.calleeReference
@@ -419,8 +436,8 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
             functionCall.transformAnnotations(transformer, data)
             functionCall.replaceLambdaArgumentInvocationKinds(session)
             functionCall.transformTypeArguments(transformer, ResolutionMode.ContextIndependent)
-            val resolvingAugmentedAssignment = data == ResolutionMode.ContextDependent.AugmentedAssignmentCallOption
-            val withTransformedArguments = if (!resolvingAugmentedAssignment) {
+            val choosingOptionForAugmentedAssignment = callResolutionMode == CallResolutionMode.OPTION_FOR_AUGMENTED_ASSIGNMENT
+            val withTransformedArguments = if (!choosingOptionForAugmentedAssignment) {
                 dataFlowAnalyzer.enterCallArguments(functionCall, functionCall.arguments)
                 // In provideDelegate mode the explicitReceiver is already resolved
                 // E.g. we have val some by someDelegate
@@ -428,7 +445,8 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
                 // at 2nd stage in provideDelegate mode we are trying to resolve someDelegate.provideDelegate(),
                 // and 'someDelegate' explicit receiver is resolved at 1st stage
                 // See also FirDeclarationsResolveTransformer.transformWrappedDelegateExpression
-                val withResolvedExplicitReceiver = if (provideDelegate) functionCall else transformExplicitReceiver(functionCall)
+                val withResolvedExplicitReceiver =
+                    if (callResolutionMode == CallResolutionMode.PROVIDE_DELEGATE) functionCall else transformExplicitReceiver(functionCall)
                 withResolvedExplicitReceiver.also {
                     it.replaceArgumentList(it.argumentList.transform(this, ResolutionMode.ContextDependent))
                     dataFlowAnalyzer.exitCallArguments()
@@ -439,9 +457,12 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
 
             val resultExpression = callResolver.resolveCallAndSelectCandidate(withTransformedArguments, data)
 
-            val completeInference = callCompleter.completeCall(resultExpression, data)
+            val completeInference = callCompleter.completeCall(
+                resultExpression, data,
+                skipEvenPartialCompletion = choosingOptionForAugmentedAssignment,
+            )
             val result = completeInference.transformToIntegerOperatorCallOrApproximateItIfNeeded(data)
-            if (!resolvingAugmentedAssignment) {
+            if (!choosingOptionForAugmentedAssignment) {
                 dataFlowAnalyzer.exitFunctionCall(result, data.forceFullCompletion)
             }
 
@@ -818,8 +839,8 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
     private fun FirFunctionCall.resolveCandidateForAssignmentOperatorCall(): FirFunctionCall {
         return transformFunctionCallInternal(
             this,
-            ResolutionMode.ContextDependent.AugmentedAssignmentCallOption,
-            provideDelegate = false
+            ResolutionMode.ContextDependent,
+            CallResolutionMode.OPTION_FOR_AUGMENTED_ASSIGNMENT
         ) as FirFunctionCall
     }
 
@@ -1641,7 +1662,7 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
     override fun transformArrayLiteral(arrayLiteral: FirArrayLiteral, data: ResolutionMode): FirStatement =
         whileAnalysing(session, arrayLiteral) {
             when (data) {
-                is ResolutionMode.ContextDependent.Default -> {
+                is ResolutionMode.ContextDependent -> {
                     // Argument in non-annotation call (unsupported) or if type of array parameter is unresolved.
                     arrayLiteral.transformChildren(transformer, data)
                     arrayLiteral