multiple-round: implement deferred symbols

by introducing Deferrable/Restorable and using KtSymbolPointer.
KtSymbolPointer handles quite a lot corner cases that we'd have to do by
ourselves if using psi directly.

TODOs:
1. Java elements are not supported by KtSymbolPointer yet. We either
   need to do it in the upstream or implement with psi in KSP.
2. KSTypeReference and KSTypeArgument are implemented upon KtType, which
   a. can change in the subsequent rounds, and
   b. is not KtSymbol and has no KtSymbolPointer.

kotlin-analysis-api/build.gradle.kts

@@ -213,3 +213,9 @@ signing {
     useInMemoryPgpKeys(signingKey, signingPassword)
     sign(extensions.getByType<PublishingExtension>().publications)
 }
+
+kotlin {
+    compilerOptions {
+        freeCompilerArgs.add("-Xcontext-receivers")
+    }
+}

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt

@@ -22,14 +22,16 @@ import com.google.devtools.ksp.AnyChanges
 import com.google.devtools.ksp.KSObjectCacheManager
 import com.google.devtools.ksp.analysisapi.providers.IncrementalKotlinDeclarationProviderFactory
 import com.google.devtools.ksp.analysisapi.providers.IncrementalKotlinPackageProviderFactory
+import com.google.devtools.ksp.impl.symbol.kotlin.Deferrable
 import com.google.devtools.ksp.impl.symbol.kotlin.KSFileImpl
 import com.google.devtools.ksp.impl.symbol.kotlin.KSFileJavaImpl
+import com.google.devtools.ksp.impl.symbol.kotlin.Restorable
 import com.google.devtools.ksp.impl.symbol.kotlin.analyze
 import com.google.devtools.ksp.processing.*
 import com.google.devtools.ksp.processing.impl.CodeGeneratorImpl
 import com.google.devtools.ksp.processing.impl.JvmPlatformInfoImpl
-import com.google.devtools.ksp.symbol.KSAnnotated
 import com.google.devtools.ksp.symbol.KSFile
+import com.google.devtools.ksp.symbol.Origin
 import com.google.devtools.ksp.toKotlinVersion
 import com.intellij.core.CoreApplicationEnvironment
 import com.intellij.core.CorePackageIndex
