Refactor map implementation and restore which-key compatibility

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.CommandPartNode
+import com.maddyhome.idea.vim.key.KeyStrokeTrie
 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.addLeafs
+import com.maddyhome.idea.vim.key.add
 import com.maddyhome.idea.vim.newapi.ij
 import com.maddyhome.idea.vim.newapi.vim
 import com.maddyhome.idea.vim.vimscript.model.datatypes.VimString
@@ -198,27 +195,23 @@ 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
       if (keyChar != KeyEvent.CHAR_UNDEFINED) {
         keyStroke = KeyStroke.getKeyStroke(keyChar)
       }
 
-      val nextNode = currentNode[keyStroke]
-
-      when (nextNode) {
-        null -> currentNode = actionsRoot
-        is CommandNode<NerdAction> -> {
-          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) }
-          }
+      keys.add(keyStroke)
+      actionsRoot.getData(keys)?.let { action ->
+        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) {
-  actionsRoot.addLeafs(default, action)
-}
-
-
-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 fun registerCommand(mapping: String, action: NerdAction) {
+  actionsRoot.add(mapping, action)
+  injector.parser.parseKeys(mapping).forEach {
+    distinctShortcuts.add(it)
   }
 }
 
+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 =
-    collectShortcuts(actionsRoot).map { RequiredShortcut(it, MappingOwner.Plugin.get(NerdTree.pluginName)) }
+  val shortcuts = distinctShortcuts.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

@@ -8,7 +8,6 @@
 
 package com.maddyhome.idea.vim.group;
 
-import com.google.common.collect.ImmutableList;
 import com.intellij.codeInsight.lookup.impl.LookupImpl;
 import com.intellij.openapi.actionSystem.*;
 import com.intellij.openapi.actionSystem.ex.ActionManagerEx;
@@ -18,17 +17,16 @@
 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;
+import com.intellij.util.containers.MultiMap;
 import com.maddyhome.idea.vim.EventFacade;
 import com.maddyhome.idea.vim.VimPlugin;
 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 void registerCommandAction(@NotNull LazyVimCommand command) {
       registerRequiredShortcut(keyStrokes, MappingOwner.IdeaVim.System.INSTANCE);
 
       for (MappingMode mappingMode : command.getModes()) {
-        Node<LazyVimCommand> node = getKeyRoot(mappingMode);
-        NodesKt.addLeafs(node, keyStrokes, command);
+        getBuiltinCommandsTrie(mappingMode).add(keyStrokes, command);
       }
     }
   }
@@ -225,53 +222,71 @@ private void registerRequiredShortcut(@NotNull List<KeyStroke> keys, MappingOwne
     return new CustomShortcutSet(shortcuts.toArray(new Shortcut[0]));
   }
 
