feat(kotlin): add callee lookup for Kotlin functions #308

Implement `lookupCallee` and `findCallees` methods in `KotlinRelatedClassProvider` to identify and return callee functions for a given Kotlin function. Added corresponding test case in `KotlinRelatedClassProviderTest`.

Author: phodal

kotlin/src/main/kotlin/cc/unitmesh/kotlin/provider/KotlinRelatedClassesProvider.kt

@@ -2,8 +2,12 @@ package cc.unitmesh.kotlin.provider
 
 import cc.unitmesh.devti.provider.RelatedClassesProvider
 import cc.unitmesh.kotlin.util.KotlinTypeResolver
+import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
+import org.jetbrains.kotlin.idea.structuralsearch.visitor.KotlinRecursiveElementWalkingVisitor
+import org.jetbrains.kotlin.psi.KtCallExpression
+import org.jetbrains.kotlin.psi.KtNamedFunction
 
 class KotlinRelatedClassProvider : RelatedClassesProvider {
     override fun lookupIO(element: PsiElement): List<PsiElement> {
@@ -13,4 +17,25 @@ class KotlinRelatedClassProvider : RelatedClassesProvider {
     override fun lookupIO(element: PsiFile): List<PsiElement> {
         return KotlinTypeResolver.resolveByElement(element).values.filterNotNull().toList()
     }
+
+    override fun lookupCallee(project: Project, element: PsiElement): List<PsiElement> {
+        return when (element) {
+            is KtNamedFunction -> findCallees(project, element)
+            else -> emptyList()
+        }
+    }
+
+    fun findCallees(project: Project, method: KtNamedFunction): List<KtNamedFunction> {
+        val calledMethods = mutableSetOf<KtNamedFunction>()
+        method.accept(object : KotlinRecursiveElementWalkingVisitor() {
+            override fun visitCallExpression(expression: KtCallExpression) {
+                val psiElement = expression.calleeExpression?.node?.psi?.reference?.resolve() ?: return
+                if (psiElement is KtNamedFunction) {
+                    calledMethods.add(psiElement)
+                }
+            }
+        })
+
+        return calledMethods.toList()
+    }
 }

kotlin/src/test/kotlin/cc/unitmesh/kotlin/provider/KotlinRelatedClassesProviderTest.kt

@@ -0,0 +1,36 @@
+package cc.unitmesh.kotlin.provider
+
+import com.intellij.psi.PsiFileFactory
+import com.intellij.testFramework.LightPlatformTestCase
+import org.intellij.lang.annotations.Language
+import org.jetbrains.kotlin.psi.KtNamedFunction
+
+class KotlinRelatedClassProviderTest : LightPlatformTestCase() {
+    @Language("Kotlin")
+    val code = """
+    fun bar() {
+        println("bar")
+    }
+
+    fun baz() {
+        println("baz")
+    }
+    
+    fun foo() {
+        bar()
+        baz()
+    }
+    """.trimIndent()
+
+    fun testShouldReturnCorrectCallees() {
+        // given
+        val file = PsiFileFactory.getInstance(project).createFileFromText("test.kt", code)
+        val lastFunction = file.children.last()
+
+        // when
+        val result = KotlinRelatedClassProvider().findCallees(project, lastFunction as KtNamedFunction)
+
+        // then
+        assertEquals(2, result.size)
+    }
+}