Introduce KeyStrokeTrie to find commands

Should also restore compatibility with idea-which-key

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

@@ -42,13 +42,10 @@ import com.maddyhome.idea.vim.extension.VimExtension
 import com.maddyhome.idea.vim.group.KeyGroup
 import com.maddyhome.idea.vim.helper.MessageHelper
 import com.maddyhome.idea.vim.helper.runAfterGotFocus
-import com.maddyhome.idea.vim.key.CommandNode
+import com.maddyhome.idea.vim.key.KeyStrokeTrie
-import com.maddyhome.idea.vim.key.CommandPartNode
 import com.maddyhome.idea.vim.key.MappingOwner
-import com.maddyhome.idea.vim.key.Node
 import com.maddyhome.idea.vim.key.RequiredShortcut
-import com.maddyhome.idea.vim.key.RootNode
+import com.maddyhome.idea.vim.key.add
-import com.maddyhome.idea.vim.key.addLeafs
 import com.maddyhome.idea.vim.newapi.ij
 import com.maddyhome.idea.vim.newapi.vim
 import com.maddyhome.idea.vim.vimscript.model.datatypes.VimString
@@ -198,6 +195,8 @@ internal class NerdTree : VimExtension {
     internal var waitForSearch = false
     internal var speedSearchListenerInstalled = false
 
+    private val keys = mutableListOf<KeyStroke>()
+
     override fun actionPerformed(e: AnActionEvent) {
       var keyStroke = getKeyStroke(e) ?: return
       val keyChar = keyStroke.keyChar
@@ -205,20 +204,14 @@ internal class NerdTree : VimExtension {
         keyStroke = KeyStroke.getKeyStroke(keyChar)
       }
 
-      val nextNode = currentNode[keyStroke]
+      keys.add(keyStroke)
-
+      actionsRoot.getData(keys)?.let { action ->
-      when (nextNode) {
+        when (action) {
-        null -> currentNode = actionsRoot
+          is NerdAction.ToIj -> Util.callAction(null, action.name, e.dataContext.vim)
-        is CommandNode<NerdAction> -> {
+          is NerdAction.Code -> e.project?.let { action.action(it, e.dataContext, e) }
-          currentNode = actionsRoot
-
-          val action = nextNode.actionHolder
-          when (action) {
-            is NerdAction.ToIj -> Util.callAction(null, action.name, e.dataContext.vim)
-            is NerdAction.Code -> e.project?.let { action.action(it, e.dataContext, e) }
-          }
         }
-        is CommandPartNode<NerdAction> -> currentNode = nextNode
+
+        keys.clear()
       }
     }
 
@@ -540,38 +533,29 @@ private fun addCommand(alias: String, handler: CommandAliasHandler) {
   VimPlugin.getCommand().setAlias(alias, CommandAlias.Call(0, -1, alias, handler))
 }
 
-private fun registerCommand(variable: String, default: String, action: NerdAction) {
+private fun registerCommand(variable: String, defaultMapping: String, action: NerdAction) {
   val variableValue = VimPlugin.getVariableService().getGlobalVariableValue(variable)
-  val mappings = if (variableValue is VimString) {
+  val mapping = if (variableValue is VimString) {
     variableValue.value
   } else {
-    default
+    defaultMapping
   }
-  actionsRoot.addLeafs(mappings, action)
+  registerCommand(mapping, action)
 }
 
-private fun registerCommand(default: String, action: NerdAction) {
+private fun registerCommand(mapping: String, action: NerdAction) {
-  actionsRoot.addLeafs(default, action)
+  actionsRoot.add(mapping, action)
-}
+  injector.parser.parseKeys(mapping).forEach {
-
+    distinctShortcuts.add(it)
-
-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.children.keys.toMutableSet()
-    res += node.children.values.map { collectShortcuts(it) }.flatten()
-    res
-  } else {
-    emptySet()
   }
 }
 
+private val actionsRoot: KeyStrokeTrie<NerdAction> = KeyStrokeTrie<NerdAction>("NERDTree")
+private val distinctShortcuts = mutableSetOf<KeyStroke>()
+
 private fun installDispatcher(project: Project) {
   val dispatcher = NerdTree.NerdDispatcher.getInstance(project)
-  val shortcuts =
+  val shortcuts = distinctShortcuts.map { RequiredShortcut(it, MappingOwner.Plugin.get(NerdTree.pluginName)) }
-    collectShortcuts(actionsRoot).map { RequiredShortcut(it, MappingOwner.Plugin.get(NerdTree.pluginName)) }
   dispatcher.registerCustomShortcutSet(
     KeyGroup.toShortcutSet(shortcuts),
     (ProjectView.getInstance(project) as ProjectViewImpl).component,

src/main/java/com/maddyhome/idea/vim/group/KeyGroup.java

@@ -18,7 +18,6 @@ import com.intellij.openapi.components.PersistentStateComponent;
 import com.intellij.openapi.components.State;
 import com.intellij.openapi.components.Storage;
 import com.intellij.openapi.diagnostic.Logger;
-import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.keymap.Keymap;
 import com.intellij.openapi.keymap.KeymapManager;
 import com.intellij.openapi.keymap.ex.KeymapManagerEx;
@@ -28,7 +27,6 @@ import com.maddyhome.idea.vim.action.VimShortcutKeyAction;
 import com.maddyhome.idea.vim.action.change.LazyVimCommand;
 import com.maddyhome.idea.vim.api.*;
 import com.maddyhome.idea.vim.command.MappingMode;
-import com.maddyhome.idea.vim.ex.ExOutputModel;
 import com.maddyhome.idea.vim.key.*;
 import com.maddyhome.idea.vim.newapi.IjNativeAction;
 import com.maddyhome.idea.vim.newapi.IjVimEditor;
@@ -199,8 +197,7 @@ public class KeyGroup extends VimKeyGroupBase implements PersistentStateComponen
       registerRequiredShortcut(keyStrokes, MappingOwner.IdeaVim.System.INSTANCE);
 
       for (MappingMode mappingMode : command.getModes()) {
-        Node<LazyVimCommand> node = getKeyRoot(mappingMode);
+        getBuiltinCommandsTrie(mappingMode).add(keyStrokes, command);
-        NodesKt.addLeafs(node, keyStrokes, command);
       }
     }
   }

src/main/java/com/maddyhome/idea/vim/key/NodesHelper.kt

@@ -1,15 +0,0 @@
-/*
- * Copyright 2003-2023 The IdeaVim authors
- *
- * Use of this source code is governed by an MIT-style
- * license that can be found in the LICENSE.txt file or at
- * https://opensource.org/licenses/MIT.
- */
-
-package com.maddyhome.idea.vim.key
-
-import com.maddyhome.idea.vim.api.injector
-
-internal fun <T> Node<T>.addLeafs(keys: String, actionHolder: T) {
-  addLeafs(injector.parser.parseKeys(keys), actionHolder)
-}

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

@@ -14,7 +14,6 @@ import com.intellij.testFramework.PlatformTestUtil
 import com.maddyhome.idea.vim.KeyHandler
 import com.maddyhome.idea.vim.api.injector
 import com.maddyhome.idea.vim.api.key
-import com.maddyhome.idea.vim.key.CommandNode
 import com.maddyhome.idea.vim.newapi.vim
 import org.jetbrains.jetCheck.Generator
 import org.jetbrains.jetCheck.ImperativeCommand
@@ -92,19 +91,23 @@ class RandomActionsPropertyTest : VimPropertyTestBase() {
 
 private class AvailableActions(private val editor: Editor) : ImperativeCommand {
   override fun performCommand(env: ImperativeCommand.Environment) {
-    val currentNode = KeyHandler.getInstance().keyHandlerState.commandBuilder.getCurrentTrie()
+    val trie = KeyHandler.getInstance().keyHandlerState.commandBuilder.getCurrentTrie()
+    val currentKeys = KeyHandler.getInstance().keyHandlerState.commandBuilder.getCurrentCommandKeys()
 
     // Note: esc is always an option
-    val possibleKeys = (currentNode.children.keys.toList() + esc).sortedBy { injector.parser.toKeyNotation(it) }
+    val possibleKeys: List<KeyStroke> = buildList {
-    println("Keys: ${possibleKeys.joinToString(", ")}")
+      add(esc)
+      trie.getTrieNode(currentKeys)?.visit { stroke, _ -> add(stroke) }
+    }.sortedBy { injector.parser.toKeyNotation(it) }
+
+//    println("Keys: ${possibleKeys.joinToString(", ")}")
     val keyGenerator = Generator.integers(0, possibleKeys.lastIndex)
       .suchThat { injector.parser.toKeyNotation(possibleKeys[it]) !in stinkyKeysList }
       .map { possibleKeys[it] }
 
     val usedKey = env.generateValue(keyGenerator, null)
-    val node = currentNode[usedKey]
+    val node = trie.getTrieNode(currentKeys + usedKey)
-
+    env.logMessage("Use command: ${injector.parser.toKeyNotation(currentKeys + usedKey)}. ${if (node?.data != null) "Action: ${node.data!!.actionId}" else ""}")
-    env.logMessage("Use command: ${injector.parser.toKeyNotation(usedKey)}. ${if (node is CommandNode) "Action: ${node.actionHolder.actionId}" else ""}")
     VimNoWriteActionTestCase.typeText(listOf(usedKey), editor, editor.project)
 
     IdeEventQueue.getInstance().flushQueue()

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

@@ -7,14 +7,12 @@
  */
 package com.maddyhome.idea.vim
 
-import com.maddyhome.idea.vim.action.change.LazyVimCommand
 import com.maddyhome.idea.vim.api.ExecutionContext
 import com.maddyhome.idea.vim.api.VimEditor
 import com.maddyhome.idea.vim.api.globalOptions
 import com.maddyhome.idea.vim.api.injector
 import com.maddyhome.idea.vim.command.Command
 import com.maddyhome.idea.vim.command.CommandFlags
-import com.maddyhome.idea.vim.command.MappingMode
 import com.maddyhome.idea.vim.command.MappingProcessor
 import com.maddyhome.idea.vim.command.OperatorArguments
 import com.maddyhome.idea.vim.diagnostic.VimLogger
@@ -23,7 +21,6 @@ import com.maddyhome.idea.vim.diagnostic.vimLogger
 import com.maddyhome.idea.vim.impl.state.toMappingMode
 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
@@ -269,7 +266,7 @@ class KeyHandler {
     editor.isReplaceCharacter = false
     editor.resetOpPending()
     keyHandlerState.partialReset(editor.mode)
-    keyHandlerState.commandBuilder.resetAll(getKeyRoot(editor.mode.toMappingMode()))
+    keyHandlerState.commandBuilder.resetAll(injector.keyGroup.getBuiltinCommandsTrie(editor.mode.toMappingMode()))
   }
 
   // TODO we should have a single reset method
@@ -277,11 +274,7 @@ class KeyHandler {
     logger.trace { "Reset is executed" }
     injector.commandLine.getActiveCommandLine()?.clearCurrentAction()
     keyHandlerState.partialReset(mode)
-    keyState.commandBuilder.resetAll(getKeyRoot(mode.toMappingMode()))
+    keyState.commandBuilder.resetAll(injector.keyGroup.getBuiltinCommandsTrie(mode.toMappingMode()))
-  }
-
-  private fun getKeyRoot(mappingMode: MappingMode): RootNode<LazyVimCommand> {
-    return injector.keyGroup.getKeyRoot(mappingMode)
   }
 
   fun updateState(keyState: KeyHandlerState) {

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

@@ -12,15 +12,19 @@ import com.maddyhome.idea.vim.command.MappingMode
 import com.maddyhome.idea.vim.extension.ExtensionHandler
 import com.maddyhome.idea.vim.key.KeyMapping
 import com.maddyhome.idea.vim.key.KeyMappingLayer
+import com.maddyhome.idea.vim.key.KeyStrokeTrie
 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): RootNode<LazyVimCommand>
+  @Suppress("DEPRECATION")
+  @Deprecated("Use getBuiltinCommandTrie", ReplaceWith("getBuiltinCommandsTrie(mappingMode)"))
+  fun getKeyRoot(mappingMode: MappingMode): com.maddyhome.idea.vim.key.CommandPartNode<LazyVimCommand>
+
+  fun getBuiltinCommandsTrie(mappingMode: MappingMode): KeyStrokeTrie<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

@@ -14,6 +14,7 @@ import com.maddyhome.idea.vim.extension.ExtensionHandler
 import com.maddyhome.idea.vim.handler.EditorActionHandlerBase
 import com.maddyhome.idea.vim.key.KeyMapping
 import com.maddyhome.idea.vim.key.KeyMappingLayer
+import com.maddyhome.idea.vim.key.KeyStrokeTrie
 import com.maddyhome.idea.vim.key.MappingInfo
 import com.maddyhome.idea.vim.key.MappingOwner
 import com.maddyhome.idea.vim.key.RequiredShortcut
@@ -29,7 +30,7 @@ abstract class VimKeyGroupBase : VimKeyGroup {
   @JvmField
   val myShortcutConflicts: MutableMap<KeyStroke, ShortcutOwnerInfo> = LinkedHashMap()
   val requiredShortcutKeys: MutableSet<RequiredShortcut> = HashSet(300)
-  val keyRoots: MutableMap<MappingMode, RootNode<LazyVimCommand>> = EnumMap(MappingMode::class.java)
+  val builtinCommands: MutableMap<MappingMode, KeyStrokeTrie<LazyVimCommand>> = EnumMap(MappingMode::class.java)
   val keyMappings: MutableMap<MappingMode, KeyMapping> = EnumMap(MappingMode::class.java)
 
   override fun removeKeyMapping(modes: Set<MappingMode>, keys: List<KeyStroke>) {
@@ -56,13 +57,19 @@ abstract class VimKeyGroupBase : VimKeyGroup {
     keyMappings.clear()
   }
 
+  @Suppress("DEPRECATION")
+  @Deprecated("Use getBuiltinCommandTrie", ReplaceWith("getBuiltinCommandsTrie(mappingMode)"))
+  override fun getKeyRoot(mappingMode: MappingMode): com.maddyhome.idea.vim.key.CommandPartNode<LazyVimCommand> =
+    RootNode(getBuiltinCommandsTrie(mappingMode))
+
   /**
-   * Returns the root of the key mapping for the given mapping mode
+   * Returns the root node of the builtin command keystroke trie
    *
    * @param mappingMode The mapping mode
-   * @return The key mapping tree root
+   * @return The root node of the builtin command trie
    */
-  override fun getKeyRoot(mappingMode: MappingMode): RootNode<LazyVimCommand> = keyRoots.getOrPut(mappingMode) { RootNode(mappingMode.name.get(0).lowercase()) }
+  override fun getBuiltinCommandsTrie(mappingMode: MappingMode): KeyStrokeTrie<LazyVimCommand> =
+    builtinCommands.getOrPut(mappingMode) { KeyStrokeTrie<LazyVimCommand>(mappingMode.name[0].lowercase()) }
 
   override fun getKeyMappingLayer(mode: MappingMode): KeyMappingLayer = getKeyMapping(mode)
 
@@ -75,6 +82,7 @@ abstract class VimKeyGroupBase : VimKeyGroup {
     for (mappingMode in mappingModes) {
       checkIdentity(mappingMode, action.id, keys)
     }
+    @Suppress("DEPRECATION")
     checkCorrectCombination(action, keys)
   }
 
@@ -236,6 +244,6 @@ abstract class VimKeyGroupBase : VimKeyGroup {
   }
 
   override fun unregisterCommandActions() {
-    keyRoots.clear()
+    builtinCommands.clear()
   }
 }

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

@@ -20,20 +20,19 @@ import com.maddyhome.idea.vim.handler.MotionActionHandler
 import com.maddyhome.idea.vim.handler.TextObjectActionHandler
 import com.maddyhome.idea.vim.helper.StrictMode
 import com.maddyhome.idea.vim.helper.noneOfEnum
-import com.maddyhome.idea.vim.key.CommandNode
+import com.maddyhome.idea.vim.key.KeyStrokeTrie
-import com.maddyhome.idea.vim.key.CommandPartNode
-import com.maddyhome.idea.vim.key.RootNode
 import org.jetbrains.annotations.TestOnly
 import javax.swing.KeyStroke
 
 class CommandBuilder private constructor(
-  private var currentCommandPartNode: CommandPartNode<LazyVimCommand>,
+  private var keyStrokeTrie: KeyStrokeTrie<LazyVimCommand>,
   private val counts: MutableList<Int>,
-  private val keyList: MutableList<KeyStroke>,
+  private val typedKeyStrokes: MutableList<KeyStroke>,
+  private val commandKeyStrokes: MutableList<KeyStroke>
 ) : Cloneable {
 
-  constructor(rootNode: RootNode<LazyVimCommand>, initialUncommittedRawCount: Int = 0)
+  constructor(keyStrokeTrie: KeyStrokeTrie<LazyVimCommand>, initialUncommittedRawCount: Int = 0)
-    : this(rootNode, mutableListOf(initialUncommittedRawCount), mutableListOf())
+    : this(keyStrokeTrie, mutableListOf(initialUncommittedRawCount), mutableListOf(), mutableListOf())
 
   private var commandState: CurrentCommandState = CurrentCommandState.NEW_COMMAND
   private var selectedRegister: Char? = null
@@ -51,7 +50,7 @@ class CommandBuilder private constructor(
     }
 
   /** Provide the typed keys for `'showcmd'` */
-  val keys: Iterable<KeyStroke> get() = keyList
+  val keys: Iterable<KeyStroke> get() = typedKeyStrokes
 
   /** Returns true if the command builder is clean and ready to start building */
   val isEmpty
@@ -167,12 +166,12 @@ class CommandBuilder private constructor(
     if (currentCount < 0) {
       currentCount = 999999999
     }
-    addKey(key)
+    addTypedKeyStroke(key)
   }
 
   fun deleteCountCharacter() {
     currentCount /= 10
-    keyList.removeAt(keyList.size - 1)
+    typedKeyStrokes.removeLast()
   }
 
   var isRegisterPending: Boolean = false
@@ -180,7 +179,7 @@ class CommandBuilder private constructor(
 
   fun startWaitingForRegister(key: KeyStroke) {
     isRegisterPending = true
-    addKey(key)
+    addTypedKeyStroke(key)
   }
 
   fun selectRegister(register: Char) {
@@ -197,9 +196,9 @@ class CommandBuilder private constructor(
    * Only public use is when entering a digraph/literal, where each key isn't handled by [CommandBuilder], but should
    * be added to the `'showcmd'` output.
    */
-  fun addKey(key: KeyStroke) {
+  fun addTypedKeyStroke(key: KeyStroke) {
     logger.trace { "added key to command builder: $key" }
-    keyList.add(key)
+    typedKeyStrokes.add(key)
   }
 
   /**
@@ -268,24 +267,26 @@ class CommandBuilder private constructor(
    * part node.
    */
   fun processKey(key: KeyStroke, processor: (EditorActionHandlerBase) -> Unit): Boolean {
-    val node = currentCommandPartNode[key]
+    commandKeyStrokes.add(key)
-    when (node) {
+    val node = keyStrokeTrie.getTrieNode(commandKeyStrokes)
-      is CommandNode -> {
+    if (node == null) {
-        logger.trace { "Found full command node ($key) - ${node.debugString}" }
+      logger.trace { "No command or part command for key sequence: ${injector.parser.toPrintableString(commandKeyStrokes)}" }
-        addKey(key)
+      commandKeyStrokes.clear()
-        processor(node.actionHolder.instance)
+      return false
-        return true
-      }
-      is CommandPartNode -> {
-        logger.trace { "Found command part node ($key) - ${node.debugString}" }
-        currentCommandPartNode = node
-        addKey(key)
-        return true
-      }
     }
 
-    logger.trace { "No command/command part node found for key: $key" }
+    addTypedKeyStroke(key)
-    return false
+
+    val command = node.data
+    if (command == null) {
+      logger.trace { "Found unfinished key sequence for ${injector.parser.toPrintableString(commandKeyStrokes)} - ${node.debugString}"}
+      return true
+    }
+
+    logger.trace { "Found command for ${injector.parser.toPrintableString(commandKeyStrokes)} - ${node.debugString}"}
+    commandKeyStrokes.clear()
+    processor(command.instance)
+    return true
   }
 
   /**
@@ -319,8 +320,8 @@ class CommandBuilder private constructor(
     //   Similarly, nmap <C-W>a <C-W>s should not try to map the second <C-W> in <C-W><C-W>
     // Note that we might still be at RootNode if we're handling a prefix, because we might be buffering keys until we
     // get a match. This means we'll still process the rest of the keys of the prefix.
-    val isMultikey = currentCommandPartNode !is RootNode
+    val isMultikey = commandKeyStrokes.isNotEmpty()
-    logger.debug { "Building multikey command: $isMultikey" }
+    logger.debug { "Building multikey command: $commandKeyStrokes" }
     return isMultikey
   }
 
@@ -332,21 +333,22 @@ class CommandBuilder private constructor(
   fun buildCommand(): Command {
     val rawCount = calculateCount0Snapshot()
     val command = Command(selectedRegister, rawCount, action!!, argument, action!!.type, action?.flags ?: noneOfEnum())
-    resetAll(currentCommandPartNode.root as RootNode<LazyVimCommand>)
+    resetAll(keyStrokeTrie)
     return command
   }
 
-  fun resetAll(rootNode: RootNode<LazyVimCommand>) {
+  fun resetAll(keyStrokeTrie: KeyStrokeTrie<LazyVimCommand>) {
     logger.trace("resetAll is executed")
-    currentCommandPartNode = rootNode
+    this.keyStrokeTrie = keyStrokeTrie
     commandState = CurrentCommandState.NEW_COMMAND
+    commandKeyStrokes.clear()
     counts.clear()
     counts.add(0)
     isRegisterPending = false
     selectedRegister = null
     action = null
     argument = null
-    keyList.clear()
+    typedKeyStrokes.clear()
     fallbackArgumentType = null
   }
 
@@ -357,13 +359,16 @@ class CommandBuilder private constructor(
    * 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>) {
+  fun resetCommandTrie(keyStrokeTrie: KeyStrokeTrie<LazyVimCommand>) {
     logger.trace("resetCommandTrieRootNode is executed")
-    currentCommandPartNode = rootNode
+    this.keyStrokeTrie = keyStrokeTrie
   }
 
   @TestOnly
-  fun getCurrentTrie(): CommandPartNode<LazyVimCommand> = currentCommandPartNode
+  fun getCurrentTrie(): KeyStrokeTrie<LazyVimCommand> = keyStrokeTrie
+
+  @TestOnly
+  fun getCurrentCommandKeys(): List<KeyStroke> = commandKeyStrokes
 
   override fun equals(other: Any?): Boolean {
     if (this === other) return true
@@ -371,12 +376,12 @@ class CommandBuilder private constructor(
 
     other as CommandBuilder
 
-    if (currentCommandPartNode != other.currentCommandPartNode) return false
+    if (keyStrokeTrie != other.keyStrokeTrie) return false
     if (counts != other.counts) return false
     if (selectedRegister != other.selectedRegister) return false
     if (action != other.action) return false
     if (argument != other.argument) return false
-    if (keyList != other.keyList) return false
+    if (typedKeyStrokes != other.typedKeyStrokes) return false
     if (commandState != other.commandState) return false
     if (expectedArgumentType != other.expectedArgumentType) return false
     if (fallbackArgumentType != other.fallbackArgumentType) return false
@@ -385,12 +390,12 @@ class CommandBuilder private constructor(
   }
 
   override fun hashCode(): Int {
-    var result = currentCommandPartNode.hashCode()
+    var result = keyStrokeTrie.hashCode()
     result = 31 * result + counts.hashCode()
     result = 31 * result + selectedRegister.hashCode()
     result = 31 * result + action.hashCode()
     result = 31 * result + argument.hashCode()
-    result = 31 * result + keyList.hashCode()
+    result = 31 * result + typedKeyStrokes.hashCode()
     result = 31 * result + commandState.hashCode()
     result = 31 * result + expectedArgumentType.hashCode()
     result = 31 * result + fallbackArgumentType.hashCode()
@@ -399,9 +404,10 @@ class CommandBuilder private constructor(
 
   public override fun clone(): CommandBuilder {
     val result = CommandBuilder(
-      currentCommandPartNode,
+      keyStrokeTrie,
       counts.toMutableList(),
-      keyList.toMutableList()
+      typedKeyStrokes.toMutableList(),
+      commandKeyStrokes.toMutableList()
     )
     result.selectedRegister = selectedRegister
     result.action = action
@@ -413,12 +419,12 @@ class CommandBuilder private constructor(
 
   override fun toString(): String {
     return "Command state = $commandState, " +
-      "key list = ${ injector.parser.toKeyNotation(keyList) }, " +
+      "key list = ${ injector.parser.toKeyNotation(typedKeyStrokes) }, " +
       "selected register = $selectedRegister, " +
       "counts = $counts, " +
       "action = $action, " +
       "argument = $argument, " +
-      "command part node - $currentCommandPartNode"
+      "command part node - $keyStrokeTrie"
   }
 
   companion object {

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

@@ -0,0 +1,156 @@
+/*
+ * Copyright 2003-2024 The IdeaVim authors
+ *
+ * Use of this source code is governed by an MIT-style
+ * license that can be found in the LICENSE.txt file or at
+ * https://opensource.org/licenses/MIT.
+ */
+
+package com.maddyhome.idea.vim.key
+
+import com.maddyhome.idea.vim.api.injector
+import javax.swing.KeyStroke
+
+/**
+ * A trie data structure for storing and retrieving values associated with sequences of keystrokes
+ *
+ * All leaves will have data, but it is not a requirement for nodes with data to have no children.
+ *
+ * @param name The name of this KeyStrokeTrie instance (for debug purposes)
+ */
+class KeyStrokeTrie<T>(private val name: String) {
+  interface TrieNode<T> {
+    val data: T?
+
+    fun visit(visitor: (KeyStroke, TrieNode<T>) -> Unit)
+
+    val debugString: String
+  }
+
+  private class TrieNodeImpl<T>(val name: String, val depth: Int, override val data: T?)
+    : TrieNode<T> {
+
+    val children = lazy { mutableMapOf<KeyStroke, TrieNodeImpl<T>>() }
+
+    override fun visit(visitor: (KeyStroke, TrieNode<T>) -> Unit) {
+      if (!children.isInitialized()) return
+      children.value.forEach { visitor(it.key, it.value) }
+    }
+
+    /**
+     * Debug helpers to dump this node and its children
+     */
+    override val debugString
+      get() = buildString { dump(this) }
+
+    private fun dump(builder: StringBuilder) {
+      builder.run {
+        append("TrieNode('")
+        append(name)
+        append("'")
+        if (data != null) {
+          append(", ")
+          append(data)
+        }
+        if (children.isInitialized() && children.value.isNotEmpty()) {
+          appendLine()
+          children.value.forEach {
+            repeat(depth + 1) { append(" ") }
+            append("'")
+            append(injector.parser.toKeyNotation(it.key))
+            append("' - ")
+            it.value.dump(this)
+            if (children.value.size > 1 || depth > 0) appendLine()
+          }
+          repeat(depth) { append(" ") }
+        }
+        append(")")
+      }
+    }
+
+    override fun toString() = "TrieNode('$name', ${children.value.size} children): $data"
+  }
+
+  private val root = TrieNodeImpl<T>("", 0, null)
+
+  fun visit(visitor: (KeyStroke, TrieNode<T>) -> Unit) {
+    // Does not visit the (empty) root node
+    root.visit(visitor)
+  }
+
+  fun add(keyStrokes: List<KeyStroke>, data: T) {
+    var current = root
+    keyStrokes.forEachIndexed { i, stroke ->
+      current = current.children.value.getOrPut(stroke) {
+        val name = current.name + injector.parser.toKeyNotation(stroke)
+        TrieNodeImpl(name, current.depth + 1, if (i == keyStrokes.lastIndex) data else null)
+      }
+    }
+  }
+
+  /**
+   * Get the data for the given key sequence if it exists
+   *
+   * @return Returns null if the key sequence does not exist, or if the data at the node is empty
+   */
+  fun getData(keyStrokes: List<KeyStroke>): T? {
+    var current = root
+    keyStrokes.forEach {
+      if (!current.children.isInitialized()) return null
+      current = current.children.value[it] ?: return null
+    }
+    return current.data
+  }
+
+  /**
+   * Get the node for the given key sequence if it exists
+   *
+   * Like [getData] but will return a node even if that node's data is empty. Will return something useful in the case
+   * of a matching sequence, or a matching prefix. If it's only a matching prefix, the [TrieNode.data] value will be
+   * null.
+   */
+  fun getTrieNode(keyStrokes: List<KeyStroke>): TrieNode<T>? {
+    var current = root
+    keyStrokes.forEach {
+      if (!current.children.isInitialized()) return null
+      current = current.children.value[it] ?: return null
+    }
+    return current
+  }
+
+  override fun toString(): String {
+    val children = if (root.children.isInitialized()) {
+      "${root.children.value.size} children"
+    }
+    else {
+      "0 children (not initialized)"
+    }
+    return "KeyStrokeTrie - '$name', $children"
+  }
+}
+
+fun <T> KeyStrokeTrie<T>.add(keys: String, data: T) {
+  add(injector.parser.parseKeys(keys), data)
+}
+
+/**
+ * Returns a map containing all keystroke sequences that start with the given prefix
+ *
+ * This only returns keystroke sequences that have associated data. A keystroke sequence without data is considered a
+ * prefix and not included in the map.
+ */
+fun <T> KeyStrokeTrie<T>.getPrefixed(prefix: List<KeyStroke>): Map<List<KeyStroke>, T> {
+  fun visitor(prefix: List<KeyStroke>, map: MutableMap<List<KeyStroke>, T>) {
+    getTrieNode(prefix)?.let { node ->
+      node.data?.let { map[prefix] = it }
+      node.visit { key, value -> visitor(prefix + key, map) }
+    }
+  }
+
+  return buildMap { visitor(prefix, this) }
+}
+
+/**
+ * Returns all keystroke sequences with associated data
+ */
+fun <T> KeyStrokeTrie<T>.getAll(): Map<List<KeyStroke>, T> = getPrefixed(emptyList())

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

@@ -8,12 +8,15 @@
 
 package com.maddyhome.idea.vim.key
 
+import com.maddyhome.idea.vim.api.VimKeyGroup
 import com.maddyhome.idea.vim.api.injector
 import javax.swing.KeyStroke
 
 /**
  * COMPATIBILITY-LAYER: Moved from common package to this one
  * Please see: https://jb.gg/zo8n0r
+ *
+ * Used by idea-which-key (latest is currently 0.10.3)
  */
 
 /**
@@ -39,36 +42,28 @@ import javax.swing.KeyStroke
  * If the command is complete, it's represented as a [CommandNode]. If this character is a part of command
  *   and the user should complete the sequence, it's [CommandPartNode]
  */
-@Suppress("GrazieInspection")
+@Deprecated("Use KeyStrokeTrie and VimKeyGroup.getBuiltinCommandsTrie instead")
-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>(override val parent: Node<T>, val actionHolder: T, private val name: String) : Node<T> {
+@Suppress("DEPRECATION")
-  override val debugString: String
+@Deprecated("Use KeyStrokeTrie and VimKeyGroup.getBuiltinCommandsTrie instead")
-    get() = toString()
+data class CommandNode<T>(val actionHolder: T) : Node<T> {
-
+  override fun toString(): String {
-  override fun toString() = "COMMAND NODE ($name - ${actionHolder.toString()})"
+    return "COMMAND NODE (${ actionHolder.toString() })"
+  }
 }
 
-/** Represents a part of the command */
+/**
-open class CommandPartNode<T>(
+ * Represents a part of the command
-  override val parent: Node<T>?,
+ *
-  internal val name: String,
+ * Vim-which-key uses this to get a map of all builtin Vim actions. Sadly, there is on Vim equivalent, so we can't
-  internal val depth: Int) : Node<T> {
+ * provide a Vim script function as an API. After retrieving with [VimKeyGroup.getKeyRoot], the node is iterated
-
+ */
-  val children = mutableMapOf<KeyStroke, Node<T>>()
+@Suppress("DEPRECATION")
-
+@Deprecated("Use KeyStrokeTrie and VimKeyGroup.getBuiltinCommandsTrie instead")
-  operator fun set(stroke: KeyStroke, node: Node<T>) {
+open class CommandPartNode<T> internal constructor(private val trieNode: KeyStrokeTrie.TrieNode<T>)
-    children[stroke] = node
+  : Node<T>, AbstractMap<KeyStroke, Node<T>>() {
-  }
-
-  operator fun get(stroke: KeyStroke): Node<T>? = children[stroke]
 
   override fun equals(other: Any?): Boolean {
     if (this === other) return true
@@ -79,57 +74,29 @@ open class CommandPartNode<T>(
 
   override fun hashCode() = super.hashCode()
 
-  override fun toString() = "COMMAND PART NODE ($name - ${children.size} children)"
+  override fun toString(): String {
+    return """
+      COMMAND PART NODE(
+      ${entries.joinToString(separator = "\n") { "    " + injector.parser.toKeyNotation(it.key) + " - " + it.value }}
+      )
+      """.trimIndent()
+  }
 
-  override val debugString
+  override val entries: Set<Map.Entry<KeyStroke, Node<T>>>
-    get() = buildString {
+    get() {
-      append("COMMAND PART NODE(")
+      return buildMap {
-      appendLine(name)
+        trieNode.visit { key, value ->
-      children.entries.forEach {
+          val node: Node<T> = if (value.data == null) {
-        repeat(depth + 1) { append(" ") }
+            CommandPartNode<T>(value)
-        append(injector.parser.toKeyNotation(it.key))
+          }
-        append(" - ")
+          else {
-        appendLine(it.value.debugString)
+            CommandNode(value.data!!)
-      }
+          }
-      repeat(depth) { append(" ") }
+          put(key, node)
-      append(")")
+        }
+      }.entries
     }
 }
 
-/** Represents a root node for the mode */
+@Suppress("DEPRECATION")
-class RootNode<T>(name: String) : CommandPartNode<T>(null, name, 0) {
+internal class RootNode<T>(trieNode: KeyStrokeTrie<T>) : CommandPartNode<T>(trieNode.getTrieNode(emptyList())!!)
-  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
-  val len = keyStrokes.size
-  // Add a child for each keystroke in the shortcut for this action
-  for (i in 0 until len) {
-    if (node !is CommandPartNode<*>) {
-      error("Error in tree constructing")
-    }
-    node = addNode(node as CommandPartNode<T>, actionHolder, keyStrokes[i], i == len - 1)
-  }
-}
-
-private fun <T> addNode(base: CommandPartNode<T>, actionHolder: T, key: KeyStroke, isLastInSequence: Boolean): Node<T> {
-  val existing = base[key]
-  if (existing != null) return existing
-
-  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/key/consumers/DigraphConsumer.kt

@@ -47,12 +47,12 @@ class DigraphConsumer : KeyConsumer {
       logger.trace("Expected argument is digraph")
       if (digraphSequence.isDigraphStart(key)) {
         digraphSequence.startDigraphSequence()
-        commandBuilder.addKey(key)
+        commandBuilder.addTypedKeyStroke(key)
         return true
       }
       if (digraphSequence.isLiteralStart(key)) {
         digraphSequence.startLiteralSequence()
-        commandBuilder.addKey(key)
+        commandBuilder.addTypedKeyStroke(key)
         return true
       }
     }
@@ -63,7 +63,7 @@ class DigraphConsumer : KeyConsumer {
       is DigraphResult.Handled -> {
         keyProcessResultBuilder.addExecutionStep { lambdaKeyState, _, _ ->
           keyHandler.setPromptCharacterEx(res.promptCharacter)
-          lambdaKeyState.commandBuilder.addKey(key)
+          lambdaKeyState.commandBuilder.addTypedKeyStroke(key)
         }
         return true
       }
@@ -87,7 +87,7 @@ class DigraphConsumer : KeyConsumer {
         }
         val stroke = res.stroke ?: return false
         keyProcessResultBuilder.addExecutionStep { lambdaKeyState, lambdaEditorState, lambdaContext ->
-          lambdaKeyState.commandBuilder.addKey(key)
+          lambdaKeyState.commandBuilder.addTypedKeyStroke(key)
           keyHandler.handleKey(lambdaEditorState, stroke, lambdaContext, lambdaKeyState)
         }
         return true

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

@@ -35,7 +35,7 @@ class RegisterConsumer : KeyConsumer {
     if (!commandBuilder.isRegisterPending) return false
 
     logger.trace("Pending mode.")
-    commandBuilder.addKey(key)
+    commandBuilder.addTypedKeyStroke(key)
 
     val chKey: Char = if (key.keyChar == KeyEvent.CHAR_UNDEFINED) 0.toChar() else key.keyChar
     handleSelectRegister(chKey, keyProcessResultBuilder)

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

@@ -23,7 +23,7 @@ data class KeyHandlerState(
   val editorCommandBuilder: CommandBuilder,
   var commandLineCommandBuilder: CommandBuilder?,
 ): Cloneable {
-  constructor() : this(MappingState(), DigraphSequence(), CommandBuilder(injector.keyGroup.getKeyRoot(MappingMode.NORMAL)), null)
+  constructor() : this(MappingState(), DigraphSequence(), CommandBuilder(injector.keyGroup.getBuiltinCommandsTrie(MappingMode.NORMAL)), null)
 
   companion object {
     private val logger = vimLogger<KeyHandlerState>()
@@ -57,7 +57,7 @@ data class KeyHandlerState(
     // argument with the search string. The command has a count of `6`. And a command such as `3:p` becomes an action to
     // process Ex entry with an argument of `.,.+2p` and a count of 3. The count is ignored by this action.
     // Note that we use the calculated count. In Vim, `2"a3"b:` transforms to `:.,.+5`, which is the same behaviour
-    commandLineCommandBuilder = CommandBuilder(injector.keyGroup.getKeyRoot(MappingMode.CMD_LINE),
+    commandLineCommandBuilder = CommandBuilder(injector.keyGroup.getBuiltinCommandsTrie(MappingMode.CMD_LINE),
       editorCommandBuilder.calculateCount0Snapshot())
   }
 
@@ -68,7 +68,7 @@ data class KeyHandlerState(
   fun partialReset(mode: Mode) {
     logger.trace("entered partialReset. mode: $mode")
     mappingState.resetMappingSequence()
-    commandBuilder.resetCommandTrieRootNode(injector.keyGroup.getKeyRoot(mode.toMappingMode()))
+    commandBuilder.resetCommandTrie(injector.keyGroup.getBuiltinCommandsTrie(mode.toMappingMode()))
   }
 
   fun reset(mode: Mode) {
@@ -77,7 +77,7 @@ data class KeyHandlerState(
     mappingState.resetMappingSequence()
 
     commandLineCommandBuilder = null
-    editorCommandBuilder.resetAll(injector.keyGroup.getKeyRoot(mode.toMappingMode()))
+    editorCommandBuilder.resetAll(injector.keyGroup.getBuiltinCommandsTrie(mode.toMappingMode()))
   }
 
   public override fun clone(): KeyHandlerState {