Better Statistician for completion

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/CompletionSession.kt

@@ -18,6 +18,7 @@ package org.jetbrains.kotlin.idea.completion
 
 import com.intellij.codeInsight.completion.*
 import com.intellij.codeInsight.completion.impl.CamelHumpMatcher
+import com.intellij.codeInsight.lookup.LookupElementDecorator
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.patterns.ElementPattern
@@ -32,6 +33,7 @@ import org.jetbrains.kotlin.idea.caches.resolve.getResolveScope
 import org.jetbrains.kotlin.idea.codeInsight.DescriptorToSourceUtilsIde
 import org.jetbrains.kotlin.idea.codeInsight.ReferenceVariantsHelper
 import org.jetbrains.kotlin.idea.core.*
+import org.jetbrains.kotlin.idea.imports.importableFqName
 import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.idea.resolve.ResolutionFacade
 import org.jetbrains.kotlin.idea.resolve.frontendService
@@ -217,6 +219,16 @@ abstract class CompletionSession(protected val configuration: CompletionSessionC
     }
 
     public fun complete(): Boolean {
+        val statisticsContext = calcContextForStatisticsInfo()
+        if (statisticsContext != null) {
+            collector.addLookupElementPostProcessor { lookupElement ->
+                // we should put data into the original element because of DecoratorCompletionStatistician
+                val unwrapped = sequence(lookupElement) { (it as? LookupElementDecorator<*>)?.delegate }.last()
+                unwrapped.putUserData(STATISTICS_INFO_CONTEXT_KEY, statisticsContext)
+                lookupElement
+            }
+        }
+
         doComplete()
         flushToResultSet()
         return !collector.isResultEmpty
@@ -242,6 +254,26 @@ abstract class CompletionSession(protected val configuration: CompletionSessionC
         return sorter
     }
 
+    protected fun calcContextForStatisticsInfo(): String? {
+        if (expectedInfos.isEmpty()) return null
+
+        var context = expectedInfos
+                .map { it.fuzzyType?.type?.constructor?.declarationDescriptor?.importableFqName }
+                .filterNotNull()
+                .singleOrNull()
+                ?.let { "expectedType=$it" }
+
+        if (context == null) {
+            context = expectedInfos
+                    .map { it.expectedName }
+                    .filterNotNull()
+                    .singleOrNull()
+                    ?.let { "expectedName=$it" }
+        }
+
+        return context
+    }
+
     protected val referenceVariants: Collection<DeclarationDescriptor> by lazy {
         if (descriptorKindFilter != null) {
             referenceVariantsHelper.getReferenceVariants(

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/CompletionUtils.kt

@@ -61,6 +61,8 @@ fun LookupElement.assignPriority(priority: ItemPriority): LookupElement {
     return this
 }
 
+val STATISTICS_INFO_CONTEXT_KEY = Key<String>("STATISTICS_INFO_CONTEXT_KEY")
+
 val NOT_IMPORTED_KEY = Key<Unit>("NOT_IMPORTED_KEY")
 
 fun LookupElement.suppressAutoInsertion() = AutoCompletionPolicy.NEVER_AUTOCOMPLETE.applyPolicy(this)

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/Statisticians.kt

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.idea.completion
 import com.intellij.codeInsight.completion.CompletionLocation
 import com.intellij.codeInsight.completion.CompletionStatistician
 import com.intellij.codeInsight.lookup.LookupElement
+import com.intellij.codeInsight.lookup.LookupElementDecorator
 import com.intellij.psi.PsiElement
 import com.intellij.psi.statistics.StatisticsInfo
 import com.intellij.psi.util.ProximityLocation
@@ -38,12 +39,15 @@ class KotlinCompletionStatistician : CompletionStatistician() {
     override fun serialize(element: LookupElement, location: CompletionLocation): StatisticsInfo? {
         val o = (element.`object` as? DeclarationLookupObject) ?: return null
 
+        assert(element !is LookupElementDecorator<*>, "LookupElementDecorator's should be unwrapped by DecoratorCompletionStatistician")
+        val context = element.getUserData(STATISTICS_INFO_CONTEXT_KEY) ?: ""
+
         if (o.descriptor != null) {
-            return KotlinStatisticsInfo.forDescriptor(o.descriptor!!)
+            return KotlinStatisticsInfo.forDescriptor(o.descriptor!!, context)
         }
         else {
             val fqName = o.importableFqName ?: return StatisticsInfo.EMPTY
-            return StatisticsInfo("", fqName.asString())
+            return StatisticsInfo(context, fqName.asString())
         }
     }
 }
@@ -67,7 +71,7 @@ object KotlinStatisticsInfo {
         parameterNameRenderingPolicy = ParameterNameRenderingPolicy.NONE
     }
 
-    fun forDescriptor(descriptor: DeclarationDescriptor): StatisticsInfo {
+    fun forDescriptor(descriptor: DeclarationDescriptor, context: String = ""): StatisticsInfo {
         if (descriptor is ClassDescriptor) {
             return descriptor.importableFqName?.let { StatisticsInfo("", it.asString()) } ?: StatisticsInfo.EMPTY
         }
@@ -80,6 +84,6 @@ object KotlinStatisticsInfo {
                                   else -> null
                               }  ?: return StatisticsInfo.EMPTY
         val signature = SIGNATURE_RENDERER.render(descriptor)
-        return StatisticsInfo("", "$containerFqName###$signature")
+        return StatisticsInfo(context, "$containerFqName###$signature")
     }
 }