From e63838703ae29b6f2ad2583d52512b29af2748cd Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Fri, 15 Sep 2017 17:39:42 +0200
Subject: [PATCH] Consistently convert local functions to UAST

---
 .../uast/kotlin/KotlinUastLanguagePlugin.kt   | 24 +++++++++++++++----
 .../uast/kotlin/expressions/LocalFunction.kt  |  4 ++--
 .../tests/AbstractKotlinRenderLogTest.kt      |  2 +-
 3 files changed, 23 insertions(+), 7 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 8cdf9a0e66f..481744cd37a 100644
--- a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinUastLanguagePlugin.kt
+++ b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinUastLanguagePlugin.kt
@@ -156,10 +156,26 @@ class KotlinUastLanguagePlugin : UastLanguagePlugin {
                         KotlinUClass.create(lightClass, givenParent)
                     }
                 }
-                is KtFunction -> el<UMethod> {
-                    val lightMethod = LightClassUtil.getLightClassMethod(original) ?: return null
-                    convertDeclaration(lightMethod, givenParent, requiredType)
-                }
+                is KtFunction ->
+                    if (original.isLocal) {
+                        el<ULambdaExpression> {
+                            if (original.name.isNullOrEmpty()) {
+                                createLocalFunctionLambdaExpression(original, givenParent)
+                            }
+                            else {
+                                val uDeclarationsExpression = createLocalFunctionDeclaration(original, givenParent)
+                                val localFunctionVar = uDeclarationsExpression.declarations.single() as KotlinLocalFunctionUVariable
+                                localFunctionVar.uastInitializer
+                            }
+                        }
+                    }
+                    else {
+                        el<UMethod> {
+                            val lightMethod = LightClassUtil.getLightClassMethod(original) ?: return null
+                            convertDeclaration(lightMethod, givenParent, requiredType)
+                        }
+                    }
+
                 is KtPropertyAccessor -> el<UMethod> {
                     val lightMethod = LightClassUtil.getLightClassAccessorMethod(original) ?: return null
                     convertDeclaration(lightMethod, givenParent, requiredType)
diff --git a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/LocalFunction.kt b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/LocalFunction.kt
index b704548a079..822e862c5e2 100644
--- a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/LocalFunction.kt
+++ b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/LocalFunction.kt
@@ -9,7 +9,7 @@ import org.jetbrains.uast.kotlin.psi.UastKotlinPsiParameter
 import org.jetbrains.uast.kotlin.psi.UastKotlinPsiVariable
 
 
-private class KotlinLocalFunctionUVariable(
+internal class KotlinLocalFunctionUVariable(
         val function: KtFunction,
         override val psi: PsiVariable,
         givenParent: UElement?
@@ -50,7 +50,7 @@ private class KotlinLocalFunctionULambdaExpression(
 
 
 fun createLocalFunctionDeclaration(function: KtFunction, parent: UElement?): UDeclarationsExpression {
-    return KotlinUDeclarationsExpression(parent).apply {
+    return KotlinUDeclarationsExpression(null, parent, function).apply {
         val functionVariable = UastKotlinPsiVariable.create(function.name.orAnonymous(), function, this)
         declarations = listOf(KotlinLocalFunctionUVariable(function, functionVariable, this))
     }
diff --git a/plugins/uast-kotlin/tests/AbstractKotlinRenderLogTest.kt b/plugins/uast-kotlin/tests/AbstractKotlinRenderLogTest.kt
index 13a43294cec..6b31d234b3d 100644
--- a/plugins/uast-kotlin/tests/AbstractKotlinRenderLogTest.kt
+++ b/plugins/uast-kotlin/tests/AbstractKotlinRenderLogTest.kt
@@ -60,7 +60,7 @@ abstract class AbstractKotlinRenderLogTest : AbstractKotlinUastTest(), RenderLog
                 if (expectedParents != null) {
                     assertNotNull("Expected to be able to convert PSI element $element", uElement)
                     val parents = generateSequence(uElement!!.uastParent) { it.uastParent }.joinToString { it.asLogString() }
-                    assertEquals("Inconsistent parents for $uElement", expectedParents, parents)
+                    assertEquals("Inconsistent parents for $uElement (converted from $element)", expectedParents, parents)
                 }
                 super.visitElement(element)
             }