Generate null-check after assignment instruction

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

GitOrigin-RevId: 35d2655f3340b604c0056dfa90059c8a645d8632

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 {
 

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;
+  }
+}

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(); }
 }