Trait interface and VariableAccess/MethodAccess implementation

rewrite-core/src/main/java/org/openrewrite/trait/SimpleTraitMatcher.java

@@ -0,0 +1,87 @@
+/*
+ * Copyright 2024 the original author or authors.
+ * <p>
+ * 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
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * 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 org.openrewrite.trait;
+
+import org.openrewrite.Cursor;
+import org.openrewrite.Incubating;
+import org.openrewrite.Tree;
+import org.openrewrite.TreeVisitor;
+import org.openrewrite.internal.lang.Nullable;
+
+import java.util.Iterator;
+import java.util.Optional;
+import java.util.stream.Stream;
+
+@Incubating(since = "8.30.0")
+public abstract class SimpleTraitMatcher<U extends Trait<?>> implements TraitMatcher<U> {
+
+    @Override
+    public Optional<U> get(Cursor cursor) {
+        return Optional.ofNullable(test(cursor));
+    }
+
+    @Override
+    public Stream<U> higher(Cursor cursor) {
+        Stream.Builder<U> stream = Stream.builder();
+        Iterator<Cursor> cursors = cursor.getPathAsCursors();
+        while (cursors.hasNext()) {
+            Cursor c = cursors.next();
+            if (c != cursor) {
+                U u = test(c);
+                if (u != null) {
+                    stream.add(u);
+                }
+            }
+        }
+        return stream.build();
+    }
+
+    @Override
+    public Stream<U> lower(Cursor cursor) {
+        Stream.Builder<U> stream = Stream.builder();
+        this.<Stream.Builder<U>>asVisitor((va, c, sb) -> {
+            sb.add(test(c));
+            return va.getTree();
+        }).visit(cursor.getValue(), stream, cursor.getParentOrThrow());
+        return stream.build();
+    }
+
+    /**
+     * This method is called on every tree. For more performant matching, traits should override
+     * and provide a narrower visitor that only calls {@link #test(Cursor)} against tree types that
+     * could potentially match.
+     *
+     * @param visitor Called for each match of the trait. The function is passed the trait, the cursor at which the
+     *                trait was found, and a context object.
+     * @param <P>     The type of the context object.
+     * @return A visitor that can be used to locate trees matching the trait.
+     */
+    @Override
+    public <P> TreeVisitor<? extends Tree, P> asVisitor(VisitFunction3<U, P> visitor) {
+        return new TreeVisitor<Tree, P>() {
+            @Override
+            public @Nullable Tree visit(@Nullable Tree tree, P p) {
+                U u = test(getCursor());
+                return u != null ?
+                        visitor.visit(u, getCursor(), p) :
+                        super.visit(tree, p);
+            }
+        };
+    }
+
+    @Nullable
+    protected abstract U test(Cursor cursor);
+}

rewrite-core/src/main/java/org/openrewrite/trait/Trait.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright 2024 the original author or authors.
+ * <p>
+ * 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
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * 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 org.openrewrite.trait;
+
+import org.openrewrite.Cursor;
+import org.openrewrite.Incubating;
+import org.openrewrite.Tree;
+
+/**
+ * A trait captures semantic information related to a syntax element.
+ *
+ * @param <T> The type of the tree that this trait is related to. When
+ *            multiple specific types of tree are possible, this should
+ *            be the lowest common super-type of all the types.
+ */
+@Incubating(since = "8.30.0")
+public interface Trait<T extends Tree> {
+    Cursor getCursor();
+
+    default T getTree() {
+        return getCursor().getValue();
+    }
+}

rewrite-core/src/main/java/org/openrewrite/trait/TraitMatcher.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2024 the original author or authors.
+ * <p>
+ * 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
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * 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 org.openrewrite.trait;
+
+import org.openrewrite.*;
+
+import java.util.Optional;
+import java.util.stream.Stream;
+
+@Incubating(since = "8.30.0")
+public interface TraitMatcher<U extends Trait<?>> {
+
+    Optional<U> get(Cursor cursor);
+
+    Stream<U> higher(Cursor cursor);
+
+    Stream<U> lower(Cursor cursor);
+
+    default Stream<U> lower(SourceFile sourceFile) {
+        return lower(new Cursor(new Cursor(null, Cursor.ROOT_VALUE), sourceFile));
+    }
+
+    /**
+     * @param visitor Called for each match of the trait. The function is passed the trait.
+     * @param <P>     The type of context object passed to the visitor.
+     * @return A visitor that can be used to inspect or modify trees matching the trait.
+     */
+    default <P> TreeVisitor<? extends Tree, P> asVisitor(VisitFunction<U> visitor) {
+        return asVisitor((u, cursor, p) -> visitor.visit(u));
+    }
+
+    /**
+     * @param visitor Called for each match of the trait. The function is passed the trait
+     *                and a context object.
+     * @param <P>     The type of context object passed to the visitor.
+     * @return A visitor that can be used to inspect or modify trees matching the trait.
+     */
+    default <P> TreeVisitor<? extends Tree, P> asVisitor(VisitFunction2<U, P> visitor) {
+        return asVisitor((u, cursor, p) -> visitor.visit(u, p));
+    }
+
+    /**
+     * @param visitor Called for each match of the trait. The function is passed the trait, the cursor at which the
+     *                trait was found, and a context object.
+     * @param <P>     The type of context object passed to the visitor.
+     * @return A visitor that can be used to inspect or modify trees matching the trait.
+     */
+    <P> TreeVisitor<? extends Tree, P> asVisitor(VisitFunction3<U, P> visitor);
+}

