diff --git a/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/DataFlowRunner.java b/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/DataFlowRunner.java
index aa08cf6ebede..3f8943eb0cbf 100644
--- a/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/DataFlowRunner.java
+++ b/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/DataFlowRunner.java
@@ -49,6 +49,9 @@ public class DataFlowRunner {
 
   private Instruction[] myInstructions;
   private final MultiMap<PsiElement, DfaMemoryState> myNestedClosures = new MultiMap<>();
+  // Closures which were registered for previous instruction and can be queried by visitor
+  // to adjust them somehow
+  private final List<DfaMemoryState> myStackTopClosures = new ArrayList<>();
   @NotNull
   private final DfaValueFactory myValueFactory;
   private final boolean myShouldCheckLimitTime;
@@ -90,6 +93,10 @@ public class DataFlowRunner {
     return Collections.singletonList(createMemoryState());
   }
 
+  List<DfaMemoryState> getStackTopClosures() {
+    return new ArrayList<>(myStackTopClosures);
+  }
+
   @NotNull
   public final RunnerResult analyzeMethod(@NotNull PsiElement psiBlock, @NotNull InstructionVisitor visitor) {
     Collection<DfaMemoryState> initialStates = createInitialStates(psiBlock, visitor);
@@ -109,6 +116,7 @@ public class DataFlowRunner {
       int endOffset = flow.getInstructionCount();
       myInstructions = flow.getInstructions();
       myNestedClosures.clear();
+      myStackTopClosures.clear();
       
       Set<Instruction> joinInstructions = ContainerUtil.newHashSet();
       for (int index = 0; index < myInstructions.length; index++) {
@@ -270,6 +278,9 @@ public class DataFlowRunner {
   @NotNull
   protected DfaInstructionState[] acceptInstruction(@NotNull InstructionVisitor visitor, @NotNull DfaInstructionState instructionState) {
     Instruction instruction = instructionState.getInstruction();
+    DfaInstructionState[] states = instruction.accept(this, instructionState.getMemoryState(), visitor);
+
+    myStackTopClosures.clear();
     PsiElement closure = DfaUtil.getClosureInside(instruction);
     if (closure instanceof PsiClass) {
       registerNestedClosures(instructionState, (PsiClass)closure);
@@ -277,7 +288,7 @@ public class DataFlowRunner {
       registerNestedClosures(instructionState, (PsiLambdaExpression)closure);
     }
 
-    return instruction.accept(this, instructionState.getMemoryState(), visitor);
+    return states;
   }
 
   private void registerNestedClosures(@NotNull DfaInstructionState instructionState, @NotNull PsiClass nestedClass) {
@@ -285,14 +296,14 @@ public class DataFlowRunner {
     for (PsiMethod method : nestedClass.getMethods()) {
       PsiCodeBlock body = method.getBody();
       if (body != null) {
-        myNestedClosures.putValue(body, state.createClosureState());
+        myNestedClosures.putValue(body, createClosureState(state));
       }
     }
     for (PsiClassInitializer initializer : nestedClass.getInitializers()) {
-      myNestedClosures.putValue(initializer.getBody(), state.createClosureState());
+      myNestedClosures.putValue(initializer.getBody(), createClosureState(state));
     }
     for (PsiField field : nestedClass.getFields()) {
-      myNestedClosures.putValue(field, state.createClosureState());
+      myNestedClosures.putValue(field, createClosureState(state));
     }
   }
   
@@ -300,10 +311,17 @@ public class DataFlowRunner {
     DfaMemoryState state = instructionState.getMemoryState();
     PsiElement body = expr.getBody();
     if (body != null) {
-      myNestedClosures.putValue(body, state.createClosureState());
+      myNestedClosures.putValue(body, createClosureState(state));
     }
   }
 
+  @NotNull
+  private DfaMemoryState createClosureState(DfaMemoryState state) {
+    DfaMemoryState closureState = state.createClosureState();
+    myStackTopClosures.add(closureState);
+    return closureState;
+  }
+
   @NotNull
   protected DfaMemoryState createMemoryState() {
     return new DfaMemoryStateImpl(myValueFactory);
diff --git a/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/DfaMemoryStateImpl.java b/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/DfaMemoryStateImpl.java
index e34694d07798..1d94c9949f35 100644
--- a/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/DfaMemoryStateImpl.java
+++ b/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/DfaMemoryStateImpl.java
@@ -53,8 +53,6 @@ public class DfaMemoryStateImpl implements DfaMemoryState {
   // dfa value id -> indices in myEqClasses list of the classes which contain the id (or wrapped)
   private final TIntObjectHashMap<int[]> myIdToEqClassesIndices;
   private final Stack<DfaValue> myStack;
-  // Closures which correspond to the stack top (do not track other closures for now)
-  private final List<DfaMemoryState> myStackTopClosures = new ArrayList<>();
   private final TLongHashSet myDistinctClasses;
   private final LinkedHashMap<DfaVariableValue,DfaVariableState> myVariableStates;
   private final Map<DfaVariableValue,DfaVariableState> myDefaultVariableStates; 
@@ -118,14 +116,9 @@ public class DfaMemoryStateImpl implements DfaMemoryState {
       copy.flushDependencies(value);
     }
     copy.emptyStack();
-    myStackTopClosures.add(copy);
     return copy;
   }
 
-  List<DfaMemoryState> getStackTopClosures() {
-    return new ArrayList<>(myStackTopClosures);
-  }
-
   public boolean equals(Object obj) {
     if (obj == this) return true;
     if (!(obj instanceof DfaMemoryStateImpl)) return false;
@@ -241,7 +234,6 @@ public class DfaMemoryStateImpl implements DfaMemoryState {
   @Override
   public DfaValue pop() {
     myCachedHash = null;
-    myStackTopClosures.clear();
     return myStack.pop();
   }
 
@@ -253,14 +245,12 @@ public class DfaMemoryStateImpl implements DfaMemoryState {
   @Override
   public void push(@NotNull DfaValue value) {
     myCachedHash = null;
-    myStackTopClosures.clear();
     myStack.push(value);
   }
 
   @Override
   public void emptyStack() {
     myCachedHash = null;
-    myStackTopClosures.clear();
     while (!myStack.isEmpty() && !(myStack.peek() instanceof DfaControlTransferValue)) {
       myStack.pop();
     }
diff --git a/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/StandardInstructionVisitor.java b/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/StandardInstructionVisitor.java
index 6d0130b73580..9b4e2984cdec 100644
--- a/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/StandardInstructionVisitor.java
+++ b/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/StandardInstructionVisitor.java
@@ -229,7 +229,7 @@ public class StandardInstructionVisitor extends InstructionVisitor {
     if (methodName == null || !OPTIONAL_METHOD_NAMES.contains(methodName)) return Collections.emptyList();
     PsiMethod method = call.resolveMethod();
     if (method == null || !TypeUtils.isOptional(method.getContainingClass())) return Collections.emptyList();
-    List<DfaMemoryState> closures = ((DfaMemoryStateImpl)memState).getStackTopClosures();
+    List<DfaMemoryState> closures = runner.getStackTopClosures();
     DfaValue[] argValues = popCallArguments(instruction, runner, memState);
     final DfaValue qualifier = popQualifier(instruction, runner, memState);
     switch (methodName) {