From 638b0beadb913f5b52b90504d5ff21d9c658fc0f Mon Sep 17 00:00:00 2001
From: Jinseong Jeon <jsjeon@google.com>
Date: Fri, 18 Jun 2021 23:14:07 -0700
Subject: [PATCH] FIR UAST: unify UMethod creation

---
 .../uast/kotlin/declarations/FirKotlinUMethod.kt   | 14 ++++----------
 1 file changed, 4 insertions(+), 10 deletions(-)

diff --git a/plugins/uast-kotlin-fir/src/org/jetbrains/uast/kotlin/declarations/FirKotlinUMethod.kt b/plugins/uast-kotlin-fir/src/org/jetbrains/uast/kotlin/declarations/FirKotlinUMethod.kt
index 24f30462959..e19d046cfc5 100644
--- a/plugins/uast-kotlin-fir/src/org/jetbrains/uast/kotlin/declarations/FirKotlinUMethod.kt
+++ b/plugins/uast-kotlin-fir/src/org/jetbrains/uast/kotlin/declarations/FirKotlinUMethod.kt
@@ -29,9 +29,9 @@ class FirKotlinUMethod(
             psi: KtLightMethod,
             givenParent: UElement?
         ): UMethod {
-            return when (psi) {
+            return when (val kotlinOrigin = psi.kotlinOrigin) {
                 is KtConstructor<*> ->
-                    FirKotlinConstructorUMethod(psi.containingClassOrObject, psi, givenParent)
+                    FirKotlinConstructorUMethod(kotlinOrigin.containingClassOrObject, psi, givenParent)
                 // TODO: FirKotlinUAnnotationMethod
                 else ->
                     FirKotlinUMethod(psi, givenParent)
@@ -48,14 +48,8 @@ class FirKotlinUMethod(
                 is KtFunction ->
                     LightClassUtil.getLightClassMethod(sourcePsi)
                 else -> null
-            } ?: return null
-            return when (sourcePsi) {
-                is KtConstructor<*> ->
-                    FirKotlinConstructorUMethod(sourcePsi.containingClassOrObject, javaPsi, sourcePsi, givenParent)
-                // TODO: FirKotlinUAnnotationMethod
-                else ->
-                    FirKotlinUMethod(javaPsi, sourcePsi, givenParent)
-            }
+            } as? KtLightMethod ?: return null
+            return create(javaPsi, givenParent)
         }
     }
 }