feat(kotlin): add companion object support

phodal · Dec 28, 2023 · b275e9f · b275e9f
1 parent 61a83d8
commit b275e9f
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 0 deletions.
diff --git a/chapi-ast-kotlin/src/main/kotlin/chapi/ast/kotlinast/KotlinBasicIdentListener.kt b/chapi-ast-kotlin/src/main/kotlin/chapi/ast/kotlinast/KotlinBasicIdentListener.kt
@@ -163,6 +163,17 @@ open class KotlinBasicIdentListener(private val fileName: String) : KotlinAstLis
         isEnteredClass.decrementAndGet()
     }
 
+    override fun enterCompanionObject(ctx: KotlinParser.CompanionObjectContext?) {
+        val classBody = ctx?.classBody() ?: return
+        classBody.classMemberDeclarations().classMemberDeclaration().forEach {
+            val propertyDeclaration = it.declaration()?.propertyDeclaration()
+            val functionDeclaration = it.declaration()?.functionDeclaration()
+
+            if (propertyDeclaration != null) currentNode.Fields += buildField(propertyDeclaration)
+            if (functionDeclaration != null) currentNode.Functions += buildFunction(functionDeclaration)
+        }
+    }
+
     override fun enterPrimaryConstructor(ctx: KotlinParser.PrimaryConstructorContext) {
         val parameters = ctx.classParameters().classParameter().map(::buildProperty)
         val fields = ctx.classParameters().classParameter().mapNotNull(::buildField)

diff --git a/chapi-ast-kotlin/src/test/kotlin/chapi/ast/kotlinast/KotlinFunctionCallTest.kt b/chapi-ast-kotlin/src/test/kotlin/chapi/ast/kotlinast/KotlinFunctionCallTest.kt
@@ -41,6 +41,26 @@ class KotlinFunctionCallTest {
         assert(codeFunction.Name == "foo")
     }
 
+    @Test
+    fun should_support_for_companion_object() {
+        val code = """
+            class A {
+                companion object {
+                    fun foo() {
+                        println("Hello, world!")
+                    }
+                }
+            }
+        """.trimIndent()
+
+        val codeContainer = KotlinAnalyser().analysis(code, "Test.kt", ParseMode.Full)
+        val dataStructures = codeContainer.DataStructures.filter { it.NodeName == "A" }
+        val codeFunction = dataStructures[0].Functions[0]
+
+        assert(codeFunction.Name == "foo")
+    }
+
+
     @Test
     fun should_success_parse_test_usecase() {
         val code = """