clean up code

src/main/kotlin/callgraph/CallGraphToolWindow.kt

@@ -2,6 +2,7 @@ package callgraph
 
 import com.intellij.ide.util.EditorHelper
 import com.intellij.psi.PsiMethod
+import java.awt.Dimension
 import java.awt.event.KeyEvent
 import java.awt.event.KeyListener
 import java.awt.geom.Point2D
@@ -106,8 +107,8 @@ class CallGraphToolWindow {
         this.showOnlyDownstreamButton.addActionListener { run(CanvasConfig.BuildType.DOWNSTREAM) }
         this.showOnlyUpstreamDownstreamButton.addActionListener { run(CanvasConfig.BuildType.UPSTREAM_DOWNSTREAM) }
         this.viewSourceCodeButton.addActionListener { viewSourceCodeHandler() }
-        this.fitGraphToViewButton.addActionListener { fitGraphToViewButtonHandler() }
-        this.fitGraphToBestRatioButton.addActionListener { fitGraphToBestRatioButtonHandler() }
+        this.fitGraphToViewButton.addActionListener { this.canvas.fitCanvasToView() }
+        this.fitGraphToBestRatioButton.addActionListener { this.canvas.fitCanvasToBestRatio() }
         this.increaseXGridButton.addActionListener { gridSizeButtonHandler(isXGrid = true, isIncrease = true) }
         this.decreaseXGridButton.addActionListener { gridSizeButtonHandler(isXGrid = true, isIncrease = false) }
         this.increaseYGridButton.addActionListener { gridSizeButtonHandler(isXGrid = false, isIncrease = true) }
@@ -118,6 +119,8 @@ class CallGraphToolWindow {
         this.canvas.addMouseListener(mouseEventHandler)
         this.canvas.addMouseMotionListener(mouseEventHandler)
         this.canvas.addMouseWheelListener(mouseEventHandler)
+        this.canvas.isVisible = false
+        this.canvasPanel.add(this.canvas)
     }
 
     fun getContent(): JPanel {
@@ -203,19 +206,24 @@ class CallGraphToolWindow {
         return this.filterAccessPrivateCheckbox.isSelected
     }
 
+    fun getCanvasSize(): Dimension = this.canvasPanel.size
+
     fun run(buildType: CanvasConfig.BuildType) {
         val project = Utils.getActiveProject()
         if (project != null) {
             Utils.runBackgroundTask(project, Runnable {
                 // set up the config object
-                val canvasConfig = CanvasConfig(project, buildType, this.canvas)
-                canvasConfig.selectedModuleName =
-                        this@CallGraphToolWindow.moduleScopeComboBox.selectedItem as String? ?: ""
-                canvasConfig.selectedDirectoryPath = this@CallGraphToolWindow.directoryScopeTextField.text
-                canvasConfig.focusedMethods = this@CallGraphToolWindow.focusedMethods
-                canvasConfig.callGraphToolWindow = this@CallGraphToolWindow
+                val canvasConfig = CanvasConfig(
+                        project,
+                        buildType,
+                        this.canvas,
+                        this@CallGraphToolWindow.moduleScopeComboBox.selectedItem as String? ?: "",
+                        this@CallGraphToolWindow.directoryScopeTextField.text,
+                        this@CallGraphToolWindow.focusedMethods,
+                        this@CallGraphToolWindow
+                )
                 // start building graph
-                setupUiBeforeRun(canvasConfig)
+                setupUiBeforeRun(buildType)
                 this@CallGraphToolWindow.canvasBuilder.build(canvasConfig)
                 setupUiAfterRun()
             })
@@ -260,10 +268,6 @@ class CallGraphToolWindow {
         }
     }
 
-    private fun fitGraphToViewButtonHandler() = this.canvas.fitCanvasToView()
-
-    private fun fitGraphToBestRatioButtonHandler() = this.canvas.fitCanvasToBestRatio()
-
     private fun gridSizeButtonHandler(isXGrid: Boolean, isIncrease: Boolean) {
         val zoomFactor = if (isIncrease) 1.25f else 1 / 1.25f
         val xZoomFactor = if (isXGrid) zoomFactor else 1.0f
@@ -279,19 +283,18 @@ class CallGraphToolWindow {
         this.focusedMethods.forEach { EditorHelper.openInEditor(it) }
     }
 
-    private fun setupUiBeforeRun(canvasConfig: CanvasConfig) {
+    private fun setupUiBeforeRun(buildType: CanvasConfig.BuildType) {
         // focus on the 'graph tab
         this.mainTabbedPanel.getComponentAt(1).isEnabled = true
         this.mainTabbedPanel.selectedIndex = 1
         // stats label
         this.statsLabel.text = "..."
         // build-type label
-        val buildTypeText = canvasConfig.buildType.label
-        when (canvasConfig.buildType) {
+        when (buildType) {
             CanvasConfig.BuildType.WHOLE_PROJECT_WITH_TEST_LIMITED,
             CanvasConfig.BuildType.WHOLE_PROJECT_WITH_TEST,
             CanvasConfig.BuildType.WHOLE_PROJECT_WITHOUT_TEST_LIMITED,
-            CanvasConfig.BuildType.WHOLE_PROJECT_WITHOUT_TEST -> this.buildTypeLabel.text = buildTypeText
+            CanvasConfig.BuildType.WHOLE_PROJECT_WITHOUT_TEST -> this.buildTypeLabel.text = buildType.label
             CanvasConfig.BuildType.MODULE_LIMITED,
             CanvasConfig.BuildType.MODULE -> {
                 val moduleName = this.moduleScopeComboBox.selectedItem as String
@@ -306,7 +309,7 @@ class CallGraphToolWindow {
             CanvasConfig.BuildType.DOWNSTREAM,
             CanvasConfig.BuildType.UPSTREAM_DOWNSTREAM -> {
                 val functionNames = this.focusedMethods.joinToString { it.name }
-                this.buildTypeLabel.text = "<html>$buildTypeText of function <b>$functionNames</b></html>"
+                this.buildTypeLabel.text = "<html>${buildType.label} of function <b>$functionNames</b></html>"
             }
         }
         // disable some checkboxes and buttons
@@ -335,18 +338,17 @@ class CallGraphToolWindow {
         // progress bar
         this.loadingProgressBar.isVisible = true
         // clear the canvas panel, ready for new graph
-        this.canvasPanel.removeAll()
+        this.canvas.isVisible = false
     }
 
     private fun setupUiAfterRun() {
+        // hide progress bar
+        this.loadingProgressBar.isVisible = false
         // show the rendered canvas
-        this.canvas.canvasPanel = this.canvasPanel
-        this.canvasPanel.add(this.canvas)
+        this.canvas.isVisible = true
         this.canvasPanel.updateUI()
         // stats label
         this.statsLabel.text = "${this.canvas.getNodesCount()} methods"
-        // hide progress bar
-        this.loadingProgressBar.isVisible = false
         // enable some checkboxes and buttons
         enableFocusedMethodButtons()
         listOf(
@@ -388,10 +390,11 @@ class CallGraphToolWindow {
     }
 
     private fun enableFocusedMethodButtons() {
-        val isEnabled = this.focusedMethods.isNotEmpty()
-        this.showOnlyUpstreamButton.isEnabled = isEnabled
-        this.showOnlyDownstreamButton.isEnabled = isEnabled
-        this.showOnlyUpstreamDownstreamButton.isEnabled = isEnabled
-        this.viewSourceCodeButton.isEnabled = isEnabled
+        listOf(
+                this.showOnlyUpstreamButton,
+                this.showOnlyDownstreamButton,
+                this.showOnlyUpstreamDownstreamButton,
+                this.viewSourceCodeButton
+        ).forEach { it.isEnabled = this.focusedMethods.isNotEmpty() }
     }
 }

src/main/kotlin/callgraph/CallGraphToolWindowProjectService.kt

@@ -1,5 +1,5 @@
 package callgraph
 
 class CallGraphToolWindowProjectService {
-    var callGraphToolWindow: CallGraphToolWindow? = null
+    lateinit var callGraphToolWindow: CallGraphToolWindow
 }

src/main/kotlin/callgraph/Canvas.kt

@@ -16,8 +16,8 @@ class Canvas(private val callGraphToolWindow: CallGraphToolWindow): JPanel() {
     private val solidLineStroke = BasicStroke(regularLineWidth)
     private val methodAccessColorMap = mapOf<String, Color>(
             PsiModifier.PUBLIC to Colors.GREEN.color,
-            PsiModifier.PROTECTED to Colors.CYAN.color,
-            PsiModifier.PACKAGE_LOCAL to Colors.LIGHT_ORANGE.color,
+            PsiModifier.PROTECTED to Colors.LIGHT_ORANGE.color,
+            PsiModifier.PACKAGE_LOCAL to Colors.BLUE.color,
             PsiModifier.PRIVATE to Colors.RED.color
     )
     private val heatMapColors = listOf(
@@ -32,9 +32,8 @@ class Canvas(private val callGraphToolWindow: CallGraphToolWindow): JPanel() {
             Colors.ORANGE.color,
             Colors.RED.color
     )
-    var canvasPanel: JPanel? = null
     var cameraOrigin = defaultCameraOrigin
-    private var graph: Graph? = null
+    private var graph = Graph()
     private var visibleNodes = setOf<Node>()
     private var visibleEdges = setOf<Edge>()
     private var nodeShapesMap = mutableMapOf<Shape, Node>()
@@ -43,9 +42,6 @@ class Canvas(private val callGraphToolWindow: CallGraphToolWindow): JPanel() {
     private var yZoomRatio = defaultZoomRatio
 
     override fun paintComponent(graphics: Graphics) {
-        if (graph == null) {
-            return
-        }
         super.paintComponent(graphics)
 
         // set up the drawing panel
@@ -108,7 +104,6 @@ class Canvas(private val callGraphToolWindow: CallGraphToolWindow): JPanel() {
         this.graph = graph
         this.visibleNodes = graph.getNodes()
         this.visibleEdges = graph.getEdges()
-        this.canvasPanel = null
         this.cameraOrigin = this.defaultCameraOrigin
         this.nodeShapesMap = mutableMapOf()
         this.hoveredNode = null
@@ -152,9 +147,9 @@ class Canvas(private val callGraphToolWindow: CallGraphToolWindow): JPanel() {
     }
 
     fun fitCanvasToView() {
-        val blueprint = this.graph!!.getNodes().associateBy({ it.id }, { it.rawLayoutPoint })
+        val blueprint = this.graph.getNodes().associateBy({ it.id }, { it.rawLayoutPoint })
         val bestFitBlueprint = Utils.fitLayoutToViewport(blueprint)
-        Utils.applyLayoutBlueprintToGraph(bestFitBlueprint, this.graph!!)
+        Utils.applyLayoutBlueprintToGraph(bestFitBlueprint, this.graph)
         this.cameraOrigin = defaultCameraOrigin
         this.xZoomRatio = defaultZoomRatio
         this.yZoomRatio = defaultZoomRatio
@@ -163,19 +158,19 @@ class Canvas(private val callGraphToolWindow: CallGraphToolWindow): JPanel() {
 
     fun fitCanvasToBestRatio() {
         // set every node coordinate to its original raw layout by GraphViz
-        this.graph!!.getNodes().forEach { it.point = it.rawLayoutPoint }
+        this.graph.getNodes().forEach { it.point = it.rawLayoutPoint }
         this.cameraOrigin = defaultCameraOrigin
         this.xZoomRatio = defaultZoomRatio
         this.yZoomRatio = defaultZoomRatio
         repaint()
     }
 
     fun getNodesCount(): Int {
-        return this.graph!!.getNodes().size
+        return this.graph.getNodes().size
     }
 
     fun filterAccessChangeHandler() {
-        this.visibleNodes = this.graph!!.getNodes()
+        this.visibleNodes = this.graph.getNodes()
                 .filter { node ->
                     val method = node.method
                     when {
@@ -187,14 +182,14 @@ class Canvas(private val callGraphToolWindow: CallGraphToolWindow): JPanel() {
                     }
                 }
                 .toSet()
-        this.visibleEdges = this.graph!!.getEdges()
+        this.visibleEdges = this.graph.getEdges()
                 .filter { this.visibleNodes.contains(it.sourceNode) && this.visibleNodes.contains(it.targetNode) }
                 .toSet()
         repaint()
     }
 
     private fun toCameraView(point: Point2D.Float): Point2D.Float {
-        val canvasSize = this.canvasPanel!!.size
+        val canvasSize = this.callGraphToolWindow.getCanvasSize()
         return Point2D.Float(
                 this.xZoomRatio * point.x * canvasSize.width - this.cameraOrigin.x,
                 this.yZoomRatio * point.y * canvasSize.height - this.cameraOrigin.y

src/main/kotlin/callgraph/CanvasBuilder.kt

@@ -62,10 +62,10 @@ class CanvasBuilder {
         val methodsToParse = Utils.getMethodsFromFiles(filesToParse)
 
         // parse method dependencies
-        canvasConfig.callGraphToolWindow?.resetProgressBar(methodsToParse.size)
+        canvasConfig.callGraphToolWindow.resetProgressBar(methodsToParse.size)
         val newDependencies = methodsToParse
                 .flatMap {
-                    canvasConfig.callGraphToolWindow?.incrementProgressBar()
+                    canvasConfig.callGraphToolWindow.incrementProgressBar()
                     Utils.getDependenciesFromMethod(it)
                 }
                 .toSet()

src/main/kotlin/callgraph/CanvasConfig.kt

@@ -3,7 +3,15 @@ package callgraph
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiMethod
 
-data class CanvasConfig(val project: Project, val buildType: BuildType, val canvas: Canvas) {
+data class CanvasConfig(
+        val project: Project,
+        val buildType: BuildType,
+        val canvas: Canvas,
+        val selectedModuleName: String,
+        val selectedDirectoryPath: String,
+        val focusedMethods: Set<PsiMethod>,
+        val callGraphToolWindow: CallGraphToolWindow
+) {
     enum class BuildType(val label: String) {
         WHOLE_PROJECT_WITH_TEST_LIMITED("Whole project (test files included), limited upstream/downstream scope"),
         WHOLE_PROJECT_WITHOUT_TEST_LIMITED("Whole project (test files excluded), limited upstream/downstream scope"),
@@ -17,9 +25,4 @@ data class CanvasConfig(val project: Project, val buildType: BuildType, val canv
         DOWNSTREAM("Downstream"),
         UPSTREAM_DOWNSTREAM("Upstream & downstream")
     }
-
-    var selectedModuleName = ""
-    var selectedDirectoryPath = ""
-    var focusedMethods = setOf<PsiMethod>()
-    var callGraphToolWindow: CallGraphToolWindow? = null
 }

src/main/kotlin/callgraph/Utils.kt

@@ -96,9 +96,8 @@ object Utils {
     fun getDependenciesFromMethod(method: PsiMethod) =
             PsiTreeUtil
                     .findChildrenOfType(method, PsiIdentifier::class.java)
-                    .map { it.context }
-                    .filter { it != null }
-                    .flatMap { it!!.references.toList() }
+                    .mapNotNull { it.context }
+                    .flatMap { it.references.toList() }
                     .map { it.resolve() }
                     .filter { it is PsiMethod }
                     .map { Dependency(method, it as PsiMethod) }
@@ -128,12 +127,11 @@ object Utils {
     }
 
     fun getMethodPackageName(psiMethod: PsiMethod): String {
-        // get class name
-        val psiClass = psiMethod.containingClass
-        val className = psiClass?.qualifiedName ?: ""
         // get package name
         val psiJavaFile = psiMethod.containingFile as PsiJavaFile
         val packageName = psiJavaFile.packageStatement?.packageName ?: ""
+        // get class name
+        val className = psiMethod.containingClass?.qualifiedName ?: ""
         return if (packageName.isBlank() || className.startsWith(packageName)) className else "$packageName.$className"
     }
 
@@ -184,7 +182,7 @@ object Utils {
                     .getToolWindow("Call Graph")
                     .activate {
                         ServiceManager.getService(project, CallGraphToolWindowProjectService::class.java)
-                                .callGraphToolWindow!!
+                                .callGraphToolWindow
                                 .clearFocusedMethods()
                                 .toggleFocusedMethod(psiElement)
                                 .run(buildType)
@@ -329,28 +327,32 @@ object Utils {
     }
 
     private fun getAverageElementDifference(elements: Set<Int>): Float {
-        return if (elements.size < 2) 0f else (elements.max()!! - elements.min()!!) / (elements.size - 1).toFloat()
+        val max = elements.max()
+        val min = elements.min()
+        return if (elements.size < 2 || max == null || min == null) 0f else (max - min) / (elements.size - 1).toFloat()
     }
 
     private fun mergeNormalizedLayouts(blueprints: List<Map<String, Point2D.Float>>): Map<String, Point2D.Float> {
         if (blueprints.isEmpty()) {
             return emptyMap()
         }
-        val blueprintHeights = blueprints
-                .associateBy(
-                        { it },
-                        { blueprint ->
-                            // set padding to the average y grid size of the previous sub-graph (but minimum 0.1)
-                            val yPoints = blueprint.values.map { it.y }
-                            yPoints.max()!! - yPoints.min()!! + normalizedGridSize
-                        }
-                )
-                .toMap()
-        val sortedBlueprints = blueprintHeights
+        val blueprintSizes = blueprints
+                .map { blueprint ->
+                    val xPoints = blueprint.values.map { it.x }
+                    val xMax = xPoints.max() ?: 0f
+                    val xMin = xPoints.min() ?: 0f
+                    val width = xMax - xMin + normalizedGridSize
+                    val yPoints = blueprint.values.map { it.y }
+                    val yMax = yPoints.max() ?: 0f
+                    val yMin = yPoints.min() ?: 0f
+                    val height = yMax - yMin + normalizedGridSize
+                    Triple(blueprint, height, width)
+                }
+        val sortedHeights = blueprintSizes.map { (_, height, _) -> height }.sortedBy { -it }
+        val sortedBlueprints = blueprintSizes
                 .toList()
-                .sortedBy { (_, height) -> -height }
-                .map { it.first }
-        val sortedHeights = blueprintHeights.values.sortedBy { -it }
+                .sortedWith(compareBy({ (_, height, _) -> -height }, { (_, _, width) -> -width }))
+                .map { (blueprint, _, _) -> blueprint }
         val xBaseline = 0.5f
         val yBaseline = 0.5f
         // put the left-most point of the first sub-graph in the view center, by using its y value as central line
@@ -362,13 +364,12 @@ object Utils {
                     // left align the graph by the left-most nodesMap, then centering the baseline
                     val minX = blueprint.values.map { it.x }.min() ?: 0f
                     //noinspection UnnecessaryLocalVariable
-                    val shiftedBlueprint = blueprint.mapValues { (_, point) ->
+                    blueprint.mapValues { (_, point) ->
                         Point2D.Float(
                                 point.x - minX + xBaseline,
                                 point.y + yOffset - yCentralLine + yBaseline
                         )
                     }
-                    shiftedBlueprint
                 }
                 .reduce { blueprintA, blueprintB -> blueprintA + blueprintB }
     }