From cff8154475134379ba66aa4654f5a245af047344 Mon Sep 17 00:00:00 2001
From: Tagir Valeev <tagir.valeev@jetbrains.com>
Date: Mon, 11 Jul 2022 17:51:07 +0200
Subject: [PATCH] [java-inference] Fix inference for case null: (case null ->
 was already supported)

Fixes IDEA-297390 Wrong `@NotNull` is inferred for switch with explicit `null` case label
Also: support switch expressions

GitOrigin-RevId: 5605eb5e4487c592597b839c26319de021ba2282
---
 .../inference/ParameterNullityInference.kt        |  5 +++--
 .../ParameterNullityInferenceFromSourceTest.java  | 15 +++++++++++++++
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/inference/ParameterNullityInference.kt b/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/inference/ParameterNullityInference.kt
index f8b8e9d2b368..3206abfe08bb 100644
--- a/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/inference/ParameterNullityInference.kt
+++ b/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/inference/ParameterNullityInference.kt
@@ -48,7 +48,7 @@ internal fun inferNotNullParameters(tree: LighterAST, parameterNames: List<Strin
           dereference(tree, expression, canBeNulls, notNulls, queue)
         }
       }
-      SWITCH_STATEMENT -> {
+      SWITCH_STATEMENT, SWITCH_EXPRESSION -> {
         queue.clear()
         val expression = JavaLightTreeUtil.findExpressionChild(tree, element)
         val hasExplicitNullCheck = findCaseLabelElementList(tree, element)
@@ -195,6 +195,7 @@ internal fun getParameterNames(tree: LighterAST, method: LighterASTNode): List<S
 
 private fun findCaseLabelElementList(tree: LighterAST, switchNode: LighterASTNode): List<LighterASTNode> {
   val codeBlock = LightTreeUtil.firstChildOfType(tree, switchNode, CODE_BLOCK) ?: return emptyList()
-  val rules = LightTreeUtil.getChildrenOfType(tree, codeBlock, SWITCH_LABELED_RULE)
+  var rules: List<LighterASTNode> = LightTreeUtil.getChildrenOfType(tree, codeBlock, SWITCH_LABELED_RULE)
+  rules += LightTreeUtil.getChildrenOfType(tree, codeBlock, SWITCH_LABEL_STATEMENT)
   return rules.mapNotNull { node -> LightTreeUtil.firstChildOfType(tree, node, CASE_LABEL_ELEMENT_LIST) }
 }
diff --git a/java/java-tests/testSrc/com/intellij/java/codeInspection/ParameterNullityInferenceFromSourceTest.java b/java/java-tests/testSrc/com/intellij/java/codeInspection/ParameterNullityInferenceFromSourceTest.java
index 967fd211304e..1cd9afe868f7 100644
--- a/java/java-tests/testSrc/com/intellij/java/codeInspection/ParameterNullityInferenceFromSourceTest.java
+++ b/java/java-tests/testSrc/com/intellij/java/codeInspection/ParameterNullityInferenceFromSourceTest.java
@@ -152,6 +152,21 @@ public class ParameterNullityInferenceFromSourceTest extends LightJavaCodeInsigh
                   "}");
   }
 
+  public void testSwitchWithPattern() {
+    assertNullity("-",
+                  "void test(String r) {\n" +
+                  "        switch (r) {\n" +
+                  "            case null:\n" +
+                  "                break;\n" +
+                  "        }\n" +
+                  "    }");
+  }
+
+  public void testSwitchExpression() {
+    assertNullity("+",
+                  "int test(String s) { return switch(s) {case \"x\"->1;default ->2;}} ");
+  }
+
   public void testUseConfiguredNullityAnnotation() {
     PsiClass clazz = myFixture.addClass("final class Foo { void foo(String s) { s.hashCode(); } }");
     PsiParameter parameter = clazz.getMethods()[0].getParameterList().getParameters()[0];