infer types in completion lookups for methods

src/main/kotlin/org/move/lang/core/completion/FuncSignature.kt

@@ -0,0 +1,50 @@
+package org.move.lang.core.completion
+
+import org.move.ide.presentation.text
+import org.move.lang.core.psi.MvFunction
+import org.move.lang.core.psi.ext.name
+import org.move.lang.core.psi.parameters
+import org.move.lang.core.types.infer.TypeFoldable
+import org.move.lang.core.types.infer.TypeFolder
+import org.move.lang.core.types.infer.TypeVisitor
+import org.move.lang.core.types.ty.Ty
+import org.move.lang.core.types.ty.TyUnit
+
+data class FuncSignature(
+    private val params: Map<String, Ty>,
+    private val retType: Ty,
+): TypeFoldable<FuncSignature> {
+
+    override fun innerFoldWith(folder: TypeFolder): FuncSignature {
+        return FuncSignature(
+            params = params.mapValues { (_, it) -> folder.fold(it) },
+            retType = folder.fold(retType)
+        )
+    }
+
+    override fun innerVisitWith(visitor: TypeVisitor): Boolean =
+        params.values.any { visitor(it) } || visitor(retType)
+
+    fun paramsText(): String {
+        return params.entries
+            .joinToString(", ", prefix = "(", postfix = ")") { (paramName, paramTy) ->
+                "$paramName: ${paramTy.text(false)}"
+            }
+    }
+
+    fun retTypeText(): String = retType.text(false)
+
+    fun retTypeSuffix(): String {
+        return if (retType is TyUnit) "" else ": ${retTypeText()}"
+    }
+
+    companion object {
+        fun fromFunction(function: MvFunction, msl: Boolean): FuncSignature {
+            val declaredType = function.declaredType(msl)
+            val params = function.parameters.zip(declaredType.paramTypes)
+                .associate { (param, paramTy) -> Pair(param.name, paramTy) }
+            val retType = declaredType.retType
+            return FuncSignature(params, retType)
+        }
+    }
+}

src/main/kotlin/org/move/lang/core/completion/LookupElements.kt

@@ -12,8 +12,7 @@ import org.move.ide.presentation.text
 import org.move.lang.core.psi.*
 import org.move.lang.core.psi.ext.*
 import org.move.lang.core.resolve.ContextScopeInfo
-import org.move.lang.core.types.infer.inference
-import org.move.lang.core.types.infer.loweredType
+import org.move.lang.core.types.infer.*
 import org.move.lang.core.types.ty.Ty
 import org.move.lang.core.types.ty.TyUnknown
 
@@ -62,13 +61,26 @@ fun MvModule.createSelfLookup(): LookupElement {
         .withBoldness(true)
 }
 
