Infix code completion

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

@@ -47,6 +47,7 @@ import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.kotlin.types.typeUtil.replaceArgumentsWithStarProjections
 import org.jetbrains.kotlin.types.checker.KotlinTypeChecker
+import org.jetbrains.kotlin.utils.addToStdlib.applyIf
 import java.util.concurrent.TimeUnit
 
 // The maximum number of completion items
@@ -159,12 +160,14 @@ data class ElementCompletionItems(val items: Sequence<CompletionItem>, val eleme
 
 /** Finds completions based on the element around the user's cursor. */
 private fun elementCompletionItems(file: CompiledFile, cursor: Int, config: CompletionConfiguration, partial: String): ElementCompletionItems {
-    val surroundingElement = completableElement(file, cursor) ?: return ElementCompletionItems(emptySequence())
-    val completions = elementCompletions(file, cursor, surroundingElement)
-
-    val matchesName = completions.filter { containsCharactersInOrder(name(it), partial, caseSensitive = false) }
-    val sorted = matchesName.takeIf { partial.length >= MIN_SORT_LENGTH }?.sortedBy { stringDistance(name(it), partial) }
-        ?: matchesName.sortedBy { if (name(it).startsWith(partial)) 0 else 1 }
+    val (surroundingElement, isGlobal) = completableElement(file, cursor) ?: return ElementCompletionItems(emptySequence())
+    val completions = elementCompletions(file, cursor, surroundingElement, isGlobal)
+        .applyIf(isGlobal) { filter { declarationIsInfix(it) } }
+        .applyIf(surroundingElement.endOffset == cursor) {
+            filter { containsCharactersInOrder(name(it), partial, caseSensitive = false) }
+        }
+    val sorted = completions.takeIf { partial.length >= MIN_SORT_LENGTH }?.sortedBy { stringDistance(name(it), partial) }
+        ?: completions.sortedBy { if (name(it).startsWith(partial)) 0 else 1 }
     val visible = sorted.filter(isVisible(file, cursor))
 
     return ElementCompletionItems(visible.map { completionItem(it, surroundingElement, file, config) }, surroundingElement)
@@ -230,28 +233,59 @@ private fun isSetter(d: DeclarationDescriptor): Boolean =
         d.name.identifier.matches(Regex("set[A-Z]\\w+")) &&
         d.valueParameters.size == 1
 
-private fun completableElement(file: CompiledFile, cursor: Int): KtElement? {
-    val el = file.parseAtPoint(cursor - 1) ?: return null
-            // import x.y.?
-    return el.findParent<KtImportDirective>()
-            // package x.y.?
-            ?: el.findParent<KtPackageDirective>()
-            // :?
-            ?: el as? KtUserType
-            ?: el.parent as? KtTypeElement
-            // .?
-            ?: el as? KtQualifiedExpression
-            ?: el.parent as? KtQualifiedExpression
-            // something::?
-            ?: el as? KtCallableReferenceExpression
-            ?: el.parent as? KtCallableReferenceExpression
-            // something.foo() with cursor in the method
-            ?: el.parent?.parent as? KtQualifiedExpression
-            // ?
-            ?: el as? KtNameReferenceExpression
+private fun isGlobalCall(el: KtElement) = el is KtBlockExpression || el is KtClassBody || el.parent is KtBinaryExpression
+
+private fun asGlobalCompletable(file: CompiledFile, cursor: Int, el: KtElement): KtElement? {
+    val psi =  file.parse.findElementAt(cursor) ?: return null
+    val element = when (val e = psi.getPrevSiblingIgnoringWhitespace() ?: psi.parent) {
+        is KtProperty -> e.children.lastOrNull()
+        is KtBinaryExpression -> el
+        else -> e
+    }
+    return element as? KtReferenceExpression
+        ?: element as? KtQualifiedExpression
+        ?: element as? KtConstantExpression
+}
+
+private fun KtElement.asKtClass(): KtElement? {
+    return this.findParent<KtImportDirective>() // import x.y.?
+        // package x.y.?
+        ?: this.findParent<KtPackageDirective>()
+        // :?
+        ?: this as? KtUserType
+        ?: this.parent as? KtTypeElement
+        // .?
+        ?: this as? KtQualifiedExpression
+        ?: this.parent as? KtQualifiedExpression
+        // something::?
+        ?: this as? KtCallableReferenceExpression
+        ?: this.parent as? KtCallableReferenceExpression
+        // something.foo() with cursor in the method
+        ?: this.parent?.parent as? KtQualifiedExpression
+        // ?
+        ?: this as? KtNameReferenceExpression
+        // x ? y (infix)
+        ?: this.parent as? KtBinaryExpression
+        // x()
+        ?: this as? KtCallExpression
+        // x (constant)
+        ?: this as? KtConstantExpression
+}
+
+private fun completableElement(file: CompiledFile, cursor: Int): Pair<KtElement, Boolean>? {
+    val parsed = file.parseAtPoint(cursor - 1) ?: return null
+    val asGlobal = isGlobalCall(parsed)
+    val el = (
+            if (asGlobal) asGlobalCompletable(file, cursor, parsed) else null
+     ) ?: parsed
+
+    return el.asKtClass()?.let {
+        Pair(it, asGlobal)
+    }
 }
 
-private fun elementCompletions(file: CompiledFile, cursor: Int, surroundingElement: KtElement): Sequence<DeclarationDescriptor> {
+@Suppress("LongMethod", "ReturnCount", "CyclomaticComplexMethod")
+private fun elementCompletions(file: CompiledFile, cursor: Int, surroundingElement: KtElement, infixCall: Boolean): Sequence<DeclarationDescriptor> {
     return when (surroundingElement) {
         // import x.y.?
         is KtImportDirective -> {
@@ -298,7 +332,8 @@ private fun elementCompletions(file: CompiledFile, cursor: Int, surroundingEleme
         // .?
         is KtQualifiedExpression -> {
             LOG.info("Completing member expression '{}'", surroundingElement.text)
-            completeMembers(file, cursor, surroundingElement.receiverExpression, surroundingElement is KtSafeQualifiedExpression)
+            val exp = if (infixCall) surroundingElement else surroundingElement.receiverExpression
+            completeMembers(file, cursor, exp, surroundingElement is KtSafeQualifiedExpression)
         }
         is KtCallableReferenceExpression -> {
             // something::?
@@ -316,8 +351,21 @@ private fun elementCompletions(file: CompiledFile, cursor: Int, surroundingEleme
         // ?
         is KtNameReferenceExpression -> {
             LOG.info("Completing identifier '{}'", surroundingElement.text)
-            val scope = file.scopeAtPoint(surroundingElement.startOffset) ?: return noResult("No scope at ${file.describePosition(cursor)}", emptySequence())
-            identifiers(scope)
+            if (infixCall) {
+                completeMembers(file, surroundingElement.startOffset, surroundingElement)
+            } else {
+                val scope = file.scopeAtPoint(surroundingElement.startOffset) ?: return noResult("No scope at ${file.describePosition(cursor)}", emptySequence())
+                identifiers(scope)
+            }
+        }
+        // x ? y (infix)
+        is KtBinaryExpression -> {
+            if (surroundingElement.operationToken == KtTokens.IDENTIFIER) {
+                completeMembers(file, cursor, surroundingElement.left!!)
+            } else emptySequence()
+        }
+        is KtCallExpression, is KtConstantExpression -> {
+            completeMembers(file, cursor, surroundingElement as KtExpression)
         }
         else -> {
             LOG.info("{} {} didn't look like a type, a member, or an identifier", surroundingElement::class.simpleName, surroundingElement.text)
@@ -370,6 +418,11 @@ private fun ClassDescriptor.getDescriptors(): Sequence<DeclarationDescriptor> {
 
 }
 
+private fun declarationIsInfix(declaration: DeclarationDescriptor): Boolean {
+    val functionDescriptor = declaration as? FunctionDescriptor ?: return false
+    return functionDescriptor.isInfix
+}
+
 private fun isCompanionOfEnum(kotlinType: KotlinType): Boolean {
     val classDescriptor = TypeUtils.getClassDescriptor(kotlinType)
     val isCompanion = DescriptorUtils.isCompanionObject(classDescriptor)

server/src/test/kotlin/org/javacs/kt/CompletionsTest.kt

@@ -65,6 +65,48 @@ class InstanceMemberTest : SingleFileTestFixture("completions", "InstanceMember.
     }
 }
 
+class InfixMethodTest : SingleFileTestFixture("completions", "InfixFunctions.kt") {
+    @Test fun `complete member function`() {
+        val completions = languageServer.textDocumentService.completion(completionParams(file, 10, 11)).get().right!!
+        val labels = completions.items.map { it.label }
+
+        assertThat(labels, hasItem(startsWith("cmpB")))
+        assertThat(labels, not(hasItem(startsWith("cmpA"))))
+    }
+
+    @Test fun `includes completion for stdlib`() {
+        val completions = languageServer.textDocumentService.completion(completionParams(file, 15, 10)).get().right!!
+        val labels = completions.items.map { it.label }
+
+        assertThat(labels, hasSize(2))
+        assertThat(labels, hasItem(startsWith("and")))
+        assertThat(labels, hasItem("andTo"))
+    }
+
+    @Test fun `complete extension function`() {
+        val completions = languageServer.textDocumentService.completion(completionParams(file, 19, 9)).get().right!!
+        val labels = completions.items.map { it.label }
+
+        assertThat(labels, hasItem("funcA"))
+    }
+
+    @Test fun `complete global scope`() {
+        val completions = languageServer.textDocumentService.completion(completionParams(file, 26, 11)).get().right!!
+        val labels = completions.items.map { it.label }
+
+        assertThat(labels, hasSize(3))
+        assertThat(labels, hasItem("ord"))
+    }
+
+    @Test fun `has global completions for binary expression`() {
+        val completions = languageServer.textDocumentService.completion(completionParams(file, 30, 7)).get().right!!
+        val labels = completions.items.map { it.label }
+
+        assertThat(labels, hasItem("funcA"))
+        assertThat(labels.filter { it.startsWith("and") }, hasSize(2))
+    }
+}
+
 class InstanceMembersJava : SingleFileTestFixture("completions", "InstanceMembersJava.kt") {
     @Test fun `convert getFileName to fileName`() {
         val completions = languageServer.textDocumentService.completion(completionParams(file, 4, 14)).get().right!!

server/src/test/resources/completions/InfixFunctions.kt

@@ -0,0 +1,32 @@
+class Q {
+    fun cmpA(): Double = 1.0
+    infix fun cmpB(x: Int): Int { return 1 }
+}
+
+infix fun Int.funcA(x: Int): Boolean = x == this
+infix fun Int.andTo(v: Int) = v
+
+private fun memberFunc() {
+    Q() cm
+}
+
+private fun stdlibFunc() {
+    val v = 1
+    v and
+}
+
+private fun extensionFunc() {
+    2 fu 3
+}
+
+enum class FOO { T, U }
+infix fun FOO.ord(n: Int) = this.ordinal == n
+
+private fun globalEnumFunc() {
+    FOO.U
+}
+
+private fun globalBinaryExpFunc() {
+    4  4
+}
+