Store extension receiver types in symbol index

Author: fwcd

server/src/main/kotlin/org/javacs/kt/completion/Completions.kt

@@ -68,7 +68,7 @@ fun completions(file: CompiledFile, cursor: Int, index: SymbolIndex, config: Com
     val elementItemLabels = elementItemList.mapNotNull { it.label }.toSet()
     val items = (
         elementItemList.asSequence()
-        + (if (!isExhaustive) indexCompletionItems(file.parse, receiver, index, partial).filter { it.label !in elementItemLabels } else emptySequence())
+        + (if (!isExhaustive) indexCompletionItems(file, cursor, receiver, index, partial).filter { it.label !in elementItemLabels } else emptySequence())
         + (if (elementItemList.isEmpty()) keywordCompletionItems(partial) else emptySequence())
     )
     val itemList = items
@@ -81,7 +81,8 @@ fun completions(file: CompiledFile, cursor: Int, index: SymbolIndex, config: Com
 }
 
 /** Finds completions in the global symbol index, for potentially unimported symbols. */
-private fun indexCompletionItems(parsedFile: KtFile, receiver: KtExpression?, index: SymbolIndex, partial: String): Sequence<CompletionItem> {
+private fun indexCompletionItems(file: CompiledFile, cursor: Int, receiver: KtExpression?, index: SymbolIndex, partial: String): Sequence<CompletionItem> {
+    val parsedFile = file.parse
     val imports = parsedFile.importDirectives
     // TODO: Deal with alias imports
     val wildcardPackages = imports
@@ -92,6 +93,8 @@ private fun indexCompletionItems(parsedFile: KtFile, receiver: KtExpression?, in
     val importedNames = imports
         .mapNotNull { it.importedFqName?.shortName() }
         .toSet()
+    val receiverType = receiver?.let { expr -> file.scopeAtPoint(cursor)?.let { file.typeOfExpression(expr, it) } }
+    val receiverTypeFqName = receiverType?.constructor?.declarationDescriptor?.fqNameSafe
 
     return index
         .query(partial, limit = MAX_COMPLETION_ITEMS)
@@ -104,6 +107,7 @@ private fun indexCompletionItems(parsedFile: KtFile, receiver: KtExpression?, in
             || it.visibility == Symbol.Visibility.PROTECTED
             || it.visibility == Symbol.Visibility.INTERNAL
         }
+        .filter { receiverTypeFqName == it.extensionReceiverType } // if both are null, it's not an extension
         .map { CompletionItem().apply {
             label = it.fqName.shortName().toString()
             kind = when (it.kind) {

server/src/main/kotlin/org/javacs/kt/index/ExtractSymbolExtensionReceiverType.kt

@@ -0,0 +1,14 @@
+package org.javacs.kt.index
+
+import org.jetbrains.kotlin.descriptors.*
+import org.jetbrains.kotlin.descriptors.impl.DeclarationDescriptorVisitorEmptyBodies
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
+
+object ExtractSymbolExtensionReceiverType : DeclarationDescriptorVisitorEmptyBodies<FqName?, Unit>() {
+    private fun convert(desc: ReceiverParameterDescriptor): FqName? = desc.value.type.constructor.declarationDescriptor?.fqNameSafe
+
+    override fun visitFunctionDescriptor(desc: FunctionDescriptor, nothing: Unit?) = desc.extensionReceiverParameter?.let(this::convert)
+
+    override fun visitVariableDescriptor(desc: VariableDescriptor, nothing: Unit?) = desc.extensionReceiverParameter?.let(this::convert)
+}

server/src/main/kotlin/org/javacs/kt/index/Symbol.kt

@@ -6,7 +6,8 @@ data class Symbol(
     // TODO: Store location (e.g. using a URI)
     val fqName: FqName,
     val kind: Kind,
-    val visibility: Visibility
+    val visibility: Visibility,
+    val extensionReceiverType: FqName?
 ) {
     enum class Kind(val rawValue: Int) {
         CLASS(0),

server/src/main/kotlin/org/javacs/kt/index/SymbolIndex.kt

@@ -8,6 +8,7 @@ import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
 import org.jetbrains.kotlin.resolve.scopes.MemberScope
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.psi2ir.intermediate.extensionReceiverType
 import org.javacs.kt.LOG
 import org.javacs.kt.progress.Progress
 import org.jetbrains.exposed.sql.Database
@@ -19,6 +20,7 @@ private object Symbols : Table() {
     val shortName = varchar("shortname", length = 80)
     val kind = integer("kind")
     val visibility = integer("visibility")
+    val extensionReceiverType = varchar("extensionreceivertype", length = 255).nullable()
 }
 
 /**
@@ -62,6 +64,7 @@ class SymbolIndex {
                                 it[shortName] = fqn.shortName().toString()
                                 it[kind] = descriptor.accept(ExtractSymbolKind, Unit).rawValue
                                 it[visibility] = descriptor.accept(ExtractSymbolVisibility, Unit).rawValue
+                                it[extensionReceiverType] = descriptor.accept(ExtractSymbolExtensionReceiverType, Unit)?.toString()
                             }
                         }
                     }
@@ -86,7 +89,8 @@ class SymbolIndex {
             .map { Symbol(
                 fqName = FqName(it[Symbols.fqName]),
                 kind = Symbol.Kind.fromRaw(it[Symbols.kind]),
-                visibility = Symbol.Visibility.fromRaw(it[Symbols.visibility])
+                visibility = Symbol.Visibility.fromRaw(it[Symbols.visibility]),
+                extensionReceiverType = it[Symbols.extensionReceiverType]?.let(::FqName)
             ) }
     }