improved type inference in plugin

lyng-idea/src/main/kotlin/net/sergeych/lyng/idea/annotators/LyngExternalAnnotator.kt

@@ -117,9 +117,18 @@ class LyngExternalAnnotator : ExternalAnnotator<LyngExternalAnnotator.Input, Lyn
         }
 
         // Parameters
-        mini.declarations.filterIsInstance<MiniFunDecl>().forEach { fn ->
-            if (fn.nameStart.source != source) return@forEach
-            fn.params.forEach { p -> putName(p.nameStart, p.name, LyngHighlighterColors.PARAMETER) }
+        fun addParams(params: List<MiniParam>) {
+            params.forEach { p ->
+                if (p.nameStart.source == source)
+                    putName(p.nameStart, p.name, LyngHighlighterColors.PARAMETER)
+            }
+        }
+        mini.declarations.forEach { d ->
+            when (d) {
+                is MiniFunDecl -> addParams(d.params)
+                is MiniClassDecl -> d.members.filterIsInstance<MiniMemberFunDecl>().forEach { addParams(it.params) }
+                else -> {}
+            }
         }
 
         // Type name segments (including generics base & args)
@@ -146,8 +155,8 @@ class LyngExternalAnnotator : ExternalAnnotator<LyngExternalAnnotator.Input, Lyn
                 null -> {}
             }
         }
-        mini.declarations.forEach { d ->
-            if (d.nameStart.source != source) return@forEach
+        fun addDeclTypeSegments(d: MiniDecl) {
+            if (d.nameStart.source != source) return
             when (d) {
                 is MiniFunDecl -> {
                     addTypeSegments(d.returnType)
@@ -161,10 +170,23 @@ class LyngExternalAnnotator : ExternalAnnotator<LyngExternalAnnotator.Input, Lyn
                 is MiniClassDecl -> {
                     d.ctorFields.forEach { addTypeSegments(it.type) }
                     d.classFields.forEach { addTypeSegments(it.type) }
+                    for (m in d.members) {
+                        when (m) {
+                            is MiniMemberFunDecl -> {
+                                addTypeSegments(m.returnType)
+                                m.params.forEach { addTypeSegments(it.type) }
+                            }
+                            is MiniMemberValDecl -> {
+                                addTypeSegments(m.type)
+                            }
+                            else -> {}
+                        }
+                    }
                 }
                 is MiniEnumDecl -> {}
             }
         }
+        mini.declarations.forEach { d -> addDeclTypeSegments(d) }
 
         ProgressManager.checkCanceled()
 
@@ -212,6 +234,13 @@ class LyngExternalAnnotator : ExternalAnnotator<LyngExternalAnnotator.Input, Lyn
                         if (d.mutable) LyngHighlighterColors.VARIABLE else LyngHighlighterColors.VALUE
 
                     is MiniFunDecl -> d.params.forEach { p -> nameRole[p.name] = LyngHighlighterColors.PARAMETER }
+                    is MiniClassDecl -> {
+                        d.members.forEach { m ->
+                            if (m is MiniMemberFunDecl) {
+                                m.params.forEach { p -> nameRole[p.name] = LyngHighlighterColors.PARAMETER }
+                            }
+                        }
+                    }
                     else -> {}
                 }
             }

lyng-idea/src/main/kotlin/net/sergeych/lyng/idea/completion/LyngCompletionContributor.kt

