Move IncrementalContextBase into common-util

(cherry picked from commit 134553499e8045d75b9602703ceb563ef167fe93)

5 files changed, 143 insertions, 129 deletions

diff --git a/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/IncrementalContextBase.kt b/common-util/src/main/kotlin/com/google/devtools/ksp/IncrementalContextBase.kt
index 632f68f8..dab746e2 100644
--- a/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/IncrementalContextBase.kt
+++ b/common-util/src/main/kotlin/com/google/devtools/ksp/IncrementalContextBase.kt
@@ -17,31 +17,18 @@
 
 package com.google.devtools.ksp
 
-import com.google.devtools.ksp.symbol.*
-import com.google.devtools.ksp.symbol.impl.findPsi
-import com.google.devtools.ksp.symbol.impl.java.KSFunctionDeclarationJavaImpl
-import com.google.devtools.ksp.symbol.impl.java.KSPropertyDeclarationJavaImpl
+import com.google.devtools.ksp.symbol.KSClassDeclaration
+import com.google.devtools.ksp.symbol.KSDeclaration
+import com.google.devtools.ksp.symbol.KSDeclarationContainer
+import com.google.devtools.ksp.symbol.KSFile
+import com.google.devtools.ksp.symbol.KSFunctionDeclaration
+import com.google.devtools.ksp.symbol.KSNode
 import com.google.devtools.ksp.visitor.KSDefaultVisitor
-import com.intellij.psi.*
-import com.intellij.psi.impl.source.PsiClassReferenceType
+import com.intellij.psi.PsiJavaFile
+import com.intellij.psi.PsiPackage
 import com.intellij.util.containers.MultiMap
-import com.intellij.util.io.DataExternalizer
-import org.jetbrains.kotlin.descriptors.CallableMemberDescriptor
-import org.jetbrains.kotlin.descriptors.ClassDescriptor
-import org.jetbrains.kotlin.resolve.descriptorUtil.getAllSuperclassesWithoutAny
-import org.jetbrains.kotlin.types.KotlinType
-import org.jetbrains.kotlin.types.typeUtil.supertypes
 import java.io.File
-import java.util.*
-
-class FileToSymbolsMap(
-    storageFile: File,
-    lookupSymbolExternalizer: DataExternalizer<LookupSymbolWrapper>
-) : PersistentMap<File, List<LookupSymbolWrapper>>(
-    storageFile,
-    FileKeyDescriptor,
-    ListExternalizer(lookupSymbolExternalizer),
-)
+import java.util.Date
 
 object symbolCollector : KSDefaultVisitor<(LookupSymbolWrapper) -> Unit, Unit>() {
     override fun defaultHandler(node: KSNode, data: (LookupSymbolWrapper) -> Unit) = Unit
@@ -523,113 +510,6 @@ abstract class IncrementalContextBase(
         }
     }
 
