Update ChangeType and ChangePackage to work with SourceFileWithReference

rewrite-core/src/main/java/org/openrewrite/SourceFileWithReferences.java

@@ -0,0 +1,66 @@
+/*
+ * 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;
+
+import lombok.AccessLevel;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import org.jspecify.annotations.Nullable;
+import org.openrewrite.trait.Reference;
+
+import java.util.*;
+
+@Incubating(since = "8.39.0")
+public interface SourceFileWithReferences extends SourceFile {
+
+    References getReferences();
+
+    @RequiredArgsConstructor(access = AccessLevel.PRIVATE)
+    @Getter
+    class References {
+        private final SourceFile sourceFile;
+        private final Set<Reference> references;
+
+        public Collection<Reference> findMatches(Reference.Matcher matcher) {
+            return findMatchesInternal(matcher, null);
+        }
+
+        public Collection<Reference> findMatches(Reference.Matcher matcher, Reference.Kind kind) {
+            return findMatchesInternal(matcher, kind);
+        }
+
+        private List<Reference> findMatchesInternal(Reference.Matcher matcher, Reference.@Nullable Kind kind) {
+            List<Reference> list = new ArrayList<>();
+            for (Reference ref : references) {
+                if ((kind == null || ref.getKind().equals(kind)) && ref.matches(matcher) ) {
+                    list.add(ref);
+                }
+            }
+            return list;
+        }
+
+        public static References build(SourceFile sourceFile) {
+            Set<Reference> references = new HashSet<>();
+            ServiceLoader<Reference.Provider> loader = ServiceLoader.load(Reference.Provider.class);
+            loader.forEach(provider -> {
+                if (provider.isAcceptable(sourceFile)) {
+                    references.addAll(provider.getReferences(sourceFile));
+                }
+            });
+            return new References(sourceFile, references);
+        }
+    }
+}

rewrite-core/src/main/java/org/openrewrite/trait/TypeReference.java → rewrite-core/src/main/java/org/openrewrite/trait/Reference.java

@@ -15,29 +15,48 @@
  */
 package org.openrewrite.trait;
 
-import org.openrewrite.Incubating;
-import org.openrewrite.SourceFile;
-import org.openrewrite.Tree;
+import org.openrewrite.*;
 
 import java.util.Set;
 
 @Incubating(since = "8.39.0")
-public interface TypeReference extends Trait<Tree> {
+public interface Reference extends Trait<Tree> {
 
-    String getName();
+    enum Kind {
+        TYPE,
+        PACKAGE
+    }
+
+    Kind getKind();
+
+    String getValue();
+
+    default boolean supportsRename() {
+        return false;
+    }
 
     default boolean matches(Matcher matcher) {
-        return matcher.matchesName(getName());
+        return matcher.matchesReference(this);
+    }
+
+    default TreeVisitor<Tree, ExecutionContext> rename(Renamer renamer, String replacement) {
+        return renamer.rename(replacement);
     }
 
     interface Provider {
 
-        Set<TypeReference> getTypeReferences(SourceFile sourceFile);
+        Set<Reference> getReferences(SourceFile sourceFile);
 
         boolean isAcceptable(SourceFile sourceFile);
     }
 
     interface Matcher {
-        boolean matchesName(String name);
+        boolean matchesReference(Reference value);
+    }
+
+    interface Renamer {
+        default TreeVisitor<Tree, ExecutionContext> rename(String replacement) {
+            throw new UnsupportedOperationException();
+        }
     }
 }

rewrite-java-test/src/test/java/org/openrewrite/java/ChangePackageTest.java

@@ -30,6 +30,7 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.openrewrite.java.Assertions.java;
+import static org.openrewrite.xml.Assertions.xml;
 
 @SuppressWarnings("ConstantConditions")
 class ChangePackageTest implements RewriteTest {
@@ -1702,4 +1703,26 @@ public enum MyEnum {
           )
         );
     }
+
+    @Test
+    void changePackageInSpringXml() {
+        rewriteRun(
+          spec -> spec.recipe(new ChangePackage("test.type", "test.test.type", true)),
+          xml(
+            """
+              <?xml version="1.0" encoding="UTF-8"?>
+              <beans xsi:schemaLocation="www.springframework.org/schema/beans">
+                <bean id="abc" class="test.type.A"/>
+              </beans>
+              """,
+            """
+              <?xml version="1.0" encoding="UTF-8"?>
+              <beans xsi:schemaLocation="www.springframework.org/schema/beans">
+                <bean id="abc" class="test.test.type.A"/>
+              </beans>
+              """
+          )
+        );
+
+    }
 }

rewrite-java-test/src/test/java/org/openrewrite/java/ChangeTypeTest.java

@@ -29,6 +29,7 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.openrewrite.java.Assertions.java;
+import static org.openrewrite.xml.Assertions.xml;
 
 @SuppressWarnings("ConstantConditions")
 class ChangeTypeTest implements RewriteTest {
@@ -2019,4 +2020,26 @@ class Letters {
         );
     }
 
