fix tool window content management

cycodehq · Feb 20, 2024 · 42b4664 · 42b4664
1 parent f30140c
commit 42b4664
Show file tree

Hide file tree

Showing 14 changed files with 97 additions and 54 deletions.
diff --git a/src/main/kotlin/com/cycode/plugin/components/toolWindow/CycodeContentTab.kt b/src/main/kotlin/com/cycode/plugin/components/toolWindow/CycodeContentTab.kt
@@ -0,0 +1,14 @@
+package com.cycode.plugin.components.toolWindow
+
+import com.cycode.plugin.components.toolWindow.components.treeView.TreeView
+import com.intellij.openapi.project.Project
+import javax.swing.JPanel
+
+class CycodeContentTab(project: Project) {
+    private val treeView = TreeView(project)
+
+    fun updateContent(rightPanel: JPanel): JPanel {
+        treeView.refreshTree()
+        return treeView.replaceRightPanel(rightPanel)
+    }
+}
diff --git a/src/main/kotlin/com/cycode/plugin/components/toolWindow/CycodeToolWindowFactory.kt b/src/main/kotlin/com/cycode/plugin/components/toolWindow/CycodeToolWindowFactory.kt
@@ -1,24 +1,21 @@
 package com.cycode.plugin.components.toolWindow
 
-import com.cycode.plugin.CycodeBundle
-import com.cycode.plugin.components.Component
 import com.cycode.plugin.components.toolWindow.components.authContentTab.AuthContentTab
 import com.cycode.plugin.components.toolWindow.components.loadingContentTab.LoadingContentTab
 import com.cycode.plugin.components.toolWindow.components.scanContentTab.ScanContentTab
 import com.cycode.plugin.icons.PluginIcons
-import com.cycode.plugin.services.CycodeService
 import com.cycode.plugin.services.cycode
 import com.cycode.plugin.services.pluginState
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.WriteAction
 import com.intellij.openapi.project.DumbAware
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.Disposer
 import com.intellij.openapi.wm.ToolWindow
 import com.intellij.openapi.wm.ToolWindowFactory
-import com.intellij.openapi.wm.ToolWindowManager
 import com.intellij.ui.content.Content
 import com.intellij.ui.content.ContentFactory
-
+import javax.swing.JPanel
 
 class CycodeToolWindowFactory : DumbAware, ToolWindowFactory {
 
@@ -27,48 +24,58 @@ class CycodeToolWindowFactory : DumbAware, ToolWindowFactory {
     }
 
     override fun createToolWindowContent(project: Project, toolWindow: ToolWindow) {
-        createLoadingTabOnly(project)
+        val service = cycode(project)
+
+        val contentTab = CycodeContentTab(project)
+        TabManager.addTab(project, contentTab)
+
+        val defaultRightPanel = LoadingContentTab().getContent(service)
+        val toolWindowContent = createToolWindowContent(contentTab.updateContent(defaultRightPanel))
+        toolWindow.contentManager.addContent(toolWindowContent)
+
+        Disposer.register(service, toolWindowContent)
     }
 
     override fun shouldBeAvailable(project: Project) = true
-}
 
-private fun getCycodeToolWindow(project: Project): ToolWindow? {
-    return ToolWindowManager.getInstance(project).getToolWindow(CycodeBundle.message("toolWindowId"))
-}
+    object TabManager {
+        private val toolWindowsTabs = mutableMapOf<Project, CycodeContentTab>()
 
-private fun replaceToolWindowContent(project: Project, content: Content) {
-    val window = getCycodeToolWindow(project) ?: return
-    window.contentManager.removeAllContents(true)
-    window.contentManager.addContent(content)
-}
+        fun addTab(project: Project, tab: CycodeContentTab) {
+            toolWindowsTabs[project] = tab
+        }
 
-private fun createToolWindowContent(project: Project, component: Component<CycodeService>): Content {
-    return ContentFactory.SERVICE.getInstance()
-        .createContent(component.getContent(cycode(project)), null, false)
-}
+        fun getTab(project: Project): CycodeContentTab? {
+            return toolWindowsTabs[project]
+        }
 
-private fun createLoadingTabOnly(project: Project) {
-    replaceToolWindowContent(project, createToolWindowContent(project, LoadingContentTab()))
-}
+        fun removeTab(project: Project) {
+            toolWindowsTabs.remove(project)
+        }
+    }
 
