Update ChangeType and ChangePackage to work with SourceFileWithReference

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

@@ -18,40 +18,40 @@
 import lombok.AccessLevel;
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
-import org.openrewrite.trait.TypeReference;
+import org.openrewrite.trait.reference.Reference;
 
 import java.util.*;
 
 @Incubating(since = "8.39.0")
-public interface SourceFileWithTypeReferences extends SourceFile {
+public interface SourceFileWithReferences extends SourceFile {
 
     TypeReferences getTypeReferences();
 
     @RequiredArgsConstructor(access = AccessLevel.PRIVATE)
     @Getter
     class TypeReferences {
         private final SourceFile sourceFile;
-        private final Set<TypeReference> typeReferences;
+        private final Set<Reference> references;
 
-        public Collection<TypeReference> findMatches(TypeReference.Matcher matcher) {
-            List<TypeReference> list = new ArrayList<>();
-            for (TypeReference ref : typeReferences) {
-                if (ref.matches(matcher)) {
+        public Collection<Reference> findMatches(Reference.MatcherMutator matcherMutator) {
+            List<Reference> list = new ArrayList<>();
+            for (Reference ref : references) {
+                if (ref.matches(matcherMutator)) {
                     list.add(ref);
                 }
             }
             return list;
         }
 
         public static TypeReferences build(SourceFile sourceFile) {
-            Set<TypeReference> typeReferences = new HashSet<>();
-            ServiceLoader<TypeReference.Provider> loader = ServiceLoader.load(TypeReference.Provider.class);
+            Set<Reference> references = new HashSet<>();
+            ServiceLoader<Reference.Provider> loader = ServiceLoader.load(Reference.Provider.class);
             loader.forEach(provider -> {
                 if (provider.isAcceptable(sourceFile)) {
-                    typeReferences.addAll(provider.getTypeReferences(sourceFile));
+                    references.addAll(provider.getTypeReferences(sourceFile));
                 }
             });
-            return new TypeReferences(sourceFile, typeReferences);
+            return new TypeReferences(sourceFile, references);
         }
     }
 }

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

@@ -15,13 +15,13 @@
  */
 package org.openrewrite;
 
-import org.openrewrite.trait.TypeReference;
+import org.openrewrite.trait.reference.Reference;
 
 import java.util.Set;
 
 public interface TypeReferenceProvider {
 
-    Set<TypeReference> getTypeReferences(SourceFile sourceFile);
+    Set<Reference> getTypeReferences(SourceFile sourceFile);
 
     boolean isAcceptable(SourceFile sourceFile);
 }

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

@@ -0,0 +1,70 @@
+/*
+ * 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.reference;
+
+import org.openrewrite.*;
+import org.openrewrite.trait.Trait;
+
+import java.util.Set;
+
+@Incubating(since = "8.39.0")
+public interface Reference extends Trait<Tree> {
+
+    enum Kind {
+        TYPE,
+        PACKAGE
+    }
+
+    Kind getKind();
+
+    String getValue();
+
+    default boolean supportsRename() {
+        return false;
+    }
+
+    default boolean matches(MatcherMutator matcher) {
+        return matcher.matchesReference(this);
+    }
+
+    default TreeVisitor<Tree, ExecutionContext>  rename(MatcherMutator renamer, String replacement) {
+        return renamer.rename(replacement);
+    }
+
+    interface Provider {
+
+        Set<Reference> getTypeReferences(SourceFile sourceFile);
+
+        boolean isAcceptable(SourceFile sourceFile);
+    }
+
+    interface MatcherMutator extends Matcher, Renamer {
+    }
+
+    interface Matcher {
+        boolean matchesReference(Reference value);
+    }
+
+    interface Renamer {
+        default TreeVisitor<Tree, ExecutionContext> rename(String replacement) {
+            return rename(replacement, false);
+        }
+
+        default TreeVisitor<Tree, ExecutionContext> rename(String replacement, boolean recursive) {
+            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,10 +23,10 @@
 import org.openrewrite.internal.ListUtils;
 import org.openrewrite.java.tree.*;
 import org.openrewrite.marker.SearchResult;
+import org.openrewrite.trait.reference.Reference;
 
 import java.nio.file.Paths;
-import java.util.IdentityHashMap;
-import java.util.Map;
+import java.util.*;
 
 import static java.util.Objects.requireNonNull;
 
@@ -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,45 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
                     }
                     stopAfterPreVisit();
                 }
+                if (tree instanceof SourceFileWithReferences) {
+                    SourceFileWithReferences cu = (SourceFileWithReferences) tree;
+                    boolean recursive = Boolean.TRUE.equals(ChangePackage.this.recursive);
+                    String recursivePackageNamePrefix = oldPackageName + ".";
+                    for (Reference ref : cu.getTypeReferences().getReferences()) {
+                        if (ref.getValue().equals(oldPackageName) || recursive && ref.getValue().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 SourceFileWithReferences;
+            }
+
+            @Override
+            public @Nullable Tree preVisit(@Nullable Tree tree, ExecutionContext ctx) {
+                if (tree instanceof JavaSourceFile) {
+                    return new JavaChangePackageVisitor().visit(tree, ctx, requireNonNull(getCursor().getParent()));
+                } else if (tree instanceof SourceFileWithReferences) {
+                    SourceFileWithReferences sourceFile = (SourceFileWithReferences) tree;
+                    SourceFileWithReferences.TypeReferences typeReferences = sourceFile.getTypeReferences();
+                    boolean recursive = Boolean.TRUE.equals(ChangePackage.this.recursive);
+                    Reference.MatcherMutator matcherMutator = new PackageMatcher(oldPackageName, recursive);
+                    Set<Reference> matches = new HashSet<>(typeReferences.findMatches(matcherMutator));
+                    return new ReferenceChangePackageVisitor(matches, matcherMutator, 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 +380,24 @@ private boolean isTargetRecursivePackageName(String packageName) {
         }
 
     }
+
+    @Value
+    @EqualsAndHashCode(callSuper = false)
+    private static class ReferenceChangePackageVisitor extends TreeVisitor<Tree, ExecutionContext> {
+        Set<Reference> matches;
+        Reference.MatcherMutator matcherMutator;
+        String newPackageName;
+
+        @Override
+        public @Nullable Tree visit(@Nullable Tree tree, ExecutionContext ctx) {
+            Tree tree1 = super.visit(tree, ctx);
+            for (Reference ref : matches) {
+                if (ref.getTree().equals(tree) && ref.supportsRename()) {
+                    return ref.rename(matcherMutator, newPackageName).visit(tree, ctx, getCursor().getParent());
+                }
+            }
+            return tree1;
+        }
+    }
+
 }