Extract KeyMappingLayer and MappingInfoLayer

src/main/java/com/maddyhome/idea/vim/KeyHandler.kt

@@ -51,7 +51,7 @@ import com.maddyhome.idea.vim.helper.commandState
 import com.maddyhome.idea.vim.helper.inNormalMode
 import com.maddyhome.idea.vim.helper.inSingleNormalMode
 import com.maddyhome.idea.vim.helper.inVisualMode
-import com.maddyhome.idea.vim.key.KeyMapping
+import com.maddyhome.idea.vim.key.KeyMappingLayer
 import com.maddyhome.idea.vim.newapi.VimActionsInitiator
 import com.maddyhome.idea.vim.newapi.ij
 import com.maddyhome.idea.vim.newapi.injector
@@ -294,7 +294,7 @@ class KeyHandler {
     // Save the unhandled key strokes until we either complete or abandon the sequence.
     LOG.trace("Add key to mapping state")
     mappingState.addKey(key)
-    val mapping = injector.keyGroup.getKeyMapping(mappingState.mappingMode)
+    val mapping = injector.keyGroup.getKeyMappingLayer(mappingState.mappingMode)
     LOG.trace { "Get keys for mapping mode. mode = " + mappingState.mappingMode }
 
     // Returns true if any of these methods handle the key. False means that the key is unrelated to mapping and should
@@ -320,7 +320,7 @@ class KeyHandler {
   private fun handleUnfinishedMappingSequence(
     editor: VimEditor,
     mappingState: MappingState,
-    mapping: KeyMapping,
+    mapping: KeyMappingLayer,
     mappingCompleted: Boolean,
   ): Boolean {
     LOG.trace("Processing unfinished mappings...")
@@ -382,12 +382,12 @@ class KeyHandler {
     editor: VimEditor,
     context: ExecutionContext,
     mappingState: MappingState,
-    mapping: KeyMapping,
+    mapping: KeyMappingLayer,
     key: KeyStroke,
   ): Boolean {
     LOG.trace("Processing complete mapping sequence...")
     // The current sequence isn't a prefix, check to see if it's a completed sequence.
-    val currentMappingInfo = mapping[mappingState.keys]
+    val currentMappingInfo = mapping.getLayer(mappingState.keys)
     var mappingInfo = currentMappingInfo
     if (mappingInfo == null) {
       LOG.trace("Haven't found any mapping info for the given sequence. Trying to apply mapping to a subsequence.")
@@ -405,7 +405,7 @@ class KeyHandler {
       mappingState.keys.forEach(Consumer { e: KeyStroke -> previouslyUnhandledKeySequence.add(e) })
       if (previouslyUnhandledKeySequence.size > 1) {
         previouslyUnhandledKeySequence.removeAt(previouslyUnhandledKeySequence.size - 1)
-        mappingInfo = mapping[previouslyUnhandledKeySequence]
+        mappingInfo = mapping.getLayer(previouslyUnhandledKeySequence)
       }
     }
     if (mappingInfo == null) {

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

@@ -340,6 +340,11 @@ public class KeyGroup implements PersistentStateComponent<Element>, VimKeyGroup
     return keyRoots.computeIfAbsent(mappingMode, (key) -> new RootNode<>());
   }
 
+  @Override
+  public @NotNull KeyMappingLayer getKeyMappingLayer(@NotNull MappingMode mode) {
+    return getKeyMapping(mode);
+  }
+
   /**
    * Registers a shortcut that is handled by KeyHandler#handleKey directly, rather than by an action
    *

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

@@ -21,6 +21,7 @@ package com.maddyhome.idea.vim.group;
 import com.maddyhome.idea.vim.common.MappingMode;
 import com.maddyhome.idea.vim.common.CommandPartNode;
 import com.maddyhome.idea.vim.key.KeyMapping;
+import com.maddyhome.idea.vim.key.KeyMappingLayer;
 import com.maddyhome.idea.vim.newapi.VimActionsInitiator;
 import org.jetbrains.annotations.NotNull;
 
@@ -28,4 +29,6 @@ public interface VimKeyGroup {
   @NotNull KeyMapping getKeyMapping(@NotNull MappingMode mode);
 
   @NotNull CommandPartNode<VimActionsInitiator> getKeyRoot(@NotNull MappingMode mappingMode);
+
+  @NotNull KeyMappingLayer getKeyMappingLayer(@NotNull MappingMode mode);
 }

src/main/java/com/maddyhome/idea/vim/key/KeyMapping.java

@@ -22,8 +22,6 @@ import com.google.common.collect.HashMultiset;
 import com.google.common.collect.Multiset;
 import com.maddyhome.idea.vim.extension.VimExtensionHandler;
 import com.maddyhome.idea.vim.helper.StringHelper;
-import com.maddyhome.idea.vim.vimscript.model.Executable;
-import com.maddyhome.idea.vim.vimscript.model.VimLContext;
 import com.maddyhome.idea.vim.vimscript.model.expressions.Expression;
 import kotlin.Pair;
 import org.jetbrains.annotations.NotNull;
@@ -39,7 +37,7 @@ import java.util.stream.Collectors;
  *
  * @author vlan
  */
-public class KeyMapping implements Iterable<List<KeyStroke>> {
+public class KeyMapping implements Iterable<List<KeyStroke>>, KeyMappingLayer {
   /**
    * Contains all key mapping for some mode.
    */
@@ -149,7 +147,8 @@ public class KeyMapping implements Iterable<List<KeyStroke>> {
       .map(o -> new Pair<>(o.getKey(), o.getValue())).collect(Collectors.toList());
   }
 
-  public boolean isPrefix(@NotNull Iterable<KeyStroke> keys) {
+  @Override
+  public boolean isPrefix(@NotNull Iterable<? extends KeyStroke> keys) {
     // Having a parameter of Iterable allows for a nicer API, because we know when a given list is immutable.
     // Perhaps we should look at changing this to a trie or something?
     assert (keys instanceof List) : "keys must be of type List<KeyStroke>";
@@ -174,4 +173,11 @@ public class KeyMapping implements Iterable<List<KeyStroke>> {
       o -> o.getValue() instanceof ToKeysMappingInfo && ((ToKeysMappingInfo)o.getValue()).getToKeys().equals(toKeys))
       .map(o -> new Pair<>(o.getKey(), o.getValue())).collect(Collectors.toList());
   }
+
+  @Nullable
+  @Override
+  public MappingInfoLayer getLayer(@NotNull Iterable<? extends KeyStroke> keys) {
+    Iterable<KeyStroke> keys1 = (Iterable<KeyStroke>)keys;
+    return get(keys1);
+  }
 }

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

@@ -0,0 +1,26 @@
+/*
+ * IdeaVim - Vim emulator for IDEs based on the IntelliJ platform
+ * Copyright (C) 2003-2022 The IdeaVim authors
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program. If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.maddyhome.idea.vim.key
+
+import javax.swing.KeyStroke
+
+interface KeyMappingLayer {
+  fun isPrefix(keys: Iterable<KeyStroke>): Boolean
+  fun getLayer(keys: Iterable<KeyStroke>): MappingInfoLayer?
+}

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

@@ -60,12 +60,12 @@ sealed class MappingInfo(
   val fromKeys: List<KeyStroke>,
   val isRecursive: Boolean,
   val owner: MappingOwner,
-) : Comparable<MappingInfo> {
+) : Comparable<MappingInfo>, MappingInfoLayer {
 
   @VimNlsSafe
-  abstract fun getPresentableString(): String
+  abstract override fun getPresentableString(): String
 
-  abstract fun execute(editor: VimEditor, context: ExecutionContext)
+  abstract override fun execute(editor: VimEditor, context: ExecutionContext)
 
   override fun compareTo(other: MappingInfo): Int {
     val size = fromKeys.size

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

@@ -0,0 +1,27 @@
+/*
+ * IdeaVim - Vim emulator for IDEs based on the IntelliJ platform
+ * Copyright (C) 2003-2022 The IdeaVim authors
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program. If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.maddyhome.idea.vim.key
+
+import com.maddyhome.idea.vim.api.ExecutionContext
+import com.maddyhome.idea.vim.api.VimEditor
+
+interface MappingInfoLayer {
+  fun getPresentableString(): String
+  fun execute(editor: VimEditor, context: ExecutionContext)
+}