From efd3214f6f3317a8a59d64821f30211ac6d9993e Mon Sep 17 00:00:00 2001
From: Knut Wannheden <knut@moderne.io>
Date: Thu, 16 Nov 2023 08:13:07 +0100
Subject: [PATCH] Add `ImportService#shortenFullyQualifiedTypeReferencesIn()`
 (#3687)

* Add `ImportService#shortenFullyQualifiedTypeReferencesIn()`

Extends `ImportService` with `shortenFullyQualifiedTypeReferencesIn()` so that other languages extending `J` can provide a customized implementation.

* Add some Javadoc
---
 ...ortenFullyQualifiedTypeReferencesTest.java |  5 ++--
 .../ShortenFullyQualifiedTypeReferences.java  | 27 ++++++++++++++++---
 .../java/service/ImportService.java           | 11 ++++++++
 3 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/rewrite-java-test/src/test/java/org/openrewrite/java/ShortenFullyQualifiedTypeReferencesTest.java b/rewrite-java-test/src/test/java/org/openrewrite/java/ShortenFullyQualifiedTypeReferencesTest.java
index 0c58a549d59..e7cbf2cff2e 100644
--- a/rewrite-java-test/src/test/java/org/openrewrite/java/ShortenFullyQualifiedTypeReferencesTest.java
+++ b/rewrite-java-test/src/test/java/org/openrewrite/java/ShortenFullyQualifiedTypeReferencesTest.java
@@ -18,6 +18,7 @@
 import org.junit.jupiter.api.Test;
 import org.openrewrite.DocumentExample;
 import org.openrewrite.ExecutionContext;
+import org.openrewrite.java.service.ImportService;
 import org.openrewrite.java.tree.J;
 import org.openrewrite.java.tree.JavaType;
 import org.openrewrite.test.RecipeSpec;
@@ -228,7 +229,7 @@ void visitSubtreeOnly() {
               @SuppressWarnings("DataFlowIssue")
               public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, ExecutionContext ctx) {
                   if (method.getSimpleName().equals("m1")) {
-                      return (J.MethodDeclaration) new ShortenFullyQualifiedTypeReferences().getVisitor().visit(method, ctx, getCursor().getParent());
+                      return (J.MethodDeclaration) ShortenFullyQualifiedTypeReferences.modifyOnly(method).visit(method, ctx, getCursor().getParent());
                   }
                   return super.visitMethodDeclaration(method, ctx);
               }
@@ -429,7 +430,7 @@ void subtree() {
               @Override
               public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, ExecutionContext ctx) {
                   if (method.getSimpleName().equals("m1")) {
-                      doAfterVisit(ShortenFullyQualifiedTypeReferences.modifyOnly(method));
+                      doAfterVisit(service(ImportService.class).shortenFullyQualifiedTypeReferencesIn(method));
                   }
                   return super.visitMethodDeclaration(method, ctx);
               }
diff --git a/rewrite-java/src/main/java/org/openrewrite/java/ShortenFullyQualifiedTypeReferences.java b/rewrite-java/src/main/java/org/openrewrite/java/ShortenFullyQualifiedTypeReferences.java
index c2ba78b55ca..f6361e80027 100644
--- a/rewrite-java/src/main/java/org/openrewrite/java/ShortenFullyQualifiedTypeReferences.java
+++ b/rewrite-java/src/main/java/org/openrewrite/java/ShortenFullyQualifiedTypeReferences.java
@@ -19,6 +19,7 @@
 import org.openrewrite.internal.lang.NonNull;
 import org.openrewrite.internal.lang.Nullable;
 import org.openrewrite.java.internal.DefaultJavaTypeSignatureBuilder;
+import org.openrewrite.java.service.ImportService;
 import org.openrewrite.java.tree.*;
 
 import java.time.Duration;
@@ -47,11 +48,31 @@ public String getDescription() {
     }
 
     @Override
-    public TreeVisitor<?, ExecutionContext> getVisitor() {
-        return getVisitor(null);
+    public JavaVisitor<ExecutionContext> getVisitor() {
+        // This wrapper is necessary so that the "correct" implementation is used when this recipe is used declaratively
+        return new JavaVisitor<ExecutionContext>() {
+            @Override
+            public @Nullable J visit(@Nullable Tree tree, ExecutionContext ctx) {
+                if (tree instanceof JavaSourceFile) {
+                    return ((JavaSourceFile) tree).service(ImportService.class).shortenFullyQualifiedTypeReferencesIn((J) tree).visit(tree, ctx);
+                }
+                return (J) tree;
+            }
+        };
     }
 
-    public static <J2 extends J> TreeVisitor<J, ExecutionContext> modifyOnly(J2 subtree) {
+    /**
+     * Returns a visitor which replaces all fully qualified references in the given subtree with simple names and adds 
+     * corresponding import statements.
+     * <p>
+     * For compatibility with other Java-based languages it is recommended to use this as a service via
+     * {@link ImportService#shortenFullyQualifiedTypeReferencesIn(J)}, as that will dispatch to the correct
+     * implementation for the language.
+     * 
+     * @see ImportService#shortenFullyQualifiedTypeReferencesIn(J)
+     * @see JavaVisitor#service(Class) 
+     */
+    public static <J2 extends J> JavaVisitor<ExecutionContext> modifyOnly(J2 subtree) {
         return getVisitor(subtree);
     }
 
diff --git a/rewrite-java/src/main/java/org/openrewrite/java/service/ImportService.java b/rewrite-java/src/main/java/org/openrewrite/java/service/ImportService.java
index c2a8a2d42ac..7ab5bd5e2ff 100644
--- a/rewrite-java/src/main/java/org/openrewrite/java/service/ImportService.java
+++ b/rewrite-java/src/main/java/org/openrewrite/java/service/ImportService.java
@@ -15,10 +15,13 @@
  */
 package org.openrewrite.java.service;
 
+import org.openrewrite.ExecutionContext;
 import org.openrewrite.Incubating;
 import org.openrewrite.internal.lang.Nullable;
 import org.openrewrite.java.AddImport;
 import org.openrewrite.java.JavaVisitor;
+import org.openrewrite.java.ShortenFullyQualifiedTypeReferences;
+import org.openrewrite.java.tree.J;
 
 @Incubating(since = "8.2.0")
 public class ImportService {
@@ -30,4 +33,12 @@ public <P> JavaVisitor<P> addImportVisitor(@Nullable String packageName,
                                                boolean onlyIfReferenced) {
         return new AddImport<>(packageName, typeName, member,  alias, onlyIfReferenced);
     }
+
+    public <J2 extends J> JavaVisitor<ExecutionContext> shortenAllFullyQualifiedTypeReferences() {
+        return new ShortenFullyQualifiedTypeReferences().getVisitor();
+    }
+
+    public <J2 extends J> JavaVisitor<ExecutionContext> shortenFullyQualifiedTypeReferencesIn(J2 subtree) {
+        return ShortenFullyQualifiedTypeReferences.modifyOnly(subtree);
+    }
 }