fix: try to six string type issue

java/src/main/kotlin/cc/unitmesh/idea/context/JavaContextCollectionUtilsKt.kt

@@ -2,10 +2,12 @@ package cc.unitmesh.idea.context
 
 import cc.unitmesh.devti.context.SimpleClassStructure
 import com.intellij.psi.*
+import com.intellij.psi.impl.source.PsiClassReferenceType
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.search.SearchScope
 import com.intellij.psi.search.searches.MethodReferencesSearch
 import com.intellij.psi.search.searches.ReferencesSearch
+import org.jetbrains.kotlin.psi.psiUtil.contains
 
 object JavaContextCollectionUtilsKt {
     fun findUsages(nameIdentifierOwner: PsiNameIdentifierOwner): List<PsiReference> {
@@ -31,31 +33,37 @@ object JavaContextCollectionUtilsKt {
      *
      * For example, if a BlogPost class includes a Comment class, and the Comment class includes a User class, then the resulting tree will be:
      *
+     * ```
      * parent: BlogPost Psi
-     *    child: I'd
-     *    child: Comment psi
-     *        child: User psi
+     *    child: id
+     *    child: Comment
+     *        child: User
+     *          child: name
+     *```
      */
-    fun dataStructure(clazz: PsiClass): HashMap<String, SimpleClassStructure> {
-        val classTree = HashMap<String, SimpleClassStructure>()
-        buildSimpleClassStructure(clazz, classTree)
-        return classTree
+    fun dataStructure(clazz: PsiClass): SimpleClassStructure {
+        val project = clazz.project
+        val searchScope = GlobalSearchScope.allScope(project) as SearchScope
+        return createSimpleStructure(clazz ,searchScope)
     }
 
-    private fun buildSimpleClassStructure(clazz: PsiClass, classTree: HashMap<String, SimpleClassStructure>) {
+    private fun createSimpleStructure(clazz: PsiClass, searchScope: SearchScope): SimpleClassStructure {
         val fields = clazz.fields
-        val children = fields.mapNotNull {
-            if (it.type is PsiClass) {
-                val childSimpleClassStructure = SimpleClassStructure(it.name, it.type.canonicalText, emptyList())
-                buildSimpleClassStructure(it.type as PsiClass, classTree)
-                childSimpleClassStructure
-            } else {
-                SimpleClassStructure(it.name, it.type.canonicalText, emptyList())
+        val children = fields.mapNotNull { field ->
+            when {
+                field.type is PsiClassReferenceType -> {
+                    val classStructure =
+                        (field.type as PsiClassType).resolve()?.let { createSimpleStructure(it, searchScope) } ?: return@mapNotNull null
+                    classStructure.builtIn = false
+                    classStructure
+                }
+
+                else -> {
+                    SimpleClassStructure(field.name, field.type.presentableText, emptyList(), builtIn = true)
+                }
             }
         }
 
-        val classSimpleClassStructure =
-            SimpleClassStructure(clazz.name ?: "Unknown", clazz.qualifiedName ?: "Unknown", children)
-        classTree[clazz.name ?: "Unknown"] = classSimpleClassStructure
+        return SimpleClassStructure(clazz.name ?: "", clazz.name ?: "", children)
     }
 }

java/src/test/kotlin/cc/unitmesh/idea/context/JavaContextCollectionUtilsKtTest.kt

@@ -16,21 +16,33 @@ class SimpleClassStructureTest : LightJavaCodeInsightFixtureTestCase() {
     Comment comment;
 }
 class Comment {
+    long id;
     User user;
 }
 class User {
+    long id;
     String name;
 }
 """
 
         val psiFile = fileFactory.createFileFromText(JavaLanguage.INSTANCE, javaCode)
         val classes = PsiTreeUtil.findChildrenOfType(psiFile, PsiClass::class.java)
 
-        val blogpost = classes.filter { it.name == "BlogPost" }.first()
-        val tree = JavaContextCollectionUtilsKt.dataStructure(blogpost)
-        TestCase.assertEquals(tree["BlogPost"]?.toString(), "class BlogPost {\n" +
+        val blogpost = classes.first { it.name == "BlogPost" }
+        val structure = JavaContextCollectionUtilsKt.dataStructure(blogpost)
+        TestCase.assertEquals(structure.children.size, 2)
+        TestCase.assertEquals(structure.toString(), "class BlogPost {\n" +
+                "  id: long\n" +
+                "  Comment: Comment\n" +
+                "}\n" +
+                "\n" +
+                "class Comment {\n" +
+                "  id: long\n" +
+                "  User: User\n" +
+                "}\n" +
+                "\n" +
+                "class User {\n" +
                 "  id: long\n" +
-                "  comment: Comment\n" +
                 "}\n")
     }
 }

src/main/kotlin/cc/unitmesh/devti/context/SimpleClassStructure.kt

@@ -1,21 +1,88 @@
 package cc.unitmesh.devti.context
 
-data class SimpleClassStructure(val fieldName: String, val fieldType: String, val children: List<SimpleClassStructure>) {
-    override fun toString(): String {
-        return buildPumlRepresentation(this)
+data class SimpleClassStructure(
+    val fieldName: String,
+    var fieldType: String,
+    val children: List<SimpleClassStructure>,
+    var builtIn: Boolean = false
+) {
+    val childPuml: MutableMap<String, String> = mutableMapOf()
+
+
+    /**
+     * Returns a PlantUML string representation of the class structure
+     * for example:
+     * ```
+     * class BlogPost {
+     *     long id;
+     *     Comment comment;
+     * }
+     * class Comment {
+     *     User user;
+     * }
+     * class User {
+     *     String name;
+     * }
+     *```
+     *
+     * will be represented as:
+     *
+     * ```puml
+     * class BlogPost {
+     *    long id;
+     *    Comment comment;
+     *}
+     *```
+     */
+    private fun simplePuml(simpleClassStructure: SimpleClassStructure): String {
+        return "class ${simpleClassStructure.fieldType} {\n" +
+                simpleClassStructure.children.joinToString("\n") { "  ${it.fieldName}: ${it.fieldType}" } +
+                "\n}\n"
     }
 
-    private fun buildPumlRepresentation(node: SimpleClassStructure, indent: String = ""): String {
-        val stringBuilder = StringBuilder()
 
-        stringBuilder.append("class ${node.fieldName} {\n")
+    /**
+     * Returns a PlantUML string representation of the class structure.
+     *
+     * This method generates a PlantUML string representation of the class structure based on the current object and its child objects. The resulting string is built using the buildPuml() method and the childPuml map. The resulting string will be a tree-like structure that shows the relationships between the classes.
+     *
+     * @return the PlantUML string representation of the class structure
+     *
+     * For example, if a BlogPost class includes a Comment class, and the Comment class includes a User class, then the resulting tree will be:
+     *
+     * ```puml
+     * class BlogPost {
+     *  id: long
+     *  comment: Comment
+     *}
+     *
+     * class Comment {
+     *   user: User
+     * }
+     *
+     * class User {
+     *  name: String
+     * }
+     *```
+     */
+    override fun toString(): String {
+        val puml = StringBuilder()
+        puml.append(simplePuml(this))
 
-        for (child in node.children) {
-            stringBuilder.append("  ${child.fieldName}: ${child.fieldType}\n")
+        createChildPuml(children)
+
+        childPuml.forEach {
+            puml.append("\n")
+            puml.append(it.value)
         }
 
-        stringBuilder.append("}\n")
+        return puml.toString()
+    }
 
-        return stringBuilder.toString()
+    private fun createChildPuml(data: List<SimpleClassStructure>) {
+        data.filter { !it.builtIn }.forEach {
+            childPuml[it.fieldType] = simplePuml(it)
+            createChildPuml(it.children)
+        }
     }
-}
+}