@@ -332,9 +334,6 @@ class KotlinSymbolProcessing(
 
     @OptIn(KtAnalysisApiInternals::class)
     fun execute() {
-        val deferredSymbols = mutableMapOf<SymbolProcessor, List<KSAnnotated>>()
-        val providers: List<SymbolProcessorProvider> = kspConfig.processorProviders
-
         // TODO: CompilerConfiguration is deprecated.
         val compilerConfiguration: CompilerConfiguration = CompilerConfiguration().apply {
             addKotlinSourceRoots(kspConfig.sourceRoots.map { it.path })
@@ -358,11 +357,13 @@ class KotlinSymbolProcessing(
 
         val kspCoreEnvironment = KSPCoreEnvironment(analysisAPISession.project as MockProject)
 
-        // TODO: deferred symbols: use PSIs; they don't change.
         // TODO: error handling, onError()
         // TODO: performance
         val project = analysisAPISession.project
         val psiManager = PsiManager.getInstance(project)
+        val providers: List<SymbolProcessorProvider> = kspConfig.processorProviders
+
+        val deferredSymbols = mutableMapOf<SymbolProcessor, List<Restorable>>()
         var finished = false
         var initialized = false
         lateinit var codeGenerator: CodeGeneratorImpl
@@ -448,17 +449,22 @@ class KotlinSymbolProcessing(
                     it.filePath in newFileNames
                 }
             }
-            // TODO: support no kotlin source input.
             val resolver = ResolverAAImpl(
                 allKSFiles,
                 newKSFiles,
-                kspConfig,
+                deferredSymbols,
                 analysisAPISession.project
             )
             ResolverAAImpl.instance = resolver
 
-            // TODO: multiple rounds
-            processors.forEach { it.process(resolver) }
+            processors.forEach {
+                deferredSymbols[it] =
+                    it.process(resolver).filter { it.origin == Origin.KOTLIN || it.origin == Origin.JAVA }
+                        .filterIsInstance<Deferrable>().mapNotNull(Deferrable::defer)
+                if (!deferredSymbols.containsKey(it) || deferredSymbols[it]!!.isEmpty()) {
+                    deferredSymbols.remove(it)
+                }
+            }
 
             if (codeGenerator.generatedFile.isEmpty()) {
                 finished = true

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/ResolverAAImpl.kt

@@ -22,6 +22,7 @@ import com.google.devtools.ksp.*
 import com.google.devtools.ksp.impl.symbol.kotlin.*
 import com.google.devtools.ksp.processing.KSBuiltIns
 import com.google.devtools.ksp.processing.Resolver
+import com.google.devtools.ksp.processing.SymbolProcessor
 import com.google.devtools.ksp.processing.impl.KSNameImpl
 import com.google.devtools.ksp.processing.impl.KSTypeReferenceSyntheticImpl
 import com.google.devtools.ksp.symbol.*
@@ -55,7 +56,7 @@ import org.jetbrains.org.objectweb.asm.Opcodes
 class ResolverAAImpl(
     val allKSFiles: List<KSFile>,
     val newKSFiles: List<KSFile>,
-    val kspConfig: KSPJvmConfig,
+    val deferredSymbols: Map<SymbolProcessor, List<Restorable>>,
     val project: Project
 ) : Resolver {
     companion object {
@@ -415,7 +416,11 @@ class ResolverAAImpl(
             file.accept(visitor, Unit)
         }
 
-        return visitor.symbols.asSequence().filter {
+        val deferred = deferredSymbols.values.flatten().mapNotNull {
+            it.restore()
+        }.toSet()
+
+        return (visitor.symbols + deferred).asSequence().filter {
             it.annotations.any {
                 it.annotationType.resolve().declaration.qualifiedName?.asString() == annotationName
             }

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/AbstractKSDeclarationImpl.kt

@@ -37,7 +37,7 @@ import org.jetbrains.kotlin.analysis.utils.printer.parentOfType
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtModifierListOwner
 
-abstract class AbstractKSDeclarationImpl(val ktDeclarationSymbol: KtDeclarationSymbol) : KSDeclaration {
+abstract class AbstractKSDeclarationImpl(val ktDeclarationSymbol: KtDeclarationSymbol) : KSDeclaration, Deferrable {
     override val origin: Origin by lazy {
         mapAAOrigin(ktDeclarationSymbol)
     }

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSClassDeclarationEnumEntryImpl.kt

@@ -127,4 +127,8 @@ class KSClassDeclarationEnumEntryImpl private constructor(private val ktEnumEntr
         // TODO: fix after .getDeclaredMemberScope() works for enum entry with no initializer.
         emptySequence()
     }
+
+    override fun defer(): Restorable? {
+        return ktEnumEntrySymbol.defer(Companion::getCached)
+    }
 }

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSClassDeclarationImpl.kt

@@ -156,6 +156,10 @@ class KSClassDeclarationImpl private constructor(internal val ktClassOrObjectSym
             }.asSequence()
         } else decls
     }
+
+    override fun defer(): Restorable? {
+        return ktClassOrObjectSymbol.defer(::getCached)
+    }
 }
 
 internal fun KtClassOrObjectSymbol.toModifiers(): Set<Modifier> {

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFileImpl.kt

@@ -36,7 +36,7 @@ import org.jetbrains.kotlin.analysis.api.symbols.KtPropertySymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtTypeAliasSymbol
 import org.jetbrains.kotlin.psi.KtFile
 
-class KSFileImpl private constructor(private val ktFileSymbol: KtFileSymbol) : KSFile {
+class KSFileImpl private constructor(private val ktFileSymbol: KtFileSymbol) : KSFile, Deferrable {
     companion object : KSObjectCache<KtFileSymbol, KSFileImpl>() {
         fun getCached(ktFileSymbol: KtFileSymbol) = cache.getOrPut(ktFileSymbol) { KSFileImpl(ktFileSymbol) }
     }
@@ -93,4 +93,15 @@ class KSFileImpl private constructor(private val ktFileSymbol: KtFileSymbol) : K
     override fun toString(): String {
         return "File: ${this.fileName}"
     }
+
+    override fun defer(): Restorable {
+        val psi = this.psi
+        return Restorable {
+            when (psi) {
+                is KtFile -> analyze { getCached(psi.getFileSymbol()) }
+                is PsiJavaFile -> null
+                else -> throw IllegalStateException("Unhandled psi file type ${psi.javaClass}")
+            }
+        }
+    }
 }

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFileJavaImpl.kt

@@ -29,7 +29,7 @@ import com.google.devtools.ksp.symbol.Location
 import com.google.devtools.ksp.symbol.Origin
 import com.intellij.psi.PsiJavaFile
 
-class KSFileJavaImpl private constructor(private val psi: PsiJavaFile) : KSFile {
+class KSFileJavaImpl private constructor(private val psi: PsiJavaFile) : KSFile, Deferrable {
     companion object : KSObjectCache<PsiJavaFile, KSFileJavaImpl>() {
         fun getCached(psi: PsiJavaFile) = cache.getOrPut(psi) { KSFileJavaImpl(psi) }
     }
@@ -63,4 +63,7 @@ class KSFileJavaImpl private constructor(private val psi: PsiJavaFile) : KSFile
     override fun toString(): String {
         return "File: ${this.fileName}"
     }
+
+    // Resolver.getSymbolsWithAnnotation never returns a java file because the latter cannot have file annotation.
+    override fun defer(): Restorable? = null
 }

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFunctionDeclarationImpl.kt

@@ -166,6 +166,10 @@ class KSFunctionDeclarationImpl private constructor(internal val ktFunctionSymbo
             super.docString
         }
     }
+
+    override fun defer(): Restorable? {
+        return ktFunctionSymbol.defer(::getCached)
+    }
 }
 
 internal fun KtFunctionLikeSymbol.toModifiers(): Set<Modifier> {

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyAccessorImpl.kt

@@ -32,7 +32,7 @@ import org.jetbrains.kotlin.psi.KtPropertyAccessor
 abstract class KSPropertyAccessorImpl(
     internal val ktPropertyAccessorSymbol: KtPropertyAccessorSymbol,
     override val receiver: KSPropertyDeclaration
-) : KSPropertyAccessor {
+) : KSPropertyAccessor, Deferrable {
 
     override val annotations: Sequence<KSAnnotation> by lazy {
         ktPropertyAccessorSymbol.annotations.asSequence()
@@ -105,6 +105,14 @@ class KSPropertySetterImpl private constructor(
     override fun toString(): String {
         return "$receiver.setter()"
     }
+
+    override fun defer(): Restorable? {
+        val other = (receiver as Deferrable).defer() ?: return null
+        return ktPropertyAccessorSymbol.defer {
+            val owner = other.restore() ?: return@defer null
+            getCached(owner as KSPropertyDeclaration, it as KtPropertySetterSymbol)
+        }
+    }
 }
 
 class KSPropertyGetterImpl private constructor(
@@ -127,4 +135,12 @@ class KSPropertyGetterImpl private constructor(
     override fun toString(): String {
         return "$receiver.getter()"
     }
+
+    override fun defer(): Restorable? {
+        val other = (receiver as Deferrable).defer() ?: return null
+        return ktPropertyAccessorSymbol.defer {
+            val owner = other.restore() ?: return@defer null
+            getCached(owner as KSPropertyDeclaration, it as KtPropertyGetterSymbol)
+        }
+    }
 }

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationImpl.kt

@@ -134,6 +134,10 @@ class KSPropertyDeclarationImpl private constructor(internal val ktPropertySymbo
     override fun <D, R> accept(visitor: KSVisitor<D, R>, data: D): R {
         return visitor.visitPropertyDeclaration(this, data)
     }
+
+    override fun defer(): Restorable? {
+        return ktPropertySymbol.defer(::getCached)
+    }
 }
 
 internal fun KtAnnotationApplication.isUseSiteTargetAnnotation(): Boolean {

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationJavaImpl.kt

@@ -64,6 +64,10 @@ class KSPropertyDeclarationJavaImpl private constructor(private val ktJavaFieldS
     override fun <D, R> accept(visitor: KSVisitor<D, R>, data: D): R {
         return visitor.visitPropertyDeclaration(this, data)
     }
+
+    override fun defer(): Restorable? {
+        TODO("Not yet implemented")
+    }
 }
 
 internal fun KtJavaFieldSymbol.toModifiers(): Set<Modifier> {

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationLocalVariableImpl.kt

@@ -51,4 +51,8 @@ class KSPropertyDeclarationLocalVariableImpl private constructor(
     override fun <D, R> accept(visitor: KSVisitor<D, R>, data: D): R {
         return visitor.visitPropertyDeclaration(this, data)
     }
+
+    override fun defer(): Restorable? {
+        return ktLocalVariableSymbol.defer(::getCached)
+    }
 }

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeAliasImpl.kt

@@ -53,4 +53,8 @@ class KSTypeAliasImpl private constructor(private val ktTypeAliasSymbol: KtTypeA
     override fun <D, R> accept(visitor: KSVisitor<D, R>, data: D): R {
         return visitor.visitTypeAlias(this, data)
     }
+
+    override fun defer(): Restorable? {
+        return ktTypeAliasSymbol.defer(::getCached)
+    }
 }

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeArgumentImpl.kt

@@ -34,7 +34,7 @@ import org.jetbrains.kotlin.analysis.api.KtTypeProjection
 class KSTypeArgumentImpl private constructor(
     private val ktTypeProjection: KtTypeProjection,
     override val parent: KSNode?
-) : KSTypeArgument {
+) : KSTypeArgument, Deferrable {
     companion object : KSObjectCache<IdKeyPair<KtTypeProjection, KSNode?>, KSTypeArgumentImpl>() {
         fun getCached(ktTypeProjection: KtTypeProjection, parent: KSNode? = null) =
             cache.getOrPut(IdKeyPair(ktTypeProjection, parent)) { KSTypeArgumentImpl(ktTypeProjection, parent) }
@@ -74,4 +74,8 @@ class KSTypeArgumentImpl private constructor(
     override fun toString(): String {
         return "$variance $type"
     }
+
+    override fun defer(): Restorable? {
+        TODO("Not yet implemented")
+    }
 }

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeArgumentLiteImpl.kt

@@ -12,7 +12,7 @@ import com.google.devtools.ksp.symbol.Origin
 import com.google.devtools.ksp.symbol.Variance
 
 class KSTypeArgumentLiteImpl private constructor(override val type: KSTypeReference, override val variance: Variance) :
-    KSTypeArgument {
+    KSTypeArgument, Deferrable {
     companion object : KSObjectCache<Pair<KSTypeReference, Variance>, KSTypeArgument>() {
         fun getCached(type: KSTypeReference, variance: Variance) = cache.getOrPut(Pair(type, variance)) {
             KSTypeArgumentLiteImpl(type, variance)
@@ -30,4 +30,8 @@ class KSTypeArgumentLiteImpl private constructor(override val type: KSTypeRefere
     override fun <D, R> accept(visitor: KSVisitor<D, R>, data: D): R {
         return visitor.visitTypeArgument(this, data)
     }
+
+    override fun defer(): Restorable? {
+        TODO("Not yet implemented")
+    }
 }

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeParameterImpl.kt

@@ -66,4 +66,8 @@ class KSTypeParameterImpl private constructor(internal val ktTypeParameterSymbol
     override fun <D, R> accept(visitor: KSVisitor<D, R>, data: D): R {
         return visitor.visitTypeParameter(this, data)
     }
+
+    override fun defer(): Restorable? {
+        return ktTypeParameterSymbol.defer(::getCached)
+    }
 }

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeReferenceImpl.kt

@@ -31,7 +31,7 @@ class KSTypeReferenceImpl private constructor(
     private val ktType: KtType,
     override val parent: KSNode?,
     private val index: Int
-) : KSTypeReference {
+) : KSTypeReference, Deferrable {
     companion object : KSObjectCache<IdKeyTriple<KtType, KSNode?, Int>, KSTypeReference>() {
         fun getCached(type: KtType, parent: KSNode? = null, index: Int = -1): KSTypeReference =
             cache.getOrPut(IdKeyTriple(type, parent, index)) { KSTypeReferenceImpl(type, parent, index) }
@@ -103,4 +103,8 @@ class KSTypeReferenceImpl private constructor(
     override fun toString(): String {
         return ktType.render()
     }
+
+    override fun defer(): Restorable? {
+        TODO("Not yet implemented")
+    }
 }

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSValueArgumentImpl.kt

@@ -32,7 +32,7 @@ import org.jetbrains.kotlin.analysis.api.annotations.KtUnsupportedAnnotationValu
 class KSValueArgumentImpl private constructor(
     private val namedAnnotationValue: KtNamedAnnotationValue,
     override val origin: Origin
-) : KSValueArgument {
+) : KSValueArgument, Deferrable {
     companion object : KSObjectCache<KtNamedAnnotationValue, KSValueArgumentImpl>() {
         fun getCached(namedAnnotationValue: KtNamedAnnotationValue, origin: Origin) =
             cache.getOrPut(namedAnnotationValue) { KSValueArgumentImpl(namedAnnotationValue, origin) }
@@ -93,4 +93,6 @@ class KSValueArgumentImpl private constructor(
         is KtConstantAnnotationValue -> this.constantValue.value
         is KtUnsupportedAnnotationValue -> null
     }
+
+    override fun defer(): Restorable = Restorable { getCached(namedAnnotationValue, origin) }
 }

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSValueParameterImpl.kt

@@ -32,7 +32,7 @@ import org.jetbrains.kotlin.psi.KtParameter
 class KSValueParameterImpl private constructor(
     private val ktValueParameterSymbol: KtValueParameterSymbol,
     override val parent: KSAnnotated
-) : KSValueParameter {
+) : KSValueParameter, Deferrable {
     companion object : KSObjectCache<KtValueParameterSymbol, KSValueParameterImpl>() {
         fun getCached(ktValueParameterSymbol: KtValueParameterSymbol, parent: KSAnnotated) =
             cache.getOrPut(ktValueParameterSymbol) { KSValueParameterImpl(ktValueParameterSymbol, parent) }
@@ -102,4 +102,11 @@ class KSValueParameterImpl private constructor(
     override fun toString(): String {
         return name?.asString() ?: "_"
     }
+
+    override fun defer(): Restorable? {
+        val other = (parent as Deferrable).defer() ?: return null
+        return ktValueParameterSymbol.defer {
+            getCached(it, other.restore() ?: return@defer null)
+        }
+    }
 }

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/util.kt

@@ -466,3 +466,21 @@ internal fun fillInDeepSubstitutor(context: KtType, substitutorBuilder: KtSubsti
         fillInDeepSubstitutor(it, substitutorBuilder)
     }
 }
+
+fun interface Restorable {
+    fun restore(): KSAnnotated?
+}
+
+fun interface Deferrable {
+    fun defer(): Restorable?
+}
+
+fun <T : KtSymbol> T.defer(restore: (T) -> KSAnnotated?): Restorable? {
+    val ptr = analyze { with(this) { this@defer.createPointer() } }
+    return Restorable {
+        analyze {
+            val restored = ptr.restoreSymbol() ?: return@analyze null
+            restore(restored as T)
+        }
+    }
+}