Support using response type keeper with KSP

CHANGELOG.md

@@ -7,6 +7,7 @@
 
  - Add explicit keep rules for RxJava `Result` types to prevent their generic information from being removed.
  - Add `allowoptimization` flags for most kept types.
+ - Support using response type keeper with KSP.
 
 **Changed**
 

gradle/libs.versions.toml

@@ -18,6 +18,7 @@ okhttp = "5.2.1"
 protobuf = "3.25.8"
 robovm = "2.3.14"
 kotlinx-serialization = "1.9.0"
+kct = "0.10.0"
 autoService = "1.1.1"
 incap = "1.0.0"
 jackson = "2.20.0"
@@ -33,6 +34,8 @@ kotlin-stdLib = { module = "org.jetbrains.kotlin:kotlin-stdlib", version.ref = "
 kotlin-gradlePlugin = { module = "org.jetbrains.kotlin:kotlin-gradle-plugin", version.ref = "kotlin" }
 kotlin-serializationPlugin = { module = "org.jetbrains.kotlin:kotlin-serialization", version.ref = "kotlin" }
 
+ksp-api = "com.google.devtools.ksp:symbol-processing-api:2.3.0"
+
 errorpronePlugin = "net.ltgt.gradle:gradle-errorprone-plugin:4.3.0"
 errorproneCore = { module = "com.google.errorprone:error_prone_core", version = "2.10.0" }
 errorproneJavac = { module = "com.google.errorprone:javac", version = "9+181-r4173-1" }
@@ -84,3 +87,5 @@ googleJavaFormat = "com.google.googlejavaformat:google-java-format:1.28.0"
 ktlint = "com.pinterest.ktlint:ktlint-cli:1.7.1"
 compileTesting = "com.google.testing.compile:compile-testing:0.23.0"
 testParameterInjector = "com.google.testparameterinjector:test-parameter-injector:1.19"
+kct-core = { module = "dev.zacsweers.kctfork:core", version.ref = "kct" }
+kct-ksp = { module = "dev.zacsweers.kctfork:ksp", version.ref = "kct" }

retrofit-response-type-keeper/build.gradle

@@ -5,11 +5,16 @@ apply plugin: 'com.vanniktech.maven.publish'
 dependencies {
   compileOnly libs.autoService.annotations
   compileOnly libs.incap.runtime
+  compileOnly libs.ksp.api
   kapt libs.autoService.compiler
   kapt libs.incap.processor
 
   testImplementation libs.junit
   testImplementation libs.compileTesting
+  testImplementation libs.kct.core
+  testImplementation libs.kct.ksp
+  testImplementation libs.ksp.api
   testImplementation libs.truth
+  testImplementation libs.testParameterInjector
   testImplementation projects.retrofit
 }

retrofit-response-type-keeper/src/main/kotlin/retrofit2/keeper/RetrofitResponseTypeKeepProcessor.kt

@@ -33,16 +33,7 @@ import net.ltgt.gradle.incap.IncrementalAnnotationProcessorType.ISOLATING
 @IncrementalAnnotationProcessor(ISOLATING)
 class RetrofitResponseTypeKeepProcessor : AbstractProcessor() {
   override fun getSupportedSourceVersion() = SourceVersion.latestSupported()
-  override fun getSupportedAnnotationTypes() = setOf(
-    "retrofit2.http.DELETE",
-    "retrofit2.http.GET",
-    "retrofit2.http.HEAD",
-    "retrofit2.http.HTTP",
-    "retrofit2.http.OPTIONS",
-    "retrofit2.http.PATCH",
-    "retrofit2.http.POST",
-    "retrofit2.http.PUT",
-  )
+  override fun getSupportedAnnotationTypes() = annotationNames
 
   override fun process(
     annotations: Set<TypeElement>,
@@ -77,12 +68,11 @@ class RetrofitResponseTypeKeepProcessor : AbstractProcessor() {
 
     for ((element, referencedTypes) in elementToReferencedTypes) {
       val typeName = element.qualifiedName.toString()
-      val outputFile = "META-INF/proguard/retrofit-response-type-keeper-$typeName.pro"
-      val rules = processingEnv.filer.createResource(CLASS_OUTPUT, "", outputFile, element)
+      val rules = processingEnv.filer.createResource(CLASS_OUTPUT, "", outputFile(typeName), element)
       rules.openWriter().buffered().use { w ->
         w.write("# $typeName\n")
         for (referencedType in referencedTypes.sorted()) {
-          w.write("-keep,allowoptimization,allowshrinking,allowobfuscation class $referencedType\n")
+          w.write("$KEEP_RULE_PREFIX $referencedType\n")
         }
       }
     }

retrofit-response-type-keeper/src/main/kotlin/retrofit2/keeper/RetrofitResponseTypeKeepSymbolProcessor.kt

@@ -0,0 +1,80 @@
+package retrofit2.keeper
+
+import com.google.auto.service.AutoService
+import com.google.devtools.ksp.processing.CodeGenerator
+import com.google.devtools.ksp.processing.Dependencies
+import com.google.devtools.ksp.processing.Resolver
+import com.google.devtools.ksp.processing.SymbolProcessor
+import com.google.devtools.ksp.processing.SymbolProcessorEnvironment
+import com.google.devtools.ksp.processing.SymbolProcessorProvider
+import com.google.devtools.ksp.symbol.KSAnnotated
+import com.google.devtools.ksp.symbol.KSClassDeclaration
+import com.google.devtools.ksp.symbol.KSFunctionDeclaration
+import com.google.devtools.ksp.symbol.KSType
+import com.google.devtools.ksp.symbol.Modifier
+
+class RetrofitResponseTypeKeepSymbolProcessor(
+  environment: SymbolProcessorEnvironment,
+) : SymbolProcessor {
+  private val codeGenerator: CodeGenerator = environment.codeGenerator
+
+  override fun process(resolver: Resolver): List<KSAnnotated> {
+    val elementToReferencedTypes = mutableMapOf<KSClassDeclaration, MutableSet<String>>()
+
+    annotationNames.flatMap { resolver.getSymbolsWithAnnotation(it) }
+      .filterIsInstance<KSFunctionDeclaration>()
+      .forEach { function ->
+        val serviceType = function.parentDeclaration as? KSClassDeclaration ?: return@forEach
+        val referenced = elementToReferencedTypes.getOrPut(serviceType, ::LinkedHashSet)
+
+        // Retrofit has special support for 'suspend fun' in Kotlin which manifests as a
+        // final Continuation parameter whose generic type is the declared return type.
+        if (function.modifiers.contains(Modifier.SUSPEND)) {
+          function.parameters.forEach {
+            it.type.resolve().recursiveParameterizedTypesTo(referenced)
+          }
+        }
+
+        val returnType = function.returnType?.resolve() ?: return@forEach
+        returnType.recursiveParameterizedTypesTo(referenced)
+      }
+
+    elementToReferencedTypes.forEach { (element, referencedTypes) ->
+      val containingFile = element.containingFile ?: return@forEach
+      val typeName = requireNotNull(element.qualifiedName).asString()
+
+      val dependencies = Dependencies(aggregating = false, containingFile)
+      codeGenerator.createNewFile(dependencies, "", outputFile(typeName), "")
+        .bufferedWriter().use { writer ->
+          writer.write("# $typeName\n")
+          for (referencedType in referencedTypes.sorted()) {
+            writer.write("$KEEP_RULE_PREFIX $referencedType\n")
+          }
+        }
+    }
+
+    return emptyList()
+  }
+
+  private fun KSType.recursiveParameterizedTypesTo(types: MutableSet<String>) {
+    val declaration = this.declaration
+    if (declaration is KSClassDeclaration) {
+      var qualifiedName = declaration.qualifiedName?.asString()
+      if (qualifiedName == "kotlin.Any") {
+        qualifiedName = "java.lang.Object"
+      }
+      qualifiedName?.let { types.add(it) }
+    }
+
+    for (typeArgument in arguments) {
+      typeArgument.type?.resolve()?.recursiveParameterizedTypesTo(types)
+    }
+  }
+
+  @Suppress("unused") // Used in service file.
+  @AutoService(SymbolProcessorProvider::class)
+  class Provider : SymbolProcessorProvider {
+    override fun create(environment: SymbolProcessorEnvironment): SymbolProcessor =
+      RetrofitResponseTypeKeepSymbolProcessor(environment)
+  }
+}

retrofit-response-type-keeper/src/main/kotlin/retrofit2/keeper/Utils.kt

@@ -0,0 +1,17 @@
+package retrofit2.keeper
+
+internal val annotationNames = setOf(
+  "retrofit2.http.DELETE",
+  "retrofit2.http.GET",
+  "retrofit2.http.HEAD",
+  "retrofit2.http.HTTP",
+  "retrofit2.http.OPTIONS",
+  "retrofit2.http.PATCH",
+  "retrofit2.http.POST",
+  "retrofit2.http.PUT",
+)
+
+internal const val KEEP_RULE_PREFIX = "-keep,allowoptimization,allowshrinking,allowobfuscation class"
+
+internal fun outputFile(typeName: String) =
+  "META-INF/proguard/retrofit-response-type-keeper-$typeName.pro"

retrofit-response-type-keeper/src/test/kotlin/retrofit2/keeper/RetrofitResponseTypeKeepProcessorTest.kt

@@ -16,147 +16,85 @@
 package retrofit2.keeper
 
 import com.google.common.truth.Truth.assertAbout
+import com.google.common.truth.Truth.assertThat
 import com.google.testing.compile.JavaFileObjects
 import com.google.testing.compile.JavaSourceSubjectFactory.javaSource
+import com.google.testing.junit.testparameterinjector.TestParameter
+import com.google.testing.junit.testparameterinjector.TestParameterInjector
+import com.tschuchort.compiletesting.KotlinCompilation
+import com.tschuchort.compiletesting.KotlinCompilation.ExitCode
+import com.tschuchort.compiletesting.SourceFile
+import com.tschuchort.compiletesting.configureKsp
+import com.tschuchort.compiletesting.kspSourcesDir
 import java.nio.charset.StandardCharsets.UTF_8
+import java.nio.file.NoSuchFileException
 import javax.tools.StandardLocation.CLASS_OUTPUT
+import kotlin.io.path.readText
+import kotlin.io.path.toPath
+import org.jetbrains.kotlin.compiler.plugin.ExperimentalCompilerApi
 import org.junit.Test
+import org.junit.runner.RunWith
 
-class RetrofitResponseTypeKeepProcessorTest {
+@RunWith(TestParameterInjector::class)
+class RetrofitResponseTypeKeepProcessorTest(
+  @param:TestParameter private val useKsp: Boolean,
+) {
+  @OptIn(ExperimentalCompilerApi::class)
   @Test
-  fun allHttpMethods() {
-    val service = JavaFileObjects.forSourceString(
-      "test.Service",
-      """
-      package test;
-      import retrofit2.*;
-      import retrofit2.http.*;
+  fun process(
+    @TestParameter(
+      "all-http-methods",
+      "nesting",
+      "kotlin-suspend",
+    ) name: String,
+  ) {
+    val rules = readResourceAsText("$name/Service.pro")
+    val generatedPath = "META-INF/proguard/retrofit-response-type-keeper-test.Service.pro"
 
-      class DeleteUser {}
-      class GetUser {}
-      class HeadUser {}
-      class HttpUser {}
-      class OptionsUser {}
-      class PatchUser {}
-      class PostUser {}
-      class PutUser {}
-
-      interface Service {
-        @DELETE("/") Call<DeleteUser> delete();
-        @GET("/") Call<GetUser> get();
-        @HEAD("/") Call<HeadUser> head();
-        @HTTP(method = "CUSTOM", path = "/") Call<HttpUser> http();
-        @OPTIONS("/") Call<OptionsUser> options();
-        @PATCH("/") Call<PatchUser> patch();
-        @POST("/") Call<PostUser> post();
-        @PUT("/") Call<PutUser> put();
+    if (useKsp) {
+      val compilation = KotlinCompilation().apply {
+        configureKsp {
+          inheritClassPath = true
+          symbolProcessorProviders += RetrofitResponseTypeKeepSymbolProcessor.Provider()
+          sources = listOf(
+            SourceFile.kotlin(
+              "Service.kt",
+              readResourceAsText("$name/Service.kt"),
+            ),
+          )
+        }
       }
-    """.trimIndent(),
-    )
+      val result = compilation.compile()
 
-    assertAbout(javaSource())
-      .that(service)
-      .processedWith(RetrofitResponseTypeKeepProcessor())
-      .compilesWithoutError()
-      .and()
-      .generatesFileNamed(
-        CLASS_OUTPUT,
-        "",
-        "META-INF/proguard/retrofit-response-type-keeper-test.Service.pro",
-      ).withStringContents(
-        UTF_8,
-        """
-        |# test.Service
-        |-keep,allowoptimization,allowshrinking,allowobfuscation class retrofit2.Call
-        |-keep,allowoptimization,allowshrinking,allowobfuscation class test.DeleteUser
-        |-keep,allowoptimization,allowshrinking,allowobfuscation class test.GetUser
-        |-keep,allowoptimization,allowshrinking,allowobfuscation class test.HeadUser
-        |-keep,allowoptimization,allowshrinking,allowobfuscation class test.HttpUser
-        |-keep,allowoptimization,allowshrinking,allowobfuscation class test.OptionsUser
-        |-keep,allowoptimization,allowshrinking,allowobfuscation class test.PatchUser
-        |-keep,allowoptimization,allowshrinking,allowobfuscation class test.PostUser
-        |-keep,allowoptimization,allowshrinking,allowobfuscation class test.PutUser
-        |
-        """.trimMargin(),
+      assertThat(result.exitCode).isEqualTo(ExitCode.OK)
+      assertThat(compilation.kspSourcesDir.resolve("resources/$generatedPath").readText())
+        .isEqualTo(rules)
+    } else {
+      val service = JavaFileObjects.forSourceString(
+        "test.Service",
+        readResourceAsText("$name/Service.java"),
       )
+      assertAbout(javaSource())
+        .that(service)
+        .processedWith(RetrofitResponseTypeKeepProcessor())
+        .compilesWithoutError()
+        .and()
+        .generatesFileNamed(
+          CLASS_OUTPUT,
+          "",
+          generatedPath,
+        ).withStringContents(
+          UTF_8,
+          rules,
+        )
+    }
   }
 
-  @Test
-  fun nesting() {
-    val service = JavaFileObjects.forSourceString(
-      "test.Service",
-      """
-      package test;
-      import retrofit2.*;
-      import retrofit2.http.*;
-
-      class One<T> {}
-      class Two<T> {}
-      class Three {}
-
-      interface Service {
-        @GET("/") Call<One<Two<Three>>> get();
-      }
-    """.trimIndent(),
-    )
-
-    assertAbout(javaSource())
-      .that(service)
-      .processedWith(RetrofitResponseTypeKeepProcessor())
-      .compilesWithoutError()
-      .and()
-      .generatesFileNamed(
-        CLASS_OUTPUT,
-        "",
-        "META-INF/proguard/retrofit-response-type-keeper-test.Service.pro",
-      ).withStringContents(
-        UTF_8,
-        """
-        |# test.Service
-        |-keep,allowoptimization,allowshrinking,allowobfuscation class retrofit2.Call
-        |-keep,allowoptimization,allowshrinking,allowobfuscation class test.One
-        |-keep,allowoptimization,allowshrinking,allowobfuscation class test.Three
-        |-keep,allowoptimization,allowshrinking,allowobfuscation class test.Two
-        |
-        """.trimMargin(),
-      )
-  }
-
-  @Test
-  fun kotlinSuspend() {
-    val service = JavaFileObjects.forSourceString(
-      "test.Service",
-      """
-      package test;
-      import kotlin.coroutines.Continuation;
-      import retrofit2.*;
-      import retrofit2.http.*;
-
-      class Body {}
-
-      interface Service {
-        @GET("/") Object get(Continuation<? extends Body> c);
-      }
-    """.trimIndent(),
-    )
-
-    assertAbout(javaSource())
-      .that(service)
-      .processedWith(RetrofitResponseTypeKeepProcessor())
-      .compilesWithoutError()
-      .and()
-      .generatesFileNamed(
-        CLASS_OUTPUT,
-        "",
-        "META-INF/proguard/retrofit-response-type-keeper-test.Service.pro",
-      ).withStringContents(
-        UTF_8,
-        """
-        |# test.Service
-        |-keep,allowoptimization,allowshrinking,allowobfuscation class java.lang.Object
-        |-keep,allowoptimization,allowshrinking,allowobfuscation class test.Body
-        |
-        """.trimMargin(),
-      )
+  private companion object {
+    fun readResourceAsText(name: String): String {
+      val resource = this::class.java.classLoader.getResource(name)
+        ?: throw NoSuchFileException("Resource $name not found.")
+      return resource.toURI().toPath().readText()
+    }
   }
 }