Incremental processing: support getSealedSubclasses from multiple files

There isn't a simple way to associate symbols with their supertypes without resolution, because of typealias. Therefore, all sealed classes / interfaces on which getSealedSubclasses is called are invalidated.
google · Apr 28, 2021 · 470c16a · 470c16a
1 parent f8dfd33
commit 470c16a
Show file tree

Hide file tree

Showing 7 changed files with 159 additions and 23 deletions.
diff --git a/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/Incremental.kt b/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/Incremental.kt
@@ -49,7 +49,17 @@ import java.nio.file.Files
 import java.nio.file.StandardCopyOption
 import java.util.*
 
-class FileToSymbolsMap(storageFile: File) : BasicMap<File, Collection<LookupSymbol>>(storageFile, FileKeyDescriptor, CollectionExternalizer(LookupSymbolExternalizer, { HashSet() })) {
+abstract class PersistentMap<K : Comparable<K>, V>(
+    storageFile: File,
+    keyDescriptor: KeyDescriptor<K>,
+    valueExternalizer: DataExternalizer<V>
+) : BasicMap<K, V>(storageFile, keyDescriptor, valueExternalizer) {
+    abstract operator fun get(key: K): V?
+    abstract operator fun set(key: K, value: V)
+    abstract fun remove(key: K)
+}
+
+class FileToSymbolsMap(storageFile: File) : PersistentMap<File, Collection<LookupSymbol>>(storageFile, FileKeyDescriptor, CollectionExternalizer(LookupSymbolExternalizer, { HashSet() })) {
     override fun dumpKey(key: File): String = key.toString()
 
     override fun dumpValue(value: Collection<LookupSymbol>): String = value.toString()
@@ -58,13 +68,13 @@ class FileToSymbolsMap(storageFile: File) : BasicMap<File, Collection<LookupSymb
         storage.append(file, listOf(symbol))
     }
 
-    operator fun get(key: File): Collection<LookupSymbol>? = storage[key]
+    override operator fun get(key: File): Collection<LookupSymbol>? = storage[key]
 
-    operator fun set(key: File, symbols: Set<LookupSymbol>) {
+    override operator fun set(key: File, symbols: Collection<LookupSymbol>) {
         storage[key] = symbols
     }
 
-    fun remove(key: File) {
+    override fun remove(key: File) {
         storage.remove(key)
     }
 
@@ -103,11 +113,11 @@ object FileExternalizer : DataExternalizer<File> {
     }
 }
 
-class FileToFilesMap(storageFile: File) : BasicMap<File, Collection<File>>(storageFile, FileKeyDescriptor, CollectionExternalizer(FileExternalizer, { HashSet() })) {
+class FileToFilesMap(storageFile: File) : PersistentMap<File, Collection<File>>(storageFile, FileKeyDescriptor, CollectionExternalizer(FileExternalizer, { HashSet() })) {
 
-    operator fun get(key: File): Collection<File>? = storage[key]
+    override operator fun get(key: File): Collection<File>? = storage[key]
 
-    operator fun set(key: File, value: Collection<File>) {
+    override operator fun set(key: File, value: Collection<File>) {
         storage[key] = value
     }
 
@@ -116,7 +126,7 @@ class FileToFilesMap(storageFile: File) : BasicMap<File, Collection<File>>(stora
     override fun dumpValue(value: Collection<File>) =
             value.dumpCollection()
 
-    fun remove(key: File) {
+    override fun remove(key: File) {
         storage.remove(key)
     }
 
@@ -160,6 +170,14 @@ class IncrementalContext(
     // Symbols defined in changed files. This is used to update symbolsMap in the end.
     private val updatedSymbols = MultiMap.createSet<File, LookupSymbol>()
 
+    // Sealed classes / interfaces on which `getSealedSubclasses` is invoked.
+    // This is used to update sealedMap in the end.
+    private val updatedSealed = MultiMap.createSet<File, LookupSymbol>()
+
+    // Sealed classes / interfaces on which `getSealedSubclasses` is invoked.
+    // This is saved across processing.
+    private val sealedMap = FileToSymbolsMap(File(options.cachesDir, "sealed"))
+
     // Symbols defined in each file. This is saved across processing.
     private val symbolsMap = FileToSymbolsMap(File(options.cachesDir, "symbols"))
 
@@ -220,14 +238,18 @@ class IncrementalContext(
         }
 
         // Add previously defined symbols in removed and modified files
-        ksFiles.filter { it.relativeFile in removed || it.relativeFile in modified }.forEach { file ->
-            symbolsMap[file.relativeFile]?.let {
+        (modified + removed).forEach { file ->
+            symbolsMap[file]?.let {
                 changedSyms.addAll(it)
             }
         }
 
+        // Invalidate all sealed classes / interfaces on which `getSealedSubclasses` was invoked.
+        // FIXME: find a better solution to deal with typealias without resolution.
+        changedSyms.addAll(sealedMap.keys.flatMap { sealedMap[it]!! })
+
         // For each changed symbol, either changed, modified or removed, invalidate files that looked them up, recursively.
-        val invalidator = DepInvalidator(lookupCache, symbolsMap, ksFiles.filter { it.relativeFile in removed || it.relativeFile in modified }.map { it.relativeFile })
+        val invalidator = DepInvalidator(lookupCache, symbolsMap, modified)
         changedSyms.forEach {
             invalidator.invalidate(it)
         }
@@ -460,28 +482,37 @@ class IncrementalContext(
         updateLookupCache(dirtyFiles, removedOutputs)
 
         // Update symbolsMap
-        if (!rebuild) {
+        fun <K: Comparable<K>, V> update(m: PersistentMap<K, Collection<V>>, u: MultiMap<K, V>) {
             // Update symbol caches from modified files.
-            updatedSymbols.keySet().forEach {
-                symbolsMap.set(it, updatedSymbols[it].toSet())
+            u.keySet().forEach {
+                m.set(it, u[it].toSet())
             }
+        }
 
+        fun <K: Comparable<K>, V> remove(m: PersistentMap<K, Collection<V>>, removedKeys: Collection<K>) {
             // Remove symbol caches from removed files.
-            removed.forEach {
-                symbolsMap.remove(it)
+            removedKeys.forEach {
+                m.remove(it)
             }
+        }
 
-            removedOutputs.forEach {
-                symbolsMap.remove(it.absoluteFile)
-            }
+        if (!rebuild) {
+            update(sealedMap, updatedSealed)
+            remove(sealedMap, removed + removedOutputs)
+
+            update(symbolsMap, updatedSymbols)
+            remove(symbolsMap, removed + removedOutputs)
         } else {
             symbolsMap.clean()
-            updatedSymbols.keySet().forEach {
-                symbolsMap.set(it, updatedSymbols[it].toSet())
-            }
+            update(symbolsMap, updatedSymbols)
+
+            sealedMap.clean()
+            update(sealedMap, updatedSealed)
         }
         symbolsMap.flush(false)
         symbolsMap.close()
+        sealedMap.flush(false)
+        sealedMap.close()
     }
 
     fun registerGeneratedFiles(newFiles: Collection<KSFile>) = closeFilesOnException {
@@ -493,6 +524,7 @@ class IncrementalContext(
             return f()
         } catch (e: Exception) {
             symbolsMap.close()
+            sealedMap.close()
             lookupCache.close()
             sourceToOutputsMap.close()
             throw e
@@ -659,6 +691,12 @@ class IncrementalContext(
         }
     }
 
+    fun recordGetSealedSubclasses(classDeclaration: KSClassDeclaration) {
+        val name = classDeclaration.simpleName.asString()
+        val scope = classDeclaration.qualifiedName?.asString()?.let { it.substring(0, Math.max(it.length - name.length - 1, 0))} ?: return
+        updatedSealed.putValue(classDeclaration.containingFile!!.relativeFile, LookupSymbol(name, scope))
+    }
+
     // Debugging and testing only.
     internal fun dumpLookupRecords(): Map<String, List<String>> {
         val map = mutableMapOf<String, List<String>>()
@@ -675,7 +713,7 @@ class IncrementalContext(
 internal class DepInvalidator(
         private val lookupCache: LookupStorage,
         private val symbolsMap: FileToSymbolsMap,
-        changedFiles: List<File>
+        changedFiles: Collection<File>
 ) {
     private val visitedSyms = mutableSetOf<LookupSymbol>()
     val visitedFiles = mutableSetOf<File>().apply { addAll(changedFiles) }

diff --git a/...ugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/kotlin/KSClassDeclarationImpl.kt b/...ugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/kotlin/KSClassDeclarationImpl.kt
@@ -56,6 +56,7 @@ class KSClassDeclarationImpl private constructor(val ktClassOrObject: KtClassOrO
 
     override fun getSealedSubclasses(): Sequence<KSClassDeclaration> {
         return if (Modifier.SEALED in modifiers) {
+            ResolverImpl.instance.incrementalContext.recordGetSealedSubclasses(this)
             descriptor.sealedSubclassesSequence()
         } else {
             emptySequence()

diff --git a/integration-tests/src/test/kotlin/com/google/devtools/ksp/test/GetSealedSubclassesIncIT.kt b/integration-tests/src/test/kotlin/com/google/devtools/ksp/test/GetSealedSubclassesIncIT.kt
@@ -0,0 +1,56 @@
+package com.google.devtools.ksp.test
+
+import org.gradle.testkit.runner.GradleRunner
+import org.junit.Assert
+import org.junit.Rule
+import org.junit.Test
+import java.io.File
+
+class GetSealedSubclassesIncIT {
+    @Rule
+    @JvmField
+    val project: TemporaryTestProject = TemporaryTestProject("sealed-subclasses", "test-processor")
+
+    @Test
+    fun testGetSealedSubclassesInc() {
+        val gradleRunner = GradleRunner.create().withProjectDir(project.root)
+
+        val expected2 = listOf(
+            "w: [ksp] Processing Impl1.kt",
+            "w: [ksp] Impl1 : []",
+            "w: [ksp] Processing Impl2.kt",
+            "w: [ksp] Impl2 : []",
+            "w: [ksp] Processing Sealed.kt",
+            "w: [ksp] Sealed : [Impl1, Impl2]",
+        )
+
+        val expected3 = listOf(
+            "w: [ksp] Processing Impl1.kt",
+            "w: [ksp] Impl1 : []",
+            "w: [ksp] Processing Impl2.kt",
+            "w: [ksp] Impl2 : []",
+            "w: [ksp] Processing Impl3.kt",
+            "w: [ksp] Impl3 : []",
+            "w: [ksp] Processing Sealed.kt",
+            "w: [ksp] Sealed : [Impl1, Impl2, Impl3]",
+        )
+
+        gradleRunner.withArguments("assemble").build().let { result ->
+            val outputs = result.output.split("\n").filter { it.startsWith("w: [ksp]")}
+            Assert.assertEquals(expected2, outputs)
+        }
+
+        File(project.root, "workload/src/main/kotlin/com/example/Impl3.kt").appendText("package com.example\n\n")
+        File(project.root, "workload/src/main/kotlin/com/example/Impl3.kt").appendText("class Impl3 : Sealed()\n")
+        gradleRunner.withArguments("assemble").build().let { result ->
+            val outputs = result.output.split("\n").filter { it.startsWith("w: [ksp]")}
+            Assert.assertEquals(expected3, outputs)
+        }
+
+        File(project.root, "workload/src/main/kotlin/com/example/Impl3.kt").delete()
+        gradleRunner.withArguments("assemble").build().let { result ->
+            val outputs = result.output.split("\n").filter { it.startsWith("w: [ksp]")}
+            Assert.assertEquals(expected2, outputs)
+        }
+    }
+}
diff --git a/...ests/src/test/resources/sealed-subclasses/test-processor/src/main/kotlin/TestProcessor.kt b/...ests/src/test/resources/sealed-subclasses/test-processor/src/main/kotlin/TestProcessor.kt
@@ -0,0 +1,32 @@
+import com.google.devtools.ksp.processing.*
+import com.google.devtools.ksp.symbol.*
+
+
+class TestProcessor(
+    val codeGenerator: CodeGenerator,
+    val logger: KSPLogger
+) : SymbolProcessor {
+    override fun process(resolver: Resolver): List<KSAnnotated> {
+        resolver.getNewFiles().forEach { f ->
+            logger.warn("Processing ${f.fileName}")
+            f.declarations.forEach {
+                if (it is KSClassDeclaration) {
+                    val subs = it.getSealedSubclasses().map { it.simpleName.asString() }.toList()
+                    logger.warn("${it.simpleName.asString()} : $subs")
+                }
+            }
+        }
+        return emptyList()
+    }
+}
+
+class TestProcessorProvider : SymbolProcessorProvider {
+    override fun create(
+        options: Map<String, String>,
+        kotlinVersion: KotlinVersion,
+        codeGenerator: CodeGenerator,
+        logger: KSPLogger
+    ): SymbolProcessor {
+        return TestProcessor(codeGenerator, logger)
+    }
+}
diff --git a/...-tests/src/test/resources/sealed-subclasses/workload/src/main/kotlin/com/example/Impl1.kt b/...-tests/src/test/resources/sealed-subclasses/workload/src/main/kotlin/com/example/Impl1.kt
@@ -0,0 +1,3 @@
+package com.example
+
+class Impl1 : Sealed()
diff --git a/...-tests/src/test/resources/sealed-subclasses/workload/src/main/kotlin/com/example/Impl2.kt b/...-tests/src/test/resources/sealed-subclasses/workload/src/main/kotlin/com/example/Impl2.kt
@@ -0,0 +1,3 @@
+package com.example
+
+class Impl2 : Sealed()
diff --git a/...tests/src/test/resources/sealed-subclasses/workload/src/main/kotlin/com/example/Sealed.kt b/...tests/src/test/resources/sealed-subclasses/workload/src/main/kotlin/com/example/Sealed.kt
@@ -0,0 +1,3 @@
+package com.example
+
+sealed class Sealed