Ensure builder resets to a root command trie node

Also refactors command nodes a bit for better debug/trace output

src/main/java/com/maddyhome/idea/vim/extension/nerdtree/NerdTree.kt

@@ -555,12 +555,13 @@ private fun registerCommand(default: String, action: NerdAction) {
 }
 
 
-private val actionsRoot: RootNode<NerdAction> = RootNode()
+private val actionsRoot: RootNode<NerdAction> = RootNode("NERDTree")
 private var currentNode: CommandPartNode<NerdAction> = actionsRoot
+
 private fun collectShortcuts(node: Node<NerdAction>): Set<KeyStroke> {
   return if (node is CommandPartNode<NerdAction>) {
-    val res = node.keys.toMutableSet()
-    res += node.values.map { collectShortcuts(it) }.flatten()
+    val res = node.children.keys.toMutableSet()
+    res += node.children.values.map { collectShortcuts(it) }.flatten()
     res
   } else {
     emptySet()

tests/property-tests/src/test/kotlin/org/jetbrains/plugins/ideavim/propertybased/RandomActionsPropertyTest.kt

@@ -95,7 +95,7 @@ private class AvailableActions(private val editor: Editor) : ImperativeCommand {
     val currentNode = KeyHandler.getInstance().keyHandlerState.commandBuilder.getCurrentTrie()
 
     // Note: esc is always an option
-    val possibleKeys = (currentNode.keys.toList() + esc).sortedBy { injector.parser.toKeyNotation(it) }
+    val possibleKeys = (currentNode.children.keys.toList() + esc).sortedBy { injector.parser.toKeyNotation(it) }
     println("Keys: ${possibleKeys.joinToString(", ")}")
     val keyGenerator = Generator.integers(0, possibleKeys.lastIndex)
       .suchThat { injector.parser.toKeyNotation(possibleKeys[it]) !in stinkyKeysList }

vim-engine/src/main/kotlin/com/maddyhome/idea/vim/KeyHandler.kt

@@ -22,9 +22,9 @@ import com.maddyhome.idea.vim.diagnostic.VimLogger
 import com.maddyhome.idea.vim.diagnostic.trace
 import com.maddyhome.idea.vim.diagnostic.vimLogger
 import com.maddyhome.idea.vim.impl.state.toMappingMode
-import com.maddyhome.idea.vim.key.CommandPartNode
 import com.maddyhome.idea.vim.key.KeyConsumer
 import com.maddyhome.idea.vim.key.KeyStack
+import com.maddyhome.idea.vim.key.RootNode
 import com.maddyhome.idea.vim.key.consumers.CharArgumentConsumer
 import com.maddyhome.idea.vim.key.consumers.CommandConsumer
 import com.maddyhome.idea.vim.key.consumers.CommandCountConsumer
@@ -282,7 +282,7 @@ class KeyHandler {
     keyState.commandBuilder.resetAll(getKeyRoot(mode.toMappingMode()))
   }
 
-  private fun getKeyRoot(mappingMode: MappingMode): CommandPartNode<LazyVimCommand> {
+  private fun getKeyRoot(mappingMode: MappingMode): RootNode<LazyVimCommand> {
     return injector.keyGroup.getKeyRoot(mappingMode)
   }
 

vim-engine/src/main/kotlin/com/maddyhome/idea/vim/api/VimKeyGroup.kt

@@ -10,17 +10,17 @@ package com.maddyhome.idea.vim.api
 import com.maddyhome.idea.vim.action.change.LazyVimCommand
 import com.maddyhome.idea.vim.command.MappingMode
 import com.maddyhome.idea.vim.extension.ExtensionHandler
-import com.maddyhome.idea.vim.key.CommandPartNode
 import com.maddyhome.idea.vim.key.KeyMapping
 import com.maddyhome.idea.vim.key.KeyMappingLayer
 import com.maddyhome.idea.vim.key.MappingInfo
 import com.maddyhome.idea.vim.key.MappingOwner
+import com.maddyhome.idea.vim.key.RootNode
 import com.maddyhome.idea.vim.key.ShortcutOwnerInfo
 import com.maddyhome.idea.vim.vimscript.model.expressions.Expression
 import javax.swing.KeyStroke
 
 interface VimKeyGroup {
-  fun getKeyRoot(mappingMode: MappingMode): CommandPartNode<LazyVimCommand>
+  fun getKeyRoot(mappingMode: MappingMode): RootNode<LazyVimCommand>
   fun getKeyMappingLayer(mode: MappingMode): KeyMappingLayer
   fun getActions(editor: VimEditor, keyStroke: KeyStroke): List<NativeAction>
   fun getKeymapConflicts(keyStroke: KeyStroke): List<NativeAction>

vim-engine/src/main/kotlin/com/maddyhome/idea/vim/api/VimKeyGroupBase.kt

@@ -12,7 +12,6 @@ import com.maddyhome.idea.vim.action.change.LazyVimCommand
 import com.maddyhome.idea.vim.command.MappingMode
 import com.maddyhome.idea.vim.extension.ExtensionHandler
 import com.maddyhome.idea.vim.handler.EditorActionHandlerBase
-import com.maddyhome.idea.vim.key.CommandPartNode
 import com.maddyhome.idea.vim.key.KeyMapping
 import com.maddyhome.idea.vim.key.KeyMappingLayer
 import com.maddyhome.idea.vim.key.MappingInfo
@@ -30,7 +29,7 @@ abstract class VimKeyGroupBase : VimKeyGroup {
   @JvmField
   val myShortcutConflicts: MutableMap<KeyStroke, ShortcutOwnerInfo> = LinkedHashMap()
   val requiredShortcutKeys: MutableSet<RequiredShortcut> = HashSet(300)
-  val keyRoots: MutableMap<MappingMode, CommandPartNode<LazyVimCommand>> = EnumMap(MappingMode::class.java)
+  val keyRoots: MutableMap<MappingMode, RootNode<LazyVimCommand>> = EnumMap(MappingMode::class.java)
   val keyMappings: MutableMap<MappingMode, KeyMapping> = EnumMap(MappingMode::class.java)
 
   override fun removeKeyMapping(modes: Set<MappingMode>, keys: List<KeyStroke>) {
@@ -63,7 +62,7 @@ abstract class VimKeyGroupBase : VimKeyGroup {
    * @param mappingMode The mapping mode
    * @return The key mapping tree root
    */
-  override fun getKeyRoot(mappingMode: MappingMode): CommandPartNode<LazyVimCommand> = keyRoots.getOrPut(mappingMode) { RootNode() }
+  override fun getKeyRoot(mappingMode: MappingMode): RootNode<LazyVimCommand> = keyRoots.getOrPut(mappingMode) { RootNode(mappingMode.name.get(0).lowercase()) }
 
   override fun getKeyMappingLayer(mode: MappingMode): KeyMappingLayer = getKeyMapping(mode)
 

vim-engine/src/main/kotlin/com/maddyhome/idea/vim/command/CommandBuilder.kt

@@ -32,14 +32,8 @@ class CommandBuilder private constructor(
   private val keyList: MutableList<KeyStroke>,
 ) : Cloneable {
 
-  constructor(
-    currentCommandPartNode: CommandPartNode<LazyVimCommand>,
-    initialUncommittedRawCount: Int = 0
-  ) : this(
-    currentCommandPartNode,
-    mutableListOf(initialUncommittedRawCount),
-    mutableListOf(),
-  )
+  constructor(rootNode: RootNode<LazyVimCommand>, initialUncommittedRawCount: Int = 0)
+    : this(rootNode, mutableListOf(initialUncommittedRawCount), mutableListOf())
 
   private var selectedRegister: Char? = null
   private var action: EditorActionHandlerBase? = null
@@ -263,13 +257,13 @@ class CommandBuilder private constructor(
     val node = currentCommandPartNode[key]
     when (node) {
       is CommandNode -> {
-        logger.trace { "Found full command node - $node ($key)" }
+        logger.trace { "Found full command node ($key) - ${node.debugString}" }
         addKey(key)
         processor(node.actionHolder.instance)
         return true
       }
       is CommandPartNode -> {
-        logger.trace { "Found command part node - $node ($key)" }
+        logger.trace { "Found command part node ($key) - ${node.debugString}" }
         currentCommandPartNode = node
         addKey(key)
         return true
@@ -290,7 +284,7 @@ class CommandBuilder private constructor(
    * @see DuplicableOperatorAction
    */
   fun convertDuplicateOperatorKeyStrokeToMotion(key: KeyStroke): KeyStroke {
-    logger.trace("convertDuplicateOperatorKeyStrokeToMotion is executed. key = $key")
+    logger.trace { "convertDuplicateOperatorKeyStrokeToMotion is executed. key = $key" }
 
     // Simple check to ensure that we're in OP_PENDING. If we don't have an action, we don't have an operator. If we
     // have an argument, we can't be in OP_PENDING
@@ -319,13 +313,13 @@ class CommandBuilder private constructor(
   fun buildCommand(): Command {
     val rawCount = calculateCount0Snapshot()
     val command = Command(selectedRegister, rawCount, action!!, argument, action!!.type, action?.flags ?: noneOfEnum())
-    resetAll(currentCommandPartNode)
+    resetAll(currentCommandPartNode.root as RootNode<LazyVimCommand>)
     return command
   }
 
-  fun resetAll(commandPartNode: CommandPartNode<LazyVimCommand>) {
+  fun resetAll(rootNode: RootNode<LazyVimCommand>) {
     logger.trace("resetAll is executed")
-    resetInProgressCommandPart(commandPartNode)
+    currentCommandPartNode = rootNode
     commandState = CurrentCommandState.NEW_COMMAND
     counts.clear()
     counts.add(0)
@@ -337,12 +331,16 @@ class CommandBuilder private constructor(
     fallbackArgumentType = null
   }
 
-  // TODO: Why do we need this to be public?
-  // Who needs to reset the current command part node without resetting the whole command builder?
-  fun resetInProgressCommandPart(commandPartNode: CommandPartNode<LazyVimCommand>) {
-    logger.trace("resetInProgressCommandPart is executed")
-    counts[counts.size - 1] = 0
-    currentCommandPartNode = commandPartNode
+  /**
+   * Change the command trie root node used to find commands for the current mode
+   *
+   * Typically, we reset the command trie root node after a command is executed, using the root node of the current
+   * mode - this is handled by [resetAll]. This function allows us to change the root node without executing a command
+   * or fully resetting the command builder, such as when switching to Op-pending while entering an operator+motion.
+   */
+  fun resetCommandTrieRootNode(rootNode: RootNode<LazyVimCommand>) {
+    logger.trace("resetCommandTrieRootNode is executed")
+    currentCommandPartNode = rootNode
   }
 
   @TestOnly

vim-engine/src/main/kotlin/com/maddyhome/idea/vim/key/Nodes.kt

@@ -40,17 +40,36 @@ import javax.swing.KeyStroke
  *   and the user should complete the sequence, it's [CommandPartNode]
  */
 @Suppress("GrazieInspection")
-interface Node<T>
+interface Node<T> {
+  val debugString: String
+  val parent: Node<T>?
+
+  val root: Node<T>
+    get() = parent?.root ?: this
+}
 
 /** Represents a complete command */
-data class CommandNode<T>(val actionHolder: T) : Node<T> {
-  override fun toString(): String {
-    return "COMMAND NODE (${ actionHolder.toString() })"
-  }
+data class CommandNode<T>(override val parent: Node<T>, val actionHolder: T, private val name: String) : Node<T> {
+  override val debugString: String
+    get() = toString()
+
+  override fun toString() = "COMMAND NODE ($name - ${actionHolder.toString()})"
 }
 
 /** Represents a part of the command */
-open class CommandPartNode<T> : Node<T>, HashMap<KeyStroke, Node<T>>() {
+open class CommandPartNode<T>(
+  override val parent: Node<T>?,
+  internal val name: String,
+  internal val depth: Int) : Node<T> {
+
+  val children = mutableMapOf<KeyStroke, Node<T>>()
+
+  operator fun set(stroke: KeyStroke, node: Node<T>) {
+    children[stroke] = node
+  }
+
+  operator fun get(stroke: KeyStroke): Node<T>? = children[stroke]
+
   override fun equals(other: Any?): Boolean {
     if (this === other) return true
     if (javaClass != other?.javaClass) return false
@@ -58,21 +77,32 @@ open class CommandPartNode<T> : Node<T>, HashMap<KeyStroke, Node<T>>() {
     return true
   }
 
-  override fun hashCode(): Int {
-    return super.hashCode()
-  }
+  override fun hashCode() = super.hashCode()
 
-  override fun toString(): String {
-    return """
-      COMMAND PART NODE(
-      ${entries.joinToString(separator = "\n") { "    " + injector.parser.toKeyNotation(it.key) + " - " + it.value }}
-      )
-      """.trimIndent()
-  }
+  override fun toString() = "COMMAND PART NODE ($name - ${children.size} children)"
+
+  override val debugString
+    get() = buildString {
+      append("COMMAND PART NODE(")
+      appendLine(name)
+      children.entries.forEach {
+        repeat(depth + 1) { append(" ") }
+        append(injector.parser.toKeyNotation(it.key))
+        append(" - ")
+        appendLine(it.value.debugString)
+      }
+      repeat(depth) { append(" ") }
+      append(")")
+    }
 }
 
 /** Represents a root node for the mode */
-class RootNode<T> : CommandPartNode<T>()
+class RootNode<T>(name: String) : CommandPartNode<T>(null, name, 0) {
+  override val debugString: String
+    get() = "ROOT NODE ($name)\n" + super.debugString
+
+  override fun toString() = "ROOT NODE ($name - ${children.size} children)"
+}
 
 fun <T> Node<T>.addLeafs(keyStrokes: List<KeyStroke>, actionHolder: T) {
   var node: Node<T> = this
@@ -90,7 +120,16 @@ private fun <T> addNode(base: CommandPartNode<T>, actionHolder: T, key: KeyStrok
   val existing = base[key]
   if (existing != null) return existing
 
-  val newNode: Node<T> = if (isLastInSequence) CommandNode(actionHolder) else CommandPartNode()
+  val childName = injector.parser.toKeyNotation(key)
+  val name = when (base) {
+    is RootNode -> base.name + "_" + childName
+    else -> base.name + childName
+  }
+  val newNode: Node<T> = if (isLastInSequence) {
+    CommandNode(base, actionHolder, name)
+  } else {
+    CommandPartNode(base, name, base.depth + 1)
+  }
   base[key] = newNode
   return newNode
 }

vim-engine/src/main/kotlin/com/maddyhome/idea/vim/state/KeyHandlerState.kt

@@ -68,7 +68,7 @@ data class KeyHandlerState(
   fun partialReset(mode: Mode) {
     logger.trace("entered partialReset. mode: $mode")
     mappingState.resetMappingSequence()
-    commandBuilder.resetInProgressCommandPart(injector.keyGroup.getKeyRoot(mode.toMappingMode()))
+    commandBuilder.resetCommandTrieRootNode(injector.keyGroup.getKeyRoot(mode.toMappingMode()))
   }
 
   fun reset(mode: Mode) {