From b437398eabad4c3b3a6db1cd8895c52dca1d88b8 Mon Sep 17 00:00:00 2001
From: "Roman.Ivanov" <Roman.Ivanov@jetbrains.com>
Date: Wed, 18 Dec 2019 16:42:09 +0700
Subject: [PATCH] RecordAugmentProvider: transform ellipsis types and simplify

GitOrigin-RevId: 4490bbff05c3958862cf470154fd34db56a65c3f
---
 .../psi/impl/RecordAugmentProvider.java       | 29 +++++++------------
 1 file changed, 10 insertions(+), 19 deletions(-)

diff --git a/java/java-psi-impl/src/com/intellij/psi/impl/RecordAugmentProvider.java b/java/java-psi-impl/src/com/intellij/psi/impl/RecordAugmentProvider.java
index 88e292b848bf..103caec34a5f 100644
--- a/java/java-psi-impl/src/com/intellij/psi/impl/RecordAugmentProvider.java
+++ b/java/java-psi-impl/src/com/intellij/psi/impl/RecordAugmentProvider.java
@@ -11,12 +11,13 @@ import com.intellij.psi.impl.source.PsiExtensibleClass;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-import java.util.*;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
-import static com.intellij.util.ObjectUtils.tryCast;
-
 public class RecordAugmentProvider extends PsiAugmentProvider {
   @NotNull
   @Override
@@ -91,37 +92,27 @@ public class RecordAugmentProvider extends PsiAugmentProvider {
     if (name == null) return null;
     String typeText = getTypeText(component, RecordAugmentProvider::hasTargetApplicableForMethod);
     if (typeText == null) return null;
-    return factory.createMethodFromText("public " + typeText + " " + name + "(){}", component.getContainingClass());
+    return factory.createMethodFromText("public " + typeText + " " + name + "(){ return " + name + "; }", component.getContainingClass());
   }
 
   @Nullable
   private static String getTypeText(@NotNull PsiRecordComponent component, Predicate<PsiAnnotation> annotationPredicate) {
+    PsiType type = component.getType();
+    if (type instanceof PsiEllipsisType) type = ((PsiEllipsisType)type).toArrayType();
     PsiTypeElement typeElement = component.getTypeElement();
     if (typeElement == null) return null;
     String annotations = Arrays.stream(component.getAnnotations())
         .filter(annotationPredicate)
         .map(annotation -> annotation.getText())
         .collect(Collectors.joining(" "));
-    return annotations + " " + typeElement.getText();
+    return annotations + " " + type.getCanonicalText();
   }
 
   private static boolean hasTargetApplicableForField(PsiAnnotation annotation) {
-    Set<PsiAnnotation.TargetType> targets = getTargets(annotation);
-    if (targets == null) return false;
-    return targets.contains(PsiAnnotation.TargetType.TYPE) || targets.contains(PsiAnnotation.TargetType.FIELD);
+    return AnnotationTargetUtil.findAnnotationTarget(annotation, PsiAnnotation.TargetType.TYPE_USE, PsiAnnotation.TargetType.FIELD) != null;
   }
 
   private static boolean hasTargetApplicableForMethod(PsiAnnotation annotation) {
-    Set<PsiAnnotation.TargetType> targets = getTargets(annotation);
-    if (targets == null) return false;
-    return targets.contains(PsiAnnotation.TargetType.TYPE) || targets.contains(PsiAnnotation.TargetType.METHOD);
-  }
-
-  private static Set<PsiAnnotation.TargetType> getTargets(PsiAnnotation annotation) {
-    PsiJavaCodeReferenceElement element = annotation.getNameReferenceElement();
-    if (element == null) return null;
-    PsiClass annotationClass = tryCast(element.resolve(), PsiClass.class);
-    if (annotationClass == null) return null;
-    return AnnotationTargetUtil.getAnnotationTargets(annotationClass);
+    return AnnotationTargetUtil.findAnnotationTarget(annotation, PsiAnnotation.TargetType.TYPE_USE, PsiAnnotation.TargetType.METHOD) != null;
   }
 }