rewrite-core/src/main/java/org/openrewrite/trait/VisitFunction.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright 2024 the original author or authors.
+ * <p>
+ * 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
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * 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 org.openrewrite.trait;
+
+import org.openrewrite.Tree;
+
+public interface VisitFunction<U extends Trait<?>> {
+    Tree visit(U data);
+}

rewrite-core/src/main/java/org/openrewrite/trait/VisitFunction2.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright 2024 the original author or authors.
+ * <p>
+ * 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
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * 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 org.openrewrite.trait;
+
+import org.openrewrite.Tree;
+
+public interface VisitFunction2<U extends Trait<?>, P> {
+    Tree visit(U data, P p);
+}

rewrite-core/src/main/java/org/openrewrite/trait/VisitFunction3.java

@@ -0,0 +1,23 @@
+/*
+ * Copyright 2024 the original author or authors.
+ * <p>
+ * 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
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * 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 org.openrewrite.trait;
+
+import org.openrewrite.Cursor;
+import org.openrewrite.Tree;
+
+public interface VisitFunction3<U extends Trait<?>, P> {
+    Tree visit(U data, Cursor cursor, P p);
+}

rewrite-core/src/main/java/org/openrewrite/trait/package-info.java

@@ -0,0 +1,19 @@
+/*
+ * Copyright 2020 the original author or authors.
+ * <p>
+ * 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
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * 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.
+ */
+@NonNullApi
+package org.openrewrite.trait;
+
+import org.openrewrite.internal.lang.NonNullApi;

rewrite-java/src/main/java/org/openrewrite/java/trait/MethodAccess.java

@@ -0,0 +1,113 @@
+/*
+ * Copyright 2024 the original author or authors.
+ * <p>
+ * 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
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * 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 org.openrewrite.java.trait;
+
+import lombok.RequiredArgsConstructor;
+import lombok.Value;
+import org.openrewrite.Cursor;
+import org.openrewrite.Incubating;
+import org.openrewrite.Tree;
+import org.openrewrite.TreeVisitor;
+import org.openrewrite.internal.lang.Nullable;
+import org.openrewrite.java.JavaVisitor;
+import org.openrewrite.java.MethodMatcher;
+import org.openrewrite.java.tree.Expression;
+import org.openrewrite.java.tree.J;
+import org.openrewrite.trait.SimpleTraitMatcher;
+import org.openrewrite.trait.Trait;
+import org.openrewrite.trait.VisitFunction3;
+
+@Incubating(since = "8.30.0")
+@Value
+public class MethodAccess implements Trait<Expression> {
+    Cursor cursor;
+
+    @RequiredArgsConstructor
+    public static class Matcher extends SimpleTraitMatcher<MethodAccess> {
+        private final MethodMatcher methodMatcher;
+
+        public Matcher(String methodPattern) {
+            this(new MethodMatcher(methodPattern));
+        }
+
+        @Override
+        public <P> TreeVisitor<? extends Tree, P> asVisitor(VisitFunction3<MethodAccess, P> visitor) {
+            return new JavaVisitor<P>() {
+                @Override
+                public J visitMethodInvocation(J.MethodInvocation method, P p) {
+                    MethodAccess methodAccess = test(getCursor());
+                    return methodAccess != null ?
+                            (J) visitor.visit(methodAccess, getCursor(), p) :
+                            super.visitMethodInvocation(method, p);
+                }
+
+                @Override
+                public J visitMethodDeclaration(J.MethodDeclaration method, P p) {
+                    MethodAccess methodAccess = test(getCursor());
+                    return methodAccess != null ?
+                            (J) visitor.visit(methodAccess, getCursor(), p) :
+                            super.visitMethodDeclaration(method, p);
+                }
+
+                @Override
+                public J visitNewClass(J.NewClass newClass, P p) {
+                    MethodAccess methodAccess = test(getCursor());
+                    return methodAccess != null ?
+                            (J) visitor.visit(methodAccess, getCursor(), p) :
+                            super.visitNewClass(newClass, p);
+                }
+
+                @Override
+                public J visitMemberReference(J.MemberReference memberRef, P p) {
+                    MethodAccess methodAccess = test(getCursor());
+                    return methodAccess != null ?
+                            (J) visitor.visit(methodAccess, getCursor(), p) :
+                            super.visitMemberReference(memberRef, p);
+                }
+            };
+        }
+
+        @Override
+        protected @Nullable MethodAccess test(Cursor cursor) {
+            Object value = cursor.getValue();
+            if (value instanceof J.MethodInvocation ||
+                value instanceof J.NewClass ||
+                value instanceof J.MemberReference) {
+                return methodMatcher.matches(((Expression) value)) ?
+                        new MethodAccess(cursor) :
+                        null;
+            }
+            if (value instanceof J.MethodDeclaration) {
+                J newClassOrClassDecl = cursor
+                        .dropParentUntil(t -> t instanceof J.ClassDeclaration ||
+                                              t instanceof J.NewClass)
+                        .getValue();
+                if (newClassOrClassDecl instanceof J.ClassDeclaration) {
+                    return methodMatcher.matches((J.MethodDeclaration) value,
+                            (J.ClassDeclaration) newClassOrClassDecl) ?
+                            new MethodAccess(cursor) :
+                            null;
+                } else if (newClassOrClassDecl instanceof J.NewClass) {
+                    return methodMatcher.matches((J.MethodDeclaration) value,
+                            (J.NewClass) newClassOrClassDecl) ?
+                            new MethodAccess(cursor) :
+                            null;
+                }
+            }
+            return null;
+        }
+    }
+}