-  private static @NotNull List<Pair<EnumSet<MappingMode>, MappingInfo>> getKeyMappingRows(@NotNull Set<? extends MappingMode> modes) {
-    final Map<ImmutableList<KeyStroke>, EnumSet<MappingMode>> actualModes = new HashMap<>();
+  private static @NotNull List<Pair<Set<MappingMode>, MappingInfo>> getKeyMappingRows(@NotNull Set<? extends MappingMode> modes) {
+    // Some map commands set a mapping for more than one mode (e.g. `map` sets for Normal, Visual, Select and
+    // Op-pending). Vim treats this as a single mapping, and when listing all maps only lists it once, with the
+    // appropriate mode indicator(s) in the first column (NVO is a space char). If the lhs mapping is changed or cleared
+    // for one of the modes, the original mapping is still a single map for the remaining modes, and the indicator
+    // changes. E.g. `map foo bar` followed by `sunmap foo` would result in `nox foo bar` in the output to `map`.
+    // Vim doesn't do automatic grouping - `nmap foo bar` followed by `omap foo bar` and `vmap foo bar` would result in
+    // 3 lines in the output to `map` - one for `n`, one for `o` and one for `v`.
+    // We store mappings separately per mode (to simplify lookup, especially when matching prefixes), but want to have
+    // the same behaviour as Vim in map output. So we store the original modes with the mapping and check they're still
+    // valid as we collect output
+    final List<Pair<Set<MappingMode>, MappingInfo>> rows = new ArrayList<>();
+    final MultiMap<List<? extends KeyStroke>, Set<MappingMode>> multiModeMappings = MultiMap.create();
     for (MappingMode mode : modes) {
       final KeyMapping mapping = VimPlugin.getKey().getKeyMapping(mode);
       for (List<? extends KeyStroke> fromKeys : mapping) {
-        final ImmutableList<KeyStroke> key = ImmutableList.copyOf(fromKeys);
-        final EnumSet<MappingMode> value = actualModes.get(key);
-        final EnumSet<MappingMode> newValue;
-        if (value != null) {
-          newValue = value.clone();
-          newValue.add(mode);
-        }
-        else {
-          newValue = EnumSet.of(mode);
-        }
-        actualModes.put(key, newValue);
-      }
-    }
-    final List<Pair<EnumSet<MappingMode>, MappingInfo>> rows = new ArrayList<>();
-    for (Map.Entry<ImmutableList<KeyStroke>, EnumSet<MappingMode>> entry : actualModes.entrySet()) {
-      final ArrayList<KeyStroke> fromKeys = new ArrayList<>(entry.getKey());
-      final EnumSet<MappingMode> mappingModes = entry.getValue();
-      if (!mappingModes.isEmpty()) {
-        final MappingMode mode = mappingModes.iterator().next();
-        final KeyMapping mapping = VimPlugin.getKey().getKeyMapping(mode);
         final MappingInfo mappingInfo = mapping.get(fromKeys);
         if (mappingInfo != null) {
-          rows.add(new Pair<>(mappingModes, mappingInfo));
+          final Set<@NotNull MappingMode> originalModes = mappingInfo.getOriginalModes();
+          if (originalModes.size() == 1) {
+            rows.add(new Pair<>(originalModes, mappingInfo));
+          }
+          else if (!multiModeMappings.get(fromKeys).contains(originalModes)) {
+            multiModeMappings.putValue(fromKeys, originalModes);
+            rows.add(new Pair<>(getModesForMapping(fromKeys, originalModes), mappingInfo));
+          }
         }
       }
     }
-    rows.sort(Comparator.comparing(Pair<EnumSet<MappingMode>, MappingInfo>::getSecond));
+    rows.sort(Comparator.comparing(Pair<Set<MappingMode>, MappingInfo>::getSecond));
     return rows;
   }
 
-  private static @NotNull @NonNls String getModesStringCode(@NotNull Set<MappingMode> modes) {
-    if (modes.equals(MappingMode.NVO)) {
-      return "";
+  private static @NotNull Set<MappingMode> getModesForMapping(@NotNull List<? extends KeyStroke> keyStrokes,
+                                                              @NotNull Set<MappingMode> originalMappingModes) {
+    final Set<MappingMode> actualModes = EnumSet.noneOf(MappingMode.class);
+    for (MappingMode mode : originalMappingModes) {
+      final MappingInfo mappingInfo = VimPlugin.getKey().getKeyMapping(mode).get(keyStrokes);
+      if (mappingInfo != null && mappingInfo.getOriginalModes() == originalMappingModes) {
+        actualModes.add(mode);
+      }
     }
-    else if (modes.contains(MappingMode.INSERT)) {
-      return "i";
+    return actualModes;
+  }
+
+  private static @NotNull @NonNls String getModesStringCode(@NotNull Set<MappingMode> modes) {
+    if (modes.equals(MappingMode.IC)) return "!";
+    if (modes.equals(MappingMode.NVO)) return " ";
+    if (modes.equals(MappingMode.C)) return "c";
+    if (modes.equals(MappingMode.I)) return "i";
+    //if (modes.equals(MappingMode.L)) return "l";
+
+    // The following modes are concatenated
+    String mode = "";
+    if (modes.containsAll(MappingMode.N)) mode += "n";
+    if (modes.containsAll(MappingMode.O)) mode += "o";
+
+    if (modes.containsAll(MappingMode.V)) {
+      mode += "v";
     }
-    else if (modes.contains(MappingMode.NORMAL)) {
-      return "n";
+    else {
+      if (modes.containsAll(MappingMode.X)) mode += "x";
+      if (modes.containsAll(MappingMode.S)) mode += "s";
     }
-    // TODO: Add more codes
-    return "";
+    return mode;
   }
 
   private @NotNull List<AnAction> getActions(@NotNull Component component, @NotNull KeyStroke keyStroke) {
@@ -336,19 +351,19 @@ public void loadState(@NotNull Element state) {
 
   @Override
   public boolean showKeyMappings(@NotNull Set<? extends MappingMode> modes, @NotNull VimEditor editor) {
-    List<Pair<EnumSet<MappingMode>, MappingInfo>> rows = getKeyMappingRows(modes);
+    List<Pair<Set<MappingMode>, MappingInfo>> rows = getKeyMappingRows(modes);
+
     final StringBuilder builder = new StringBuilder();
-    for (Pair<EnumSet<MappingMode>, MappingInfo> row : rows) {
+    for (Pair<Set<MappingMode>, MappingInfo> row : rows) {
       MappingInfo mappingInfo = row.getSecond();
-      builder.append(StringsKt.padEnd(getModesStringCode(row.getFirst()), 2, ' '));
-      builder.append(" ");
-      builder.append(StringsKt.padEnd(VimInjectorKt.getInjector().getParser().toKeyNotation(mappingInfo.getFromKeys()), 11, ' '));
-      builder.append(" ");
-      builder.append(mappingInfo.isRecursive() ? " " : "*");
-      builder.append(" ");
+      builder.append(StringsKt.padEnd(getModesStringCode(row.getFirst()), 3, ' '));
+      builder.append(StringsKt.padEnd(VimInjectorKt.getInjector().getParser().toKeyNotation(mappingInfo.getFromKeys()) + " ", 12, ' '));
+      builder.append(mappingInfo.isRecursive() ? " " : "*");  // Or `&` if script-local mappings being recursive
+      builder.append(" ");  // Should be `@` if it's a buffer-local mapping
       builder.append(mappingInfo.getPresentableString());
       builder.append("\n");
     }
+
     VimOutputPanel outputPanel = injector.getOutputPanel().getOrCreate(editor, injector.getExecutionContextManager().getEditorExecutionContext(editor));
     outputPanel.addText(builder.toString(), true);
     outputPanel.show();