From c835d651ed7ae94459dcef1eaafa6928414cd3ac Mon Sep 17 00:00:00 2001
From: Anna Kozlova <anna.kozlova@jetbrains.com>
Date: Fri, 24 Sep 2021 19:13:25 +0200
Subject: [PATCH] [java] local can be final: special treat for switch branches
 (IDEA-229646)

GitOrigin-RevId: cdd20359dbf92ed5a032b86075fc840a0c09601f
---
 .../localCanBeFinal/LocalCanBeFinal.java      | 29 ++++++++++++++-----
 .../localCanBeFinal/SwitchBraches.java        | 18 ++++++++++++
 .../codeInspection/LocalCanBeFinalTest.java   |  5 ++++
 3 files changed, 45 insertions(+), 7 deletions(-)
 create mode 100644 java/java-tests/testData/inspection/localCanBeFinal/SwitchBraches.java

diff --git a/java/java-analysis-impl/src/com/intellij/codeInspection/localCanBeFinal/LocalCanBeFinal.java b/java/java-analysis-impl/src/com/intellij/codeInspection/localCanBeFinal/LocalCanBeFinal.java
index f911c68bef34..89ba35cbb606 100644
--- a/java/java-analysis-impl/src/com/intellij/codeInspection/localCanBeFinal/LocalCanBeFinal.java
+++ b/java/java-analysis-impl/src/com/intellij/codeInspection/localCanBeFinal/LocalCanBeFinal.java
@@ -1,4 +1,4 @@
-// Copyright 2000-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+// Copyright 2000-2021 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package com.intellij.codeInspection.localCanBeFinal;
 
 import com.intellij.codeInspection.*;
@@ -125,14 +125,29 @@ public class LocalCanBeFinal extends AbstractBaseJavaLocalInspectionTool impleme
           return;
         }
         super.visitCodeBlock(block);
+        Set<PsiVariable> declared = getDeclaredVariables(block);
         PsiElement anchor = block;
-        if (block.getParent() instanceof PsiSwitchStatement) {
+        if (block.getParent() instanceof PsiSwitchBlock) {
           anchor = block.getParent();
+
+          //special case: switch legs
+          Set<PsiReferenceExpression> writeRefs = 
+            SyntaxTraverser.psiTraverser().withRoot(block)
+              .filter(PsiReferenceExpression.class)
+              .filter(ref -> PsiUtil.isOnAssignmentLeftHand(ref)).toSet();
+
+          for (PsiReferenceExpression ref : writeRefs) {
+            PsiElement resolve = ref.resolve();
+            if (resolve instanceof PsiVariable && declared.contains(resolve) && ((PsiVariable)resolve).hasInitializer()) {
+              declared.remove(resolve);
+            }
+          }
         }
+
         int from = flow.getStartOffset(anchor);
         int end = flow.getEndOffset(anchor);
         List<PsiVariable> ssa = ControlFlowUtil.getSSAVariables(flow, from, end, true);
-        HashSet<PsiElement> declared = getDeclaredVariables(block);
+       
         for (PsiVariable psiVariable : ssa) {
           if (declared.contains(psiVariable)) {
             result.add(psiVariable);
@@ -180,8 +195,8 @@ public class LocalCanBeFinal extends AbstractBaseJavaLocalInspectionTool impleme
         }
       }
 
-      private HashSet<PsiElement> getDeclaredVariables(PsiCodeBlock block) {
-        final HashSet<PsiElement> result = new HashSet<>();
+      private Set<PsiVariable> getDeclaredVariables(PsiCodeBlock block) {
+        final HashSet<PsiVariable> result = new HashSet<>();
         PsiElement[] children = block.getChildren();
         for (PsiElement child : children) {
           child.accept(new JavaElementVisitor() {
@@ -192,7 +207,7 @@ public class LocalCanBeFinal extends AbstractBaseJavaLocalInspectionTool impleme
             @Override public void visitDeclarationStatement(PsiDeclarationStatement statement) {
               PsiElement[] declaredElements = statement.getDeclaredElements();
               for (PsiElement declaredElement : declaredElements) {
-                if (declaredElement instanceof PsiVariable) result.add(declaredElement);
+                if (declaredElement instanceof PsiVariable) result.add((PsiVariable)declaredElement);
               }
             }
 
@@ -207,7 +222,7 @@ public class LocalCanBeFinal extends AbstractBaseJavaLocalInspectionTool impleme
               final PsiElement[] declaredElements = declarationStatement.getDeclaredElements();
               for (final PsiElement declaredElement : declaredElements) {
                 if (declaredElement instanceof PsiVariable) {
-                  result.add(declaredElement);
+                  result.add((PsiVariable)declaredElement);
                 }
               }
             }
diff --git a/java/java-tests/testData/inspection/localCanBeFinal/SwitchBraches.java b/java/java-tests/testData/inspection/localCanBeFinal/SwitchBraches.java
new file mode 100644
index 000000000000..24bb460694b2
--- /dev/null
+++ b/java/java-tests/testData/inspection/localCanBeFinal/SwitchBraches.java
@@ -0,0 +1,18 @@
+class Test {
+  void f(int i) {
+    switch (i) {
+      case 1:
+        int <warning descr="Variable 'j' can have 'final' modifier">j</warning> = 0;
+        long <warning descr="Variable 'k' can have 'final' modifier">k</warning>;
+        k = 0;
+        long l = 0;
+        break;
+      case 2:
+        k = 2;
+        l = 3;
+        break;
+      default:
+        break;
+    }
+  }
+}
\ No newline at end of file
diff --git a/java/java-tests/testSrc/com/intellij/java/codeInspection/LocalCanBeFinalTest.java b/java/java-tests/testSrc/com/intellij/java/codeInspection/LocalCanBeFinalTest.java
index a1e9477239cb..223150283567 100644
--- a/java/java-tests/testSrc/com/intellij/java/codeInspection/LocalCanBeFinalTest.java
+++ b/java/java-tests/testSrc/com/intellij/java/codeInspection/LocalCanBeFinalTest.java
@@ -44,6 +44,11 @@ public class LocalCanBeFinalTest extends LightJavaCodeInsightFixtureTestCase {
     myTool.REPORT_VARIABLES = true;
     doTest();
   }
+  public void testSwitchBraches() {
+    myTool.REPORT_PARAMETERS = false;
+    myTool.REPORT_VARIABLES = true;
+    doTest();
+  }
   public void testIncompleteAssignment() {
     myTool.REPORT_PARAMETERS = true;
     myTool.REPORT_VARIABLES = true;