FIR/UAST: commonize base UMethod

plugins/uast-kotlin-base/src/org/jetbrains/uast/kotlin/declarations/BaseKotlinUMethod.kt

@@ -0,0 +1,84 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.uast.kotlin
+
+import com.intellij.psi.PsiFile
+import com.intellij.psi.PsiMethod
+import com.intellij.psi.PsiNameIdentifierOwner
+import org.jetbrains.kotlin.asJava.elements.KtLightElement
+import org.jetbrains.kotlin.asJava.elements.KtLightMethod
+import org.jetbrains.kotlin.asJava.elements.isGetter
+import org.jetbrains.kotlin.asJava.elements.isSetter
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.uast.*
+
+abstract class BaseKotlinUMethod(
+    psi: PsiMethod,
+    final override val sourcePsi: KtDeclaration?,
+    givenParent: UElement?
+) : KotlinAbstractUElement(givenParent), UMethod, UAnchorOwner, PsiMethod by psi {
+
+    override val psi: PsiMethod = unwrap<UMethod, PsiMethod>(psi)
+
+    override val javaPsi = psi
+
+    override fun getSourceElement() = sourcePsi ?: this
+
+    private val kotlinOrigin = getKotlinMemberOrigin(psi.originalElement) ?: sourcePsi
+
+    override fun getContainingFile(): PsiFile? {
+        kotlinOrigin?.containingFile?.let { return it }
+        return unwrapFakeFileForLightClass(psi.containingFile)
+    }
+
+    override fun getNameIdentifier() = UastLightIdentifier(psi, kotlinOrigin)
+
+    override val uAnnotations: List<UAnnotation> by lz {
+        psi.annotations
+            .mapNotNull { (it as? KtLightElement<*, *>)?.kotlinOrigin as? KtAnnotationEntry }
+            .map { baseResolveProviderService.baseKotlinConverter.convertAnnotation(it, this) }
+    }
+
+    protected val receiverTypeReference by lz {
+        when (sourcePsi) {
+            is KtCallableDeclaration -> sourcePsi
+            is KtPropertyAccessor -> sourcePsi.property
+            else -> null
+        }?.receiverTypeReference
+    }
+
+    override val uastAnchor: UIdentifier? by lz {
+        val identifierSourcePsi = when (val sourcePsi = sourcePsi) {
+            is PsiNameIdentifierOwner -> sourcePsi.nameIdentifier
+            is KtObjectDeclaration -> sourcePsi.getObjectKeyword()
+            is KtPropertyAccessor -> sourcePsi.namePlaceholder
+            else -> sourcePsi?.navigationElement
+        }
+        KotlinUIdentifier(nameIdentifier, identifierSourcePsi, this)
+    }
+
+    override val uastBody: UExpression? by lz {
+        if (kotlinOrigin?.canAnalyze() != true) return@lz null // EA-137193
+        val bodyExpression = when (sourcePsi) {
+            is KtFunction -> sourcePsi.bodyExpression
+            is KtPropertyAccessor -> sourcePsi.bodyExpression
+            is KtProperty -> when {
+                psi is KtLightMethod && psi.isGetter -> sourcePsi.getter?.bodyExpression
+                psi is KtLightMethod && psi.isSetter -> sourcePsi.setter?.bodyExpression
+                else -> null
+            }
+            else -> null
+        } ?: return@lz null
+
+        wrapExpressionBody(this, bodyExpression)
+    }
+
+    override val returnTypeReference: UTypeReferenceExpression? by lz {
+        (sourcePsi as? KtCallableDeclaration)?.typeReference?.let {
+            KotlinUTypeReferenceExpression(it, this) { javaPsi.returnType ?: UastErrorType }
+        }
+    }
+}

plugins/uast-kotlin-fir/src/org/jetbrains/uast/kotlin/declarations/FirKotlinUMethod.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.uast.kotlin
 
 import com.intellij.psi.PsiMethod
-import com.intellij.psi.PsiNameIdentifierOwner
 import org.jetbrains.kotlin.asJava.LightClassUtil
 import org.jetbrains.kotlin.asJava.elements.*
 import org.jetbrains.kotlin.psi.*
