From c49e231d56af4ebfef31a5e636c9fa2814129e48 Mon Sep 17 00:00:00 2001
From: Bas Leijdekkers <basleijdekkers@gmail.com>
Date: Sat, 25 Jan 2020 16:32:08 +0100
Subject: [PATCH] SSR: match enum constant argument lists and initializers
 (IDEA-231441)

GitOrigin-RevId: c7ba6b448e4d9bf385f0ce68cdfd40179655aa4d
---
 .../impl/matcher/JavaMatchingVisitor.java     | 26 ++++++++++++++++---
 .../StructuralSearchTest.java                 | 12 ++++++---
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/java/structuralsearch-java/src/com/intellij/structuralsearch/impl/matcher/JavaMatchingVisitor.java b/java/structuralsearch-java/src/com/intellij/structuralsearch/impl/matcher/JavaMatchingVisitor.java
index 2adfde2f9cb9..f2b757ba6192 100644
--- a/java/structuralsearch-java/src/com/intellij/structuralsearch/impl/matcher/JavaMatchingVisitor.java
+++ b/java/structuralsearch-java/src/com/intellij/structuralsearch/impl/matcher/JavaMatchingVisitor.java
@@ -5,6 +5,7 @@ import com.intellij.dupLocator.iterators.ArrayBackedNodeIterator;
 import com.intellij.dupLocator.iterators.NodeIterator;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.*;
+import com.intellij.psi.impl.light.LightElement;
 import com.intellij.psi.javadoc.PsiDocComment;
 import com.intellij.psi.javadoc.PsiDocTag;
 import com.intellij.psi.tree.IElementType;
@@ -348,23 +349,37 @@ public class JavaMatchingVisitor extends JavaElementVisitor {
 
   @Override
   public void visitField(PsiField field) {
+    final PsiField other = getElement(PsiField.class);
+    if (other == null) return;
     final PsiDocComment comment = field.getDocComment();
-    final PsiField other = (PsiField)myMatchingVisitor.getElement();
     if (comment != null && !myMatchingVisitor.setResult(myMatchingVisitor.match(comment, other))) return;
     if (!myMatchingVisitor.setResult(checkHierarchy(other, field))) return;
     super.visitField(field);
   }
 
+  @Override
+  public void visitEnumConstant(PsiEnumConstant enumConstant) {
+    final PsiEnumConstant other = getElement(PsiEnumConstant.class);
+    if (other == null) return;
+    final PsiExpressionList argumentList = enumConstant.getArgumentList();
+    if (argumentList != null && !myMatchingVisitor.setResult(myMatchingVisitor.matchSons(argumentList, other.getArgumentList()))) return;
+    final PsiEnumConstantInitializer enumConstantInitializer = enumConstant.getInitializingClass();
+    if (enumConstantInitializer != null &&
+        !myMatchingVisitor.setResult(myMatchingVisitor.match(enumConstantInitializer, other.getInitializingClass()))) return;
+    super.visitEnumConstant(enumConstant);
+  }
+
   @Override
   public void visitAnonymousClass(PsiAnonymousClass clazz) {
-    final PsiAnonymousClass other = (PsiAnonymousClass)myMatchingVisitor.getElement();
+    final PsiAnonymousClass other = getElement(PsiAnonymousClass.class);
+    if (other == null) return;
     final PsiElement classReference = clazz.getBaseClassReference();
     final boolean isTypedVar = myMatchingVisitor.getMatchContext().getPattern().isTypedVar(classReference);
 
     if (myMatchingVisitor.setResult((isTypedVar || myMatchingVisitor.match(clazz.getBaseClassReference(), other.getBaseClassReference())) &&
                                     myMatchingVisitor.matchSons(clazz.getArgumentList(), other.getArgumentList()) &&
                                     matchClasses(clazz, other)) && isTypedVar) {
-      myMatchingVisitor.setResult(matchType(classReference, other.getBaseClassReference()));
+      myMatchingVisitor.setResult(classReference instanceof LightElement || matchType(classReference, other.getBaseClassReference()));
     }
   }
 
@@ -1783,4 +1798,9 @@ public class JavaMatchingVisitor extends JavaElementVisitor {
     }
     return myMatchingVisitor.setResult(aClass.isInstance(other)) ? aClass.cast(other) : null;
   }
+
+  private <T extends PsiElement> T getElement(Class<T> aClass) {
+    final PsiElement other = myMatchingVisitor.getElement();
+    return myMatchingVisitor.setResult(aClass.isInstance(other)) ? aClass.cast(other) : null;
+  }
 }
diff --git a/platform/structuralsearch/testSource/com/intellij/structuralsearch/StructuralSearchTest.java b/platform/structuralsearch/testSource/com/intellij/structuralsearch/StructuralSearchTest.java
index 50a752ae9e3c..05347f4f0cc0 100644
--- a/platform/structuralsearch/testSource/com/intellij/structuralsearch/StructuralSearchTest.java
+++ b/platform/structuralsearch/testSource/com/intellij/structuralsearch/StructuralSearchTest.java
@@ -1,4 +1,4 @@
-// Copyright 2000-2019 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-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.
 package com.intellij.structuralsearch;
 
 import com.intellij.openapi.fileTypes.LanguageFileType;
@@ -2260,9 +2260,15 @@ public class StructuralSearchTest extends StructuralSearchTestCase {
 
   public void testFindEnums() {
     String s = "class Foo {} class Bar {} enum X {}";
-    String s2 = "enum 'x {}";
 
-    assertEquals(1, findMatchesCount(s,s2));
+    assertEquals(1, findMatchesCount(s, "enum 'x {}"));
+
+    String in = "enum E {" +
+                "  A(1), B(2), C(3)" +
+                "}";
+    assertEquals(1, findMatchesCount(in, "enum '_E { 'A(2) }"));
+    assertEquals(0, findMatchesCount(in, "enum '_E { 'A('_x{0,0}) }"));
+    assertEquals(0, findMatchesCount(in, "enum '_E { 'A(2) {} }"));
   }
 
   public void testFindDeclaration() {