+}
 
-private fun createAuthTabOnly(project: Project) {
-    replaceToolWindowContent(project, createToolWindowContent(project, AuthContentTab()))
+private fun replaceToolWindowRightPanel(project: Project, panel: JPanel) {
+    val contentTab = CycodeToolWindowFactory.TabManager.getTab(project) ?: return
+    contentTab.updateContent(panel)
 }
 
-private fun createScanTabOnly(project: Project) {
-    replaceToolWindowContent(project, createToolWindowContent(project, ScanContentTab()))
+private fun createToolWindowContent(component: JPanel): Content {
+    return ContentFactory.SERVICE.getInstance().createContent(component, null, false)
 }
 
+
 fun updateToolWindowState(project: Project) {
     val pluginState = pluginState()
+    val service = cycode(project)
+
     ApplicationManager.getApplication().invokeLater {
         WriteAction.run<Error> {
             if (pluginState.cliAuthed) {
-                createScanTabOnly(project)
+                replaceToolWindowRightPanel(project, ScanContentTab().getContent(service))
             } else {
-                createAuthTabOnly(project)
+                replaceToolWindowRightPanel(project, AuthContentTab().getContent(service))
             }
         }
     }

diff --git a/...otlin/com/cycode/plugin/components/toolWindow/components/scanContentTab/ScanContentTab.kt b/...otlin/com/cycode/plugin/components/toolWindow/components/scanContentTab/ScanContentTab.kt
@@ -3,7 +3,6 @@ package com.cycode.plugin.components.toolWindow.components.scanContentTab
 import com.cycode.plugin.CycodeBundle
 import com.cycode.plugin.components.Component
 import com.cycode.plugin.components.common.createClickableLabel
-import com.cycode.plugin.components.toolWindow.components.scanContentTab.components.treeView.TreeView
 import com.cycode.plugin.services.CycodeService
 import com.intellij.util.ui.JBUI
 import java.awt.GridBagConstraints
@@ -13,7 +12,7 @@ import javax.swing.JPanel
 
 class ScanContentTab : Component<CycodeService>() {
     override fun getContent(service: CycodeService): JPanel {
-        val rightPanel = JPanel().apply {
+        return JPanel().apply {
             layout = GridBagLayout()
             add(add(JPanel().apply {
                 add(createClickableLabel(CycodeBundle.message("scanTabTitleLabel")))
@@ -62,7 +61,5 @@ class ScanContentTab : Component<CycodeService>() {
                 anchor = GridBagConstraints.NORTHWEST
             })
         }
-
-        return TreeView(service.project, rightPanel)
     }
 }
diff --git a/...b/components/treeView/TreeCellRenderer.kt → ...w/components/treeView/TreeCellRenderer.kt b/...b/components/treeView/TreeCellRenderer.kt → ...w/components/treeView/TreeCellRenderer.kt
@@ -1,6 +1,6 @@
-package com.cycode.plugin.components.toolWindow.components.scanContentTab.components.treeView
+package com.cycode.plugin.components.toolWindow.components.treeView
 
-import com.cycode.plugin.components.toolWindow.components.scanContentTab.components.treeView.nodes.AbstractNode
+import com.cycode.plugin.components.toolWindow.components.treeView.nodes.AbstractNode
 import com.intellij.ui.ColoredTreeCellRenderer
 import com.intellij.ui.JBColor
 import com.intellij.ui.SimpleTextAttributes

diff --git a/...ontentTab/components/treeView/TreeView.kt → ...oolWindow/components/treeView/TreeView.kt b/...ontentTab/components/treeView/TreeView.kt → ...oolWindow/components/treeView/TreeView.kt
@@ -1,4 +1,4 @@
-package com.cycode.plugin.components.toolWindow.components.scanContentTab.components.treeView
+package com.cycode.plugin.components.toolWindow.components.treeView
 
 import com.cycode.plugin.CycodeBundle
 import com.cycode.plugin.cli.CliResult
@@ -7,7 +7,7 @@ import com.cycode.plugin.cli.models.scanResult.DetectionBase
 import com.cycode.plugin.cli.models.scanResult.ScanResultBase
 import com.cycode.plugin.cli.models.scanResult.sca.ScaDetection
 import com.cycode.plugin.cli.models.scanResult.secret.SecretDetection
-import com.cycode.plugin.components.toolWindow.components.scanContentTab.components.treeView.nodes.*
+import com.cycode.plugin.components.toolWindow.components.treeView.nodes.*
 import com.cycode.plugin.icons.PluginIcons
 import com.cycode.plugin.services.scanResults
 import com.intellij.openapi.project.Project
@@ -27,16 +27,23 @@ import javax.swing.tree.TreeSelectionModel
 
 const val DIFFERENCE_BETWEEN_SCA_LINE_NUMBERS = 1
 
-class TreeView(val project: Project, defaultRightPane: Component) : JPanel(GridLayout(1, 0)), TreeSelectionListener {
+class TreeView(
+    val project: Project, defaultRightPane: Component? = null
+) : JPanel(GridLayout(1, 0)), TreeSelectionListener {
     private val tree: Tree
+
+    // dummyRootNode is a workaround to allow us to hide the root node of the tree
+    private val dummyRootNode = createNode(DummyNode())
     private val rootNodes: RootNodes = RootNodes()
+
+    private val splitPane: JSplitPane = JSplitPane(JSplitPane.HORIZONTAL_SPLIT)
+
     private val scanResults = scanResults(project)
 
     init {
-        val top = createNode(DummyNode())
-        createNodes(top)
+        createNodes(dummyRootNode)
 
-        tree = Tree(top)
+        tree = Tree(dummyRootNode)
         tree.setRootVisible(false)
         tree.setCellRenderer(TreeCellRenderer())
 
@@ -45,17 +52,17 @@ class TreeView(val project: Project, defaultRightPane: Component) : JPanel(GridL
 
         val treeView = JBScrollPane(tree)
 
-        val splitPane = JSplitPane(JSplitPane.HORIZONTAL_SPLIT)
         flattenJSplitPane(splitPane)
-
         splitPane.leftComponent = treeView
-        splitPane.rightComponent = defaultRightPane
+        splitPane.resizeWeight = 0.5
 
         val minimumSize = Dimension(400, 100)
-        defaultRightPane.minimumSize = minimumSize
         treeView.minimumSize = minimumSize
 
-        splitPane.resizeWeight = 0.5
+        if (defaultRightPane != null) {
+            splitPane.rightComponent = defaultRightPane
+            defaultRightPane.minimumSize = minimumSize
+        }
 
         add(splitPane)
     }
@@ -168,6 +175,18 @@ class TreeView(val project: Project, defaultRightPane: Component) : JPanel(GridL
         createDetectionNodes(CliScanType.Sca, scaDetections.result, ::createScaDetectionNode)
     }
 
+    fun replaceRightPanel(newRightPanel: Component): TreeView {
+        splitPane.rightComponent = newRightPanel
+        return this
+    }
+
+    fun refreshTree() {
+        // TODO(MarshalX): is possible to optimize this to only update the nodes that have changed
+        dummyRootNode.removeAllChildren()
+        createNodes(dummyRootNode)
+        tree.updateUI()
+    }
+
     private fun createNodes(top: DefaultMutableTreeNode) {
         rootNodes.createNodes(top)
         createSecretDetectionNodes()

diff --git a/...components/treeView/nodes/AbstractNode.kt → ...components/treeView/nodes/AbstractNode.kt b/...components/treeView/nodes/AbstractNode.kt → ...components/treeView/nodes/AbstractNode.kt
@@ -1,4 +1,4 @@
-package com.cycode.plugin.components.toolWindow.components.scanContentTab.components.treeView.nodes
+package com.cycode.plugin.components.toolWindow.components.treeView.nodes
 
 import javax.swing.Icon
 

diff --git a/...mponents/treeView/nodes/DetectionNodes.kt → ...mponents/treeView/nodes/DetectionNodes.kt b/...mponents/treeView/nodes/DetectionNodes.kt → ...mponents/treeView/nodes/DetectionNodes.kt
@@ -1,4 +1,4 @@
-package com.cycode.plugin.components.toolWindow.components.scanContentTab.components.treeView.nodes
+package com.cycode.plugin.components.toolWindow.components.treeView.nodes
 
 import com.cycode.plugin.cli.models.scanResult.sca.ScaDetection
 import com.cycode.plugin.cli.models.scanResult.secret.SecretDetection

diff --git a/...ab/components/treeView/nodes/DummyNode.kt → ...ow/components/treeView/nodes/DummyNode.kt b/...ab/components/treeView/nodes/DummyNode.kt → ...ow/components/treeView/nodes/DummyNode.kt
@@ -1,4 +1,4 @@
-package com.cycode.plugin.components.toolWindow.components.scanContentTab.components.treeView.nodes
+package com.cycode.plugin.components.toolWindow.components.treeView.nodes
 
 import com.cycode.plugin.CycodeBundle
 import javax.swing.Icon

diff --git a/...Tab/components/treeView/nodes/FileNode.kt → ...dow/components/treeView/nodes/FileNode.kt b/...Tab/components/treeView/nodes/FileNode.kt → ...dow/components/treeView/nodes/FileNode.kt
@@ -1,4 +1,4 @@
-package com.cycode.plugin.components.toolWindow.components.scanContentTab.components.treeView.nodes
+package com.cycode.plugin.components.toolWindow.components.treeView.nodes
 
 import javax.swing.Icon
 

diff --git a/...ab/components/treeView/nodes/RootNodes.kt → ...ow/components/treeView/nodes/RootNodes.kt b/...ab/components/treeView/nodes/RootNodes.kt → ...ow/components/treeView/nodes/RootNodes.kt
@@ -1,4 +1,4 @@
-package com.cycode.plugin.components.toolWindow.components.scanContentTab.components.treeView.nodes
+package com.cycode.plugin.components.toolWindow.components.treeView.nodes
 
 import com.cycode.plugin.CycodeBundle
 import com.cycode.plugin.cli.CliScanType

diff --git a/...components/treeView/nodes/ScanTypeNode.kt → ...components/treeView/nodes/ScanTypeNode.kt b/...components/treeView/nodes/ScanTypeNode.kt → ...components/treeView/nodes/ScanTypeNode.kt
@@ -1,4 +1,4 @@
-package com.cycode.plugin.components.toolWindow.components.scanContentTab.components.treeView.nodes
+package com.cycode.plugin.components.toolWindow.components.treeView.nodes
 
 import javax.swing.Icon
 

diff --git a/...entTab/components/treeView/nodes/utils.kt → ...Window/components/treeView/nodes/utils.kt b/...entTab/components/treeView/nodes/utils.kt → ...Window/components/treeView/nodes/utils.kt
@@ -1,4 +1,4 @@
-package com.cycode.plugin.components.toolWindow.components.scanContentTab.components.treeView.nodes
+package com.cycode.plugin.components.toolWindow.components.treeView.nodes
 
 import javax.swing.tree.DefaultMutableTreeNode
 

diff --git a/...anContentTab/components/treeView/utils.kt → ...s/toolWindow/components/treeView/utils.kt b/...anContentTab/components/treeView/utils.kt → ...s/toolWindow/components/treeView/utils.kt
@@ -1,4 +1,4 @@
-package com.cycode.plugin.components.toolWindow.components.scanContentTab.components.treeView
+package com.cycode.plugin.components.toolWindow.components.treeView
 
 import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.fileEditor.OpenFileDescriptor

diff --git a/src/main/kotlin/com/cycode/plugin/services/CycodeService.kt b/src/main/kotlin/com/cycode/plugin/services/CycodeService.kt
@@ -1,8 +1,10 @@
 package com.cycode.plugin.services
 
 import com.cycode.plugin.CycodeBundle
+import com.cycode.plugin.components.toolWindow.CycodeToolWindowFactory
 import com.cycode.plugin.components.toolWindow.updateToolWindowState
 import com.cycode.plugin.utils.CycodeNotifier
+import com.intellij.openapi.Disposable
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.progress.ProgressIndicator
@@ -12,7 +14,7 @@ import com.intellij.psi.PsiDocumentManager
 
 
 @Service(Service.Level.PROJECT)
-class CycodeService(val project: Project) {
+class CycodeService(val project: Project) : Disposable {
     private val cliService = cli(project)
     private val cliDownloadService = cliDownload()
 
@@ -108,4 +110,8 @@ class CycodeService(val project: Project) {
 
         startPathScaScan(projectRoot)
     }
+
+    override fun dispose() {
+        CycodeToolWindowFactory.TabManager.removeTab(project)
+    }
 }