+    @Test
+    void changeTypeInSpringXml() {
+        rewriteRun(
+          spec -> spec.recipe(new ChangeType("test.type.A", "test.type.B", true)),
+          xml(
+            """
+              <?xml version="1.0" encoding="UTF-8"?>
+              <beans xsi:schemaLocation="www.springframework.org/schema/beans">
+                <bean id="abc" class="test.type.A"/>
+              </beans>
+              """,
+            """
+              <?xml version="1.0" encoding="UTF-8"?>
+              <beans xsi:schemaLocation="www.springframework.org/schema/beans">
+                <bean id="abc" class="test.type.B"/>
+              </beans>
+              """
+          )
+        );
+
+    }
+
 }

rewrite-java/src/main/java/org/openrewrite/java/ChangePackage.java

@@ -23,8 +23,10 @@
 import org.openrewrite.internal.ListUtils;
 import org.openrewrite.java.tree.*;
 import org.openrewrite.marker.SearchResult;
+import org.openrewrite.trait.Reference;
 
 import java.nio.file.Paths;
+import java.util.HashMap;
 import java.util.IdentityHashMap;
 import java.util.Map;
 
@@ -83,11 +85,12 @@ public Validated<Object> validate() {
 
     @Override
     public TreeVisitor<?, ExecutionContext> getVisitor() {
-        JavaIsoVisitor<ExecutionContext> condition = new JavaIsoVisitor<ExecutionContext>() {
+        TreeVisitor<?, ExecutionContext> condition = new TreeVisitor<Tree, ExecutionContext>() {
             @Override
-            public @Nullable J preVisit(J tree, ExecutionContext ctx) {
+            public @Nullable Tree preVisit(@Nullable Tree tree, ExecutionContext ctx) {
+                stopAfterPreVisit();
                 if (tree instanceof JavaSourceFile) {
-                    JavaSourceFile cu = (JavaSourceFile) requireNonNull(tree);
+                    JavaSourceFile cu = (JavaSourceFile) tree;
                     if (cu.getPackageDeclaration() != null) {
                         String original = cu.getPackageDeclaration().getExpression()
                                 .printTrimmed(getCursor()).replaceAll("\\s", "");
@@ -111,16 +114,48 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
                             }
                         }
                     }
-                    stopAfterPreVisit();
+                } else if (tree instanceof SourceFileWithReferences) {
+                    SourceFileWithReferences cu = (SourceFileWithReferences) tree;
+                    boolean recursive = Boolean.TRUE.equals(ChangePackage.this.recursive);
+                    String recursivePackageNamePrefix = oldPackageName + ".";
+                    for (Reference ref : cu.getReferences().getReferences()) {
+                        if (ref.getValue().equals(oldPackageName) || recursive && ref.getValue().startsWith(recursivePackageNamePrefix)) {
+                            return SearchResult.found(cu);
+                        }
+                    }
                 }
-                return super.preVisit(tree, ctx);
+                return tree;
             }
         };
 
-        return Preconditions.check(condition, new ChangePackageVisitor());
+        return Preconditions.check(condition, new TreeVisitor<Tree, ExecutionContext>() {
+            @Override
+            public boolean isAcceptable(SourceFile sourceFile, ExecutionContext ctx) {
+                return sourceFile instanceof JavaSourceFile || sourceFile instanceof SourceFileWithReferences;
+            }
+
+            @Override
+            public @Nullable Tree preVisit(@Nullable Tree tree, ExecutionContext ctx) {
+                stopAfterPreVisit();
+                if (tree instanceof JavaSourceFile) {
+                    return new JavaChangePackageVisitor().visit(tree, ctx, requireNonNull(getCursor().getParent()));
+                } else if (tree instanceof SourceFileWithReferences) {
+                    SourceFileWithReferences sourceFile = (SourceFileWithReferences) tree;
+                    SourceFileWithReferences.References references = sourceFile.getReferences();
+                    boolean recursive = Boolean.TRUE.equals(ChangePackage.this.recursive);
+                    PackageMatcher matcher = new PackageMatcher(oldPackageName, recursive);
+                    Map<Tree, Reference> matches = new HashMap<>();
+                    for (Reference ref : references.findMatches(matcher)) {
+                        matches.put(ref.getTree(), ref);
+                    }
+                    return new ReferenceChangePackageVisitor(matches, matcher, newPackageName).visit(tree, ctx, requireNonNull(getCursor().getParent()));
+                }
+                return tree;
+            }
+        });
     }
 
-    private class ChangePackageVisitor extends JavaVisitor<ExecutionContext> {
+    private class JavaChangePackageVisitor extends JavaVisitor<ExecutionContext> {
         private static final String RENAME_TO_KEY = "renameTo";
         private static final String RENAME_FROM_KEY = "renameFrom";
 
@@ -349,4 +384,22 @@ private boolean isTargetRecursivePackageName(String packageName) {
         }
 
     }
+
+    @Value
+    @EqualsAndHashCode(callSuper = false)
+    private static class ReferenceChangePackageVisitor extends TreeVisitor<Tree, ExecutionContext> {
+        Map<Tree, Reference> matches;
+        Reference.Renamer renamer;
+        String newPackageName;
+
+        @Override
+        public @Nullable Tree preVisit(@Nullable Tree tree, ExecutionContext ctx) {
+            Reference reference = matches.get(tree);
+            if (reference != null && reference.supportsRename()) {
+                return reference.rename(renamer, newPackageName).visit(tree, ctx, getCursor().getParent());
+            }
+            return tree;
+        }
+    }
+
 }