diff --git a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosticCompilerTestFE10TestdataTestGenerated.java b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosticCompilerTestFE10TestdataTestGenerated.java
index 2161486..60ac660 100644
--- a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosticCompilerTestFE10TestdataTestGenerated.java
+++ b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosticCompilerTestFE10TestdataTestGenerated.java
@@ -16650,6 +16650,12 @@
                 }
 
                 @Test
+                @TestMetadata("considerContractsOfIncompleteCalls.kt")
+                public void testConsiderContractsOfIncompleteCalls() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/builderInference/considerContractsOfIncompleteCalls.kt");
+                }
+
+                @Test
                 @TestMetadata("equalityChecksOnIntegerTypes.kt")
                 public void testEqualityChecksOnIntegerTypes() throws Exception {
                     runTest("compiler/testData/diagnostics/tests/inference/builderInference/equalityChecksOnIntegerTypes.kt");
diff --git a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated.java b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated.java
index cb3b4b1..7e30c74 100644
--- a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated.java
+++ b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated.java
@@ -16650,6 +16650,12 @@
                 }
 
                 @Test
+                @TestMetadata("considerContractsOfIncompleteCalls.kt")
+                public void testConsiderContractsOfIncompleteCalls() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/builderInference/considerContractsOfIncompleteCalls.kt");
+                }
+
+                @Test
                 @TestMetadata("equalityChecksOnIntegerTypes.kt")
                 public void testEqualityChecksOnIntegerTypes() throws Exception {
                     runTest("compiler/testData/diagnostics/tests/inference/builderInference/equalityChecksOnIntegerTypes.kt");
diff --git a/compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirLightTreeOldFrontendDiagnosticsTestGenerated.java b/compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirLightTreeOldFrontendDiagnosticsTestGenerated.java
index 95d3d1b..a7f0424 100644
--- a/compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirLightTreeOldFrontendDiagnosticsTestGenerated.java
+++ b/compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirLightTreeOldFrontendDiagnosticsTestGenerated.java
@@ -16644,6 +16644,12 @@
                 }
 
                 @Test
+                @TestMetadata("considerContractsOfIncompleteCalls.kt")
+                public void testConsiderContractsOfIncompleteCalls() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/builderInference/considerContractsOfIncompleteCalls.kt");
+                }
+
+                @Test
                 @TestMetadata("equalityChecksOnIntegerTypes.kt")
                 public void testEqualityChecksOnIntegerTypes() throws Exception {
                     runTest("compiler/testData/diagnostics/tests/inference/builderInference/equalityChecksOnIntegerTypes.kt");
diff --git a/compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirPsiOldFrontendDiagnosticsTestGenerated.java b/compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirPsiOldFrontendDiagnosticsTestGenerated.java
index 524b317..d792e6f 100644
--- a/compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirPsiOldFrontendDiagnosticsTestGenerated.java
+++ b/compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirPsiOldFrontendDiagnosticsTestGenerated.java
@@ -16650,6 +16650,12 @@
                 }
 
                 @Test
