From 76b9ba05fdd807de778281bb51dfe66007fa595f Mon Sep 17 00:00:00 2001
From: Brian Norman <brian.norman@jetbrains.com>
Date: Fri, 26 Jan 2024 09:18:55 -0600
Subject: [PATCH] [FIR] Lambda reassigned arguments should not have contracts
 applied

When applying function contracts to arguments, make sure the argument
variable hasn't been reassigned within a lambda. Contracts can only be
applied to unchanged variables, otherwise outdated type statements could
be added to the variable.

^KT-63151 Fixed
---
 .../fir/resolve/dfa/FirDataFlowAnalyzer.kt    | 35 +++++++++++++------
 .../variables/reassignedInRhs.fir.kt          |  2 +-
 2 files changed, 26 insertions(+), 11 deletions(-)

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 10d0e96f0a1..47e85b43cd6 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
@@ -516,10 +516,8 @@ abstract class FirDataFlowAnalyzer(
         }
 
         // Only consider the LHS variable if it has not been reassigned in the RHS.
-        val leftOperandVariable = variableStorage.getOrCreateIfReal(flow, leftOperand).takeIf {
-            val variable = variableStorage.getRealVariableWithoutUnwrappingAlias(lhsExitFlow, leftOperand)
-            variable == null || logicSystem.isSameValueIn(lhsExitFlow, flow, variable)
-        }
+        val leftOperandVariable = variableStorage.getOrCreateIfReal(flow, leftOperand)
+            .takeIf { isSameValueIn(lhsExitFlow, leftOperand, flow) }
         val rightOperandVariable = variableStorage.getOrCreateIfReal(flow, rightOperand)
         if (leftOperandVariable == null && rightOperandVariable == null) return
         val expressionVariable = variableStorage.createSynthetic(expression)
@@ -818,7 +816,7 @@ abstract class FirDataFlowAnalyzer(
 
     fun exitQualifiedAccessExpression(qualifiedAccessExpression: FirQualifiedAccessExpression) {
         graphBuilder.exitQualifiedAccessExpression(qualifiedAccessExpression).mergeIncomingFlow { _, flow ->
-            processConditionalContract(flow, qualifiedAccessExpression)
+            processConditionalContract(flow, qualifiedAccessExpression, callArgsExit = null)
         }
     }
 
@@ -909,8 +907,10 @@ abstract class FirDataFlowAnalyzer(
 
     fun exitFunctionCall(functionCall: FirFunctionCall, callCompleted: Boolean) {
         context.variableAssignmentAnalyzer.exitFunctionCall(callCompleted)
-        graphBuilder.exitFunctionCall(functionCall, callCompleted).mergeIncomingFlow { _, flow ->
-            processConditionalContract(flow, functionCall)
+        val node = graphBuilder.exitFunctionCall(functionCall, callCompleted)
+        node.mergeIncomingFlow { _, flow ->
+            val callArgsExit = node.previousNodes.singleOrNull { it is FunctionCallArgumentsExitNode }
+            processConditionalContract(flow, functionCall, callArgsExit?.flow)
         }
     }
 
@@ -958,7 +958,11 @@ abstract class FirDataFlowAnalyzer(
         }
     }
 
-    private fun processConditionalContract(flow: MutableFlow, qualifiedAccess: FirStatement) {
+    private fun processConditionalContract(
+        flow: MutableFlow,
+        qualifiedAccess: FirStatement,
+        callArgsExit: PersistentFlow?,
+    ) {
         // contracts has no effect on non-body resolve stages
         if (!components.transformer.baseTransformerPhase.isBodyResolve) return
 
@@ -982,7 +986,13 @@ abstract class FirDataFlowAnalyzer(
         if (conditionalEffects.isEmpty()) return
 
         val arguments = qualifiedAccess.orderedArguments(callee) ?: return
-        val argumentVariables = Array(arguments.size) { i -> arguments[i]?.let { variableStorage.getOrCreateIfReal(flow, it) } }
+        val argumentVariables = Array(arguments.size) { i ->
+            arguments[i]?.let { argument ->
+                variableStorage.getOrCreateIfReal(flow, argument)
+                    // Only apply contract information to argument if it has not been reassigned in a lambda.
+                    .takeIf { callArgsExit == null || isSameValueIn(callArgsExit, argument, flow) }
+            }
+        }
         if (argumentVariables.all { it == null }) return
 
         val typeParameters = callee.typeParameters
@@ -1043,7 +1053,7 @@ abstract class FirDataFlowAnalyzer(
                     logicSystem.recordNewAssignment(flow, variable, context.newAssignmentIndex())
                 }
             }
-            processConditionalContract(flow, assignment)
+            processConditionalContract(flow, assignment, callArgsExit = null)
         }
     }
 
@@ -1463,6 +1473,11 @@ abstract class FirDataFlowAnalyzer(
         }
     }
 
+    private fun isSameValueIn(other: PersistentFlow, fir: FirElement, original: MutableFlow): Boolean {
+        val variable = variableStorage.getRealVariableWithoutUnwrappingAlias(other, fir)
+        return variable == null || logicSystem.isSameValueIn(other, original, variable)
+    }
+
     private fun MutableFlow.addImplication(statement: Implication) {
         logicSystem.addImplication(this, statement)
     }
diff --git a/compiler/testData/diagnostics/tests/smartCasts/variables/reassignedInRhs.fir.kt b/compiler/testData/diagnostics/tests/smartCasts/variables/reassignedInRhs.fir.kt
index 79563038161..8ca2ffee730 100644
--- a/compiler/testData/diagnostics/tests/smartCasts/variables/reassignedInRhs.fir.kt
+++ b/compiler/testData/diagnostics/tests/smartCasts/variables/reassignedInRhs.fir.kt
@@ -46,6 +46,6 @@ fun safeCall() {
 fun contractFunction() {
     var x: String? = ""
     if (isNotNullAlsoCall(x) { x = null }) {
-        x.length // bad (#2 == true => x#0 != null; but this is x#1 = null)
+        x<!UNSAFE_CALL!>.<!>length // bad (#2 == true => x#0 != null; but this is x#1 = null)
     }
 }