From a1d1b473de3e44b223df747ca65cec7e33e9641d Mon Sep 17 00:00:00 2001
From: Phodal Huang <h@phodal.com>
Date: Thu, 29 Sep 2022 14:39:08 +0800
Subject: [PATCH] feat: make structure view works

---
 src/main/kotlin/com/feakin/intellij/FkFile.kt |  2 +-
 .../structure/FkStructureViewElement.kt       | 52 ++++++++++++++-----
 .../structure/FkStructureViewModel.kt         |  6 ---
 3 files changed, 41 insertions(+), 19 deletions(-)

diff --git a/src/main/kotlin/com/feakin/intellij/FkFile.kt b/src/main/kotlin/com/feakin/intellij/FkFile.kt
index 99015b8..0356c27 100644
--- a/src/main/kotlin/com/feakin/intellij/FkFile.kt
+++ b/src/main/kotlin/com/feakin/intellij/FkFile.kt
@@ -12,4 +12,4 @@ class FkFile(viewProvider: FileViewProvider) : PsiFileBase(viewProvider, FkLangu
     override fun toString(): String {
         return "Feakin"
     }
-}
\ No newline at end of file
+}
diff --git a/src/main/kotlin/com/feakin/intellij/structure/FkStructureViewElement.kt b/src/main/kotlin/com/feakin/intellij/structure/FkStructureViewElement.kt
index dffe0ec..ca63d2c 100644
--- a/src/main/kotlin/com/feakin/intellij/structure/FkStructureViewElement.kt
+++ b/src/main/kotlin/com/feakin/intellij/structure/FkStructureViewElement.kt
@@ -1,7 +1,10 @@
 package com.feakin.intellij.structure
 
+import com.feakin.intellij.FkFile
+import com.feakin.intellij.FkIcons
 import com.feakin.intellij.psi.FeakinContextDeclaration
 import com.feakin.intellij.psi.FeakinContextMapDeclaration
+import com.feakin.intellij.psi.FeakinDeclaration
 import com.intellij.ide.projectView.PresentationData
 import com.intellij.ide.structureView.StructureViewTreeElement
 import com.intellij.ide.util.treeView.TreeAnchorizer
@@ -13,9 +16,8 @@ import com.intellij.openapi.editor.DefaultLanguageHighlighterColors
 import com.intellij.pom.Navigatable
 import com.intellij.psi.NavigatablePsiElement
 import com.intellij.psi.PsiElement
-import com.intellij.psi.util.PsiTreeUtil
 
-class FkStructureViewElement(private val element: NavigatablePsiElement) : StructureViewTreeElement {
+class FkStructureViewElement(private val element: PsiElement) : StructureViewTreeElement {
     companion object {
         private val log: Logger = logger<FkStructureViewElement>()
     }
@@ -37,12 +39,11 @@ class FkStructureViewElement(private val element: NavigatablePsiElement) : Struc
     override fun getPresentation(): ItemPresentation {
         val struct = PresentationData()
         struct.setIcon(null)
-        log.warn(psi?.javaClass?.name)
 
         when (val psi = psi) {
             is FeakinContextMapDeclaration -> {
                 val structNameDeclaration = psi.contextMapName
-                struct.setAttributesKey(DefaultLanguageHighlighterColors.KEYWORD)
+                struct.setAttributesKey(DefaultLanguageHighlighterColors.IDENTIFIER)
                 struct.presentableText = structNameDeclaration.text
                 return struct
             }
@@ -54,27 +55,54 @@ class FkStructureViewElement(private val element: NavigatablePsiElement) : Struc
                 return struct
             }
 
+            is FkFile -> {
+                struct.setAttributesKey(DefaultLanguageHighlighterColors.KEYWORD)
+                struct.setIcon(FkIcons.FILE)
+                struct.presentableText = psi.originalFile.name
+                return struct
+            }
+
             else -> {
-                val presentation = element.presentation
-                return presentation ?: PresentationData()
+                return PresentationData()
             }
         }
     }
 
     override fun getChildren(): Array<TreeElement> {
-        log.warn(psi?.javaClass?.name)
-        return when (val psi = psi) {
+        val treeElements = ArrayList<TreeElement>()
+
+        psi?.children?.forEach { child ->
+            val nodes = findContentNode(child).filterNotNull()
+            nodes.forEach {
+                treeElements.add(FkStructureViewElement(it))
+            }
+        }
+
+        return treeElements.toTypedArray()
+    }
+
+    // todo: make it better.
+    private fun findContentNode(psi: PsiElement?): List<PsiElement?> {
+        when (psi) {
+            is FeakinDeclaration -> {
+                val res = ArrayList<PsiElement?>()
+                psi.children.forEach {
+                    res.addAll(findContentNode(it).filterNotNull())
+                }
+                return res
+            }
+
             is FeakinContextMapDeclaration -> {
-                PsiTreeUtil.getChildrenOfTypeAsList(psi, FeakinContextMapDeclaration::class.java)
+                return listOf(psi)
             }
 
             is FeakinContextDeclaration -> {
-                PsiTreeUtil.getChildrenOfTypeAsList(psi, FeakinContextDeclaration::class.java)
+                return listOf(psi)
             }
 
             else -> {
-                emptyList()
+                return emptyList()
             }
-        }.map(::FkStructureViewElement).toTypedArray()
+        }
     }
 }
\ No newline at end of file
diff --git a/src/main/kotlin/com/feakin/intellij/structure/FkStructureViewModel.kt b/src/main/kotlin/com/feakin/intellij/structure/FkStructureViewModel.kt
index 201dcee..c98bb71 100644
--- a/src/main/kotlin/com/feakin/intellij/structure/FkStructureViewModel.kt
+++ b/src/main/kotlin/com/feakin/intellij/structure/FkStructureViewModel.kt
@@ -33,15 +33,9 @@ class FkStructureViewModel(editor: Editor?, psiFile: FkFile) :
             is FeakinContextMapDeclaration,
             is FeakinContextDeclaration,
             is FeakinAggregateDeclaration,
-
-            is FeakinContextName,
-            is FeakinContextMapName,
-            is FeakinAggregateName,
             -> true
 
             else -> false
         }
     }
-
-
 }
\ No newline at end of file