DfaMemoryStateImpl#uniteClasses do not corrupt distinctClasses when returning false

In NaN handling memstate may still be alive even if classes to unite are already distinct. In this case we uniteClasses processing was incomplete resulting in corrupted memory state.

java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/DfaMemoryStateImpl.java

@@ -519,6 +519,8 @@ public class DfaMemoryStateImpl implements DfaMemoryState {
   }
 
   private boolean uniteClasses(int c1Index, int c2Index) {
+    if (!myDistinctClasses.unite(c1Index, c2Index)) return false;
+
     EqClass c1 = myEqClasses.get(c1Index);
     EqClass c2 = myEqClasses.get(c2Index);
 
@@ -555,9 +557,6 @@ public class DfaMemoryStateImpl implements DfaMemoryState {
       addToMap(c, c1Index);
     }
 
-    if (!myDistinctClasses.unite(c1Index, c2Index)) {
-      return false;
-    }
     myEqClasses.set(c2Index, null);
     checkInvariants();
 

java/java-tests/testData/inspection/dataFlow/fixture/LoopDoubleComparisonNotComplex.java

@@ -0,0 +1,26 @@
+import java.util.List;
+
+class LoopDoubleComparisonNotComplex {
+  private static Node prob(List<Node> nodes) {
+    Node maxNode = null;
+    double maxProbability = 0.0;
+
+    for (Node node : nodes) {
+      double probability = probability(node);
+      if (probability > maxProbability) {
+        maxProbability = probability;
+        maxNode = node;
+      } else if (probability == maxProbability) {
+      }
+    }
+
+    return <warning descr="Expression 'maxNode' might evaluate to null but is returned by the method which is not declared as @Nullable">maxNode</warning>;
+  }
+
+  interface Node {}
+
+  static double probability(Node node) {
+    return 0.0;
+  }
+
+}

java/java-tests/testSrc/com/intellij/java/codeInspection/DataFlowInspectionTest.java

@@ -585,4 +585,5 @@ public class DataFlowInspectionTest extends DataFlowInspectionTestCase {
   public void testNullableGetterInLoop() { doTest(); }
   public void testNullabilityBasics() { doTest(); }
   public void testReassignedVarInLoop() { doTest(); }
+  public void testLoopDoubleComparisonNotComplex() { doTest(); }
 }