Preliminary declaration visitor for estimating local variable's predictability for smart casts

Predictability estimation algorithm is completely new, but backward compatibility should present.
A large set of tests. Some updated tests.
Smart casts allowed for captured variables if they are not modified in closure #KT-9051 Fixed
Also #KT-8643 Fixed
Also #KT-7976 Fixed
Correct handling of lambda arguments in functions #KT-9143 Fixed

compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt

@@ -310,3 +310,5 @@ public fun SearchScope.contains(element: PsiElement): Boolean = PsiSearchScopeUt
 
 public fun <E : PsiElement> E.createSmartPointer(): SmartPsiElementPointer<E> =
         SmartPointerManager.getInstance(getProject()).createSmartPsiElementPointer(this)
+
+public fun PsiElement.before(element: PsiElement) = textRange.endOffset <= element.textRange.startOffset

compiler/frontend/src/org/jetbrains/kotlin/resolve/BindingContext.java

@@ -40,6 +40,7 @@ import org.jetbrains.kotlin.types.DeferredType;
 import org.jetbrains.kotlin.types.JetType;
 import org.jetbrains.kotlin.types.expressions.CaptureKind;
 import org.jetbrains.kotlin.types.expressions.JetTypeInfo;
+import org.jetbrains.kotlin.types.expressions.PreliminaryDeclarationVisitor;
 import org.jetbrains.kotlin.util.Box;
 import org.jetbrains.kotlin.util.slicedMap.*;
 
