Trait interface and VariableAccess/MethodAccess implementation (#…

…4309)

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, sb) -> {
+            sb.add(test(va.getCursor()));
+            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(VisitFunction2<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, 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,89 @@
+/*
+ * 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;
+
+/**
+ * A trait matcher builds {@link Trait} instances when they match the criteria for that trait.
+ * Required constructor arguments in implementing classes represent the minimum information
+ * needed to meaningfully search for this trait. Optional further filtering criteria should
+ * be expressed as setter methods in the builder style (lacking a 'set' prefix and returning
+ * the matcher instance).
+ *
+ * @param <U> The type of {@link Trait} that this matcher builds.
+ */
+@Incubating(since = "8.30.0")
+public interface TraitMatcher<U extends Trait<?>> {
+
+    /**
+     * Tests whether a tree at the cursor matches the trait, and if so, returns
+     * a trait instance containing the semantic information represented by the tree.
+     *
+     * @param cursor The starting point of the search.
+     * @return Optionally a trait instance if the tree pointed at by the cursor
+     * matches the trait criteria.
+     */
+    Optional<U> get(Cursor cursor);
+
+    /**
+     * Looks up the cursor stack (ancestors) for trees that match the trait.
+     *
+     * @param cursor The starting point of the search.
+     * @return A stream of trees that are ancestors of the cursor that match the trait.
+     */
+    Stream<U> higher(Cursor cursor);
+
+    /**
+     * Looks down the syntax tree starting at the cursor stack (descendents)
+     * for trees that match the trait.
+     *
+     * @param cursor The starting point of the search.
+     * @return A stream of trees that are descendents of the cursor that match the trait.
+     */
+    Stream<U> lower(Cursor cursor);
+
+    /**
+     * Looks down the syntax tree starting at the source file (root LST)
+     * for trees that match the trait.
+     *
+     * @param sourceFile A whole source file.
+     * @return A stream of all trees in a source file that match the trait.
+     */
+    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, 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.
+     */
+    <P> TreeVisitor<? extends Tree, P> asVisitor(VisitFunction2<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/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,97 @@
+/*
+ * 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.java.tree.JavaType;
+import org.openrewrite.java.tree.MethodCall;
+import org.openrewrite.trait.SimpleTraitMatcher;
+import org.openrewrite.trait.Trait;
+import org.openrewrite.trait.VisitFunction2;
+
+import java.util.function.Predicate;
+
+@Incubating(since = "8.30.0")
+@Value
+public class MethodAccess implements Trait<MethodCall> {
+    Cursor cursor;
+
+    @RequiredArgsConstructor
+    public static class Matcher extends SimpleTraitMatcher<MethodAccess> {
+        private final MethodMatcher methodMatcher;
+        private Predicate<@Nullable JavaType> returnsTest = m -> true;
+
+        public Matcher(String methodPattern) {
+            this(new MethodMatcher(methodPattern));
+        }
+
+        public Matcher returns(Predicate<@Nullable JavaType> returnsTest) {
+            this.returnsTest = returnsTest;
+            return this;
+        }
+
+        @Override
+        public <P> TreeVisitor<? extends Tree, P> asVisitor(VisitFunction2<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, p) :
+                            super.visitMethodInvocation(method, p);
+                }
+
+                @Override
+                public J visitNewClass(J.NewClass newClass, P p) {
+                    MethodAccess methodAccess = test(getCursor());
+                    return methodAccess != null ?
+                            (J) visitor.visit(methodAccess, 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, p) :
+                            super.visitMemberReference(memberRef, p);
+                }
+            };
+        }
+
+        @Override
+        protected @Nullable MethodAccess test(Cursor cursor) {
+            Object value = cursor.getValue();
+            JavaType.Method methodType = ((MethodCall) value).getMethodType();
+            JavaType returnType = methodType == null ? null : methodType.getReturnType();
+
+            return methodMatcher.matches(((Expression) value)) &&
+                   returnsTest.test(returnType) ?
+                    new MethodAccess(cursor) :
+                    null;
+        }
+    }
+}