Add inspection for 'for-each' statements to suggest using 'var'

Author: KalbarczykDev

src/main/java/dev/kalbarczyk/canbevar/VarCanBeUsedInspection.java

@@ -3,11 +3,11 @@
 import com.intellij.codeInspection.*;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.*;
+import com.intellij.psi.util.PsiUtil;
 import org.jetbrains.annotations.NotNull;
 
 public final class VarCanBeUsedInspection extends AbstractBaseJavaLocalInspectionTool {
 
-    //TODO: add inspection for the for each
 
     @Override
     public @NotNull PsiElementVisitor buildVisitor(final @NotNull ProblemsHolder holder, final boolean isOnTheFly) {
@@ -19,16 +19,12 @@ public void visitLocalVariable(final @NotNull PsiLocalVariable variable) {
                 var initializer = variable.getInitializer();
 
                 if (initializer == null) return;
-                if ("var".equals(typeElement.getText())) return;
-
-                //lombok feature (final + var = val)
-                if ("val".equals(typeElement.getText())) return;
-
+                if ("var".equals(typeElement.getText()) || "val".equals(typeElement.getText())) return;
                 if (type.equalsToText("var")) return;
 
-                if (!(initializer.getType() instanceof PsiClassType)) return;
-
                 var initType = initializer.getType();
+                if (initType == null) return;
+                if (!(initType instanceof PsiClassType) && !(initType instanceof PsiArrayType)) return;
 
                 if (initType.equals(type)) {
                     holder.registerProblem(
@@ -40,10 +36,43 @@ public void visitLocalVariable(final @NotNull PsiLocalVariable variable) {
                 }
 
             }
-        };
 
+            @Override
+            public void visitForeachStatement(@NotNull PsiForeachStatement statement) {
+                var iteratedValue = statement.getIteratedValue();
+                var parameter = statement.getIterationParameter();
+
+                var typeElement = parameter.getTypeElement();
+                if (typeElement == null) return;
+                if ("var".equals(typeElement.getText()) || "val".equals(typeElement.getText())) return;
+
+                var iteratedType = iteratedValue != null ? iteratedValue.getType() : null;
+                if (iteratedType == null) return;
+
+                var parameterType = parameter.getType();
+
+                PsiType elementType = null;
+                if (iteratedType instanceof PsiClassType classType) {
+                    elementType = PsiUtil.extractIterableTypeParameter(classType, false);
+                } else if (iteratedType instanceof PsiArrayType arrayType) {
+                    elementType = arrayType.getComponentType();
+                }
+
+                if (elementType == null) return;
+
+                if (elementType.equals(parameterType)) {
+                    holder.registerProblem(
+                            typeElement,
+                            "Explicit type in 'for-each' can be replaced with 'var'",
+                            ProblemHighlightType.WEAK_WARNING,
+                            new ReplaceWithVarQuickFix()
+                    );
+                }
+            }
+        };
     }
 
+
     private final static class ReplaceWithVarQuickFix implements LocalQuickFix {
         @Override
         public void applyFix(final @NotNull Project project, final @NotNull ProblemDescriptor problemDescriptor) {