feat(java): add method caller and called method lookup #29

- Add method caller and called method lookup functionality to JavaTestHelper. - Implement findCallers and findCallees methods to retrieve callers and callees of a given method.
phodal · Jul 2, 2024 · 6b46c00 · 6b46c00
1 parent 75b58de
commit 6b46c00
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 5 deletions.
diff --git a/core/src/main/kotlin/com/phodal/shirecore/provider/variable/PsiContextVariable.kt b/core/src/main/kotlin/com/phodal/shirecore/provider/variable/PsiContextVariable.kt
@@ -59,6 +59,10 @@ enum class PsiContextVariable(val variableName: String, val description: String?
      * codeSmell
      */
     CODE_SMELL("codeSmell", "Include psi error and warning"),
+
+    METHOD_CALLER("methodCaller", "The method that initiates the current call"),
+
+    CALLED_METHOD("calledMethod", "The method that is being called by the current method"),
     ;
 
     companion object {

diff --git a/...java/src/main/kotlin/com/phodal/shirelang/java/variable/JavaPsiContextVariableProvider.kt b/...java/src/main/kotlin/com/phodal/shirelang/java/variable/JavaPsiContextVariableProvider.kt
@@ -2,6 +2,7 @@ package com.phodal.shirelang.java.variable
 
 import com.intellij.openapi.editor.Editor
 import com.intellij.psi.*
+import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.testIntegration.TestFinderHelper
 import com.phodal.shirecore.provider.context.LanguageToolchainProvider
 import com.phodal.shirecore.provider.context.ToolchainPrepareContext
@@ -57,6 +58,14 @@ class JavaPsiContextVariableProvider : PsiContextVariableProvider {
             }
 
             PsiContextVariable.CODE_SMELL -> CodeSmellBuilder.collectElementProblemAsSting(psiElement, project, editor)
+            PsiContextVariable.METHOD_CALLER -> {
+                if (psiElement !is PsiMethod) return ""
+                return JavaTestHelper.findCallers(psiElement).joinToString("\n") { it.text }
+            }
+            PsiContextVariable.CALLED_METHOD -> {
+                if (psiElement !is PsiMethod) return ""
+                return JavaTestHelper.findCallees(psiElement).joinToString("\n") { it.text }
+            }
         }
     }
 }

diff --git a/languages/java/src/main/kotlin/com/phodal/shirelang/java/variable/JavaTestHelper.kt b/languages/java/src/main/kotlin/com/phodal/shirelang/java/variable/JavaTestHelper.kt
@@ -4,12 +4,17 @@ import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiMethod
 import com.intellij.psi.PsiMethodCallExpression
+import com.intellij.psi.PsiReference
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.search.PsiShortNamesCache
 import com.intellij.psi.search.searches.ClassInheritorsSearch
+import com.intellij.psi.search.searches.MethodReferencesSearch
+import com.intellij.psi.search.searches.ReferencesSearch
 import com.intellij.psi.util.*
 import com.phodal.shirecore.search.CodeNamingTokenizer
 import com.phodal.shirecore.search.TfIdf
+import java.util.function.Consumer
+
 
 object JavaTestHelper {
     fun lookupUnderTestMethod(project: Project, psiElement: PsiElement): String {
@@ -69,15 +74,12 @@ object JavaTestHelper {
             val testMethods = mutableListOf<PsiMethod>()
             val scope = GlobalSearchScope.projectScope(project)
 
-            // 获取所有文件名中包含 'Test' 的类
-            PsiShortNamesCache.getInstance(project).getAllClassNames()
+            PsiShortNamesCache.getInstance(project).allClassNames
                 .filter { it.contains("Test") }
                 .forEach { className ->
-                    // 获取所有匹配的类
                     PsiShortNamesCache.getInstance(project).getClassesByName(className, scope)
                         .filter { it.containingFile.name.endsWith("Test.java") }
                         .forEach { psiClass ->
-                            // 收集所有方法
                             testMethods.addAll(psiClass.methods)
                         }
                 }
@@ -86,4 +88,52 @@ object JavaTestHelper {
         }
 
         return cachedValue.value
-    }}
+    }
+
+    /**
+     * Finds all the callers of a given method.
+     *
+     * @param method the method for which callers need to be found
+     * @return a list of PsiMethod objects representing the callers of the given method
+     */
+    fun findCallers(method: PsiMethod): List<PsiMethod> {
+        val callers: MutableList<PsiMethod> = ArrayList()
+
+        val references = ReferencesSearch.search(method).findAll()
+
+        for (reference in references) {
+            val element = reference.element
+
+            if (element is PsiMethodCallExpression) {
+                val callerMethod = element.resolveMethod()
+                if (callerMethod != null) {
+                    callers.add(callerMethod)
+                }
+            }
+        }
+
+        return callers
+    }
+
+    /**
+     * Finds all the methods called by the given method.
+     *
+     * @param method the method for which callees need to be found
+     * @return a list of PsiMethod objects representing the methods called by the given method
+     */
+    fun findCallees(method: PsiMethod): List<PsiMethod> {
+        val callees: MutableList<PsiMethod> = ArrayList()
+
+        MethodReferencesSearch.search(method).forEach(Consumer { reference: PsiReference ->
+            val element = reference.element
+            if (element is PsiMethodCallExpression) {
+                val resolvedMethod = element.resolveMethod()
+                if (resolvedMethod != null) {
+                    callees.add(resolvedMethod)
+                }
+            }
+        })
+
+        return callees
+    }
+}