From 8235f95686b89d2a74bcd06748e52f5572b47651 Mon Sep 17 00:00:00 2001
From: Tagir Valeev <Tagir.Valeev@jetbrains.com>
Date: Tue, 11 Feb 2020 15:53:26 +0700
Subject: [PATCH] Generate null-check after assignment instruction

Fixes IDEA-232554 False positive @nullable method returns non-null only

GitOrigin-RevId: 35d2655f3340b604c0056dfa90059c8a645d8632
---
 .../dataFlow/ControlFlowAnalyzer.java         |  3 ++-
 .../NullableNotNullAssignmentInReturn.java    | 24 +++++++++++++++++++
 .../DataFlowInspection8Test.java              |  1 +
 3 files changed, 27 insertions(+), 1 deletion(-)
 create mode 100644 java/java-tests/testData/inspection/dataFlow/fixture/NullableNotNullAssignmentInReturn.java

diff --git a/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/ControlFlowAnalyzer.java b/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/ControlFlowAnalyzer.java
index a7273f4a03db..9fb537a9f711 100644
--- a/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/ControlFlowAnalyzer.java
+++ b/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/ControlFlowAnalyzer.java
@@ -237,6 +237,7 @@ public class ControlFlowAnalyzer extends JavaElementVisitor {
     }
 
     addInstruction(new AssignInstruction(rExpr, myFactory.createValue(lExpr)));
+    addNullCheck(expression);
 
     finishElement(expression);
   }
@@ -815,7 +816,7 @@ public class ControlFlowAnalyzer extends JavaElementVisitor {
     PsiExpression returnValue = statement.getReturnValue();
 
     if (myExpressionBlockContext != null) {
-      // We treat return inside switch expression (which is disallowed syntax) as break-with-value
+      // We treat return inside switch expression (which is disallowed syntax) as yield
       myExpressionBlockContext.generateReturn(returnValue, this);
     } else {
 
diff --git a/java/java-tests/testData/inspection/dataFlow/fixture/NullableNotNullAssignmentInReturn.java b/java/java-tests/testData/inspection/dataFlow/fixture/NullableNotNullAssignmentInReturn.java
new file mode 100644
index 000000000000..029f3023b714
--- /dev/null
+++ b/java/java-tests/testData/inspection/dataFlow/fixture/NullableNotNullAssignmentInReturn.java
@@ -0,0 +1,24 @@
+import org.jetbrains.annotations.Nullable;
+class Hello {
+  @Nullable
+  String something;
+  @Nullable
+  private String getAndCacheSomething() {
+    if (something != null) {
+      return something;
+    }
+    return something = getSomething();
+  }
+  @Nullable
+  private String getAndCacheSomething2() {
+    if (something != null) {
+      return something;
+    }
+    something = getSomething();
+    return something;
+  }
+  @Nullable
+  String getSomething() {
+    return null;
+  }
+}
\ No newline at end of file
diff --git a/java/java-tests/testSrc/com/intellij/java/codeInspection/DataFlowInspection8Test.java b/java/java-tests/testSrc/com/intellij/java/codeInspection/DataFlowInspection8Test.java
index 1831c1b7ab3b..69b86dd6e921 100644
--- a/java/java-tests/testSrc/com/intellij/java/codeInspection/DataFlowInspection8Test.java
+++ b/java/java-tests/testSrc/com/intellij/java/codeInspection/DataFlowInspection8Test.java
@@ -276,4 +276,5 @@ public class DataFlowInspection8Test extends DataFlowInspectionTestCase {
   }
   public void testInlineLambdaFromLocal() { doTest(); }
   public void testAllowRequireNonNullInCtor() { doTest(); }
+  public void testNullableNotNullAssignmentInReturn() { doTest(); }
 }
\ No newline at end of file