Preliminary: support smart casts on variables in closures

No feature support yet
So #KT-14486 In Progress

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/smartcasts/DataFlowValueFactory.kt

@@ -40,6 +40,7 @@ import org.jetbrains.kotlin.resolve.scopes.receivers.ImplicitReceiver
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
 import org.jetbrains.kotlin.resolve.scopes.receivers.TransientReceiver
 import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.expressions.AssignedVariablesSearcher.Writer
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingUtils
 import org.jetbrains.kotlin.types.expressions.PreliminaryDeclarationVisitor
 import org.jetbrains.kotlin.types.isError
@@ -296,18 +297,40 @@ object DataFlowValueFactory {
             false
     }
 
+    private fun hasNoWritersInClosures(
+            variableContainingDeclaration: DeclarationDescriptor,
+            writers: Set<Writer>,
+            bindingContext: BindingContext
+    ): Boolean {
+        return writers.none { (_, writerDeclaration) ->
+            val writerDescriptor = writerDeclaration?.let {
+                ControlFlowInformationProvider.getDeclarationDescriptorIncludingConstructors(bindingContext, it)
+            }
+            writerDeclaration != null && variableContainingDeclaration != writerDescriptor
+        }
+    }
+
+    private fun isAccessedInsideClosureAfterAllWriters(
+            writers: Set<Writer>,
+            accessElement: KtElement
+    ): Boolean {
+        val parent = ControlFlowInformationProvider.getElementParentDeclaration(accessElement) ?: return false
+        return writers.none { (assignment) -> !assignment.before(parent) }
+    }
+
     private fun isAccessedBeforeAllClosureWriters(
             variableContainingDeclaration: DeclarationDescriptor,
-            writers: Set<KtDeclaration?>,
+            writers: Set<Writer>,
             bindingContext: BindingContext,
             accessElement: KtElement
     ): Boolean {
         // All writers should be before access element, with the exception:
         // writer which is the same with declaration site does not count
-        writers.filterNotNull().forEach { writer ->
+        writers.mapNotNull { it.declaration }.forEach { writerDeclaration ->
-            val writerDescriptor = ControlFlowInformationProvider.getDeclarationDescriptorIncludingConstructors(bindingContext, writer)
+            val writerDescriptor = ControlFlowInformationProvider.getDeclarationDescriptorIncludingConstructors(
-            // Access is after some writer
+                    bindingContext, writerDeclaration)
-            if (variableContainingDeclaration != writerDescriptor && !accessElement.before(writer)) {
+            // Access is after some writerDeclaration
+            if (variableContainingDeclaration != writerDescriptor && !accessElement.before(writerDeclaration)) {
                 return false
             }
         }
@@ -353,7 +376,16 @@ object DataFlowValueFactory {
 
         // If access element is inside closure: captured
         val variableContainingDeclaration = variableDescriptor.containingDeclaration
-        if (isAccessedInsideClosure(variableContainingDeclaration, bindingContext, accessElement)) return CAPTURED_VARIABLE
+        if (isAccessedInsideClosure(variableContainingDeclaration, bindingContext, accessElement)) {
+            // stable iff we have no writers in closures AND this closure is AFTER all writers
+            return if (hasNoWritersInClosures(variableContainingDeclaration, writers, bindingContext) &&
+                       isAccessedInsideClosureAfterAllWriters(writers, accessElement)) {
+                STABLE_VARIABLE
+            }
+            else {
+                CAPTURED_VARIABLE
+            }
+        }
 
         // Otherwise, stable iff considered position is BEFORE all writers except declarer itself
         return if (isAccessedBeforeAllClosureWriters(variableContainingDeclaration, writers, bindingContext, accessElement))

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/AssignedVariablesSearcher.kt

@@ -25,9 +25,11 @@ import org.jetbrains.kotlin.psi.*
 
 abstract class AssignedVariablesSearcher: KtTreeVisitorVoid() {
 
-    private val assignedNames: SetMultimap<Name, KtDeclaration?> = LinkedHashMultimap.create()
+    data class Writer(val assignment: KtBinaryExpression, val declaration: KtDeclaration?)
 
-    open fun writers(variableDescriptor: VariableDescriptor): MutableSet<KtDeclaration?> = assignedNames[variableDescriptor.name]
+    private val assignedNames: SetMultimap<Name, Writer> = LinkedHashMultimap.create()
+
+    open fun writers(variableDescriptor: VariableDescriptor): MutableSet<Writer> = assignedNames[variableDescriptor.name]
 
     fun hasWriters(variableDescriptor: VariableDescriptor) = writers(variableDescriptor).isNotEmpty()
 
@@ -53,7 +55,7 @@ abstract class AssignedVariablesSearcher: KtTreeVisitorVoid() {
         if (binaryExpression.operationToken === KtTokens.EQ) {
             val left = KtPsiUtil.deparenthesize(binaryExpression.left)
             if (left is KtNameReferenceExpression) {
-                assignedNames.put(left.getReferencedNameAsName(), currentDeclaration)
+                assignedNames.put(left.getReferencedNameAsName(), Writer(binaryExpression, currentDeclaration))
             }
         }
         super.visitBinaryExpression(binaryExpression)

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/PreliminaryDeclarationVisitor.kt

@@ -28,7 +28,7 @@ import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
 
 class PreliminaryDeclarationVisitor(val declaration: KtDeclaration): AssignedVariablesSearcher() {
 
-    override fun writers(variableDescriptor: VariableDescriptor): MutableSet<KtDeclaration?> {
+    override fun writers(variableDescriptor: VariableDescriptor): MutableSet<Writer> {
         lazyTrigger
         return super.writers(variableDescriptor)
     }

compiler/testData/diagnostics/tests/smartCasts/varnotnull/capturedInClosureModifiedBefore.kt

@@ -0,0 +1,60 @@
+// !LANGUAGE: +CapturedInClosureSmartCasts
+
+fun run(f: () -> Unit) = f()
+
+fun foo(s: String?) {
+    var x: String? = null
+    if (s != null) {
+        x = s
+    }
+    if (x != null) {
+        run {
+            <!DEBUG_INFO_SMARTCAST!>x<!>.hashCode()
+        }
+    }
+}
+
+fun bar(s: String?) {
+    var x = s
+    if (x != null) {
+        run {
+            <!DEBUG_INFO_SMARTCAST!>x<!>.hashCode()
+        }
+    }
+}
+
+fun baz(s: String?) {
+    var x = s
+    if (x != null) {
+        run {
+            <!SMARTCAST_IMPOSSIBLE!>x<!>.hashCode()
+        }
+        run {
+            <!SMARTCAST_IMPOSSIBLE!>x<!>.hashCode()
+            x = null
+        }
+    }
+}
+
+fun gaz(s: String?) {
+    var x = s
+    if (x != null) {
+        run {
+            <!SMARTCAST_IMPOSSIBLE!>x<!>.hashCode()
+            x = null
+        }
+        run {
+            <!SMARTCAST_IMPOSSIBLE!>x<!>.hashCode()
+        }
+    }
+}
+
+fun gav(s: String?) {
+    var x = s
+    if (x != null) {
+        run {
+            <!SMARTCAST_IMPOSSIBLE!>x<!>.hashCode()
+        }
+        x = null
+    }
+}

compiler/testData/diagnostics/tests/smartCasts/varnotnull/capturedInClosureModifiedBefore.txt

@@ -0,0 +1,8 @@
+package
+
+public fun bar(/*0*/ s: kotlin.String?): kotlin.Unit
+public fun baz(/*0*/ s: kotlin.String?): kotlin.Unit
+public fun foo(/*0*/ s: kotlin.String?): kotlin.Unit
+public fun gav(/*0*/ s: kotlin.String?): kotlin.Unit
+public fun gaz(/*0*/ s: kotlin.String?): kotlin.Unit
+public fun run(/*0*/ f: () -> kotlin.Unit): kotlin.Unit

compiler/testData/diagnostics/tests/smartCasts/varnotnull/capturedInClosureOff.kt

@@ -0,0 +1,15 @@
+// !LANGUAGE: -CapturedInClosureSmartCasts
+
+fun run(f: () -> Unit) = f()
+
+fun foo(s: String?) {
+    var x: String? = null
+    if (s != null) {
+        x = s
+    }
+    if (x != null) {
+        run {
+            <!DEBUG_INFO_SMARTCAST!>x<!>.hashCode()
+        }
+    }
+}

compiler/testData/diagnostics/tests/smartCasts/varnotnull/capturedInClosureOff.txt

@@ -0,0 +1,4 @@
+package
+
+public fun foo(/*0*/ s: kotlin.String?): kotlin.Unit
+public fun run(/*0*/ f: () -> kotlin.Unit): kotlin.Unit

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

@@ -21309,6 +21309,18 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTest {
                     doTest(fileName);
                 }
 
+                @TestMetadata("capturedInClosureModifiedBefore.kt")
+                public void testCapturedInClosureModifiedBefore() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/smartCasts/varnotnull/capturedInClosureModifiedBefore.kt");
+                    doTest(fileName);
+                }
+
+                @TestMetadata("capturedInClosureOff.kt")
+                public void testCapturedInClosureOff() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/smartCasts/varnotnull/capturedInClosureOff.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("doWhileWithBreak.kt")
                 public void testDoWhileWithBreak() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/smartCasts/varnotnull/doWhileWithBreak.kt");

compiler/util/src/org/jetbrains/kotlin/config/LanguageVersionSettings.kt

@@ -56,6 +56,7 @@ enum class LanguageFeature(
     DeprecatedFieldForInvisibleCompanionObject(KOTLIN_1_2),
     SafeCastCheckBoundSmartCasts(KOTLIN_1_2),
     BooleanElvisBoundSmartCasts(KOTLIN_1_2),
+    CapturedInClosureSmartCasts(KOTLIN_1_2),
 
     // Experimental features