-fun MvNamedElement.getLookupElementBuilder(structAsType: Boolean = false): LookupElementBuilder {
+fun MvNamedElement.getLookupElementBuilder(
+    completionCtx: CompletionContext,
+    subst: Substitution = emptySubstitution,
+    structAsType: Boolean = false
+): LookupElementBuilder {
     val lookupElementBuilder = this.createLookupElementWithIcon()
+    val msl = completionCtx.isMsl()
     return when (this) {
-        is MvFunction -> lookupElementBuilder
-            .withTailText(this.signatureText)
-            .withTypeText(this.outerFileName)
-
+        is MvFunction -> {
+            val signature = FuncSignature.fromFunction(this, msl).substitute(subst)
+            if (completionCtx.contextElement is MvMethodOrField) {
+                lookupElementBuilder
+                    .withTailText(signature.paramsText())
+                    .withTypeText(signature.retTypeText())
+            } else {
+                lookupElementBuilder
+                    .withTailText("${signature.paramsText()}${signature.retTypeSuffix()}")
+                    .withTypeText(this.outerFileName)
+            }
+        }
         is MvSpecFunction -> lookupElementBuilder
             .withTailText(this.parameters.joinToSignature())
             .withTypeText(this.returnType?.type?.text ?: "()")
@@ -84,21 +96,23 @@ fun MvNamedElement.getLookupElementBuilder(structAsType: Boolean = false): Looku
                 .withTypeText(this.containingFile?.name)
         }
 
-        is MvStructField -> lookupElementBuilder
-            .withTypeText(this.typeAnnotation?.type?.text)
-
+        is MvStructField -> {
+            val fieldTy = this.type?.loweredType(msl)?.substitute(subst) ?: TyUnknown
+            lookupElementBuilder
+                .withTypeText(fieldTy.text(false))
+        }
         is MvConst -> {
-            val msl = this.isMslOnlyItem
+//            val msl = this.isMslOnlyItem
             val constTy = this.type?.loweredType(msl) ?: TyUnknown
             lookupElementBuilder
                 .withTypeText(constTy.text(true))
         }
 
         is MvBindingPat -> {
-            val msl = this.isMslOnlyItem
-            val inference = this.inference(msl)
+//            val msl = this.isMslOnlyItem
+            val bindingInference = this.inference(msl)
             // race condition sometimes happens, when file is too big, inference is not finished yet
-            val ty = inference?.getPatTypeOrUnknown(this) ?: TyUnknown
+            val ty = bindingInference?.getPatTypeOrUnknown(this) ?: TyUnknown
             lookupElementBuilder
                 .withTypeText(ty.text(true))
         }
@@ -114,19 +128,27 @@ data class CompletionContext(
     val contextElement: MvElement,
     val contextScopeInfo: ContextScopeInfo,
     val expectedTy: Ty? = null,
-)
+) {
+    fun isMsl(): Boolean = contextScopeInfo.isMslScope
+}
 
 
 fun MvNamedElement.createLookupElement(
     completionContext: CompletionContext,
+    subst: Substitution = emptySubstitution,
     structAsType: Boolean = false,
     priority: Double = DEFAULT_PRIORITY,
     insertHandler: InsertHandler<LookupElement> = DefaultInsertHandler(completionContext),
 ): LookupElement {
-    val builder = this.getLookupElementBuilder(structAsType)
-        .withInsertHandler(insertHandler)
-        .withPriority(priority)
-    val props = getLookupElementProperties(this, completionContext)
+    val builder =
+        this.getLookupElementBuilder(
+            completionContext,
+            subst = subst,
+            structAsType = structAsType
+        )
+            .withInsertHandler(insertHandler)
+            .withPriority(priority)
+    val props = getLookupElementProperties(this, subst, completionContext)
     return builder.toMvLookupElement(properties = props)
 }
 

src/main/kotlin/org/move/lang/core/completion/MvLookupElement.kt

@@ -57,63 +57,30 @@ data class LookupElementProperties(
     val typeHasAllRequiredAbilities: Boolean = false,
 )
 
-fun getLookupElementProperties(element: MvNamedElement, context: CompletionContext): LookupElementProperties {
+fun getLookupElementProperties(
+    element: MvNamedElement,
+    subst: Substitution,
+    context: CompletionContext
+): LookupElementProperties {
     var props = LookupElementProperties()
     val expectedTy = context.expectedTy
     if (expectedTy != null) {
-        val msl = context.contextScopeInfo.isMslScope
-        val contextElement = context.contextElement
-        val itemTy =
-            when {
-                contextElement is MvStructDotField -> {
-                    val receiverTy = contextElement.inferReceiverTy(msl)
-                    when (element) {
-                        is MvFunction -> {
-                            // infer method return type using known self type
-                            val declaredFuncTy =
-                                element.declaredType(msl).substitute(element.tyInfers) as TyFunction
-                            val declaredSelfTy = declaredFuncTy.paramTypes.first()
-                            val actualSelfTy =
-                                autoborrow(receiverTy, declaredSelfTy)
-                                    ?: error("unreachable, references always compatible")
-
-                            val inferenceCtx = InferenceContext(msl)
-                            inferenceCtx.combineTypes(actualSelfTy, expectedTy)
-                            inferenceCtx.resolveTypeVarsIfPossible(declaredFuncTy.retType)
-                        }
-                        is MvStructField -> {
-                            // infer field type using struct type
-                            val struct = element.structItem
-                            val tyVars = struct.tyInfers
-                            val declaredStructTy = struct.declaredType(msl).substitute(tyVars) as TyStruct
-
-                            val inferenceCtx = InferenceContext(msl)
-                            inferenceCtx.combineTypes(receiverTy, declaredStructTy)
-
-                            val fieldTy = element.type?.loweredType(msl)?.substitute(tyVars) ?: TyUnknown
-                            inferenceCtx.resolveTypeVarsIfPossible(fieldTy)
-                        }
-                        is MvFunctionLike -> {
-                            // TODO: spec functions inference
-                            element.declaredType(msl).retType
-                        }
-                        else -> TyUnknown
-                    }
-                }
-                else -> when (element) {
-                    is MvFunctionLike -> element.declaredType(msl).retType
-                    is MvStruct -> element.declaredType(msl)
-                    is MvConst -> element.type?.loweredType(msl) ?: TyUnknown
-                    is MvBindingPat -> {
-                        val inference = element.inference(msl)
-                        // sometimes type inference won't be able to catch up with the completion, and this line crashes,
-                        // so changing to infallible getPatTypeOrUnknown()
-                        inference?.getPatTypeOrUnknown(element) ?: TyUnknown
-                    }
-                    is MvStructField -> element.type?.loweredType(msl) ?: TyUnknown
-                    else -> TyUnknown
+        val msl = context.isMsl()
+        val declaredTy =
+            when (element) {
+                is MvFunctionLike -> element.declaredType(msl).retType
+                is MvStruct -> element.declaredType(msl)
+                is MvConst -> element.type?.loweredType(msl) ?: TyUnknown
+                is MvBindingPat -> {
+                    val inference = element.inference(msl)
+                    // sometimes type inference won't be able to catch up with the completion, and this line crashes,
+                    // so changing to infallible getPatTypeOrUnknown()
+                    inference?.getPatTypeOrUnknown(element) ?: TyUnknown
                 }
+                is MvStructField -> element.type?.loweredType(msl) ?: TyUnknown
+                else -> TyUnknown
             }
+        val itemTy = declaredTy.substitute(subst)
 
         // NOTE: it is required for the TyInfer.TyVar to always have a different underlying unification table
         val isCompat = isCompatible(expectedTy, itemTy, msl) && compatAbilities(expectedTy, itemTy, msl)

src/main/kotlin/org/move/lang/core/completion/providers/ImportsCompletionProvider.kt

@@ -1,5 +1,6 @@
 package org.move.lang.core.completion.providers
 
+import com.intellij.codeInsight.completion.BasicInsertHandler
 import com.intellij.codeInsight.completion.CompletionParameters
 import com.intellij.codeInsight.completion.CompletionResultSet
 import com.intellij.patterns.ElementPattern
@@ -59,8 +60,11 @@ object ImportsCompletionProvider: MvCompletionProvider() {
         val completionContext = CompletionContext(itemImport, contextScopeInfo)
         processModuleItems(referredModule, ns, vs, contextScopeInfo) {
             result.addElement(
-                it.element.createLookupElement(completionContext, structAsType = true)
-//                it.element.createLookupElement(BasicInsertHandler(), structAsType = true)
+                it.element.createLookupElement(
+                    completionContext,
+                    insertHandler = BasicInsertHandler(),
+                    structAsType = true
+                )
             )
             false
         }

src/main/kotlin/org/move/lang/core/completion/providers/MethodOrFieldCompletionProvider.kt

@@ -8,10 +8,17 @@ import com.intellij.psi.PsiElement
 import com.intellij.util.ProcessingContext
 import org.jetbrains.annotations.VisibleForTesting
 import org.move.lang.core.completion.*
-import org.move.lang.core.psi.*
+import org.move.lang.core.psi.MvNamedElement
 import org.move.lang.core.psi.ext.*
+import org.move.lang.core.psi.refItemScopes
+import org.move.lang.core.psi.tyInfers
 import org.move.lang.core.resolve.ContextScopeInfo
 import org.move.lang.core.resolve.letStmtScope
+import org.move.lang.core.types.infer.InferenceContext
+import org.move.lang.core.types.infer.Substitution
+import org.move.lang.core.types.infer.substitute
+import org.move.lang.core.types.ty.TyFunction
+import org.move.lang.core.types.ty.TyReference
 import org.move.lang.core.types.ty.TyStruct
 import org.move.lang.core.types.ty.knownOrNull
 import org.move.lang.core.withParent
@@ -50,20 +57,31 @@ object MethodOrFieldCompletionProvider: MvCompletionProvider() {
         if (structTy != null) {
             getFieldVariants(element, structTy, msl)
                 .forEach { (_, field) ->
-                    result.addElement(field.createLookupElement(ctx))
+                    val lookupElement = field.createLookupElement(
+                        ctx,
+                        subst = structTy.substitution
+                    )
+                    result.addElement(lookupElement)
                 }
         }
         getMethodVariants(element, receiverTy, msl)
             .forEach { (_, function) ->
-                val lookupElement = function.createLookupElement(ctx)
-//                val lookupProperties = lookupProperties(function, context = ctx)
-//                val builder = lookupElement
-//                    .withTailText("")
-//                    .withTypeText("")
-//                    .withInsertHandler(DefaultInsertHandler())
-//                val mvLookupElement = builder.withPriority(DEFAULT_PRIORITY).toMvLookupElement(lookupProperties)
+                // TODO: can instantiation of TyFunction every time be avoided here? is it slow?
+                val subst = function.tyInfers
+                val declaredFuncTy = function.declaredType(msl).substitute(subst) as TyFunction
+                val declaredSelfTy = declaredFuncTy.paramTypes.first()
+                val autoborrowedReceiverTy =
+                    TyReference.autoborrow(receiverTy, declaredSelfTy)
+                        ?: error("unreachable, references always compatible")
+
+                val inferenceCtx = InferenceContext(msl)
+                inferenceCtx.combineTypes(declaredSelfTy, autoborrowedReceiverTy)
+
+                val lookupElement = function.createLookupElement(
+                    ctx,
+                    subst = inferenceCtx.resolveTypeVarsIfPossible(subst)
+                )
                 result.addElement(lookupElement)
             }
     }
-
 }

src/main/kotlin/org/move/lang/core/completion/providers/ModulesCompletionProvider.kt

@@ -75,13 +75,9 @@ object ModulesCompletionProvider: MvCompletionProvider() {
                 candidate.element.createLookupElement(
                     completionCtx,
                     structAsType = Namespace.TYPE in importContext.namespaces,
-                    priority = UNIMPORTED_ITEM_PRIORITY
+                    priority = UNIMPORTED_ITEM_PRIORITY,
+                    insertHandler = ImportInsertHandler(parameters, candidate)
                 )
-//            candidate.element.createLookupElement(
-//                ImportInsertHandler(parameters, candidate),
-//                structAsType = Namespace.TYPE in importContext.namespaces,
-//                priority = UNIMPORTED_ITEM_PRIORITY,
-//            )
             result.addElement(lookupElement)
         }
     }

src/main/kotlin/org/move/lang/core/completion/providers/StructFieldsCompletionProvider.kt

@@ -39,7 +39,7 @@ object StructFieldsCompletionProvider: MvCompletionProvider() {
 
         if (element is MvBindingPat) element = element.parent as MvElement
 
-        val completionCtx = CompletionContext(element, ContextScopeInfo.msl())
+        val completionCtx = CompletionContext(element, ContextScopeInfo.default())
         when (element) {
             is MvStructPatField -> {
                 val structPat = element.structPat

src/main/kotlin/org/move/lang/core/resolve/NameResolution.kt

@@ -3,9 +3,11 @@ package org.move.lang.core.resolve
 import com.intellij.psi.search.GlobalSearchScope
 import org.move.lang.MoveFile
 import org.move.lang.core.psi.*
+import org.move.lang.core.psi.NamedItemScope.MAIN
 import org.move.lang.core.psi.NamedItemScope.VERIFY
 import org.move.lang.core.psi.ext.*
 import org.move.lang.core.resolve.LetStmtScope.EXPR_STMT
+import org.move.lang.core.resolve.LetStmtScope.NONE
 import org.move.lang.core.resolve.ref.MvReferenceElement
 import org.move.lang.core.resolve.ref.Namespace
 import org.move.lang.core.resolve.ref.Visibility
@@ -31,6 +33,7 @@ data class ContextScopeInfo(
 
     companion object {
         /// really does not affect anything, created just to allow creating CompletionContext everywhere
+        fun default(): ContextScopeInfo = ContextScopeInfo(setOf(MAIN), NONE)
         fun msl(): ContextScopeInfo = ContextScopeInfo(setOf(VERIFY), EXPR_STMT)
     }
 }

src/test/kotlin/org/move/lang/completion/lookups/BuiltInFunctionLookupTest.kt

@@ -44,7 +44,7 @@ class BuiltInFunctionLookupTest: MvTestBase() {
         val moduleElement = myFixture.findElementInEditor<MvModule>()
         val lookup =
             moduleElement.builtinFunctions().single { it.name == name }.let {
-                it.createLookupElement(CompletionContext(it, ContextScopeInfo.msl()))
+                it.createLookupElement(CompletionContext(it, ContextScopeInfo.default()))
             }
         checkLookupPresentation(
             lookup,

src/test/kotlin/org/move/lang/completion/lookups/LookupElementTest.kt

@@ -93,7 +93,7 @@ class LookupElementTest: MvTestBase() {
                   //^ 
             }
         }        
-    """, typeText = "u8", isBold = true
+    """, typeText = "u8"
     )
 
     fun `test generic method`() = checkMethodOrFieldProvider(
@@ -130,7 +130,7 @@ class LookupElementTest: MvTestBase() {
         val element = myFixture.findElementInEditor<T>() as? MvNamedElement
             ?: error("Marker `^` should point to the MvNamedElement")
 
-        val completionCtx = CompletionContext(element, ContextScopeInfo.msl())
+        val completionCtx = CompletionContext(element, ContextScopeInfo.default())
         val lookup = element.createLookupElement(completionCtx)
         checkLookupPresentation(
             lookup,