[java, psi] introduce new PSI element: PsiForeachPatternStatement and extract base interface for foreach stmt

GitOrigin-RevId: d57f03553d360bf18bcb3a5a0a1f0c57fea58f91

java/java-psi-api/src/com/intellij/psi/JavaElementVisitor.java

@@ -161,6 +161,10 @@ public abstract class JavaElementVisitor extends PsiElementVisitor {
     visitStatement(statement);
   }
 
+  public void visitForeachPatternStatement(@NotNull PsiForeachPatternStatement statement) {
+    visitStatement(statement);
+  }
+
   public void visitForStatement(@NotNull PsiForStatement statement) {
     visitStatement(statement);
   }

java/java-psi-api/src/com/intellij/psi/PsiForeachDeclarationElement.java

@@ -1,9 +0,0 @@
-// Copyright 2000-2022 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
-package com.intellij.psi;
-
-
-import org.jetbrains.annotations.ApiStatus;
-
-@ApiStatus.NonExtendable
-public interface PsiForeachDeclarationElement extends PsiElement {
-}

java/java-psi-api/src/com/intellij/psi/PsiForeachPatternStatement.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2000-2009 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.intellij.psi;
+
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * Represents a Java  {@code for} statement with a pattern.
+ */
+public interface PsiForeachPatternStatement extends PsiForeachStatementBase {
+  /**
+   * @return pattern used in the foreach, e. g. {@code for (Rec(var x): recs) }
+   */
+  @NotNull
+  PsiPattern getIterationPattern();
+}

java/java-psi-api/src/com/intellij/psi/PsiForeachStatement.java

@@ -16,49 +16,18 @@
 package com.intellij.psi;
 
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
 
 /**
  * Represents a Java enhanced {@code for}   statement.
  *
  * @author dsl
  */