@@ -116,7 +116,7 @@ class LyngCompletionContributor : CompletionContributor() {
             if (memberDotPos != null && engineItems.isEmpty()) {
                 if (DEBUG_COMPLETION) log.info("[LYNG_DEBUG] Fallback: engine returned 0 in member context; trying local inference")
                 // Build imported modules from text (lenient) + stdlib; avoid heavy MiniAst here
-                val fromText = extractImportsFromText(text)
+                val fromText = DocLookupUtils.extractImportsFromText(text)
                 val imported = LinkedHashSet<String>().apply {
                     fromText.forEach { add(it) }
                     add("lyng.stdlib")
@@ -176,7 +176,7 @@ class LyngCompletionContributor : CompletionContributor() {
             // In member context, ensure stdlib extension-like methods (e.g., String.re) are present
             if (memberDotPos != null) {
                 val existing = engineItems.map { it.name }.toMutableSet()
-                val fromText = extractImportsFromText(text)
+                val fromText = DocLookupUtils.extractImportsFromText(text)
                 val imported = LinkedHashSet<String>().apply {
                     fromText.forEach { add(it) }
                     add("lyng.stdlib")
@@ -249,7 +249,7 @@ class LyngCompletionContributor : CompletionContributor() {
             // If in member context and engine items are suspiciously sparse, try to enrich via local inference + offerMembers
             if (memberDotPos != null && engineItems.size < 3) {
                 if (DEBUG_COMPLETION) log.info("[LYNG_DEBUG] Engine produced only ${engineItems.size} items in member context — trying enrichment")
-                val fromText = extractImportsFromText(text)
+                val fromText = DocLookupUtils.extractImportsFromText(text)
                 val imported = LinkedHashSet<String>().apply {
                     fromText.forEach { add(it) }
                     add("lyng.stdlib")
@@ -410,13 +410,18 @@ class LyngCompletionContributor : CompletionContributor() {
                 for (name in keys) {
                     val list = map[name] ?: continue
                     // Choose a representative for display:
-                    // 1) Prefer a method with a known return type
-                    // 2) Else any method
-                    // 3) Else the first variant
+                    // 1) Prefer a method with return type AND parameters
+                    // 2) Prefer a method with parameters
+                    // 3) Prefer a method with return type
+                    // 4) Else any method
+                    // 5) Else the first variant
                     val rep =
-                        list.asSequence()
-                            .filterIsInstance<MiniMemberFunDecl>()
-                            .firstOrNull { it.returnType != null }
+                        list.asSequence().filterIsInstance<MiniMemberFunDecl>()
+                            .firstOrNull { it.returnType != null && it.params.isNotEmpty() }
+                            ?: list.asSequence().filterIsInstance<MiniMemberFunDecl>()
+                                .firstOrNull { it.params.isNotEmpty() }
+                            ?: list.asSequence().filterIsInstance<MiniMemberFunDecl>()
+                                .firstOrNull { it.returnType != null }
                             ?: list.firstOrNull { it is MiniMemberFunDecl }
                             ?: list.first()
                     when (rep) {
@@ -603,32 +608,10 @@ class LyngCompletionContributor : CompletionContributor() {
             }
         }
 
-        // Lenient textual import extractor (duplicated from QuickDoc privately)
-        private fun extractImportsFromText(text: String): List<String> {
-            val result = LinkedHashSet<String>()
-            val re = Regex("^\\s*import\\s+([a-zA-Z_][a-zA-Z0-9_]*(?:\\.[a-zA-Z_][a-zA-Z0-9_]*)*)", RegexOption.MULTILINE)
-            re.findAll(text).forEach { m ->
-                val raw = m.groupValues.getOrNull(1)?.trim().orEmpty()
-                if (raw.isNotEmpty()) {
-                    val canon = if (raw.startsWith("lyng.")) raw else "lyng.$raw"
-                    result.add(canon)
-                }
-            }
-            return result.toList()
-        }
 
         private fun typeOf(t: MiniTypeRef?): String {
-            return when (t) {
-                null -> ""
-                is MiniTypeName -> t.segments.lastOrNull()?.name?.let { ": $it" } ?: ""
-                is MiniGenericType -> {
-                    val base = typeOf(t.base).removePrefix(": ")
-                    val args = t.args.joinToString(",") { typeOf(it).removePrefix(": ") }
-                    ": ${base}<${args}>"
-                }
-                is MiniFunctionType -> ": (fn)"
-                is MiniTypeVar -> ": ${t.name}"
-            }
+            val s = DocLookupUtils.typeOf(t)
+            return if (s.isEmpty()) "" else ": $s"
         }
     }
 }

lyng-idea/src/main/kotlin/net/sergeych/lyng/idea/docs/LyngDocumentationProvider.kt

@@ -561,16 +561,9 @@ class LyngDocumentationProvider : AbstractDocumentationProvider() {
         return sb.toString()
     }
 
-    private fun typeOf(t: MiniTypeRef?): String = when (t) {
-        is MiniTypeName -> ": ${t.segments.joinToString(".") { it.name }}${if (t.nullable) "?" else ""}"
-        is MiniGenericType -> {
-            val base = typeOf(t.base).removePrefix(": ")
-            val args = t.args.joinToString(", ") { typeOf(it).removePrefix(": ") }
-            ": ${base}<${args}>${if (t.nullable) "?" else ""}"
-        }
-        is MiniFunctionType -> ": (..) -> ..${if (t.nullable) "?" else ""}"
-        is MiniTypeVar -> ": ${t.name}${if (t.nullable) "?" else ""}"
-        null -> ": Object?"
+    private fun typeOf(t: MiniTypeRef?): String {
+        val s = DocLookupUtils.typeOf(t)
+        return if (s.isEmpty()) (if (t == null) ": Object?" else "") else ": $s"
     }
 
     private fun signatureOf(fn: MiniFunDecl): String {