diff --git a/chapi-ast-rust/src/main/kotlin/chapi/ast/rustast/RustAstBaseListener.kt b/chapi-ast-rust/src/main/kotlin/chapi/ast/rustast/RustAstBaseListener.kt index 2b0f49ea..80e0bf3e 100644 --- a/chapi-ast-rust/src/main/kotlin/chapi/ast/rustast/RustAstBaseListener.kt +++ b/chapi-ast-rust/src/main/kotlin/chapi/ast/rustast/RustAstBaseListener.kt @@ -29,6 +29,10 @@ open class RustAstBaseListener(private val fileName: String) : RustParserBaseLis private var structMap = mutableMapOf() + private var currentModule: String = "" + /// for testing + private var lastModule: String = "" + /** * localVars will store all local variables in the current scope @@ -82,6 +86,15 @@ open class RustAstBaseListener(private val fileName: String) : RustParserBaseLis Imports = imports } + override fun enterModule(ctx: RustParser.ModuleContext?) { + currentModule = (ctx?.identifier()?.text ?: return) + } + + override fun exitModule(ctx: RustParser.ModuleContext?) { + lastModule = currentModule + currentModule = "" + } + override fun enterUseDeclaration(ctx: RustParser.UseDeclarationContext?) { val importString: List> = buildToPath(ctx) @@ -150,6 +163,7 @@ open class RustAstBaseListener(private val fileName: String) : RustParserBaseLis val codeStruct = CodeDataStruct( NodeName = structName, + Module = currentModule, Package = codeContainer.PackageName, Annotations = annotation, Fields = buildFields(ctx.structFields()), @@ -270,6 +284,7 @@ open class RustAstBaseListener(private val fileName: String) : RustParserBaseLis val codeStruct = CodeDataStruct( NodeName = enumName, + Module = currentModule, Package = codeContainer.PackageName, Annotations = annotation, Fields = buildEnumFields(ctx.enumItems()), @@ -367,6 +382,7 @@ open class RustAstBaseListener(private val fileName: String) : RustParserBaseLis } else { currentNode = CodeDataStruct( NodeName = nodeName, + Module = currentModule, Package = codeContainer.PackageName, Position = buildPosition(ctx ?: return) ) @@ -417,6 +433,7 @@ open class RustAstBaseListener(private val fileName: String) : RustParserBaseLis return listOf( CodeDataStruct().apply { NodeName = fileName.substringBeforeLast('.') + Module = if (lastModule == "tests") lastModule else "" Type = DataStructType.OBJECT Package = codeContainer.PackageName FilePath = codeContainer.FullName diff --git a/chapi-ast-rust/src/test/kotlin/chapi/ast/rustast/RustFullIdentListenerTest.kt b/chapi-ast-rust/src/test/kotlin/chapi/ast/rustast/RustFullIdentListenerTest.kt index db9b546d..1e2a84dd 100644 --- a/chapi-ast-rust/src/test/kotlin/chapi/ast/rustast/RustFullIdentListenerTest.kt +++ b/chapi-ast-rust/src/test/kotlin/chapi/ast/rustast/RustFullIdentListenerTest.kt @@ -421,4 +421,23 @@ class RustFullIdentListenerTest { assertEquals("actix_web::HttpServer", secondFunction.FunctionCalls[0].NodeName) assertEquals("run", secondFunction.FunctionCalls[0].FunctionName) } + + @Test + fun should_handle_test_mod() { + val code = """ + #[cfg(test)] + mod tests { + use super::*; + + #[test] + fn test_add() { + assert_eq!(add(1, 2), 3); + } + } + """.trimIndent() + + val codeContainer = RustAnalyser().analysis(code, "lib.rs") + val codeDataStruct = codeContainer.DataStructures[0] + assertEquals("tests", codeDataStruct.Module) + } } diff --git a/chapi-domain/src/main/kotlin/chapi/domain/core/CodeDataStruct.kt b/chapi-domain/src/main/kotlin/chapi/domain/core/CodeDataStruct.kt index b1dd5d0b..77bb2d82 100644 --- a/chapi-domain/src/main/kotlin/chapi/domain/core/CodeDataStruct.kt +++ b/chapi-domain/src/main/kotlin/chapi/domain/core/CodeDataStruct.kt @@ -27,6 +27,7 @@ data class CodeDataStruct( * for Rust, if is a variable, function, it will be naming to `default` */ var NodeName: String = "", + // var Module: String = "", var Type: DataStructType = DataStructType.EMPTY, var Package: String = "",