+                @TestMetadata("considerContractsOfIncompleteCalls.kt")
+                public void testConsiderContractsOfIncompleteCalls() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/builderInference/considerContractsOfIncompleteCalls.kt");
+                }
+
+                @Test
                 @TestMetadata("equalityChecksOnIntegerTypes.kt")
                 public void testEqualityChecksOnIntegerTypes() throws Exception {
                     runTest("compiler/testData/diagnostics/tests/inference/builderInference/equalityChecksOnIntegerTypes.kt");
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt
index 1970f78..5da3a6e 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt
@@ -20,6 +20,7 @@
 import org.jetbrains.kotlin.fir.references.toResolvedPropertySymbol
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
+import org.jetbrains.kotlin.fir.resolve.calls.candidate
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.*
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutorByMap
@@ -30,7 +31,10 @@
 import org.jetbrains.kotlin.fir.scopes.impl.toConeType
 import org.jetbrains.kotlin.fir.scopes.unsubstitutedScope
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.*
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.StandardClassIds
@@ -903,6 +907,10 @@
 
     private fun FirStatement.orderedArguments(callee: FirFunction): Array<out FirExpression?>? {
         fun FirQualifiedAccessExpression.firstReceiver(): FirExpression? {
+            val candidate = candidate()
+            if (candidate != null) {
+                return candidate.chosenExtensionReceiverExpression() ?: candidate.dispatchReceiverExpression()
+            }
             return extensionReceiver ?: dispatchReceiver
         }
 
@@ -914,7 +922,7 @@
 
         return when (this) {
             is FirFunctionCall -> {
-                val argumentToParameter = resolvedArgumentMapping ?: return null
+                val argumentToParameter = resolvedArgumentMapping ?: candidate()?.argumentMapping ?: return null
                 val parameterToArgument = argumentToParameter.entries.associate { it.value to it.key.unwrapArgument() }
                 Array(callee.valueParameters.size + 1) { i ->
                     if (i > 0) parameterToArgument[callee.valueParameters[i - 1]] else receiver
diff --git a/compiler/testData/diagnostics/tests/inference/builderInference/considerContractsOfIncompleteCalls.fir.kt b/compiler/testData/diagnostics/tests/inference/builderInference/considerContractsOfIncompleteCalls.fir.kt
new file mode 100644
index 0000000..654a487
--- /dev/null
+++ b/compiler/testData/diagnostics/tests/inference/builderInference/considerContractsOfIncompleteCalls.fir.kt
@@ -0,0 +1,24 @@
+// WITH_STDLIB
+import kotlin.contracts.ExperimentalContracts
+import kotlin.contracts.contract
+
+class Controller<T> {
+    fun yield(t: T): Boolean = true
+}
+
+@OptIn(ExperimentalContracts::class)
+fun <R, B : Any> Controller<R>.ensureNotNull(value: B?, shift: () -> R): B {
+    contract { returns() implies (value != null) }
+    return value!!
+}
+
+fun <S> generate(g: suspend Controller<S>.() -> Unit): S = TODO()
+
+fun bar(x: Int) {}
+
+fun foo(x: Int?) {
+    generate {
+        ensureNotNull(x) { "" }
+        bar(x)
+    }
+}
diff --git a/compiler/testData/diagnostics/tests/inference/builderInference/considerContractsOfIncompleteCalls.kt b/compiler/testData/diagnostics/tests/inference/builderInference/considerContractsOfIncompleteCalls.kt
new file mode 100644
index 0000000..15741b7
--- /dev/null
+++ b/compiler/testData/diagnostics/tests/inference/builderInference/considerContractsOfIncompleteCalls.kt
@@ -0,0 +1,24 @@
+// WITH_STDLIB
+import kotlin.contracts.ExperimentalContracts
+import kotlin.contracts.contract
+
+class Controller<T> {
+    fun yield(t: T): Boolean = true
+}
+
+@OptIn(ExperimentalContracts::class)
+fun <R, B : Any> Controller<R>.ensureNotNull(value: B?, shift: () -> R): B {
+    contract { returns() implies (value != null) }
+    return value!!
+}
+
+fun <S> generate(g: suspend Controller<S>.() -> Unit): S = TODO()
+
+fun bar(x: Int) {}
+
+fun foo(x: Int?) {
+    generate {
+        ensureNotNull(x) { "" }
+        bar(<!DEBUG_INFO_SMARTCAST!>x<!>)
+    }
+}
diff --git a/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java b/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java
index 678aa44..81de371 100644
--- a/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java
+++ b/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java
@@ -16650,6 +16650,12 @@
                 }
 
                 @Test
+                @TestMetadata("considerContractsOfIncompleteCalls.kt")
+                public void testConsiderContractsOfIncompleteCalls() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/builderInference/considerContractsOfIncompleteCalls.kt");
+                }
+
+                @Test
                 @TestMetadata("equalityChecksOnIntegerTypes.kt")
                 public void testEqualityChecksOnIntegerTypes() throws Exception {
                     runTest("compiler/testData/diagnostics/tests/inference/builderInference/equalityChecksOnIntegerTypes.kt");