feat: <java> add enum support

phodal · Aug 19, 2022 · 3fa19ef · 3fa19ef
1 parent c1fc5fa
commit 3fa19ef
Show file tree

Hide file tree

Showing 4 changed files with 101 additions and 1 deletion.
diff --git a/chapi-ast-java/src/main/kotlin/chapi/ast/javaast/JavaBasicIdentListener.kt b/chapi-ast-java/src/main/kotlin/chapi/ast/javaast/JavaBasicIdentListener.kt
@@ -75,6 +75,19 @@ open class JavaBasicIdentListener(fileName: String) : JavaAstListener() {
         return implements
     }
 
+    open fun buildImplements(ctx: JavaParser.EnumDeclarationContext): Array<String> {
+        var implements = arrayOf<String>()
+
+        val typeText = ctx.typeList().text
+        for (imp in imports) {
+            if (imp.Source.endsWith(".$typeText")) {
+                implements += typeText
+            }
+        }
+
+        return implements
+    }
+
     override fun enterInterfaceDeclaration(ctx: JavaParser.InterfaceDeclarationContext?) {
         hasEnterClass = true
         currentNode = CodeDataStruct(
@@ -186,6 +199,30 @@ open class JavaBasicIdentListener(fileName: String) : JavaAstListener() {
         }
     }
 
+    override fun enterEnumDeclaration(ctx: JavaParser.EnumDeclarationContext?) {
+        hasEnterClass = true
+        currentNode.Type = DataStructType.ENUM
+        currentNode.Package = codeContainer.PackageName
+
+        if (ctx!!.identifier() != null) {
+            currentNode.NodeName = ctx.identifier().text
+        }
+
+        if (ctx.IMPLEMENTS() != null) {
+            currentNode.Implements = buildImplements(ctx)
+        }
+
+        currentFunction = CodeFunction()
+    }
+
+    override fun exitEnumBodyDeclarations(ctx: JavaParser.EnumBodyDeclarationsContext?) {
+        hasEnterClass = false
+        if (currentNode.NodeName != "") {
+            classNodes += currentNode
+        }
+        currentNode = CodeDataStruct()
+    }
+
     fun getNodeInfo(): CodeContainer {
         codeContainer.DataStructures = classNodes
         return codeContainer

diff --git a/chapi-ast-java/src/main/kotlin/chapi/ast/javaast/JavaFullIdentListener.kt b/chapi-ast-java/src/main/kotlin/chapi/ast/javaast/JavaFullIdentListener.kt
@@ -760,4 +760,52 @@ open class JavaFullIdentListener(fileName: String, val classes: Array<String>) :
         codeContainer.DataStructures = classNodes
         return codeContainer
     }
+
+    override fun enterEnumDeclaration(ctx: JavaParser.EnumDeclarationContext?) {
+        currentType = DataStructType.ENUM
+        currentNode.Type = currentType
+        currentNode.Position = buildPosition(ctx!!)
+
+        hasEnterClass = true
+        buildEnumExtension(ctx, currentNode)
+
+        lastNode = currentNode
+        classNodeStack.push(currentNode)
+    }
+
+    override fun exitEnumDeclaration(ctx: JavaParser.EnumDeclarationContext?) {
+        classNodeStack.pop()
+        if (classNodeStack.count() == 0) {
+            this.exitBody()
+        } else {
+            updateStructMethods()
+        }
+    }
+
+    private fun buildEnumExtension(ctx: JavaParser.EnumDeclarationContext?, classNode: CodeDataStruct) {
+        classNode.Package = codeContainer.PackageName
+        classNode.FilePath = codeContainer.FullName
+
+        if (ctx!!.identifier() != null) {
+            currentClz = ctx.identifier().text
+            classNode.NodeName = currentClz
+        }
+
+        currentClzExtend = ""
+
+        if (ctx.IMPLEMENTS() != null) {
+            classNode.Implements = buildEnumImplements(ctx)
+        }
+    }
+
+    open fun buildEnumImplements(ctx: JavaParser.EnumDeclarationContext): Array<String> {
+        var implements = arrayOf<String>()
+        val type = ctx.typeList()
+        var target = this.warpTargetFullType(type.text).targetType
+        if (target == "") {
+            target = type.text
+        }
+        implements += target
+        return implements
+    }
 }
diff --git a/chapi-ast-java/src/test/kotlin/chapi/ast/javaast/JavaBasicIdentListenerTest.kt b/chapi-ast-java/src/test/kotlin/chapi/ast/javaast/JavaBasicIdentListenerTest.kt
@@ -225,4 +225,18 @@ internal class JavaBasicIdentListenerTest {
         assertEquals(codeFunction.Extension.jsonObject["IsReturnNull"], JsonPrimitive("true"))
         assertEquals(codeFunction.isReturnNull(), true)
     }
+
+    @Test
+    fun shouldEnum() {
+        val code = """
+            package chapi.ast.javaast;
+            public enum AbcEnum {
+                A;
+            }
+        """
+
+        val codeFile = JavaAnalyser().identBasicInfo(code, "basic")
+        assertEquals(codeFile.DataStructures.size, 1)
+        assertEquals(codeFile.DataStructures[0].NodeName, "AbcEnum")
+    }
 }
diff --git a/chapi-domain/src/main/kotlin/chapi/domain/core/CodeDataStruct.kt b/chapi-domain/src/main/kotlin/chapi/domain/core/CodeDataStruct.kt
@@ -16,7 +16,8 @@ enum class DataStructType (val structType: String) {
     CREATOR_CLASS("CreatorClass"),
     ABSTRACT_CLASS("AbstractClass"),
     // for scala, Rust
-    TRAIT("Trait")
+    TRAIT("Trait"),
+    ENUM("enum")
 }
 
 @Serializable