From 11c0d924a6d937b7a2545a020e9c5b1d96f90d74 Mon Sep 17 00:00:00 2001
From: Andrey Breslav <abreslav@abreslav-mac.labs.intellij.net>
Date: Sun, 12 Jun 2011 15:32:50 +0400
Subject: [PATCH] One more case for nullability: comparison to a non-null
 expression

---
 .../jet/lang/types/JetTypeInferrer.java       | 51 +++++++++++++------
 idea/testData/checker/Nullability.jet         | 17 +++++++
 2 files changed, 52 insertions(+), 16 deletions(-)

diff --git a/idea/src/org/jetbrains/jet/lang/types/JetTypeInferrer.java b/idea/src/org/jetbrains/jet/lang/types/JetTypeInferrer.java
index f80d4c7aaa2..abba6650322 100644
--- a/idea/src/org/jetbrains/jet/lang/types/JetTypeInferrer.java
+++ b/idea/src/org/jetbrains/jet/lang/types/JetTypeInferrer.java
@@ -1381,24 +1381,39 @@ public class JetTypeInferrer {
                              || operationToken == JetTokens.EXCLEQEQEQ) {
                         JetExpression left = expression.getLeft();
                         JetExpression right = expression.getRight();
-                        if (left instanceof JetConstantExpression) {
+                        if (right == null) return;
+
+                        if (!(left instanceof JetSimpleNameExpression)) {
                             JetExpression tmp = left;
                             left = right;
                             right = tmp;
+
+                            if (!(left instanceof JetSimpleNameExpression)) {
+                                return;
+                            }
                         }
 
-                        if (!(left instanceof JetSimpleNameExpression)) {
-                            return;
-                        }
-                        JetSimpleNameExpression nameExpression = (JetSimpleNameExpression) left;
-                        DeclarationDescriptor declarationDescriptor = trace.getBindingContext().resolveReferenceExpression(nameExpression);
-                        if (!(declarationDescriptor instanceof VariableDescriptor)) {
-                            return;
-                        }
-                        VariableDescriptor variableDescriptor = (VariableDescriptor) declarationDescriptor;
+                        VariableDescriptor variableDescriptor = getVariableDescriptorFromSimpleName(left);
 
                         // TODO : validate that DF makes sense for this variable: local, val, internal w/backing field, etc
 
+                        // Comparison to a non-null expression
+                        JetType rhsType = trace.getBindingContext().getExpressionType(right);
+                        if (rhsType != null && !rhsType.isNullable()) {
+                            extendDataFlowWithNullComparison(operationToken, variableDescriptor, !conditionValue);
+                            return;
+                        }
+
+                        VariableDescriptor rightVariable = getVariableDescriptorFromSimpleName(right);
+                        if (rightVariable != null) {
+                            JetType lhsType = trace.getBindingContext().getExpressionType(left);
+                            if (lhsType != null && !lhsType.isNullable()) {
+                                extendDataFlowWithNullComparison(operationToken, rightVariable, !conditionValue);
+                                return;
+                            }
+                        }
+
+                        // Comparison to 'null'
                         if (!(right instanceof JetConstantExpression)) {
                             return;
                         }
@@ -1407,12 +1422,16 @@ public class JetTypeInferrer {
                             return;
                         }
 
-                        if (operationToken == JetTokens.EQEQ || operationToken == JetTokens.EQEQEQ) {
-                            result[0] = dataFlowInfo.equalsToNull(variableDescriptor, !conditionValue);
-                        }
-                        else if (operationToken == JetTokens.EXCLEQ || operationToken == JetTokens.EXCLEQEQEQ) {
-                            result[0] = dataFlowInfo.equalsToNull(variableDescriptor, conditionValue);
-                        }
+                        extendDataFlowWithNullComparison(operationToken, variableDescriptor, conditionValue);
+                    }
+                }
+
+                private void extendDataFlowWithNullComparison(IElementType operationToken, VariableDescriptor variableDescriptor, boolean equalsToNull) {
+                    if (operationToken == JetTokens.EQEQ || operationToken == JetTokens.EQEQEQ) {
+                        result[0] = dataFlowInfo.equalsToNull(variableDescriptor, !equalsToNull);
+                    }
+                    else if (operationToken == JetTokens.EXCLEQ || operationToken == JetTokens.EXCLEQEQEQ) {
+                        result[0] = dataFlowInfo.equalsToNull(variableDescriptor, equalsToNull);
                     }
                 }
 
diff --git a/idea/testData/checker/Nullability.jet b/idea/testData/checker/Nullability.jet
index 5109d94f3c3..770ed4c746b 100644
--- a/idea/testData/checker/Nullability.jet
+++ b/idea/testData/checker/Nullability.jet
@@ -255,3 +255,20 @@ fun f7(s : String?, t : String?) {
     t?.get(0)
   }
 }
+
+fun f8(b : String?, a : String) {
+  b?.get(0)
+  if (b == a) {
+    b.get(0);
+  }
+  b?.get(0)
+  if (a == b) {
+    b.get(0)
+  }
+  if (a != b) {
+    b?.get(0)
+  }
+  else {
+    b.get(0)
+  }
+}