From aae8cd2a7ccf8d16e5f049bce769d8e997f564e8 Mon Sep 17 00:00:00 2001
From: Brian Norman <brian.norman@jetbrains.com>
Date: Tue, 2 Jan 2024 15:06:17 -0600
Subject: [PATCH] [FIR] Local variable assignment must be propagated before
 loops

When performing lookahead for local variable assignments, make sure
assignments taking place within loops are being propagated before loops.
This makes sure smartcasts within non-inline declarations before the
loop are disallowed.

^KT-63867 Fixed
---
 .../dfa/FirLocalVariableAssignmentAnalyzer.kt |  3 ++
 .../lambdaWithCallInPlace.fir.kt              | 39 +++++++++++++++++--
 .../lambdaWithCallInPlace.kt                  | 33 ++++++++++++++++
 3 files changed, 72 insertions(+), 3 deletions(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirLocalVariableAssignmentAnalyzer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirLocalVariableAssignmentAnalyzer.kt
index 1d763c56039..d17ad4cfacf 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirLocalVariableAssignmentAnalyzer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirLocalVariableAssignmentAnalyzer.kt
@@ -333,6 +333,9 @@ internal class FirLocalVariableAssignmentAnalyzer {
             override fun visitLoop(loop: FirLoop, data: MiniCfgData) {
                 val entry = data.flow
                 val assignedInside = visitElementWithLexicalScope(loop, data)
+                // Now that the inner variables have been discarded, the rest can be propagated to prevent smartcasts
+                // in declarations that came before this loop.
+                entry.recordAssignments(assignedInside)
                 // All forks in the loop should have the same set of variables assigned later, equal to the set
                 // at the start of the loop.
                 data.flow.recordAssignments(assignedInside)
diff --git a/compiler/testData/diagnostics/tests/smartCasts/lambdasWithContracts/lambdaWithCallInPlace.fir.kt b/compiler/testData/diagnostics/tests/smartCasts/lambdasWithContracts/lambdaWithCallInPlace.fir.kt
index 21dc9e562d4..c1d2ce2459e 100644
--- a/compiler/testData/diagnostics/tests/smartCasts/lambdasWithContracts/lambdaWithCallInPlace.fir.kt
+++ b/compiler/testData/diagnostics/tests/smartCasts/lambdasWithContracts/lambdaWithCallInPlace.fir.kt
@@ -468,7 +468,7 @@ fun test48() {
     var x: Any = materialize()
     runWithoutContract {
         while (x is String) {
-            x.length
+            <!SMARTCAST_IMPOSSIBLE!>x<!>.length
             x = 10
         }
     }
@@ -536,7 +536,7 @@ fun test55() {
     runWithoutContract {
         for (i in 1..3) {
             require(x is String)
-            x.length
+            <!SMARTCAST_IMPOSSIBLE!>x<!>.length
             x = 10
         }
     }
@@ -619,4 +619,37 @@ fun test63() {
         x?.length ?: -1
         x = null
     }
-}
\ No newline at end of file
+}
+
+fun test64() {
+    var x: Any = materialize()
+    require(x is String)
+    runWithoutContract {
+        <!SMARTCAST_IMPOSSIBLE!>x<!>.length
+    }
+    for (i in 1..3) {
+        x = 10
+    }
+}
+
+fun test65() {
+    var x: Any = materialize()
+    require(x is String)
+    atLeastOnce {
+        x.length
+    }
+    for (i in 1..3) {
+        x = 10
+    }
+}
+
+fun test66() {
+    var x: Any = materialize()
+    require(x is String)
+    exactlyOnce {
+        x.length
+    }
+    for (i in 1..3) {
+        x = 10
+    }
+}
diff --git a/compiler/testData/diagnostics/tests/smartCasts/lambdasWithContracts/lambdaWithCallInPlace.kt b/compiler/testData/diagnostics/tests/smartCasts/lambdasWithContracts/lambdaWithCallInPlace.kt
index 7a0c77f74a0..06381824edb 100644
--- a/compiler/testData/diagnostics/tests/smartCasts/lambdasWithContracts/lambdaWithCallInPlace.kt
+++ b/compiler/testData/diagnostics/tests/smartCasts/lambdasWithContracts/lambdaWithCallInPlace.kt
@@ -620,3 +620,36 @@ fun test63() {
         x = null
     }
 }
+
+fun test64() {
+    var x: Any = materialize()
+    require(x is String)
+    runWithoutContract {
+        <!SMARTCAST_IMPOSSIBLE!>x<!>.length
+    }
+    for (i in 1..3) {
+        x = 10
+    }
+}
+
+fun test65() {
+    var x: Any = materialize()
+    require(x is String)
+    atLeastOnce {
+        <!SMARTCAST_IMPOSSIBLE!>x<!>.length
+    }
+    for (i in 1..3) {
+        x = 10
+    }
+}
+
+fun test66() {
+    var x: Any = materialize()
+    require(x is String)
+    exactlyOnce {
+        <!SMARTCAST_IMPOSSIBLE!>x<!>.length
+    }
+    for (i in 1..3) {
+        x = 10
+    }
+}