@@ -150,6 +151,7 @@ public interface BindingContext {
     WritableSlice<JetElement, Boolean> UNREACHABLE_CODE = Slices.createSimpleSetSlice();
 
     WritableSlice<VariableDescriptor, CaptureKind> CAPTURED_IN_CLOSURE = new BasicWritableSlice<VariableDescriptor, CaptureKind>(DO_NOTHING);
+    WritableSlice<JetDeclaration, PreliminaryDeclarationVisitor> PRELIMINARY_VISITOR = new BasicWritableSlice<JetDeclaration, PreliminaryDeclarationVisitor>(DO_NOTHING);
 
     WritableSlice<CallableMemberDescriptor, Boolean> NEED_SYNTHETIC_ACCESSOR = new BasicWritableSlice<CallableMemberDescriptor, Boolean>(DO_NOTHING);
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/BodyResolver.java

@@ -46,6 +46,7 @@ import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue;
 import org.jetbrains.kotlin.types.*;
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingContext;
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingServices;
+import org.jetbrains.kotlin.types.expressions.PreliminaryDeclarationVisitor;
 import org.jetbrains.kotlin.types.expressions.ValueParameterResolver;
 import org.jetbrains.kotlin.types.expressions.typeInfoFactory.TypeInfoFactoryKt;
 import org.jetbrains.kotlin.util.Box;
@@ -521,6 +522,8 @@ public class BodyResolver {
         if (!classDescriptor.getConstructors().isEmpty()) {
             JetExpression body = anonymousInitializer.getBody();
             if (body != null) {
+                PreliminaryDeclarationVisitor.Companion.createForDeclaration(
+                        (JetDeclaration) anonymousInitializer.getParent().getParent(), trace);
                 expressionTypingServices.getType(scopeForInitializers, body, NO_EXPECTED_TYPE, outerDataFlowInfo, trace);
             }
             processModifiersOnInitializer(anonymousInitializer, scopeForInitializers);
@@ -582,6 +585,7 @@ public class BodyResolver {
     ) {
         computeDeferredType(propertyDescriptor.getReturnType());
 
+        PreliminaryDeclarationVisitor.Companion.createForDeclaration(property, trace);
         JetExpression initializer = property.getInitializer();
         LexicalScope propertyScope = getScopeForProperty(c, property);
         if (parentScope == null) {
@@ -784,6 +788,7 @@ public class BodyResolver {
             @Nullable Function1<LexicalScope, DataFlowInfo> beforeBlockBody,
             @NotNull CallChecker callChecker
     ) {
+        PreliminaryDeclarationVisitor.Companion.createForDeclaration(function, trace);
         LexicalScope innerScope = FunctionDescriptorUtil.getFunctionInnerScope(scope, functionDescriptor, trace);
         List<JetParameter> valueParameters = function.getValueParameters();
         List<ValueParameterDescriptor> valueParameterDescriptors = functionDescriptor.getValueParameters();

compiler/frontend/src/org/jetbrains/kotlin/resolve/DescriptorResolver.java

@@ -55,6 +55,7 @@ import org.jetbrains.kotlin.storage.StorageManager;
 import org.jetbrains.kotlin.types.*;
 import org.jetbrains.kotlin.types.checker.JetTypeChecker;
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingServices;
+import org.jetbrains.kotlin.types.expressions.PreliminaryDeclarationVisitor;
 
 import java.util.*;
 
@@ -865,6 +866,7 @@ public class DescriptorResolver {
                             new Function0<JetType>() {
                                 @Override
                                 public JetType invoke() {
+                                    PreliminaryDeclarationVisitor.Companion.createForDeclaration(variable, trace);
                                     JetType initializerType = resolveInitializerType(scope, variable.getInitializer(), dataFlowInfo, trace);
                                     setConstantForVariableIfNeeded(variableDescriptor, scope, variable, dataFlowInfo, initializerType, trace);
                                     return transformAnonymousTypeIfNeeded(variableDescriptor, variable, initializerType, trace);

compiler/frontend/src/org/jetbrains/kotlin/resolve/FunctionDescriptorResolver.kt

@@ -50,6 +50,7 @@ import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.kotlin.types.checker.JetTypeChecker
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingServices
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingUtils
+import org.jetbrains.kotlin.types.expressions.PreliminaryDeclarationVisitor
 import java.util.*
 
 class FunctionDescriptorResolver(
@@ -123,6 +124,7 @@ class FunctionDescriptorResolver(
         }
         else if (function.hasBody()) {
             DeferredType.createRecursionIntolerant(storageManager, trace) {
+                PreliminaryDeclarationVisitor.createForDeclaration(function, trace);
                 val type = expressionTypingServices.getBodyExpressionType(trace, scope, dataFlowInfo, function, functionDescriptor)
                 transformAnonymousTypeIfNeeded(functionDescriptor, function, type, trace)
             }

compiler/frontend/src/org/jetbrains/kotlin/resolve/ScriptBodyResolver.java

@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.types.JetType;
 import org.jetbrains.kotlin.types.expressions.CoercionStrategy;
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingContext;
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingServices;
+import org.jetbrains.kotlin.types.expressions.PreliminaryDeclarationVisitor;
 
 import java.util.Map;
 
@@ -63,6 +64,7 @@ public class ScriptBodyResolver {
                 DataFlowInfo.EMPTY,
                 NO_EXPECTED_TYPE
         );
+        PreliminaryDeclarationVisitor.Companion.createForDeclaration(script, trace);
         JetType returnType = expressionTypingServices.getBlockReturnedType(script.getBlockExpression(), CoercionStrategy.NO_COERCION, context).getType();
         if (returnType == null) {
             returnType = ErrorUtils.createErrorType("getBlockReturnedType returned null");

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

@@ -17,28 +17,36 @@
 package org.jetbrains.kotlin.resolve.calls.smartcasts;
 
 import com.intellij.openapi.util.Pair;
+import com.intellij.psi.PsiElement;
 import com.intellij.psi.tree.IElementType;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.JetNodeTypes;
 import org.jetbrains.kotlin.descriptors.*;
 import org.jetbrains.kotlin.descriptors.impl.LocalVariableDescriptor;
+import org.jetbrains.kotlin.descriptors.impl.SyntheticFieldDescriptor;
 import org.jetbrains.kotlin.lexer.JetTokens;
 import org.jetbrains.kotlin.psi.*;
+import org.jetbrains.kotlin.psi.psiUtil.PsiUtilsKt;
 import org.jetbrains.kotlin.resolve.BindingContext;
-import org.jetbrains.kotlin.resolve.BindingContextUtils;
 import org.jetbrains.kotlin.resolve.DescriptorUtils;
 import org.jetbrains.kotlin.resolve.calls.callUtil.CallUtilKt;
 import org.jetbrains.kotlin.resolve.calls.context.ResolutionContext;
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall;
 import org.jetbrains.kotlin.resolve.descriptorUtil.DescriptorUtilsKt;
+import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValue.Kind;
 import org.jetbrains.kotlin.resolve.scopes.receivers.*;
 import org.jetbrains.kotlin.types.JetType;
 import org.jetbrains.kotlin.types.TypeUtils;
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingUtils;
+import org.jetbrains.kotlin.types.expressions.PreliminaryDeclarationVisitor;
+
+import java.util.Set;
 
 import static org.jetbrains.kotlin.builtins.KotlinBuiltIns.isNullableNothing;
+import static org.jetbrains.kotlin.resolve.BindingContext.DECLARATION_TO_DESCRIPTOR;
 import static org.jetbrains.kotlin.resolve.BindingContext.REFERENCE_TARGET;
+import static org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValue.Kind.*;
 
 /**
  * This class is intended to create data flow values for different kind of expressions.
@@ -84,7 +92,7 @@ public class DataFlowValueFactory {
             // fun <T : Any?> foo(x: T) = x!!.hashCode() // there no way in type system to denote that `x!!` is not nullable
             return new DataFlowValue(expression,
                                      type,
-                                     DataFlowValue.Kind.OTHER,
+                                     OTHER,
                                      Nullability.NOT_NULL);
         }
 
@@ -98,7 +106,7 @@ public class DataFlowValueFactory {
     @NotNull
     public static DataFlowValue createDataFlowValue(@NotNull ThisReceiver receiver) {
         JetType type = receiver.getType();
-        return new DataFlowValue(receiver, type, DataFlowValue.Kind.STABLE_VALUE, getImmanentNullability(type));
+        return new DataFlowValue(receiver, type, STABLE_VALUE, getImmanentNullability(type));
     }
 
     @NotNull
@@ -119,7 +127,7 @@ public class DataFlowValueFactory {
         if (receiverValue instanceof TransientReceiver || receiverValue instanceof ScriptReceiver) {
             // SCRIPT: smartcasts data flow
             JetType type = receiverValue.getType();
-            return new DataFlowValue(receiverValue, type, DataFlowValue.Kind.STABLE_VALUE, getImmanentNullability(type));
+            return new DataFlowValue(receiverValue, type, STABLE_VALUE, getImmanentNullability(type));
         }
         else if (receiverValue instanceof ClassReceiver || receiverValue instanceof ExtensionReceiver) {
             return createDataFlowValue((ThisReceiver) receiverValue);
@@ -139,14 +147,16 @@ public class DataFlowValueFactory {
     }
 
     @NotNull
-    public static DataFlowValue createDataFlowValue(
+    public static DataFlowValue createDataFlowValueForProperty(
+            @NotNull JetProperty property,
             @NotNull VariableDescriptor variableDescriptor,
             @NotNull BindingContext bindingContext,
             @Nullable ModuleDescriptor usageContainingModule
     ) {
         JetType type = variableDescriptor.getType();
         return new DataFlowValue(variableDescriptor, type,
-                                 variableKind(variableDescriptor, usageContainingModule, bindingContext),
+                                 variableKind(variableDescriptor, usageContainingModule,
+                                              bindingContext, property),
                                  getImmanentNullability(type));
     }
 
@@ -157,17 +167,17 @@ public class DataFlowValueFactory {
 
     private static class IdentifierInfo {
         public final Object id;
-        public final DataFlowValue.Kind kind;
+        public final Kind kind;
         public final boolean isPackage;
 
-        private IdentifierInfo(Object id, DataFlowValue.Kind kind, boolean isPackage) {
+        private IdentifierInfo(Object id, Kind kind, boolean isPackage) {
             this.id = id;
             this.kind = kind;
             this.isPackage = isPackage;
         }
     }
 
-    private static final IdentifierInfo NO_IDENTIFIER_INFO = new IdentifierInfo(null, DataFlowValue.Kind.OTHER, false) {
+    private static final IdentifierInfo NO_IDENTIFIER_INFO = new IdentifierInfo(null, OTHER, false) {
         @Override
         public String toString() {
             return "NO_IDENTIFIER_INFO";
@@ -175,18 +185,18 @@ public class DataFlowValueFactory {
     };
 
     @NotNull
-    private static IdentifierInfo createInfo(Object id, DataFlowValue.Kind kind) {
+    private static IdentifierInfo createInfo(Object id, Kind kind) {
         return new IdentifierInfo(id, kind, false);
     }
 
     @NotNull
     private static IdentifierInfo createStableInfo(Object id) {
-        return createInfo(id, DataFlowValue.Kind.STABLE_VALUE);
+        return createInfo(id, STABLE_VALUE);
     }
 
     @NotNull
     private static IdentifierInfo createPackageOrClassInfo(Object id) {
-        return new IdentifierInfo(id, DataFlowValue.Kind.STABLE_VALUE, true);
+        return new IdentifierInfo(id, STABLE_VALUE, true);
     }
 
     @NotNull
@@ -199,9 +209,9 @@ public class DataFlowValueFactory {
         }
         return createInfo(Pair.create(receiverInfo.id, selectorInfo.id),
                           receiverInfo.kind.isStable() && selectorInfo.kind.isStable()
-                          ? DataFlowValue.Kind.STABLE_VALUE
+                          ? STABLE_VALUE
                           // x.y can never be a local variable
-                          : DataFlowValue.Kind.OTHER);
+                          : OTHER);
     }
 
     @NotNull
@@ -275,8 +285,10 @@ public class DataFlowValueFactory {
                     resolvedCall != null ? getIdForImplicitReceiver(resolvedCall.getDispatchReceiver(), simpleNameExpression) : null;
 
             VariableDescriptor variableDescriptor = (VariableDescriptor) declarationDescriptor;
-            return combineInfo(receiverInfo, createInfo(variableDescriptor,
-                                                        variableKind(variableDescriptor, usageModuleDescriptor, bindingContext)));
+            return combineInfo(receiverInfo,
+                               createInfo(variableDescriptor,
+                                          variableKind(variableDescriptor, usageModuleDescriptor,
+                                                       bindingContext, simpleNameExpression)));
         }
         if (declarationDescriptor instanceof PackageViewDescriptor || declarationDescriptor instanceof ClassDescriptor) {
             return createPackageOrClassInfo(declarationDescriptor);
@@ -311,18 +323,92 @@ public class DataFlowValueFactory {
         return NO_IDENTIFIER_INFO;
     }
 
-    public static DataFlowValue.Kind variableKind(
+    @NotNull
+    private static DeclarationDescriptor getVariableContainingDeclaration(@NotNull VariableDescriptor variableDescriptor) {
+        DeclarationDescriptor containingDeclarationDescriptor = variableDescriptor.getContainingDeclaration();
+        if (containingDeclarationDescriptor instanceof ConstructorDescriptor
+            && ((ConstructorDescriptor) containingDeclarationDescriptor).isPrimary()) {
+            // This code is necessary just because JetClassInitializer has no associated descriptor in trace
+            // Because of it we have to use class itself instead of initializer,
+            // otherwise we could not find this descriptor inside isAccessedInsideClosure below
+            containingDeclarationDescriptor = containingDeclarationDescriptor.getContainingDeclaration();
+            assert containingDeclarationDescriptor != null : "No containing declaration for primary constructor";
+        }
+        return containingDeclarationDescriptor;
+    }
+
+    private static boolean isAccessedInsideClosure(
+            @NotNull DeclarationDescriptor variableContainingDeclaration,
+            @NotNull BindingContext bindingContext,
+            @NotNull JetElement accessElement
+    ) {
+        PsiElement parent = accessElement.getParent();
+        while (parent != null) {
+            // We are inside some declaration
+            if (parent instanceof JetDeclarationWithBody || parent instanceof JetClassOrObject) {
+                DeclarationDescriptor descriptor = bindingContext.get(DECLARATION_TO_DESCRIPTOR, parent);
+                if (variableContainingDeclaration.equals(descriptor)) {
+                    // Access is at the same declaration: not in closure
+                    break;
+                }
+                else {
+                    // Access is lower than parent: in closure
+                    return true;
+                }
+            }
+            parent = parent.getParent();
+        }
+        return false;
+    }
+
+    private static boolean isAccessedBeforeAllClosureWriters(
+            @NotNull DeclarationDescriptor variableContainingDeclaration,
+            @NotNull Set<JetDeclaration> writers,
+            @NotNull BindingContext bindingContext,
+            @NotNull JetElement accessElement
+    ) {
+        // All writers should be before access element, with the exception:
+        // writer which is the same with declaration site does not count
+        for (JetDeclaration writer : writers) {
+            DeclarationDescriptor writerDescriptor = bindingContext.get(DECLARATION_TO_DESCRIPTOR, writer);
+            // Access is after some writer
+            if (!variableContainingDeclaration.equals(writerDescriptor) && !PsiUtilsKt.before(accessElement, writer)) {
+                return false;
+            }
+        }
+        // Access is before all writers
+        return true;
+    }
+
+    private static Kind variableKind(
             @NotNull VariableDescriptor variableDescriptor,
             @Nullable ModuleDescriptor usageModule,
-            @NotNull BindingContext bindingContext
+            @NotNull BindingContext bindingContext,
+            @NotNull JetElement accessElement
     ) {
-        if (isStableVariable(variableDescriptor, usageModule)) return DataFlowValue.Kind.STABLE_VALUE;
+        if (isStableValue(variableDescriptor, usageModule)) return STABLE_VALUE;
         boolean isLocalVar = variableDescriptor.isVar() && variableDescriptor instanceof LocalVariableDescriptor;
-        if (!isLocalVar) return DataFlowValue.Kind.OTHER;
-        if (BindingContextUtils.isVarCapturedInClosure(bindingContext, variableDescriptor)) {
-            return DataFlowValue.Kind.UNPREDICTABLE_VARIABLE;
-        }
-        return DataFlowValue.Kind.PREDICTABLE_VARIABLE;
+        if (!isLocalVar) return OTHER;
+        if (variableDescriptor instanceof SyntheticFieldDescriptor) return OTHER;
+
+        // Local variable classification: PREDICTABLE or UNPREDICTABLE
+        PreliminaryDeclarationVisitor preliminaryVisitor =
+                PreliminaryDeclarationVisitor.Companion.getVisitorByVariable(variableDescriptor, bindingContext);
+        // A case when we just analyse an expression alone: counts as unpredictable
+        if (preliminaryVisitor == null) return UNPREDICTABLE_VARIABLE;
+
+        // Analyze who writes variable
+        // If there is no writer: predictable
+        Set<JetDeclaration> writers = preliminaryVisitor.writers(variableDescriptor);
+        if (writers.isEmpty()) return PREDICTABLE_VARIABLE;
+
+        // If access element is inside closure: unpredictable
+        DeclarationDescriptor variableContainingDeclaration = getVariableContainingDeclaration(variableDescriptor);
+        if (isAccessedInsideClosure(variableContainingDeclaration, bindingContext, accessElement)) return UNPREDICTABLE_VARIABLE;
+
+        // Otherwise, predictable iff considered position is BEFORE all writers except declarer itself
+        if (isAccessedBeforeAllClosureWriters(variableContainingDeclaration, writers, bindingContext, accessElement)) return PREDICTABLE_VARIABLE;
+        else return UNPREDICTABLE_VARIABLE;
     }
 
     /**
@@ -339,7 +425,7 @@ public class DataFlowValueFactory {
      * @param usageModule a module with a considered usage place, or null if it's not known (not recommended)
      * @return true if variable is stable, false otherwise
      */
-    public static boolean isStableVariable(
+    public static boolean isStableValue(
             @NotNull VariableDescriptor variableDescriptor,
             @Nullable ModuleDescriptor usageModule
     ) {

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

@@ -16,25 +16,50 @@
 
 package org.jetbrains.kotlin.types.expressions
 
+import com.google.common.collect.LinkedHashMultimap
+import com.google.common.collect.SetMultimap
+import org.jetbrains.kotlin.descriptors.VariableDescriptor
 import org.jetbrains.kotlin.lexer.JetTokens
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.psi.JetBinaryExpression
-import org.jetbrains.kotlin.psi.JetNameReferenceExpression
-import org.jetbrains.kotlin.psi.JetPsiUtil
-import org.jetbrains.kotlin.psi.JetTreeVisitorVoid
-import java.util.*
+import org.jetbrains.kotlin.psi.*
 
 abstract class AssignedVariablesSearcher: JetTreeVisitorVoid() {
 
-    protected val assignedNames: MutableSet<Name> = LinkedHashSet()
+    private val assignedNames: SetMultimap<Name, JetDeclaration?> = LinkedHashMultimap.create()
+
+    public open fun writers(variableDescriptor: VariableDescriptor) = assignedNames[variableDescriptor.name]
+
+    public fun hasWriters(variableDescriptor: VariableDescriptor) = writers(variableDescriptor).isNotEmpty()
+
+    private var currentDeclaration: JetDeclaration? = null
+
+    override fun visitDeclaration(declaration: JetDeclaration) {
+        val previous = currentDeclaration
+        if (declaration is JetDeclarationWithBody || declaration is JetClassOrObject) {
+            currentDeclaration = declaration
+        }
+        else if (declaration is JetClassInitializer) {
+            // Go to class declaration: init -> body -> class
+            currentDeclaration = declaration.parent.parent as JetDeclaration
+        }
+        super.visitDeclaration(declaration)
+        currentDeclaration = previous
+    }
+
+    override fun visitFunctionLiteralExpression(functionLiteralExpression: JetFunctionLiteralExpression) {
+        val previous = currentDeclaration
+        currentDeclaration = functionLiteralExpression.functionLiteral
+        super.visitFunctionLiteralExpression(functionLiteralExpression)
+        currentDeclaration = previous
+    }
 
     override fun visitBinaryExpression(binaryExpression: JetBinaryExpression) {
         if (binaryExpression.operationToken === JetTokens.EQ) {
             val left = JetPsiUtil.deparenthesize(binaryExpression.left)
             if (left is JetNameReferenceExpression) {
-                assignedNames += left.getReferencedNameAsName()
+                assignedNames.put(left.getReferencedNameAsName(), currentDeclaration)
             }
         }
+        super.visitBinaryExpression(binaryExpression)
     }
-
 }

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/ExpressionTypingVisitorForStatements.java

@@ -145,8 +145,8 @@ public class ExpressionTypingVisitorForStatements extends ExpressionTypingVisito
             // We can comment first part of this condition to take them into account, like here: var s: String? = "xyz"
             // In this case s will be not-nullable until it is changed
             if (property.getTypeReference() == null && type != null) {
-                DataFlowValue variableDataFlowValue = DataFlowValueFactory.createDataFlowValue(
-                        propertyDescriptor, context.trace.getBindingContext(),
+                DataFlowValue variableDataFlowValue = DataFlowValueFactory.createDataFlowValueForProperty(
+                        property, propertyDescriptor, context.trace.getBindingContext(),
                         DescriptorUtils.getContainingModuleOrNull(scope.getOwnerDescriptor()));
                 DataFlowValue initializerDataFlowValue = DataFlowValueFactory.createDataFlowValue(initializer, type, context);
                 // We cannot say here anything new about initializerDataFlowValue

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

@@ -0,0 +1,65 @@
+/*
+ * Copyright 2010-2015 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.types.expressions
+
+import org.jetbrains.kotlin.descriptors.VariableDescriptor
+import org.jetbrains.kotlin.psi.JetDeclaration
+import org.jetbrains.kotlin.psi.JetExpression
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.BindingTrace
+import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
+
+class PreliminaryDeclarationVisitor(val declaration: JetDeclaration): AssignedVariablesSearcher() {
+
+    override fun writers(variableDescriptor: VariableDescriptor): MutableSet<JetDeclaration?> {
+        lazyTrigger
+        return super.writers(variableDescriptor)
+    }
+
+    private val lazyTrigger by lazy {
+        declaration.accept(this)
+    }
+
+    companion object {
+
+        fun createForExpression(expression: JetExpression, trace: BindingTrace) {
+            expression.getStrictParentOfType<JetDeclaration>()?.let { createForDeclaration(it, trace) }
+        }
+
+        fun createForDeclaration(declaration: JetDeclaration, trace: BindingTrace) {
+            // TODO: find top-most declaration (but not class!!!)
+            // TODO: check if already exists
+            val visitor = PreliminaryDeclarationVisitor(declaration)
+            // Can this declaration be synthetic? If yes, it would be better not to record it
+            trace.record(BindingContext.PRELIMINARY_VISITOR, declaration, visitor);
+        }
+
+        fun getVisitorByVariable(variableDescriptor: VariableDescriptor, bindingContext: BindingContext): PreliminaryDeclarationVisitor? {
+            // Search for preliminary visitor of parent descriptor
+            val containingDescriptor = variableDescriptor.containingDeclaration
+            var currentDeclaration: JetDeclaration? =
+                    DescriptorToSourceUtils.descriptorToDeclaration(containingDescriptor) as? JetDeclaration ?: return null
+            var preliminaryVisitor = bindingContext.get(BindingContext.PRELIMINARY_VISITOR, currentDeclaration)
+            while (preliminaryVisitor == null && currentDeclaration != null) {
+                currentDeclaration = currentDeclaration.getStrictParentOfType()
+                preliminaryVisitor = bindingContext.get(BindingContext.PRELIMINARY_VISITOR, currentDeclaration)
+            }
+            return preliminaryVisitor
+        }
+    }
+}

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

@@ -36,7 +36,7 @@ class PreliminaryLoopVisitor private constructor() : AssignedVariablesSearcher()
             // Only predictable variables are under interest here
             val id = value.id
             if (value.kind == DataFlowValue.Kind.PREDICTABLE_VARIABLE && id is LocalVariableDescriptor) {
-                if (assignedNames.contains(id.name)) {
+                if (hasWriters(id)) {
                     valueSetToClear.add(value)
                 }
             }