New LST models to support use-site annotation and deprecate `Annotati…

…onUseSite` marker (#549) * New LST models to support use-site annotation and deprecate AnnotationUseSite marker * Remove legacy `typealias` handling from `KotlinPrinter` --------- Co-authored-by: Knut Wannheden <[email protected]>
openrewrite · Dec 19, 2023 · 3804358 · 3804358
1 parent d38c9b8
commit 3804358
Show file tree

Hide file tree

Showing 7 changed files with 277 additions and 45 deletions.
diff --git a/src/main/java/org/openrewrite/kotlin/Assertions.java b/src/main/java/org/openrewrite/kotlin/Assertions.java
@@ -23,7 +23,6 @@
 import org.openrewrite.internal.lang.Nullable;
 import org.openrewrite.java.search.FindMissingTypes;
 import org.openrewrite.java.tree.*;
-import org.openrewrite.kotlin.marker.AnnotationUseSite;
 import org.openrewrite.kotlin.marker.Extension;
 import org.openrewrite.kotlin.marker.IndexedAccess;
 import org.openrewrite.kotlin.tree.K;
@@ -470,7 +469,7 @@ private boolean isAllowedToHaveNullType(J.Identifier ident) {
             return inPackageDeclaration() || inImport() || isClassName()
                    || isMethodName() || isMethodInvocationName() || isFieldAccess(ident) || isBeingDeclared(ident) || isParameterizedType(ident)
                    || isNewClass(ident) || isTypeParameter() || isMemberReference(ident) || isCaseLabel() || isLabel() || isAnnotationField(ident)
-                   || isInJavaDoc(ident) || isWhenLabel();
+                   || isInJavaDoc(ident) || isWhenLabel() || isUseSite();
         }
 
         private boolean inPackageDeclaration() {
@@ -544,6 +543,11 @@ private boolean isWhenLabel() {
             return getCursor().getParentTreeCursor().getValue() instanceof K.WhenBranch;
         }
 
+        private boolean isUseSite() {
+            Tree value = getCursor().getParentTreeCursor().getValue();
+            return value instanceof K.AnnotationType || value instanceof K.MultiAnnotationType;
+        }
+
         private boolean isLabel() {
             return getCursor().firstEnclosing(J.Label.class) != null;
         }
@@ -557,7 +561,7 @@ private boolean isAnnotationField(J.Identifier ident) {
         private boolean isValidated(J.Identifier i) {
             J j = getCursor().dropParentUntil(it -> it instanceof J).getValue();
             // TODO: replace with AnnotationUseSite tree.
-            return !j.getMarkers().findFirst(AnnotationUseSite.class).isPresent() && !(j instanceof K.KReturn);
+            return !(j instanceof K.KReturn);
         }
 
         private boolean isValidated(J.MethodInvocation mi) {

diff --git a/src/main/java/org/openrewrite/kotlin/KotlinVisitor.java b/src/main/java/org/openrewrite/kotlin/KotlinVisitor.java
@@ -323,6 +323,24 @@ public J visitUnary(K.Unary unary, P p) {
         return u;
     }
 
+    public J visitAnnotationType(K.AnnotationType annotationType, P p) {
+        K.AnnotationType at = annotationType;
+        at = at.withPrefix(visitSpace(at.getPrefix(), Space.Location.ANNOTATION_PREFIX, p));
+        at = at.withMarkers(visitMarkers(at.getMarkers(), p));
+        at = at.getPadding().withUseSite(visitRightPadded(at.getPadding().getUseSite(), JRightPadded.Location.ANNOTATION_ARGUMENT, p));
+        at = at.withCallee(visitAndCast(at.getCallee(), p));
+        return at;
+    }
+
+    public J visitMultiAnnotationType(K.MultiAnnotationType multiAnnotationType, P p) {
+        K.MultiAnnotationType mat = multiAnnotationType;
+        mat = mat.withPrefix(visitSpace(mat.getPrefix(), Space.Location.ANNOTATION_PREFIX, p));
+        mat = mat.withMarkers(visitMarkers(mat.getMarkers(), p));
+        mat = mat.getPadding().withUseSite(visitRightPadded(mat.getPadding().getUseSite(), JRightPadded.Location.ANNOTATION_ARGUMENT, p));
+        mat = mat.withAnnotations(visitContainer(mat.getAnnotations(), p));
+        return mat;
+    }
+
     public J visitWhen(K.When when, P p) {
         K.When w = when;
         w = w.withPrefix(visitSpace(w.getPrefix(), KSpace.Location.WHEN_PREFIX, p));
@@ -417,10 +435,7 @@ public <T> JRightPadded<T> visitRightPadded(@Nullable JRightPadded<T> right, KRi
     @Override
     public <M extends Marker> M visitMarker(Marker marker, P p) {
         Marker m = super.visitMarker(marker, p);
-        if (m instanceof AnnotationUseSite) {
-            AnnotationUseSite acs = (AnnotationUseSite) marker;
-            m = acs.withPrefix(visitSpace(acs.getPrefix(), KSpace.Location.ANNOTATION_CALL_SITE_PREFIX, p));
-        } else if (marker instanceof TypeReferencePrefix) {
+        if (marker instanceof TypeReferencePrefix) {
             TypeReferencePrefix tr = (TypeReferencePrefix) marker;
             m = tr.withPrefix(visitSpace(tr.getPrefix(), KSpace.Location.TYPE_REFERENCE_PREFIX, p));
         }

diff --git a/src/main/java/org/openrewrite/kotlin/internal/KotlinPrinter.java b/src/main/java/org/openrewrite/kotlin/internal/KotlinPrinter.java
@@ -417,6 +417,28 @@ public J visitUnary(K.Unary unary, PrintOutputCapture<P> p) {
         return unary;
     }
 
+    @Override
+    public J visitAnnotationType(K.AnnotationType annotationType, PrintOutputCapture<P> p) {
+        beforeSyntax(annotationType, Space.Location.ANNOTATION_PREFIX, p);
+        visitRightPadded(annotationType.getPadding().getUseSite(), p);
+        p.append(":");
+        visit(annotationType.getCallee(), p);
+        return annotationType;
+    }
+
+    @Override
+    public J visitMultiAnnotationType(K.MultiAnnotationType multiAnnotationType, PrintOutputCapture<P> p) {
+        beforeSyntax(multiAnnotationType, Space.Location.ANNOTATION_PREFIX, p);
+        visitRightPadded(multiAnnotationType.getPadding().getUseSite(), p);
+
+        if (!(multiAnnotationType.getUseSite() instanceof J.Empty)) {
+            p.append(":");
+        }
+
+        delegate.visitContainer("[", multiAnnotationType.getAnnotations(), JContainer.Location.TYPE_PARAMETERS, "", "]",  p);
+        return multiAnnotationType;
+    }
+
     @Override
     public J visitWhen(K.When when, PrintOutputCapture<P> p) {
         beforeSyntax(when, KSpace.Location.WHEN_PREFIX, p);
@@ -469,22 +491,6 @@ public J visitAnnotation(J.Annotation annotation, PrintOutputCapture<P> p) {
             String afterArgs = ")";
             String delimiter = ",";
 
-            AnnotationUseSite useSite = annotation.getMarkers().findFirst(AnnotationUseSite.class).orElse(null);
-            if (useSite != null) {
-                kotlinPrinter.visitSpace(useSite.getPrefix(), KSpace.Location.ANNOTATION_CALL_SITE_PREFIX, p);
-                p.append(":");
-
-                if (!useSite.isImplicitBracket()) {
-                    beforeArgs = "[";
-                    afterArgs = "]";
-                } else {
-                    beforeArgs = "";
-                    afterArgs = "";
-                }
-
-                delimiter = "";
-            }
-
             visitContainer(beforeArgs, annotation.getPadding().getArguments(), JContainer.Location.ANNOTATION_ARGUMENTS, delimiter, afterArgs, p);
             afterSyntax(annotation, p);
             return annotation;
@@ -1129,14 +1135,6 @@ public J visitWildcard(J.Wildcard wildcard, PrintOutputCapture<P> p) {
 
         @Override
         public J visitVariableDeclarations(J.VariableDeclarations multiVariable, PrintOutputCapture<P> p) {
-            // TypeAliases are converted into a J.VariableDeclaration to re-use complex recipes like RenameVariable and ChangeType.
-            // However, a type alias has different syntax and is printed separately to reduce code complexity in visitVariableDeclarations.
-            // This is a temporary solution until K.TypeAlias is added to the model, and RenameVariable is revised to operator from a J.Identifier.
-            if (multiVariable.getLeadingAnnotations().stream().anyMatch(it -> "typealias".equals(it.getSimpleName()))) {
-                visitTypeAlias(multiVariable, p);
-                return multiVariable;
-            }
-
             beforeSyntax(multiVariable, Space.Location.VARIABLE_DECLARATIONS_PREFIX, p);
 
             visit(multiVariable.getLeadingAnnotations(), p);

diff --git a/src/main/java/org/openrewrite/kotlin/internal/KotlinTreeParserVisitor.java b/src/main/java/org/openrewrite/kotlin/internal/KotlinTreeParserVisitor.java
@@ -1732,32 +1732,44 @@ public J visitKtFile(KtFile file, ExecutionContext data) {
 
     @Override
     public J visitAnnotation(KtAnnotation annotation, ExecutionContext data) {
-        if (annotation.getUseSiteTarget() == null) {
-            throw new UnsupportedOperationException("TODO, Some cases we don't know");
+        Expression target;
+
+        if (annotation.getUseSiteTarget() != null) {
+            target = (J.Identifier) annotation.getUseSiteTarget().accept(this, data);
+        } else {
+            target = new J.Empty(randomId(), Space.EMPTY, Markers.EMPTY);
         }
-        List<KtAnnotationEntry> annotationEntries = annotation.getEntries();
-        List<JRightPadded<Expression>> rpAnnotations = new ArrayList<>(annotationEntries.size());
 
+        List<KtAnnotationEntry> annotationEntries = annotation.getEntries();
+        List<JRightPadded<J.Annotation>> rpAnnotations = new ArrayList<>(annotationEntries.size());
+        J.Annotation anno = null;
         for (KtAnnotationEntry ktAnnotationEntry : annotationEntries) {
-            J.Annotation anno = (J.Annotation) ktAnnotationEntry.accept(this, data);
+            anno = (J.Annotation) ktAnnotationEntry.accept(this, data);
             anno = anno.withMarkers(anno.getMarkers().addIfAbsent(new AnnotationConstructor(randomId())));
             rpAnnotations.add(padRight(anno, Space.EMPTY));
         }
 
-        PsiElement maybeLBracket = findFirstChild(annotation, anno -> anno.getNode().getElementType() == KtTokens.LBRACKET);
+        PsiElement maybeLBracket = findFirstChild(annotation, an -> an.getNode().getElementType() == KtTokens.LBRACKET);
         boolean isImplicitBracket = maybeLBracket == null;
         Space beforeLBracket = isImplicitBracket ? Space.EMPTY : prefix(maybeLBracket);
 
         if (!isImplicitBracket) {
             rpAnnotations = ListUtils.mapLast(rpAnnotations,
-                    rp -> rp.withAfter(prefix(findFirstChild(annotation, anno -> anno.getNode().getElementType() == KtTokens.RBRACKET))));
+                    rp -> rp.withAfter(prefix(findFirstChild(annotation, an -> an.getNode().getElementType() == KtTokens.RBRACKET))));
+        }
+
+        NameTree annotationType;
+        if (isImplicitBracket) {
+            annotationType = new K.AnnotationType(randomId(), Space.EMPTY, Markers.EMPTY, padRight(target, suffix(annotation.getUseSiteTarget())), anno);
+        } else {
+            annotationType = new K.MultiAnnotationType(randomId(), Space.EMPTY, Markers.EMPTY, padRight(target, suffix(annotation.getUseSiteTarget())), JContainer.build(beforeLBracket, rpAnnotations, Markers.EMPTY));
         }
 
         return mapType(new J.Annotation(randomId(),
                 Space.EMPTY,
-                Markers.EMPTY.addIfAbsent(new AnnotationUseSite(randomId(), suffix(annotation.getUseSiteTarget()), isImplicitBracket)),
-                (NameTree) annotation.getUseSiteTarget().accept(this, data),
-                JContainer.build(beforeLBracket, rpAnnotations, Markers.EMPTY)
+                Markers.EMPTY,
+                annotationType,
+                null
         ));
     }
 
@@ -1773,16 +1785,20 @@ public J visitAnnotationEntry(@NotNull KtAnnotationEntry annotationEntry, Execut
         }
 
         if (isUseSite) {
-            nameTree = (NameTree) annotationEntry.getUseSiteTarget().accept(this, data);
-            markers = markers.addIfAbsent(new AnnotationUseSite(randomId(), prefix(findFirstChild(annotationEntry, p -> p.getNode().getElementType() == KtTokens.COLON)), true));
-            J.Annotation argAnno = new J.Annotation(
+            J.Annotation callee = new J.Annotation(
                     randomId(),
                     Space.EMPTY,
                     Markers.EMPTY.addIfAbsent(new AnnotationConstructor(randomId())),
                     (NameTree) requireNonNull(annotationEntry.getCalleeExpression()).accept(this, data),
                     annotationEntry.getValueArgumentList() != null ? mapValueArguments(annotationEntry.getValueArgumentList(), data) : null
             );
-            args = JContainer.build(Space.EMPTY, singletonList(padRight(argAnno, Space.EMPTY)), Markers.EMPTY);
+
+            nameTree = new K.AnnotationType(randomId(),
+                    Space.EMPTY,
+                    Markers.EMPTY,
+                    padRight(convertToExpression(annotationEntry.getUseSiteTarget().accept(this, data)),
+                            prefix(findFirstChild(annotationEntry, p -> p.getNode().getElementType() == KtTokens.COLON))),
+                    callee);
         } else {
             nameTree = (NameTree) requireNonNull(annotationEntry.getCalleeExpression()).accept(this, data);
             if (annotationEntry.getValueArgumentList() != null) {

diff --git a/src/main/java/org/openrewrite/kotlin/marker/AnnotationUseSite.java b/src/main/java/org/openrewrite/kotlin/marker/AnnotationUseSite.java
@@ -22,6 +22,7 @@
 
 import java.util.UUID;
 
+@Deprecated
 @Value
 @With
 public class AnnotationUseSite implements Marker {

diff --git a/src/main/java/org/openrewrite/kotlin/tree/K.java b/src/main/java/org/openrewrite/kotlin/tree/K.java
@@ -2149,4 +2149,168 @@ public K.Unary withOperator(JLeftPadded<K.Unary.Type> operator) {
         }
     }
 
+    @FieldDefaults(makeFinal = true, level = AccessLevel.PRIVATE)
+    @EqualsAndHashCode(callSuper = false, onlyExplicitlyIncluded = true)
+    @RequiredArgsConstructor
+    @AllArgsConstructor(access = AccessLevel.PRIVATE)
+    final class AnnotationType implements K, NameTree {
+        @Nullable
+        @NonFinal
+        transient WeakReference<K.AnnotationType.Padding> padding;
+
+        @Getter
+        @With
+        @EqualsAndHashCode.Include
+        UUID id;
+
+        @Getter
+        @With
+        Space prefix;
+
+        @Getter
+        @With
+        Markers markers;
+
+        JRightPadded<Expression> useSite;
+
+        @Getter
+        @With
+        J.Annotation callee;
+
+        public Expression getUseSite() {
+            return useSite.getElement();
+        }
+
+        @Override
+        public @Nullable JavaType getType() {
+            // use site has no type
+            return null;
+        }
+
+        @Override
+        public <T extends J> T withType(@Nullable JavaType type) {
+            return (T) this;
+        }
+
+        @Override
+        public <P> J acceptKotlin(KotlinVisitor<P> v, P p) {
+            return v.visitAnnotationType(this, p);
+        }
+
+        public K.AnnotationType.Padding getPadding() {
+            K.AnnotationType.Padding p;
+            if (this.padding == null) {
+                p = new K.AnnotationType.Padding(this);
+                this.padding = new WeakReference<>(p);
+            } else {
+                p = this.padding.get();
+                if (p == null || p.t != this) {
+                    p = new K.AnnotationType.Padding(this);
+                    this.padding = new WeakReference<>(p);
+                }
+            }
+            return p;
+        }
+
+        @RequiredArgsConstructor
+        public static class Padding {
+            private final K.AnnotationType t;
+
+            public JRightPadded<Expression> getUseSite() {
+                return t.useSite;
+            }
+
+            public K.AnnotationType withUseSite(JRightPadded<Expression> useSite) {
+                return t.useSite == useSite ? t : new K.AnnotationType(t.id,
+                        t.prefix,
+                        t.markers,
+                        useSite,
+                        t.callee
+                );
+            }
+        }
+    }
+
+    @FieldDefaults(makeFinal = true, level = AccessLevel.PRIVATE)
+    @EqualsAndHashCode(callSuper = false, onlyExplicitlyIncluded = true)
+    @RequiredArgsConstructor
+    @AllArgsConstructor(access = AccessLevel.PRIVATE)
+    final class MultiAnnotationType implements K, NameTree {
+        @Nullable
+        @NonFinal
+        transient WeakReference<K.MultiAnnotationType.Padding> padding;
+
+        @Getter
+        @With
+        @EqualsAndHashCode.Include
+        UUID id;
+
+        @Getter
+        @With
+        Space prefix;
+
+        @Getter
+        @With
+        Markers markers;
+
+        JRightPadded<Expression> useSite;
+
+        @Getter
+        @With
+        JContainer<J.Annotation> annotations;
+
+        public Expression getUseSite() {
+            return useSite.getElement();
+        }
+
+        @Override
+        public @Nullable JavaType getType() {
+            // use site has no type
+            return null;
+        }
+
+        @Override
+        public <T extends J> T withType(@Nullable JavaType type) {
+            return (T) this;
+        }
+
+        @Override
+        public <P> J acceptKotlin(KotlinVisitor<P> v, P p) {
+            return v.visitMultiAnnotationType(this, p);
+        }
+
+        public K.MultiAnnotationType.Padding getPadding() {
+            K.MultiAnnotationType.Padding p;
+            if (this.padding == null) {
+                p = new K.MultiAnnotationType.Padding(this);
+                this.padding = new WeakReference<>(p);
+            } else {
+                p = this.padding.get();
+                if (p == null || p.t != this) {
+                    p = new K.MultiAnnotationType.Padding(this);
+                    this.padding = new WeakReference<>(p);
+                }
+            }
+            return p;
+        }
+
+        @RequiredArgsConstructor
+        public static class Padding {
+            private final K.MultiAnnotationType t;
+
+            public JRightPadded<Expression> getUseSite() {
+                return t.useSite;
+            }
+
+            public K.MultiAnnotationType withUseSite(JRightPadded<Expression> useSite) {
+                return t.useSite == useSite ? t : new K.MultiAnnotationType(t.id,
+                        t.prefix,
+                        t.markers,
+                        useSite,
+                        t.annotations
+                );
+            }
+        }
+    }
+
 }
-Original file line number
+Diff line change
@@ Expand Up / @@ -22,6 +22,7 @@ @@
     import java.util.UUID;
+    @Deprecated
     @Value
     @With
     public class AnnotationUseSite implements Marker {
@@ Expand Down @@