enums and enum variants in structure view

pontem-network · Nov 8, 2024 · 6c04aba · 6c04aba
1 parent 7b565bf
commit 6c04aba
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 1 deletion.
diff --git a/src/main/kotlin/org/move/ide/structureView/MvStructureViewTreeElement.kt b/src/main/kotlin/org/move/ide/structureView/MvStructureViewTreeElement.kt
@@ -48,11 +48,13 @@ class MvStructureViewTreeElement(val element: NavigatablePsiElement): StructureV
                 listOf(
                     element.consts(),
                     element.structs(),
+                    element.enumList,
                     element.allFunctions(),
                     element.specFunctions(),
                 ).flatten()
             }
-            is MvStruct -> element.namedFields
+            is MvFieldsOwner -> element.namedFields
+            is MvEnum -> element.variants
             else -> emptyList()
         }
         return items.map { MvStructureViewTreeElement(it) }.toTypedArray()

diff --git a/src/main/kotlin/org/move/lang/core/psi/ext/MvEnum.kt b/src/main/kotlin/org/move/lang/core/psi/ext/MvEnum.kt
@@ -1,6 +1,8 @@
 package org.move.lang.core.psi.ext
 
+import com.intellij.ide.projectView.PresentationData
 import com.intellij.lang.ASTNode
+import com.intellij.navigation.ItemPresentation
 import com.intellij.psi.PsiElement
 import com.intellij.psi.stubs.IStubElementType
 import org.move.ide.MoveIcons
@@ -41,5 +43,15 @@ abstract class MvEnumMixin: MvStubbedNamedElementImpl<MvEnumStub>,
             return ItemQualName(this, moduleFQName.address, moduleFQName.itemName, itemName)
         }
 
+    override fun getPresentation(): ItemPresentation? {
+        val enumName = this.name ?: return null
+        return PresentationData(
+            enumName,
+            this.locationString(true),
+            MoveIcons.STRUCT,
+            null
+        )
+    }
+
     override val abilitiesList: MvAbilitiesList? get() = abilitiesListList.firstOrNull()
 }
diff --git a/src/main/kotlin/org/move/lang/core/psi/ext/MvEnumVariant.kt b/src/main/kotlin/org/move/lang/core/psi/ext/MvEnumVariant.kt
@@ -1,6 +1,8 @@
 package org.move.lang.core.psi.ext
 
+import com.intellij.ide.projectView.PresentationData
 import com.intellij.lang.ASTNode
+import com.intellij.navigation.ItemPresentation
 import com.intellij.psi.stubs.IStubElementType
 import org.move.ide.MoveIcons
 import org.move.lang.core.psi.MvEnum
@@ -20,4 +22,25 @@ abstract class MvEnumVariantMixin: MvStubbedNamedElementImpl<MvEnumVariantStub>,
     constructor(stub: MvEnumVariantStub, nodeType: IStubElementType<*, *>): super(stub, nodeType)
 
     override fun getIcon(flags: Int): Icon = MoveIcons.STRUCT
+
+    override fun getPresentation(): ItemPresentation? {
+        val variant = this
+        val variantName = this.name ?: return null
+        val presentationText = buildString {
+            append(variantName)
+            val fields = variant.tupleFields
+            if (fields != null) {
+                append('(')
+                val xs = fields.tupleFieldDeclList.map { it.type.text }
+                append(xs.joinToString(", "))
+                append(')')
+            }
+        }
+        return PresentationData(
+            presentationText,
+            this.locationString(true),
+            MoveIcons.ENUM_VARIANT,
+            null
+        )
+    }
 }
diff --git a/src/test/kotlin/org/move/ide/structureView/StructureViewTest.kt b/src/test/kotlin/org/move/ide/structureView/StructureViewTest.kt
@@ -43,6 +43,7 @@ address 0x1 {
         struct Struct2 {
             field2: u8
         }
+        enum S { One, Two, Three { val: u8 }, Four(u8, u8)}
         public fun call_pub() {}
         public(friend) fun call_pub_friend() {}
         entry fun call_entry() {}
@@ -59,6 +60,12 @@ address 0x1 {
    field1: u8
   -Struct2
    field2: u8
+  -S
+   One
+   Two
+   -Three
+    val: u8
+   Four(u8, u8)
   call_pub()
   call_pub_friend()
   call_entry()