From 498c40548bead507e90b6cb67d4c2ff66bbedd74 Mon Sep 17 00:00:00 2001
From: Toshiaki Kameyama <kmymtsak@gmail.com>
Date: Mon, 4 May 2020 19:00:00 +0900
Subject: [PATCH] Structure view: add visibility sorter

#KT-36444 Fixed
---
 .../KotlinStructureViewElement.kt             | 25 ++++++++++++++-----
 .../structureView/KotlinStructureViewModel.kt | 15 +++++++++--
 2 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/structureView/KotlinStructureViewElement.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/structureView/KotlinStructureViewElement.kt
index 54e660ef86d..3a0f625a027 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/structureView/KotlinStructureViewElement.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/structureView/KotlinStructureViewElement.kt
@@ -44,9 +44,9 @@ class KotlinStructureViewElement(
                 KotlinStructureElementPresentation(isInherited, element, countDescriptor())
             }
 
-    var isPublic
+    var visibility
             by AssignableLazyProperty {
-                isPublic(countDescriptor())
+                Visibility(countDescriptor())
             }
         private set
 
@@ -54,7 +54,7 @@ class KotlinStructureViewElement(
         if (element !is KtElement) {
             // Avoid storing descriptor in fields
             kotlinPresentation = KotlinStructureElementPresentation(isInherited, element, descriptor)
-            isPublic = isPublic(descriptor)
+            visibility = Visibility(descriptor)
         }
     }
 
@@ -103,9 +103,6 @@ class KotlinStructureViewElement(
         return result
     }
 
-    private fun isPublic(descriptor: DeclarationDescriptor?) =
-        (descriptor as? DeclarationDescriptorWithVisibility)?.visibility == Visibilities.PUBLIC
-
     private fun countDescriptor(): DeclarationDescriptor? {
         val element = element
         return when {
@@ -120,6 +117,22 @@ class KotlinStructureViewElement(
             }
         }
     }
+
+    class Visibility(descriptor: DeclarationDescriptor?) {
+        private val visibility = (descriptor as? DeclarationDescriptorWithVisibility)?.visibility
+
+        val isPublic: Boolean
+            get() = visibility == Visibilities.PUBLIC
+
+        val accessLevel: Int?
+            get() = when {
+                visibility == Visibilities.PUBLIC -> 1
+                visibility == Visibilities.INTERNAL -> 2
+                visibility == Visibilities.PROTECTED -> 3
+                visibility?.let { Visibilities.isPrivate(it) } == true -> 4
+                else -> null
+            }
+    }
 }
 
 private class AssignableLazyProperty<in R, T : Any>(val init: () -> T) : ReadWriteProperty<R, T> {
diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/structureView/KotlinStructureViewModel.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/structureView/KotlinStructureViewModel.kt
index 337c7d5ec4a..9a26d7ae547 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/structureView/KotlinStructureViewModel.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/structureView/KotlinStructureViewModel.kt
@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.idea.structureView
 import com.intellij.ide.structureView.StructureViewModel
 import com.intellij.ide.structureView.StructureViewModelBase
 import com.intellij.ide.structureView.StructureViewTreeElement
+import com.intellij.ide.structureView.impl.java.VisibilitySorter
 import com.intellij.ide.util.treeView.smartTree.*
 import com.intellij.openapi.editor.Editor
 import com.intellij.psi.PsiElement
@@ -22,7 +23,7 @@ class KotlinStructureViewModel(ktFile: KtFile, editor: Editor?) :
     StructureViewModel.ElementInfoProvider {
 
     init {
-        withSorters(Sorter.ALPHA_SORTER)
+        withSorters(KotlinVisibilitySorter, Sorter.ALPHA_SORTER)
     }
 
     override fun isSuitable(element: PsiElement?): Boolean = element is KtDeclaration &&
@@ -51,9 +52,19 @@ class KotlinStructureViewModel(ktFile: KtFile, editor: Editor?) :
     }
 }
 
+object KotlinVisibilitySorter : VisibilitySorter() {
+    override fun getComparator() = Comparator<Any> { a1, a2 -> a1.accessLevel() - a2.accessLevel() }
+
+    private fun Any.accessLevel() = (this as? KotlinStructureViewElement)?.visibility?.accessLevel ?: Int.MAX_VALUE
+
+    override fun getName() = ID
+
+    const val ID = "KOTLIN_VISIBILITY_SORTER"
+}
+
 object PublicElementsFilter : Filter {
     override fun isVisible(treeNode: TreeElement): Boolean {
-        return (treeNode as? KotlinStructureViewElement)?.isPublic ?: true
+        return (treeNode as? KotlinStructureViewElement)?.visibility?.isPublic ?: true
     }
 
     override fun getPresentation(): ActionPresentation {