-    // Record a *leaf* type reference. This doesn't address type arguments.
-    private fun recordLookup(ref: PsiClassReferenceType, def: PsiClass) {
-        val psiFile = ref.reference.containingFile as? PsiJavaFile ?: return
-        // A type parameter doesn't have qualified name.
-        //
-        // Note that bounds of type parameters, or other references in classes,
-        // are not addressed recursively here. They are recorded in other places
-        // with more contexts, when necessary.
-        def.qualifiedName?.let { recordLookup(psiFile, it) }
-    }
-
-    // Record a type reference, including its type arguments.
-    fun recordLookup(ref: PsiType) {
-        when (ref) {
-            is PsiArrayType -> recordLookup(ref.componentType)
-            is PsiClassReferenceType -> {
-                val def = ref.resolve() ?: return
-                recordLookup(ref, def)
-                // in case the corresponding KotlinType is passed through ways other than KSTypeReferenceJavaImpl
-                ref.typeArguments().forEach {
-                    if (it is PsiType) {
-                        recordLookup(it)
-                    }
-                }
-            }
-            is PsiWildcardType -> ref.bound?.let { recordLookup(it) }
-        }
-    }
-
-    // Record all references to super types (if they are written in Java) of a given type,
-    // in its type hierarchy.
-    fun recordLookupWithSupertypes(kotlinType: KotlinType) {
-        (listOf(kotlinType) + kotlinType.supertypes()).mapNotNull {
-            it.constructor.declarationDescriptor?.findPsi() as? PsiClass
-        }.forEach {
-            it.superTypes.forEach {
-                recordLookup(it)
-            }
-        }
-    }
-
-    // Record all type references in a Java field.
-    private fun recordLookupForJavaField(psi: PsiField) {
-        recordLookup(psi.type)
-    }
-
-    // Record all type references in a Java method.
-    private fun recordLookupForJavaMethod(psi: PsiMethod) {
-        psi.parameterList.parameters.forEach {
-            recordLookup(it.type)
-        }
-        psi.returnType?.let { recordLookup(it) }
-        psi.typeParameters.forEach {
-            it.bounds.mapNotNull { it as? PsiType }.forEach {
-                recordLookup(it)
-            }
-        }
-    }
-
-    // Record all type references in a KSDeclaration
-    fun recordLookupForDeclaration(declaration: KSDeclaration) {
-        when (declaration) {
-            is KSPropertyDeclarationJavaImpl -> recordLookupForJavaField(declaration.psi)
-            is KSFunctionDeclarationJavaImpl -> recordLookupForJavaMethod(declaration.psi)
-        }
-    }
-
-    // Record all type references in a CallableMemberDescriptor
-    fun recordLookupForCallableMemberDescriptor(descriptor: CallableMemberDescriptor) {
-        val psi = descriptor.findPsi()
-        when (psi) {
-            is PsiMethod -> recordLookupForJavaMethod(psi)
-            is PsiField -> recordLookupForJavaField(psi)
-        }
-    }
-
-    // Record references from all declared functions in the type hierarchy of the given class.
-    // TODO: optimization: filter out inaccessible members
-    fun recordLookupForGetAllFunctions(descriptor: ClassDescriptor) {
-        recordLookupForGetAll(descriptor) {
-            it.methods.forEach {
-                recordLookupForJavaMethod(it)
-            }
-        }
-    }
-
-    // Record references from all declared fields in the type hierarchy of the given class.
-    // TODO: optimization: filter out inaccessible members
-    fun recordLookupForGetAllProperties(descriptor: ClassDescriptor) {
-        recordLookupForGetAll(descriptor) {
-            it.fields.forEach {
-                recordLookupForJavaField(it)
-            }
-        }
-    }
-
-    fun recordLookupForGetAll(descriptor: ClassDescriptor, doChild: (PsiClass) -> Unit) {
-        (descriptor.getAllSuperclassesWithoutAny() + descriptor).mapNotNull {
-            it.findPsi() as? PsiClass
-        }.forEach { psiClass ->
-            psiClass.superTypes.forEach {
-                recordLookup(it)
-            }
-            doChild(psiClass)
-        }
-    }
-
     fun recordGetSealedSubclasses(classDeclaration: KSClassDeclaration) {
         val name = classDeclaration.simpleName.asString()
         val scope = classDeclaration.qualifiedName?.asString()
diff --git a/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/IncrementalWrapperBase.kt b/common-util/src/main/kotlin/com/google/devtools/ksp/IncrementalWrapperBase.kt
index 88accbec..88accbec 100644
--- a/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/IncrementalWrapperBase.kt
+++ b/common-util/src/main/kotlin/com/google/devtools/ksp/IncrementalWrapperBase.kt
diff --git a/common-util/src/main/kotlin/com/google/devtools/ksp/PersistentMap.kt b/common-util/src/main/kotlin/com/google/devtools/ksp/PersistentMap.kt
index 18a440ae..a6159b52 100644
--- a/common-util/src/main/kotlin/com/google/devtools/ksp/PersistentMap.kt
+++ b/common-util/src/main/kotlin/com/google/devtools/ksp/PersistentMap.kt
@@ -83,3 +83,12 @@ class FileToFilesMap(
     FileKeyDescriptor,
     ListExternalizer(FileExternalizer)
 )
+
+class FileToSymbolsMap(
+    storageFile: File,
+    lookupSymbolExternalizer: DataExternalizer<LookupSymbolWrapper>
+) : PersistentMap<File, List<LookupSymbolWrapper>>(
+    storageFile,
+    FileKeyDescriptor,
+    ListExternalizer(lookupSymbolExternalizer),
+)
diff --git a/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/IncrementalContext.kt b/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/IncrementalContext.kt
index a6cfc6b8..5918b259 100644
--- a/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/IncrementalContext.kt
+++ b/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/IncrementalContext.kt
@@ -1,10 +1,24 @@
 package com.google.devtools.ksp
 
+import com.google.devtools.ksp.symbol.KSDeclaration
+import com.google.devtools.ksp.symbol.impl.findPsi
+import com.google.devtools.ksp.symbol.impl.java.KSFunctionDeclarationJavaImpl
+import com.google.devtools.ksp.symbol.impl.java.KSPropertyDeclarationJavaImpl
+import com.intellij.psi.PsiArrayType
+import com.intellij.psi.PsiClass
+import com.intellij.psi.PsiField
+import com.intellij.psi.PsiJavaFile
+import com.intellij.psi.PsiMethod
+import com.intellij.psi.PsiType
+import com.intellij.psi.PsiWildcardType
+import com.intellij.psi.impl.source.PsiClassReferenceType
 import com.intellij.util.containers.MultiMap
 import com.intellij.util.io.DataExternalizer
 import com.intellij.util.io.IOUtil
 import org.jetbrains.kotlin.container.ComponentProvider
 import org.jetbrains.kotlin.container.get
+import org.jetbrains.kotlin.descriptors.CallableMemberDescriptor
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.incremental.IncrementalCompilationContext
 import org.jetbrains.kotlin.incremental.LookupStorage
 import org.jetbrains.kotlin.incremental.LookupSymbol
@@ -14,6 +28,9 @@ import org.jetbrains.kotlin.incremental.components.Position
 import org.jetbrains.kotlin.incremental.components.ScopeKind
 import org.jetbrains.kotlin.incremental.storage.FileToPathConverter
 import org.jetbrains.kotlin.incremental.update
+import org.jetbrains.kotlin.resolve.descriptorUtil.getAllSuperclassesWithoutAny
+import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.typeUtil.supertypes
 import java.io.DataInput
 import java.io.DataOutput
 import java.io.File
@@ -65,6 +82,113 @@ class IncrementalContext(
         }
         return map
     }
