feat(goast): add receiver mapping for method calls and test case

Adds a new private map to store receiver names for method calls in GoFullIdentListener. This allows for resolving the correct receiver when analyzing method declarations. A new test case is included to verify the functionality.
phodal · Nov 9, 2024 · 420c2d4 · 420c2d4
1 parent ee1006d
commit 420c2d4
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 0 deletions.
diff --git a/chapi-ast-go/src/main/kotlin/chapi/ast/goast/GoFullIdentListener.kt b/chapi-ast-go/src/main/kotlin/chapi/ast/goast/GoFullIdentListener.kt
@@ -15,6 +15,11 @@ class GoFullIdentListener(var fileName: String) : GoAstListener() {
     private var structMap = mutableMapOf<String, CodeDataStruct>()
     private var localVars = mutableMapOf<String, String>()
 
+    /**
+     * Only for resolve receiver for method call
+     */
+    private var receiverForCall = mutableMapOf<String, String>()
+
     private var currentFunction = CodeFunction(IsConstructor = false)
 
     override fun enterImportDecl(ctx: GoParser.ImportDeclContext?) {
@@ -87,6 +92,12 @@ class GoFullIdentListener(var fileName: String) : GoAstListener() {
     }
 
     override fun enterMethodDecl(ctx: GoParser.MethodDeclContext?) {
+        val receiverName = this.getStructNameFromReceiver(ctx!!.receiver()?.parameters())
+        if (ctx.receiver() != null) {
+            val text = ctx.receiver().parameters().parameterDecl()[0].identifierList().text
+            receiverForCall[text] = receiverName
+        }
+
         currentFunction = CodeFunction(
             Name = ctx!!.IDENTIFIER().text,
             MultipleReturns = buildReturnTypeFromSignature(ctx.signature()),
@@ -307,6 +318,11 @@ class GoFullIdentListener(var fileName: String) : GoAstListener() {
                 first.text
             }
         }
+
+        if (receiverForCall.containsKey(nodeName)) {
+            return receiverForCall[nodeName]!!
+        }
+
         return nodeName
     }
 

diff --git a/chapi-ast-go/src/test/kotlin/chapi/ast/goast/GoFullIdentListenerTest.kt b/chapi-ast-go/src/test/kotlin/chapi/ast/goast/GoFullIdentListenerTest.kt
@@ -1,6 +1,7 @@
 package chapi.ast.goast
 
 import chapi.domain.core.DataStructType
+import org.intellij.lang.annotations.Language
 import org.junit.jupiter.api.Test
 import kotlin.test.assertEquals
 
@@ -295,4 +296,32 @@ func (a *Animal) Move() {
         assertEquals(codeFile.DataStructures[0].Functions[0].LocalVariables[2].TypeValue, "b")
         assertEquals(codeFile.DataStructures[0].Functions[0].LocalVariables[3].TypeValue, "c")
     }
+
+    @Test
+    internal fun shouldSuccessGetSqlOfNode() {
+        @Language("Go")
+        val code= """
+package dao
+
+import (
+	"database/sql"
+)
+
+func (d *Dao) QueryBuglyProjectList() (projectList []string, err error) {
+	var (
+		rows *sql.Rows
+	)
+	sql := "select DISTINCT project_name from bugly_projects"
+	if rows, err = d.db.Raw(sql).Rows(); err != nil {
+		return
+	}
+}
+"""
+        val codeFile = GoAnalyser().analysis(code, "")
+        val functionCalls = codeFile.DataStructures[0].Functions[0].FunctionCalls
+        println(functionCalls)
+
+        assertEquals(functionCalls.size, 1)
+        assertEquals(functionCalls[0].NodeName, "Dao")
+    }
 }