From 91459bbd6a3d4724dab1b1cde2429be22db95296 Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Tue, 12 Sep 2017 14:51:45 +0200
Subject: [PATCH] Enable caching of converted UElements for Kotlin

---
 .../uast/kotlin/KotlinUastLanguagePlugin.kt   | 30 ++++++++++++-------
 1 file changed, 20 insertions(+), 10 deletions(-)

diff --git a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinUastLanguagePlugin.kt b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinUastLanguagePlugin.kt
index f88af3a9d7d..cbc062c8d4a 100644
--- a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinUastLanguagePlugin.kt
+++ b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinUastLanguagePlugin.kt
@@ -40,6 +40,7 @@ import org.jetbrains.uast.kotlin.declarations.KotlinUMethod
 import org.jetbrains.uast.kotlin.expressions.*
 import org.jetbrains.uast.kotlin.psi.UastKotlinPsiParameter
 import org.jetbrains.uast.kotlin.psi.UastKotlinPsiVariable
+import java.lang.ref.WeakReference
 
 interface KotlinUastBindingContextProviderService {
     fun getBindingContext(element: KtElement): BindingContext
@@ -59,8 +60,7 @@ class KotlinUastLanguagePlugin : UastLanguagePlugin {
     }
 
     override fun convertElement(element: PsiElement, parent: UElement?, requiredType: Class<out UElement>?): UElement? {
-        return convertDeclaration(element, parent.toCallback(), requiredType)
-               ?: KotlinConverter.convertPsiElement(element, parent.toCallback(), requiredType)
+        return convertDeclarationOrElement(element, parent.toCallback(), requiredType)
     }
     
     override fun convertElementWithParent(element: PsiElement, requiredType: Class<out UElement>?): UElement? {
@@ -78,8 +78,24 @@ class KotlinUastLanguagePlugin : UastLanguagePlugin {
             else
                 return convertElementWithParent(parentUnwrapped, null)
         }
-        return convertDeclaration(element, parentCallback, requiredType)
-               ?: KotlinConverter.convertPsiElement(element, parentCallback, requiredType)
+        return convertDeclarationOrElement(element, parentCallback, requiredType)
+    }
+
+    private fun convertDeclarationOrElement(element: PsiElement, parentCallback: (() -> UElement?)?, requiredType: Class<out UElement>?): UElement? {
+        if (element is UElement) return element
+
+        if (element.isValid) {
+            element.getUserData(KOTLIN_CACHED_UELEMENT_KEY)?.get()?.let { cachedUElement ->
+                return if (requiredType == null || requiredType.isInstance(cachedUElement)) cachedUElement else null
+            }
+        }
+
+        val uElement = convertDeclaration(element, parentCallback, requiredType)
+                        ?: KotlinConverter.convertPsiElement(element, parentCallback, requiredType)
+        if (uElement != null) {
+            element.putUserData(KOTLIN_CACHED_UELEMENT_KEY, WeakReference(uElement))
+        }
+        return uElement
     }
 
     override fun getMethodCallExpression(
@@ -125,12 +141,6 @@ class KotlinUastLanguagePlugin : UastLanguagePlugin {
     private fun convertDeclaration(element: PsiElement,
                                    parentCallback: (() -> UElement?)?,
                                    requiredType: Class<out UElement>?): UElement? {
-        if (element is UElement) return element
-
-        if (element.isValid) element.getUserData(KOTLIN_CACHED_UELEMENT_KEY)?.let { ref ->
-            ref.get()?.let { return it }
-        }
-
         fun <P : PsiElement> build(ctor: (P, UElement?) -> UElement): () -> UElement? {
             return fun(): UElement? {
                 val parent = if (parentCallback == null) null else (parentCallback() ?: return null)