Fix: do not use KaFileSymbol across rounds

google · Sep 17, 2024 · d43b61b · d43b61b
1 parent 2878157
commit d43b61b
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 22 deletions.
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt
@@ -547,6 +547,8 @@ class KotlinSymbolProcessing(
                     }
                 }
 
+                val allKSFilesPointers = allDirtyKSFiles.filterIsInstance<Deferrable>().map { it.defer() }
+
                 // Drop caches
                 KotlinGlobalModificationService.getInstance(project).publishGlobalSourceModuleStateModification()
                 KaSessionProvider.getInstance(project).clearCaches()
@@ -565,21 +567,10 @@ class KotlinSymbolProcessing(
                     codeGenerator.generatedFile.filter { it.extension.lowercase() == "kt" },
                     codeGenerator.generatedFile.filter { it.extension.lowercase() == "java" },
                 )
-                // Now that caches are dropped, KtSymbols and KS* are invalid. They need to be re-created from PSI.
-                allDirtyKSFiles = allDirtyKSFiles.map {
-                    when (it) {
-                        is KSFileImpl -> {
-                            val ktFile = it.ktFileSymbol.psi!! as KtFile
-                            analyze { KSFileImpl.getCached(ktFile.symbol) }
-                        }
-
-                        is KSFileJavaImpl -> {
-                            KSFileJavaImpl.getCached(it.psi)
-                        }
-
-                        else -> throw IllegalArgumentException("Unknown KSFile implementation: $it")
-                    }
-                } + newKSFiles
+                // Now that caches are dropped, KtSymbols and KS* are invalid. They need to be restored from deferred.
+                // Do not replace `!!` with `?.`. Implementations of KSFile in KSP2 must implement Deferrable and
+                // return non-null.
+                allDirtyKSFiles = allKSFilesPointers.map { it!!.restore() as KSFile } + newKSFiles
                 incrementalContext.registerGeneratedFiles(newKSFiles)
                 codeGenerator.closeFiles()
             }

diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFileImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFileImpl.kt
@@ -93,12 +93,9 @@ class KSFileImpl private constructor(internal val ktFileSymbol: KaFileSymbol) :
     }
 
     override fun defer(): Restorable {
-        val psi = this.psi
+        val ptr = analyze { ktFileSymbol.createPointer() }
         return Restorable {
-            when (psi) {
-                is KtFile -> analyze { getCached(psi.symbol) }
-                else -> throw IllegalStateException("Unhandled psi file type ${psi.javaClass}")
-            }
+            analyze { getCached(ptr.restoreSymbol() as KaFileSymbol) }
         }
     }
 }
diff --git a/...analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFileJavaImpl.kt b/...analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFileJavaImpl.kt
@@ -64,6 +64,11 @@ class KSFileJavaImpl private constructor(val psi: PsiJavaFile) : KSFile, Deferra
         return "File: ${this.fileName}"
     }
 
-    // Resolver.getSymbolsWithAnnotation never returns a java file because the latter cannot have file annotation.
-    override fun defer(): Restorable? = null
+    // Although Resolver.getSymbolsWithAnnotation never returns a java file because the latter cannot have file
+    // annotations, this is used internally to restore files across rounds.
+    override fun defer(): Restorable {
+        return Restorable {
+            analyze { getCached(psi) }
+        }
+    }
 }