Overriding LookupElement.getPsiElement() for all our LookupElement's because some IDEA features may use it

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

@@ -223,8 +223,7 @@ abstract class CompletionSession(protected val configuration: CompletionSessionC
         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.putUserDataDeep(STATISTICS_INFO_CONTEXT_KEY, statisticsContext)
                 lookupElement
             }
         }

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

@@ -48,6 +48,26 @@ import org.jetbrains.kotlin.types.typeUtil.TypeNullability
 import org.jetbrains.kotlin.types.typeUtil.nullability
 import java.util.*
 
+@tailRecursive
+fun <T : Any> LookupElement.putUserDataDeep(key: Key<T>, value: T?) {
+    if (this is LookupElementDecorator<*>) {
+        getDelegate().putUserDataDeep(key, value)
+    }
+    else {
+        putUserData(key, value)
+    }
+}
+
+@tailRecursive
+fun <T : Any> LookupElement.getUserDataDeep(key: Key<T>): T? {
+    if (this is LookupElementDecorator<*>) {
+        return getDelegate().getUserDataDeep(key)
+    }
+    else {
+        return getUserData(key)
+    }
+}
+
 enum class ItemPriority {
     DEFAULT,
     BACKING_FIELD,

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

@@ -68,7 +68,7 @@ public class KotlinCompletionCharFilter() : CharFilter() {
         }
 
         if (!lookup.isSelectionTouched()) {
-            currentItem?.putUserData(JUST_TYPING_PREFIX, lookup.itemPattern(currentItem))
+            currentItem?.putUserDataDeep(JUST_TYPING_PREFIX, lookup.itemPattern(currentItem))
         }
 
         return when (c) {

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

@@ -28,6 +28,7 @@ import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.PropertyDescriptor
 import org.jetbrains.kotlin.idea.completion.handlers.*
+import org.jetbrains.kotlin.idea.core.completion.DeclarationLookupObject
 import org.jetbrains.kotlin.idea.resolve.ResolutionFacade
 import org.jetbrains.kotlin.types.JetType
 import java.util.*
@@ -196,6 +197,13 @@ class LookupElementsCollector(
             result = postProcessor(result)
         }
 
+        val psiElement = (result.`object` as? DeclarationLookupObject)?.psiElement
+        if (psiElement != null) {
+            result = object : LookupElementDecorator<LookupElement>(result) {
+                override fun getPsiElement() = psiElement
+            }
+        }
+
         elements.getOrPut(prefixMatcher) { ArrayList() }.add(result)
     }
 
@@ -203,7 +211,7 @@ class LookupElementsCollector(
     private fun isJustTyping(context: InsertionContext, element: LookupElement): Boolean {
         if (!completionParameters.isAutoPopup()) return false
         val insertedText = context.getDocument().getText(TextRange(context.getStartOffset(), context.getTailOffset()))
-        return insertedText == element.getUserData(KotlinCompletionCharFilter.JUST_TYPING_PREFIX)
+        return insertedText == element.getUserDataDeep(KotlinCompletionCharFilter.JUST_TYPING_PREFIX)
     }
 
     public fun addElements(elements: Iterable<LookupElement>, prefixMatcher: PrefixMatcher = defaultPrefixMatcher, notImported: Boolean = false) {

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

@@ -19,7 +19,6 @@ 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
@@ -39,8 +38,7 @@ 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) ?: ""
+        val context = element.getUserDataDeep(STATISTICS_INFO_CONTEXT_KEY) ?: ""
 
         if (o.descriptor != null) {
             return KotlinStatisticsInfo.forDescriptor(o.descriptor!!.original, context)