5b226a688094a655c50708436034bde64823ae91
[idea/community.git] / plugins / git4idea / src / git4idea / ui / branch / dashboard / BranchesTreeModel.kt
1 // Copyright 2000-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
2 package git4idea.ui.branch.dashboard
3
4 import com.intellij.openapi.actionSystem.DataKey
5 import com.intellij.util.ThreeState
6 import git4idea.i18n.GitBundle.message
7 import git4idea.repo.GitRepository
8 import java.util.*
9 import javax.swing.tree.DefaultMutableTreeNode
10
11 internal val GIT_BRANCHES = DataKey.create<Set<BranchInfo>>("GitBranchKey")
12 internal val GIT_BRANCH_FILTERS = DataKey.create<List<String>>("GitBranchFilterKey")
13
14 internal data class BranchInfo(val branchName: String,
15                                val isLocal: Boolean,
16                                val isCurrent: Boolean,
17                                var isFavorite: Boolean,
18                                val repositories: List<GitRepository>) {
19   var isMy: ThreeState = ThreeState.UNSURE
20   override fun toString() = branchName
21 }
22
23 internal data class BranchNodeDescriptor(val type: NodeType,
24                                          val branchInfo: BranchInfo? = null,
25                                          val displayName: String? = branchInfo?.branchName,
26                                          val parent: BranchNodeDescriptor? = null) {
27   override fun toString(): String {
28     val suffix = branchInfo?.branchName ?: displayName
29     return if (suffix != null) "$type:$suffix" else "$type"
30   }
31
32   fun getDisplayText() = displayName ?: branchInfo?.branchName
33 }
34
35 internal enum class NodeType {
36   ROOT, LOCAL_ROOT, REMOTE_ROOT, BRANCH, GROUP_NODE, HEAD_NODE
37 }
38
39 internal class BranchTreeNode(nodeDescriptor: BranchNodeDescriptor) : DefaultMutableTreeNode(nodeDescriptor) {
40
41   fun getTextRepresentation(): String {
42     val nodeDescriptor = userObject as? BranchNodeDescriptor ?: return super.toString()
43     return when (nodeDescriptor.type) {
44       NodeType.LOCAL_ROOT -> message("group.Git.Local.Branch.title")
45       NodeType.REMOTE_ROOT -> message("group.Git.Remote.Branch.title")
46       NodeType.HEAD_NODE -> message("group.Git.HEAD.Branch.Filter.title")
47       else -> nodeDescriptor.getDisplayText() ?: super.toString()
48     }
49   }
50
51   fun getNodeDescriptor() = userObject as BranchNodeDescriptor
52
53   override fun equals(other: Any?): Boolean {
54     if (this === other) return true
55     if (javaClass != other?.javaClass) return false
56     if (other !is BranchTreeNode) return false
57
58     return Objects.equals(this.userObject, other.userObject)
59   }
60
61   override fun hashCode() = Objects.hash(userObject)
62 }
63
64 internal class NodeDescriptorsModel(private val localRootNodeDescriptor: BranchNodeDescriptor,
65                                     private val remoteRootNodeDescriptor: BranchNodeDescriptor) {
66   /**
67    * Parent node descriptor to direct children map
68    */
69   private val branchNodeDescriptors = hashMapOf<BranchNodeDescriptor, MutableSet<BranchNodeDescriptor>>()
70
71   fun clear() = branchNodeDescriptors.clear()
72
73   fun getChildrenForParent(parent: BranchNodeDescriptor): Set<BranchNodeDescriptor> =
74     branchNodeDescriptors.getOrDefault(parent, emptySet())
75
76   fun populateFrom(branches: Sequence<BranchInfo>, useGrouping: Boolean) {
77     branches.forEach { branch -> populateFrom(branch, useGrouping) }
78   }
79
80   private fun populateFrom(br: BranchInfo, useGrouping: Boolean) {
81     val branch = with(br) { BranchInfo(branchName, isLocal, isCurrent, isFavorite, repositories) }
82     var curParent: BranchNodeDescriptor = if (branch.isLocal) localRootNodeDescriptor else remoteRootNodeDescriptor
83
84     if (!useGrouping) {
85       addChild(curParent, BranchNodeDescriptor(NodeType.BRANCH, branch, parent = curParent))
86       return
87     }
88
89     val iter = branch.branchName.split("/").iterator()
90
91     while (iter.hasNext()) {
92       val branchNamePart = iter.next()
93       val groupNode = iter.hasNext()
94       val nodeType = if (groupNode) NodeType.GROUP_NODE else NodeType.BRANCH
95       val branchInfo = if (nodeType == NodeType.BRANCH) branch else null
96
97       val branchNodeDescriptor = BranchNodeDescriptor(nodeType, branchInfo, displayName = branchNamePart, parent = curParent)
98       addChild(curParent, branchNodeDescriptor)
99       curParent = branchNodeDescriptor
100     }
101   }
102
103   private fun addChild(parent: BranchNodeDescriptor, child: BranchNodeDescriptor) {
104     val directChildren = branchNodeDescriptors.computeIfAbsent(parent) { sortedSetOf(BRANCH_TREE_NODE_COMPARATOR) }
105     directChildren.add(child)
106     branchNodeDescriptors[parent] = directChildren
107   }
108 }
109
110 internal val BRANCH_TREE_NODE_COMPARATOR = Comparator<BranchNodeDescriptor> { d1, d2 ->
111   val b1 = d1.branchInfo
112   val b2 = d2.branchInfo
113   val displayText1 = d1.getDisplayText()
114   val displayText2 = d2.getDisplayText()
115   val b1GroupNode = d1.type == NodeType.GROUP_NODE
116   val b2GroupNode = d2.type == NodeType.GROUP_NODE
117   val b1Current = b1 != null && b1.isCurrent
118   val b2Current = b2 != null && b2.isCurrent
119   val b1Favorite = b1 != null && b1.isFavorite
120   val b2Favorite = b2 != null && b2.isFavorite
121   fun compareByDisplayTextOrType() =
122     if (displayText1 != null && displayText2 != null) displayText1.compareTo(displayText2) else d1.type.compareTo(d2.type)
123
124   when {
125     b1Current && b2Current -> compareByDisplayTextOrType()
126     b1Current -> -1
127     b2Current -> 1
128     b1Favorite && b2Favorite -> compareByDisplayTextOrType()
129     b1Favorite -> -1
130     b2Favorite -> 1
131     b1GroupNode && b2GroupNode -> compareByDisplayTextOrType()
132     b1GroupNode -> -1
133     b2GroupNode -> 1
134     else -> compareByDisplayTextOrType()
135   }
136 }