+
+    // Record a type reference, including its type arguments.
+    fun recordLookup(ref: PsiType) {
+        // Record a *leaf* type reference. This doesn't address type arguments.
+        fun recordLookup(ref: PsiClassReferenceType, def: PsiClass) {
+            val psiFile = ref.reference.containingFile as? PsiJavaFile ?: return
+            // A type parameter doesn't have qualified name.
+            //
+            // Note that bounds of type parameters, or other references in classes,
+            // are not addressed recursively here. They are recorded in other places
+            // with more contexts, when necessary.
+            def.qualifiedName?.let { recordLookup(psiFile, it) }
+        }
+
+        when (ref) {
+            is PsiArrayType -> recordLookup(ref.componentType)
+            is PsiClassReferenceType -> {
+                val def = ref.resolve() ?: return
+                recordLookup(ref, def)
+                // in case the corresponding KotlinType is passed through ways other than KSTypeReferenceJavaImpl
+                ref.typeArguments().forEach {
+                    if (it is PsiType) {
+                        recordLookup(it)
+                    }
+                }
+            }
+            is PsiWildcardType -> ref.bound?.let { recordLookup(it) }
+        }
+    }
+
+    // Record all type references in a Java field.
+    private fun recordLookupForJavaField(psi: PsiField) {
+        recordLookup(psi.type)
+    }
+
+    // Record all type references in a Java method.
+    private fun recordLookupForJavaMethod(psi: PsiMethod) {
+        psi.parameterList.parameters.forEach {
+            recordLookup(it.type)
+        }
+        psi.returnType?.let { recordLookup(it) }
+        psi.typeParameters.forEach {
+            it.bounds.mapNotNull { it as? PsiType }.forEach {
+                recordLookup(it)
+            }
+        }
+    }
+
+    // Record all references to super types (if they are written in Java) of a given type,
+    // in its type hierarchy.
+    fun recordLookupWithSupertypes(kotlinType: KotlinType) {
+        (listOf(kotlinType) + kotlinType.supertypes()).mapNotNull {
+            it.constructor.declarationDescriptor?.findPsi() as? PsiClass
+        }.forEach {
+            it.superTypes.forEach {
+                recordLookup(it)
+            }
+        }
+    }
+
+    // Record all type references in a CallableMemberDescriptor
+    fun recordLookupForCallableMemberDescriptor(descriptor: CallableMemberDescriptor) {
+        val psi = descriptor.findPsi()
+        when (psi) {
+            is PsiMethod -> recordLookupForJavaMethod(psi)
+            is PsiField -> recordLookupForJavaField(psi)
+        }
+    }
+
+    // Record references from all declared functions in the type hierarchy of the given class.
+    // TODO: optimization: filter out inaccessible members
+    fun recordLookupForGetAllFunctions(descriptor: ClassDescriptor) {
+        recordLookupForGetAll(descriptor) {
+            it.methods.forEach {
+                recordLookupForJavaMethod(it)
+            }
+        }
+    }
+
+    // Record references from all declared fields in the type hierarchy of the given class.
+    // TODO: optimization: filter out inaccessible members
+    fun recordLookupForGetAllProperties(descriptor: ClassDescriptor) {
+        recordLookupForGetAll(descriptor) {
+            it.fields.forEach {
+                recordLookupForJavaField(it)
+            }
+        }
+    }
+
+    private fun recordLookupForGetAll(descriptor: ClassDescriptor, doChild: (PsiClass) -> Unit) {
+        (descriptor.getAllSuperclassesWithoutAny() + descriptor).mapNotNull {
+            it.findPsi() as? PsiClass
+        }.forEach { psiClass ->
+            psiClass.superTypes.forEach {
+                recordLookup(it)
+            }
+            doChild(psiClass)
+        }
+    }
+
+    // Record all type references in a KSDeclaration
+    fun recordLookupForDeclaration(declaration: KSDeclaration) {
+        when (declaration) {
+            is KSPropertyDeclarationJavaImpl -> recordLookupForJavaField(declaration.psi)
+            is KSFunctionDeclarationJavaImpl -> recordLookupForJavaMethod(declaration.psi)
+        }
+    }
 }
 
 class LookupTrackerWrapperImpl(val lookupTracker: LookupTracker) : LookupTrackerWrapper {
diff --git a/test-utils/build.gradle.kts b/test-utils/build.gradle.kts
index e2b7db7e..17802f49 100644
--- a/test-utils/build.gradle.kts
+++ b/test-utils/build.gradle.kts
@@ -52,6 +52,7 @@ dependencies {
 
     implementation(kotlin("stdlib", kotlinBaseVersion))
     implementation(project(":common-deps"))
+    implementation(project(":common-util"))
     implementation("org.jetbrains.kotlin:kotlin-compiler:$kotlinBaseVersion")
     implementation("org.jetbrains.kotlin:kotlin-compiler-internal-test-framework:$kotlinBaseVersion")
     implementation("org.jetbrains.kotlin:kotlin-scripting-compiler:$kotlinBaseVersion")