Fix rewrite at slice exception for callable references inside flatMap

#KT-40254 Fixed

compiler/fir/analysis-tests/tests/org/jetbrains/kotlin/fir/FirOldFrontendDiagnosticsTestGenerated.java

@@ -2803,6 +2803,11 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirOldFronte
                 runTest("compiler/testData/diagnostics/tests/callableReference/resolve/resolveReferenceAgainstKFunctionAndKPrpoerty.kt");
             }
 
+            @TestMetadata("resolveTwoReferencesAgainstGenerics.kt")
+            public void testResolveTwoReferencesAgainstGenerics() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/callableReference/resolve/resolveTwoReferencesAgainstGenerics.kt");
+            }
+
             @TestMetadata("valVsFun.kt")
             public void testValVsFun() throws Exception {
                 runTest("compiler/testData/diagnostics/tests/callableReference/resolve/valVsFun.kt");

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/KotlinCallCompleter.kt

@@ -101,11 +101,8 @@ class KotlinCallCompleter(
             return candidates
         }
 
-        val newAtoms = mutableMapOf<KotlinResolutionCandidate, ResolvedLambdaAtom>()
-        for ((candidate, atom) in lambdas.entries) {
-            val newAtom = kotlinConstraintSystemCompleter.prepareLambdaAtomForFactoryPattern(atom, candidate, candidate)
-            newAtoms[candidate] = newAtom
-            candidate.addResolvedKtPrimitive(newAtom)
+        val newAtoms = lambdas.mapValues { (candidate, atom) ->
+            kotlinConstraintSystemCompleter.prepareLambdaAtomForFactoryPattern(atom, candidate, candidate)
         }
 
         val diagnosticHolderForLambda = KotlinDiagnosticsHolder.SimpleHolder()

compiler/testData/diagnostics/tests/callableReference/resolve/resolveTwoReferencesAgainstGenerics.fir.kt

@@ -0,0 +1,10 @@
+// !LANGUAGE: +NewInference
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+// WITH_RUNTIME
+
+
+fun <T> List<Option<T>>.flatten(): List<T> = <!AMBIGUITY!>flatMap<!> { <!UNRESOLVED_REFERENCE!>it<!>.<!INAPPLICABLE_CANDIDATE!>fold<!>(::emptyList, <!UNRESOLVED_REFERENCE!>::listOf<!>) }
+
+class Option<out T> {
+    fun <R> fold(ifEmpty: () -> R, ifSome: (T) -> R): R = TODO()
+}

compiler/testData/diagnostics/tests/callableReference/resolve/resolveTwoReferencesAgainstGenerics.kt

@@ -0,0 +1,10 @@
+// !LANGUAGE: +NewInference
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+// WITH_RUNTIME
+
+
+fun <T> List<Option<T>>.flatten(): List<T> = flatMap { it.fold(::emptyList, ::listOf) }
+
+class Option<out T> {
+    fun <R> fold(ifEmpty: () -> R, ifSome: (T) -> R): R = TODO()
+}

compiler/testData/diagnostics/tests/callableReference/resolve/resolveTwoReferencesAgainstGenerics.txt

@@ -0,0 +1,11 @@
+package
+
+public fun </*0*/ T> kotlin.collections.List<Option<T>>.flatten(): kotlin.collections.List<T>
+
+public final class Option</*0*/ out T> {
+    public constructor Option</*0*/ out T>()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public final fun </*0*/ R> fold(/*0*/ ifEmpty: () -> R, /*1*/ ifSome: (T) -> R): R
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

compiler/tests/org/jetbrains/kotlin/checkers/DiagnosticsTestGenerated.java

@@ -2810,6 +2810,11 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTestWithFirVali
                     runTest("compiler/testData/diagnostics/tests/callableReference/resolve/resolveReferenceAgainstKFunctionAndKPrpoerty.kt");
                 }
 
+                @TestMetadata("resolveTwoReferencesAgainstGenerics.kt")
+                public void testResolveTwoReferencesAgainstGenerics() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/callableReference/resolve/resolveTwoReferencesAgainstGenerics.kt");
+                }
+
                 @TestMetadata("valVsFun.kt")
                 public void testValVsFun() throws Exception {
                     runTest("compiler/testData/diagnostics/tests/callableReference/resolve/valVsFun.kt");

compiler/tests/org/jetbrains/kotlin/checkers/javac/DiagnosticsUsingJavacTestGenerated.java

@@ -2805,6 +2805,11 @@ public class DiagnosticsUsingJavacTestGenerated extends AbstractDiagnosticsUsing
                     runTest("compiler/testData/diagnostics/tests/callableReference/resolve/resolveReferenceAgainstKFunctionAndKPrpoerty.kt");
                 }
 
+                @TestMetadata("resolveTwoReferencesAgainstGenerics.kt")
+                public void testResolveTwoReferencesAgainstGenerics() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/callableReference/resolve/resolveTwoReferencesAgainstGenerics.kt");
+                }
+
                 @TestMetadata("valVsFun.kt")
                 public void testValVsFun() throws Exception {
                     runTest("compiler/testData/diagnostics/tests/callableReference/resolve/valVsFun.kt");