@@ -17,36 +16,14 @@ import org.jetbrains.uast.*
 
 open class FirKotlinUMethod(
     psi: PsiMethod,
-    final override val sourcePsi: KtDeclaration?,
+    sourcePsi: KtDeclaration?,
     givenParent: UElement?
-) : KotlinAbstractUElement(givenParent), UMethod, UAnchorOwner, PsiMethod by psi {
+) : BaseKotlinUMethod(psi, sourcePsi, givenParent) {
     constructor(
         psi: KtLightMethod,
         givenParent: UElement?
     ) : this(psi, getKotlinMemberOrigin(psi), givenParent)
 
-    override val psi: PsiMethod = unwrap<UMethod, PsiMethod>(psi)
-
-    override val javaPsi = psi
-
-    override fun getSourceElement() = sourcePsi
-
-    private val kotlinOrigin = getKotlinMemberOrigin(psi.originalElement) ?: sourcePsi
-
-    override val uAnnotations: List<UAnnotation>
-        get() {
-            // TODO: Not yet implemented
-            return emptyList()
-        }
-
-    private val receiverTypeReference by lz {
-        when (sourcePsi) {
-            is KtCallableDeclaration -> sourcePsi
-            is KtPropertyAccessor -> sourcePsi.property
-            else -> null
-        }?.receiverTypeReference
-    }
-
     override val uastParameters: List<UParameter> by lz {
         val lightParams = psi.parameterList.parameters
         val receiverTypeReference =
@@ -57,38 +34,6 @@ open class FirKotlinUMethod(
         uParameters
     }
 
-    override val uastAnchor: UIdentifier? by lz {
-        val identifierSourcePsi = when (val sourcePsi = sourcePsi) {
-            is PsiNameIdentifierOwner -> sourcePsi.nameIdentifier
-            is KtObjectDeclaration -> sourcePsi.getObjectKeyword()
-            is KtPropertyAccessor -> sourcePsi.namePlaceholder
-            else -> sourcePsi?.navigationElement
-        }
-        KotlinUIdentifier(nameIdentifier, identifierSourcePsi, this)
-    }
-
-    override val uastBody: UExpression? by lz {
-        if (kotlinOrigin?.canAnalyze() != true) return@lz null // EA-137193
-        val bodyExpression = when (sourcePsi) {
-            is KtFunction -> sourcePsi.bodyExpression
-            is KtPropertyAccessor -> sourcePsi.bodyExpression
-            is KtProperty -> when {
-                psi is KtLightMethod && psi.isGetter -> sourcePsi.getter?.bodyExpression
-                psi is KtLightMethod && psi.isSetter -> sourcePsi.setter?.bodyExpression
-                else -> null
-            }
-            else -> null
-        } ?: return@lz null
-
-        wrapExpressionBody(this, bodyExpression)
-    }
-
-    override val returnTypeReference: UTypeReferenceExpression? by lz {
-        (sourcePsi as? KtCallableDeclaration)?.typeReference?.let {
-            KotlinUTypeReferenceExpression(it, this) { javaPsi.returnType ?: UastErrorType }
-        }
-    }
-
     companion object {
         fun create(
             psi: KtLightMethod,

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUMethod.kt

@@ -24,50 +24,19 @@ import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 import org.jetbrains.kotlin.utils.SmartList
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import org.jetbrains.uast.*
-import org.jetbrains.uast.java.internal.JavaUElementWithComments
 import org.jetbrains.uast.kotlin.*
 import org.jetbrains.uast.kotlin.psi.UastFakeLightMethod
 import org.jetbrains.uast.kotlin.psi.UastKotlinPsiParameter
 
 open class KotlinUMethod(
     psi: PsiMethod,
-    final override val sourcePsi: KtDeclaration?,
+    sourcePsi: KtDeclaration?,
     givenParent: UElement?
-) : KotlinAbstractUElement(givenParent), UMethodTypeSpecific, UAnchorOwner, JavaUElementWithComments, PsiMethod by psi {
+) : BaseKotlinUMethod(psi, sourcePsi, givenParent) {
-
+    constructor(
-    constructor(psi: KtLightMethod, givenParent: UElement?) : this(psi, getKotlinMemberOrigin(psi), givenParent)
+        psi: KtLightMethod,
-
+        givenParent: UElement?
-    override val comments: List<UComment>
+    ) : this(psi, getKotlinMemberOrigin(psi), givenParent)
-        get() = super<KotlinAbstractUElement>.comments
-
-    override val psi: PsiMethod = unwrap<UMethod, PsiMethod>(psi)
-
-    override val javaPsi = psi
-
-    override fun getSourceElement() = sourcePsi ?: this
-
-    private val kotlinOrigin = getKotlinMemberOrigin(psi.originalElement) ?: sourcePsi
-
-    override fun getContainingFile(): PsiFile? {
-        kotlinOrigin?.containingFile?.let { return it }
-        return unwrapFakeFileForLightClass(psi.containingFile)
-    }
-
-    override fun getNameIdentifier() = UastLightIdentifier(psi, kotlinOrigin)
-
-    override val uAnnotations: List<UAnnotation> by lz {
-        psi.annotations
-            .mapNotNull { (it as? KtLightElement<*, *>)?.kotlinOrigin as? KtAnnotationEntry }
-            .map { baseResolveProviderService.baseKotlinConverter.convertAnnotation(it, this) }
-    }
-
-    private val receiver by lz {
-        when (sourcePsi) {
-            is KtCallableDeclaration -> sourcePsi
-            is KtPropertyAccessor -> sourcePsi.property
-            else -> null
-        }?.receiverTypeReference
-    }
 
     override val uastParameters by lz {
 
@@ -78,52 +47,13 @@ open class KotlinUMethod(
         }
 
         val lightParams = psi.parameterList.parameters
-        val receiver = receiver ?: return@lz lightParams.map { KotlinUParameter(it, parameterOrigin(it), this) }
+        val receiver = receiverTypeReference ?: return@lz lightParams.map { KotlinUParameter(it, parameterOrigin(it), this) }
         val receiverLight = lightParams.firstOrNull() ?: return@lz emptyList<UParameter>()
         val uParameters = SmartList<UParameter>(KotlinReceiverUParameter(receiverLight, receiver, this))
         lightParams.drop(1).mapTo(uParameters) { KotlinUParameter(it, parameterOrigin(it), this) }
         uParameters
     }
 
-    override val uastAnchor by lazy {
-        KotlinUIdentifier(
-            nameIdentifier,
-            sourcePsi.let { sourcePsi ->
-                when (sourcePsi) {
-                    is PsiNameIdentifierOwner -> sourcePsi.nameIdentifier
-                    is KtObjectDeclaration -> sourcePsi.getObjectKeyword()
-                    is KtPropertyAccessor -> sourcePsi.namePlaceholder
-                    else -> sourcePsi?.navigationElement
-                }
-            },
-            this
-        )
-    }
-
-    override val uastBody by lz {
-        if (kotlinOrigin?.canAnalyze() != true) return@lz null // EA-137193
-        val bodyExpression = when (kotlinOrigin) {
-            is KtFunction -> kotlinOrigin.bodyExpression
-            is KtPropertyAccessor -> kotlinOrigin.bodyExpression
-            is KtProperty -> when {
-                psi is KtLightMethod && psi.isGetter -> kotlinOrigin.getter?.bodyExpression
-                psi is KtLightMethod && psi.isSetter -> kotlinOrigin.setter?.bodyExpression
-                else -> null
-            }
-            else -> null
-        } ?: return@lz null
-
-        wrapExpressionBody(this, bodyExpression)
-    }
-
-    override val returnTypeReference: UTypeReferenceExpression? by lz {
-        (sourcePsi as? KtCallableDeclaration)?.typeReference?.let {
-            KotlinUTypeReferenceExpression(it, this) { javaPsi.returnType ?: UastErrorType }
-        }
-    }
-
-    override fun equals(other: Any?) = other is KotlinUMethod && psi == other.psi
-
     companion object {
         fun create(psi: KtLightMethod, containingElement: UElement?): KotlinUMethod {
             val kotlinOrigin = psi.kotlinOrigin