-public interface PsiForeachStatement extends PsiLoopStatement {
+public interface PsiForeachStatement extends PsiForeachStatementBase {
   /**
    * Returns the variable containing the iteration parameter of the statement.
    *
    * @return the iteration parameter instance.
    */
-  @Nullable
+  @NotNull
   PsiParameter getIterationParameter();
-
-  /**
-   * @return either the iteration parameter (PsiParameter) of the statement OR pattern (PsiPattern).
-   */
-  @NotNull
-  PsiForeachDeclarationElement getIterationDeclaration();
-
-  /**
-   * Returns the expression representing the sequence over which the iteration is performed.
-   *
-   * @return the iterated value expression instance, or null if the statement is incomplete.
-   */
-  @Nullable
-  PsiExpression getIteratedValue();
-
-  /**
-   * Returns the opening parenthesis enclosing the statement header.
-   *
-   * @return the opening parenthesis.
-   */
-  @NotNull
-  PsiJavaToken getLParenth();
-
-  /**
-   * Returns the closing parenthesis enclosing the statement header.
-   *
-   * @return the closing parenthesis, or null if the statement is incomplete.
-   */
-  @Nullable
-  PsiJavaToken getRParenth();
 }

java/java-psi-api/src/com/intellij/psi/PsiForeachStatementBase.java

@@ -0,0 +1,33 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.psi;
+
+import org.jetbrains.annotations.ApiStatus;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+@ApiStatus.NonExtendable
+public interface PsiForeachStatementBase extends PsiLoopStatement {
+  /**
+   * Returns the expression representing the sequence over which the iteration is performed.
+   *
+   * @return the iterated value expression instance, or null if the statement is incomplete.
+   */
+  @Nullable
+  PsiExpression getIteratedValue();
+
+  /**
+   * Returns the opening parenthesis enclosing the statement header.
+   *
+   * @return the opening parenthesis.
+   */
+  @NotNull
+  PsiJavaToken getLParenth();
+
+  /**
+   * Returns the closing parenthesis enclosing the statement header.
+   *
+   * @return the closing parenthesis, or null if the statement is incomplete.
+   */
+  @Nullable
+  PsiJavaToken getRParenth();
+}

java/java-psi-api/src/com/intellij/psi/PsiParameter.java

@@ -10,7 +10,7 @@ import org.jetbrains.annotations.Nullable;
 /**
  * Represents the parameter of a Java method, foreach (enhanced for) statement or catch block.
  */
-public interface PsiParameter extends PsiVariable, JvmParameter, PsiJvmModifiersOwner, PsiForeachDeclarationElement {
+public interface PsiParameter extends PsiVariable, JvmParameter, PsiJvmModifiersOwner {
   /**
    * The empty array of PSI parameters which can be reused to avoid unnecessary allocations.
    */

java/java-psi-api/src/com/intellij/psi/PsiPattern.java

@@ -4,6 +4,6 @@ package com.intellij.psi;
 /**
  * Represents pattern which is used in {@code instanceof} expressions or switch case labels.
  */
-public interface PsiPattern extends PsiCaseLabelElement, PsiForeachDeclarationElement {
+public interface PsiPattern extends PsiCaseLabelElement {
   PsiPattern[] EMPTY = new PsiPattern[]{};
 }

java/java-psi-impl/src/com/intellij/lang/java/parser/StatementParser.java

@@ -371,18 +371,18 @@ public class StatementParser {
     if (isRecordPatternInForEach(builder)) {
       myParser.getPatternParser().parsePattern(builder);
       if (builder.getTokenType() == JavaTokenType.COLON) {
-        return parseForEachFromColon(builder, statement);
+        return parseForEachFromColon(builder, statement, JavaElementType.FOREACH_PATTERN_STATEMENT);
       }
       error(builder, JavaPsiBundle.message("expected.colon"));
       // recovery: just skip everything until ')'
       while (true) {
         IElementType tokenType = builder.getTokenType();
         if (tokenType == null) {
-          done(statement, JavaElementType.FOREACH_STATEMENT);
+          done(statement, JavaElementType.FOREACH_PATTERN_STATEMENT);
           return statement;
         }
         if (tokenType == JavaTokenType.RPARENTH) {
-          return parserForEachFromRparenth(builder, statement);
+          return parserForEachFromRparenth(builder, statement, JavaElementType.FOREACH_PATTERN_STATEMENT);
         }
         builder.advanceLexer();
       }
@@ -396,7 +396,7 @@ public class StatementParser {
     }
     else {
       afterParenth.drop();
-      return parseForEachFromColon(builder, statement);
+      return parseForEachFromColon(builder, statement, JavaElementType.FOREACH_STATEMENT);
     }
   }
 
@@ -483,22 +483,22 @@ public class StatementParser {
   }
 
   @NotNull
-  private PsiBuilder.Marker parseForEachFromColon(PsiBuilder builder, PsiBuilder.Marker statement) {
+  private PsiBuilder.Marker parseForEachFromColon(PsiBuilder builder, PsiBuilder.Marker statement, IElementType foreachStatement) {
     builder.advanceLexer();
 
     if (myParser.getExpressionParser().parse(builder) == null) {
       error(builder, JavaPsiBundle.message("expected.expression"));
     }
 
-    return parserForEachFromRparenth(builder, statement);
+    return parserForEachFromRparenth(builder, statement, foreachStatement);
   }
 
-  private PsiBuilder.Marker parserForEachFromRparenth(PsiBuilder builder, PsiBuilder.Marker statement) {
+  private PsiBuilder.Marker parserForEachFromRparenth(PsiBuilder builder, PsiBuilder.Marker statement, IElementType forEachType) {
     if (expectOrError(builder, JavaTokenType.RPARENTH, "expected.rparen") && parseStatement(builder) == null) {
       error(builder, JavaPsiBundle.message("expected.statement"));
     }
 
-    done(statement, JavaElementType.FOREACH_STATEMENT);
+    done(statement, forEachType);
     return statement;
   }
 

java/java-psi-impl/src/com/intellij/psi/controlFlow/ControlFlowAnalyzer.java

@@ -713,6 +713,15 @@ final class ControlFlowAnalyzer extends JavaElementVisitor {
 
   @Override
   public void visitForeachStatement(@NotNull PsiForeachStatement statement) {
+    handleForeach(statement);
+  }
+
+  @Override
+  public void visitForeachPatternStatement(@NotNull PsiForeachPatternStatement statement) {
+    handleForeach(statement);
+  }
+
+  private void handleForeach(@NotNull PsiForeachStatementBase statement) {
     startElement(statement);
     final PsiStatement body = statement.getBody();
     myStartStatementStack.pushStatement(body == null ? statement : body, false);
@@ -727,14 +736,13 @@ final class ControlFlowAnalyzer extends JavaElementVisitor {
     myCurrentFlow.addInstruction(instruction);
     addElementOffsetLater(statement, false);
 
-    PsiForeachDeclarationElement iterationDeclaration = statement.getIterationDeclaration();
-    if (iterationDeclaration instanceof PsiParameter) {
-      final PsiParameter iterationParameter = (PsiParameter)iterationDeclaration;
+    if (statement instanceof PsiForeachStatement) {
+      final PsiParameter iterationParameter = ((PsiForeachStatement)statement).getIterationParameter();
       if (myPolicy.isParameterAccepted(iterationParameter)) {
         generateWriteInstruction(iterationParameter);
       }
-    } else if (iterationDeclaration instanceof PsiPattern) {
-      PsiPattern pattern = (PsiPattern)iterationDeclaration;
+    } else if (statement instanceof PsiForeachPatternStatement) {
+      PsiPattern pattern = ((PsiForeachPatternStatement)statement).getIterationPattern();
       processPattern(pattern);
     }
     if (body != null) {

java/java-psi-impl/src/com/intellij/psi/impl/source/tree/JavaElementType.java

@@ -115,6 +115,7 @@ public interface JavaElementType {
   IElementType WHILE_STATEMENT = new JavaCompositeElementType("WHILE_STATEMENT", () -> new PsiWhileStatementImpl());
   IElementType FOR_STATEMENT = new JavaCompositeElementType("FOR_STATEMENT", () -> new PsiForStatementImpl());
   IElementType FOREACH_STATEMENT = new JavaCompositeElementType("FOREACH_STATEMENT", () -> new PsiForeachStatementImpl());
+  IElementType FOREACH_PATTERN_STATEMENT = new JavaCompositeElementType("FOREACH_PATTERN_STATEMENT", () -> new PsiForeachPatternStatementImpl());
   IElementType DO_WHILE_STATEMENT = new JavaCompositeElementType("DO_WHILE_STATEMENT", () -> new PsiDoWhileStatementImpl());
   IElementType SWITCH_STATEMENT = new JavaCompositeElementType("SWITCH_STATEMENT", () -> new PsiSwitchStatementImpl());
   IElementType SWITCH_EXPRESSION = new JavaCompositeElementType("SWITCH_EXPRESSION", () -> new PsiSwitchExpressionImpl());

java/java-psi-impl/src/com/intellij/psi/impl/source/tree/java/PsiForeachPatternStatementImpl.java

@@ -0,0 +1,137 @@
+// Copyright 2000-2022 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.psi.impl.source.tree.java;
+
+import com.intellij.lang.ASTNode;
+import com.intellij.openapi.diagnostic.Logger;
+import com.intellij.psi.*;
+import com.intellij.psi.impl.PsiImplUtil;
+import com.intellij.psi.impl.source.Constants;
+import com.intellij.psi.impl.source.tree.ChildRole;
+import com.intellij.psi.scope.PsiScopeProcessor;
+import com.intellij.psi.tree.ChildRoleBase;
+import com.intellij.psi.tree.IElementType;
+import com.intellij.psi.util.PsiTreeUtil;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.Objects;
+
+public class PsiForeachPatternStatementImpl extends PsiLoopStatementImpl implements PsiForeachPatternStatement, Constants {
+  private static final Logger LOG = Logger.getInstance(PsiForeachPatternStatementImpl.class);
+  public PsiForeachPatternStatementImpl() {
+    super(FOREACH_PATTERN_STATEMENT);
+  }
+
+  @Override
+  public PsiExpression getIteratedValue() {
+    return (PsiExpression) findChildByRoleAsPsiElement(ChildRole.FOR_ITERATED_VALUE);
+  }
+
+  @Override
+  public PsiStatement getBody() {
+    return (PsiStatement) findChildByRoleAsPsiElement(ChildRole.LOOP_BODY);
+  }
+
+  @Override
+  @NotNull
+  public PsiJavaToken getLParenth() {
+    return (PsiJavaToken)Objects.requireNonNull(findChildByRoleAsPsiElement(ChildRole.LPARENTH));
+  }
+
+  @Override
+  public PsiJavaToken getRParenth() {
+    return (PsiJavaToken) findChildByRoleAsPsiElement(ChildRole.RPARENTH);
+  }
+
+  @Override
+  public ASTNode findChildByRole(int role) {
+    LOG.assertTrue(ChildRole.isUnique(role));
+
+    switch(role) {
+      case ChildRole.LOOP_BODY:
+        return PsiImplUtil.findStatementChild(this);
+
+      case ChildRole.FOR_ITERATED_VALUE:
+        return findChildByType(EXPRESSION_BIT_SET);
+
+      case ChildRole.FOR_KEYWORD:
+        return getFirstChildNode();
+
+      case ChildRole.LPARENTH:
+        return findChildByType(LPARENTH);
+
+      case ChildRole.RPARENTH:
+        return findChildByType(RPARENTH);
+
+      case ChildRole.COLON:
+        return findChildByType(COLON);
+
+      default:
+        return null;
+    }
+  }
+
+  @Override
+  public int getChildRole(@NotNull ASTNode child) {
+    LOG.assertTrue(child.getTreeParent() == this);
+
+    IElementType i = child.getElementType();
+    if (i == FOR_KEYWORD) {
+      return ChildRole.FOR_KEYWORD;
+    }
+    else if (i == LPARENTH) {
+      return ChildRole.LPARENTH;
+    }
+    else if (i == RPARENTH) {
+      return ChildRole.RPARENTH;
+    }
+    else if (i == COLON) {
+      return ChildRole.COLON;
+    }
+    else {
+      if (EXPRESSION_BIT_SET.contains(child.getElementType())) {
+        return ChildRole.FOR_ITERATED_VALUE;
+      }
+      else if (child.getPsi() instanceof PsiStatement) {
+        return ChildRole.LOOP_BODY;
+      }
+      else {
+        return ChildRoleBase.NONE;
+      }
+    }
+  }
+
+  @Override
+  public String toString() {
+    return "PsiForeachPatternStatement";
+  }
+
+  @Override
+  public boolean processDeclarations(@NotNull PsiScopeProcessor processor, @NotNull ResolveState state, PsiElement lastParent, @NotNull PsiElement place) {
+    processor.handleEvent(PsiScopeProcessor.Event.SET_DECLARATION_HOLDER, this);
+    if (lastParent == null || lastParent.getParent() != this || lastParent == getIteratedValue())
+      // Parent element should not see our vars
+      return true;
+
+    PsiPattern pattern = getIterationPattern();
+    if (pattern instanceof PsiDeconstructionPattern) {
+      PsiDeconstructionPattern deconstructionPattern = (PsiDeconstructionPattern)pattern;
+      return deconstructionPattern.processDeclarations(processor, state, lastParent, place);
+    }
+    return false;
+  }
+
+  @Override
+  public void accept(@NotNull PsiElementVisitor visitor) {
+    if (visitor instanceof JavaElementVisitor) {
+      ((JavaElementVisitor)visitor).visitForeachPatternStatement(this);
+    }
+    else {
+      visitor.visitElement(this);
+    }
+  }
+
+  @Override
+  public @NotNull PsiPattern getIterationPattern() {
+    return Objects.requireNonNull(PsiTreeUtil.getChildOfType(this, PsiPattern.class));
+  }
+}

java/java-psi-impl/src/com/intellij/psi/impl/source/tree/java/PsiForeachStatementImpl.java

@@ -26,26 +26,10 @@ public class PsiForeachStatementImpl extends PsiLoopStatementImpl implements Psi
     super(FOREACH_STATEMENT);
   }
 
-  @Override
-  @Nullable
-  public PsiParameter getIterationParameter() {
-    PsiParameter parameter = (PsiParameter)findChildByRoleAsPsiElement(ChildRole.FOR_ITERATION_PARAMETER);
-    if (parameter == null) {
-      LOG.error("getIterationParameter is used when forEach element contains pattern. Migrate to getIterationDeclaration()");
-      return null;
-    }
-    return parameter;
-  }
-
   @Override
   @NotNull
-  public PsiForeachDeclarationElement getIterationDeclaration() {
-    PsiParameter parameter = (PsiParameter)findChildByRoleAsPsiElement(ChildRole.FOR_ITERATION_PARAMETER);
-    if (parameter != null) {
-      return parameter;
-    } else {
-      return Objects.requireNonNull(PsiTreeUtil.getChildOfType(this, PsiPattern.class));
-    }
+  public PsiParameter getIterationParameter() {
+    return (PsiParameter)Objects.requireNonNull(findChildByRoleAsPsiElement(ChildRole.FOR_ITERATION_PARAMETER));
   }
 
   @Override
@@ -61,7 +45,7 @@ public class PsiForeachStatementImpl extends PsiLoopStatementImpl implements Psi
   @Override
   @NotNull
   public PsiJavaToken getLParenth() {
-    return (PsiJavaToken) findChildByRoleAsPsiElement(ChildRole.LPARENTH);
+    return (PsiJavaToken)Objects.requireNonNull(findChildByRoleAsPsiElement(ChildRole.LPARENTH));
   }
 
   @Override
@@ -145,16 +129,7 @@ public class PsiForeachStatementImpl extends PsiLoopStatementImpl implements Psi
       // Parent element should not see our vars
       return true;
 
-    PsiForeachDeclarationElement iterationDeclaration = getIterationDeclaration();
-    if (iterationDeclaration instanceof PsiParameter) {
-      PsiParameter parameter = (PsiParameter)iterationDeclaration;
-      return processor.execute(parameter, state);
-    } else if (iterationDeclaration instanceof PsiDeconstructionPattern) {
-      PsiDeconstructionPattern deconstructionPattern = (PsiDeconstructionPattern)iterationDeclaration;
-      return deconstructionPattern.processDeclarations(processor, state, lastParent, place);
-    } {
-      return false;
-    }
+    return processor.execute(getIterationParameter(), state);
   }
 
   @Override

java/java-tests/testData/psi/parser-partial/statements/ForEach1.txt

@@ -1,5 +1,5 @@
 PsiJavaFile:ForEach1.java
-  PsiForeachStatement
+  PsiForeachPatternStatement
     PsiKeyword:for('for')
     PsiWhiteSpace(' ')
     PsiJavaToken:LPARENTH('(')

java/java-tests/testData/psi/parser-partial/statements/ForEachIncomplete1.txt

@@ -1,5 +1,5 @@
 PsiJavaFile:ForEachIncomplete1.java
-  PsiForeachStatement
+  PsiForeachPatternStatement
     PsiKeyword:for('for')
     PsiWhiteSpace(' ')
     PsiJavaToken:LPARENTH('(')

java/java-tests/testData/psi/parser-partial/statements/ForEachIncomplete2.txt

@@ -1,5 +1,5 @@
 PsiJavaFile:ForEachIncomplete2.java
-  PsiForeachStatement
+  PsiForeachPatternStatement
     PsiKeyword:for('for')
     PsiWhiteSpace(' ')
     PsiJavaToken:LPARENTH('(')

java/java-tests/testData/psi/parser-partial/statements/ForEachIncomplete3.txt

@@ -1,5 +1,5 @@
 PsiJavaFile:ForEachIncomplete3.java
-  PsiForeachStatement
+  PsiForeachPatternStatement
     PsiKeyword:for('for')
     PsiWhiteSpace(' ')
     PsiJavaToken:LPARENTH('(')

java/java-tests/testData/psi/parser-partial/statements/ForEachIncomplete4.txt

@@ -1,5 +1,5 @@
 PsiJavaFile:ForEachIncomplete4.java
-  PsiForeachStatement
+  PsiForeachPatternStatement
     PsiKeyword:for('for')
     PsiWhiteSpace(' ')
     PsiJavaToken:LPARENTH('(')

java/java-tests/testData/psi/parser-partial/statements/ForEachIncomplete5.txt

@@ -1,5 +1,5 @@
 PsiJavaFile:ForEachIncomplete5.java
-  PsiForeachStatement
+  PsiForeachPatternStatement
     PsiKeyword:for('for')
     PsiWhiteSpace(' ')
     PsiJavaToken:LPARENTH('(')