Update ChangeType and ChangePackage to work with SourceFileWithReference

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

@@ -15,9 +15,7 @@
  */
 package org.openrewrite.trait;
 
-import org.openrewrite.Incubating;
-import org.openrewrite.SourceFile;
-import org.openrewrite.Tree;
+import org.openrewrite.*;
 
 import java.util.Set;
 
@@ -26,6 +24,14 @@ public interface TypeReference extends Trait<Tree> {
 
     String getName();
 
+    default boolean supportsRename() {
+        return false;
+    }
+
+    default TreeVisitor<Tree, ExecutionContext> renameTo(String name) {
+        throw new UnsupportedOperationException();
+    }
+
     default boolean matches(Matcher matcher) {
         return matcher.matchesName(getName());
     }

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,13 +23,13 @@
 import org.openrewrite.internal.ListUtils;
 import org.openrewrite.java.tree.*;
 import org.openrewrite.marker.SearchResult;
+import org.openrewrite.trait.TypeReference;
 
 import java.nio.file.Paths;
+import java.util.HashMap;
 import java.util.IdentityHashMap;
 import java.util.Map;
 
-import static java.util.Objects.requireNonNull;
-
 /**
  * A recipe that will rename a package name in package statements, imports, and fully-qualified types (see: NOTE).
  * <p>
@@ -83,11 +83,11 @@ 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(Tree tree, ExecutionContext ctx) {
                 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", "");
@@ -113,14 +113,52 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
                     }
                     stopAfterPreVisit();
                 }
+                if (tree instanceof SourceFileWithTypeReferences) {
+                    SourceFileWithTypeReferences cu = (SourceFileWithTypeReferences) tree;
+                    boolean recursive = Boolean.TRUE.equals(ChangePackage.this.recursive);
+                    String recursivePackageNamePrefix = oldPackageName + ".";
+                    for (TypeReference ref : cu.getTypeReferences().getTypeReferences()) {
+                        if (ref.getName().equals(oldPackageName) || recursive && ref.getName().startsWith(recursivePackageNamePrefix)) {
+                            return SearchResult.found(cu);
+                        }
+                    }
+                    stopAfterPreVisit();
+                }
                 return super.preVisit(tree, ctx);
             }
         };
 
-        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 SourceFileWithTypeReferences;
+            }
+
+            @Override
+            public @Nullable Tree preVisit(@Nullable Tree tree, ExecutionContext ctx) {
+                if (tree instanceof JavaSourceFile) {
+                    return new JavaChangePackageVisitor().visit(tree, ctx);
+                } else if (tree instanceof SourceFileWithTypeReferences) {
+                    SourceFileWithTypeReferences sourceFile = (SourceFileWithTypeReferences) tree;
+                    SourceFileWithTypeReferences.TypeReferences typeReferences = sourceFile.getTypeReferences();
+                    boolean recursive = Boolean.TRUE.equals(ChangePackage.this.recursive);
+                    String recursivePackageNamePrefix = oldPackageName + ".";
+                    Map<Tree, TypeReference> matches = new HashMap<>();
+                    for (TypeReference ref : typeReferences.getTypeReferences()) {
+                        if (ref.supportsRename()) {
+                            if (ref.getName().equals(oldPackageName) || recursive && ref.getName().startsWith(recursivePackageNamePrefix)) {
+                                matches.put(tree, ref);
+                            }
+                        }
+                    }
+                    return new TypeReferenceChangePackageVisitor(matches, oldPackageName, newPackageName).visit(tree, ctx, 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 +387,23 @@ private boolean isTargetRecursivePackageName(String packageName) {
         }
 
     }
+
+    @Value
+    @EqualsAndHashCode(callSuper = false)
+    private static class TypeReferenceChangePackageVisitor extends TreeVisitor<Tree, ExecutionContext> {
+        Map<Tree, TypeReference> matches;
+        String oldPackageName;
+        String newPackageName;
+
+        @Override
+        public @Nullable Tree visit(@Nullable Tree tree, ExecutionContext ctx) {
+            Tree tree1 = super.visit(tree, ctx);
+            TypeReference ref = matches.get(tree);
+            if (ref != null) {
+                return ref.renameTo(ref.getName().replace(oldPackageName, newPackageName)).visit(tree, ctx, getCursor().getParent());
+            }
+            return tree1;
+        }
+    }
+
 }

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

@@ -24,6 +24,7 @@
 import org.openrewrite.java.tree.*;
 import org.openrewrite.marker.Markers;
 import org.openrewrite.marker.SearchResult;
+import org.openrewrite.trait.TypeReference;
 
 import java.nio.file.Path;
 import java.nio.file.Paths;
@@ -81,22 +82,45 @@ public String getDescription() {
     public TreeVisitor<?, ExecutionContext> getVisitor() {
         TreeVisitor<?, ExecutionContext> condition = new TreeVisitor<Tree, ExecutionContext>() {
             @Override
-            public @Nullable Tree visit(@Nullable Tree tree, ExecutionContext ctx) {
+            public @Nullable Tree preVisit(@Nullable Tree tree, ExecutionContext ctx) {
                 if (tree instanceof JavaSourceFile) {
-                    JavaSourceFile cu = (JavaSourceFile) requireNonNull(tree);
+                    JavaSourceFile cu = (JavaSourceFile) tree;
                     if (!Boolean.TRUE.equals(ignoreDefinition) && containsClassDefinition(cu, oldFullyQualifiedTypeName)) {
                         return SearchResult.found(cu);
                     }
                     return new UsesType<>(oldFullyQualifiedTypeName, true).visitNonNull(cu, ctx);
                 }
+                if (tree instanceof SourceFileWithTypeReferences) {
+                    SourceFileWithTypeReferences cu = (SourceFileWithTypeReferences) tree;
+                    return new UsesType<>(oldFullyQualifiedTypeName, true).visitNonNull(cu, ctx);
+                }
                 return tree;
             }
         };
 
-        return Preconditions.check(condition, new ChangeTypeVisitor(oldFullyQualifiedTypeName, newFullyQualifiedTypeName, ignoreDefinition));
+        return Preconditions.check(condition, new TreeVisitor<Tree, ExecutionContext>() {
+            @Override
+            public boolean isAcceptable(SourceFile sourceFile, ExecutionContext ctx) {
+                return sourceFile instanceof JavaSourceFile || sourceFile instanceof SourceFileWithTypeReferences;
+            }
+
+            @Override
+            public @Nullable Tree preVisit(@Nullable Tree tree, ExecutionContext ctx) {
+                if (tree instanceof JavaSourceFile) {
+                    return new JavaChangeTypeVisitor(oldFullyQualifiedTypeName, newFullyQualifiedTypeName, ignoreDefinition).visit(tree, ctx);
+                } else if (tree instanceof SourceFileWithTypeReferences) {
+                    SourceFileWithTypeReferences sourceFile = (SourceFileWithTypeReferences) tree;
+                    SourceFileWithTypeReferences.TypeReferences typeReferences = sourceFile.getTypeReferences();
+                    TypeMatcher matcher = new TypeMatcher(oldFullyQualifiedTypeName);
+                    Set<TypeReference> matches = new HashSet<>(typeReferences.findMatches(matcher));
+                    return new TypeReferenceChangeTypeVisitor(matches, newFullyQualifiedTypeName).visit(tree, ctx, getCursor().getParent());
+                }
+                return tree;
+            }
+        });
     }
 
-    private static class ChangeTypeVisitor extends JavaVisitor<ExecutionContext> {
+    private static class JavaChangeTypeVisitor extends JavaVisitor<ExecutionContext> {
         private final JavaType.Class originalType;
         private final JavaType targetType;
 
@@ -109,7 +133,7 @@ private static class ChangeTypeVisitor extends JavaVisitor<ExecutionContext> {
         private final Map<JavaType, JavaType> oldNameToChangedType = new IdentityHashMap<>();
         private final Set<String> topLevelClassnames = new HashSet<>();
 
-        private ChangeTypeVisitor(String oldFullyQualifiedTypeName, String newFullyQualifiedTypeName, @Nullable Boolean ignoreDefinition) {
+        private JavaChangeTypeVisitor(String oldFullyQualifiedTypeName, String newFullyQualifiedTypeName, @Nullable Boolean ignoreDefinition) {
             this.originalType = JavaType.ShallowClass.build(oldFullyQualifiedTypeName);
             this.targetType = JavaType.buildType(newFullyQualifiedTypeName);
             this.ignoreDefinition = ignoreDefinition;
@@ -527,6 +551,24 @@ private boolean hasNoConflictingImport(@Nullable JavaSourceFile sf) {
         }
     }
 
+    @Value
+    @EqualsAndHashCode(callSuper = false)
+    private static class TypeReferenceChangeTypeVisitor extends TreeVisitor<Tree, ExecutionContext> {
+        Set<TypeReference> matches;
+        String newFullyQualifiedName;
+
+        @Override
+        public @Nullable Tree preVisit(@Nullable Tree tree, ExecutionContext ctx) {
+            Tree tree1 = super.visit(tree, ctx);
+            for (TypeReference ref : matches) {
+                if (ref.getTree().equals(tree) && ref.supportsRename()) {
+                    return ref.renameTo(newFullyQualifiedName).visit(tree, ctx, getCursor().getParent());
+                }
+            }
+            return tree1;
+        }
+    }
+
     private static class ChangeClassDefinition extends JavaIsoVisitor<ExecutionContext> {
         private final JavaType.Class originalType;
         private final JavaType.Class targetType;
@@ -579,7 +621,7 @@ private boolean updatePath(JavaSourceFile sf, String oldPath, String newPath) {
         }
 
         @Override
-        public  J.@Nullable Package visitPackage(J.Package pkg, ExecutionContext ctx) {
+        public J.@Nullable Package visitPackage(J.Package pkg, ExecutionContext ctx) {
             Boolean updatePackage = getCursor().pollNearestMessage("UPDATE_PACKAGE");
             if (updatePackage != null && updatePackage) {
                 String original = pkg.getExpression().printTrimmed(getCursor()).replaceAll("\\s", "");

rewrite-xml/src/main/java/org/openrewrite/xml/trait/SpringTypeReference.java

@@ -17,9 +17,7 @@
 
 import lombok.Value;
 import org.jspecify.annotations.Nullable;
-import org.openrewrite.Cursor;
-import org.openrewrite.SourceFile;
-import org.openrewrite.Tree;
+import org.openrewrite.*;
 import org.openrewrite.trait.SimpleTraitMatcher;
 import org.openrewrite.trait.TypeReference;
 import org.openrewrite.xml.XPathMatcher;
@@ -52,6 +50,27 @@ public String getName() {
         throw new IllegalArgumentException("getTree() must be an Xml.Attribute or Xml.Tag: " + getTree().getClass());
     }
 
+    @Override
+    public boolean supportsRename() {
+        return true;
+    }
+
+    @Override
+    public TreeVisitor<Tree, ExecutionContext> renameTo(String name) {
+        return new TreeVisitor<Tree, ExecutionContext>() {
+            @Override
+            public @Nullable Tree visit(@Nullable Tree tree, ExecutionContext ctx) {
+                if (tree instanceof Xml.Attribute) {
+                    return ((Xml.Attribute) tree).withValue(((Xml.Attribute) tree).getValue().withValue(name));
+                }
+                if (tree instanceof Xml.Tag) {
+                    return ((Xml.Tag) tree).withValue(name);
+                }
+                return super.visit(tree, ctx);
+            }
+        };
+    }
+
     static class Matcher extends SimpleTraitMatcher<SpringTypeReference> {
         private final Pattern typeReference = Pattern.compile("(?:[a-zA-Z_][a-zA-Z0-9_]*\\.)+[A-Z*][a-zA-Z0-9_]*(?:<[a-zA-Z0-9_,?<> ]*>)?");
         private final XPathMatcher classXPath = new XPathMatcher("//@class");