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