From 6c04aba35932fbc4144954f24724b250a3946006 Mon Sep 17 00:00:00 2001
From: Maksim Kurnikov <maxim.kurnikov@gmail.com>
Date: Fri, 8 Nov 2024 01:36:25 +0100
Subject: [PATCH] enums and enum variants in structure view

---
 .../MvStructureViewTreeElement.kt             |  4 +++-
 .../org/move/lang/core/psi/ext/MvEnum.kt      | 12 ++++++++++
 .../move/lang/core/psi/ext/MvEnumVariant.kt   | 23 +++++++++++++++++++
 .../ide/structureView/StructureViewTest.kt    |  7 ++++++
 4 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/src/main/kotlin/org/move/ide/structureView/MvStructureViewTreeElement.kt b/src/main/kotlin/org/move/ide/structureView/MvStructureViewTreeElement.kt
index 5ff359216..dfdf55214 100644
--- 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
index 67091550f..013251990 100644
--- 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()
 }
\ No newline at end of file
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
index 5d5912f22..37ea8ea9e 100644
--- 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
+        )
+    }
 }
\ No newline at end of file
diff --git a/src/test/kotlin/org/move/ide/structureView/StructureViewTest.kt b/src/test/kotlin/org/move/ide/structureView/StructureViewTest.kt
index 30078e2ce..bac268d04 100644
--- 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()