[VIM-2929]: Reset the key stack in case of exception during the execu…

…tion
JetBrains · Jan 2, 2024 · b307c7d · b307c7d
1 parent 47d4445
commit b307c7d
Show file tree

Hide file tree

Showing 7 changed files with 147 additions and 26 deletions.
diff --git a/src/main/java/com/maddyhome/idea/vim/group/MacroGroup.kt b/src/main/java/com/maddyhome/idea/vim/group/MacroGroup.kt
@@ -63,27 +63,33 @@ internal class MacroGroup : VimMacroBase() {
     try {
       myPotemkinProgress.text2 = if (isInternalMacro) "Executing internal macro" else ""
       val runnable = runnable@{
-        // Handle one keystroke then queue up the next key
-        for (i in 0 until total) {
-          myPotemkinProgress.fraction = (i + 1).toDouble() / total
-          while (keyStack.hasStroke()) {
-            val key = keyStack.feedStroke()
+        try {
+          // Handle one keystroke then queue up the next key
+          for (i in 0 until total) {
             try {
-              myPotemkinProgress.checkCanceled()
-            } catch (e: ProcessCanceledException) {
-              return@runnable
+              myPotemkinProgress.fraction = (i + 1).toDouble() / total
+              while (keyStack.hasStroke()) {
+                val key = keyStack.feedStroke()
+                try {
+                  myPotemkinProgress.checkCanceled()
+                } catch (e: ProcessCanceledException) {
+                  return@runnable
+                }
+                ProgressManager.getInstance().executeNonCancelableSection {
+                  // Prevent autocompletion during macros.
+                  // See https://github.com/JetBrains/ideavim/pull/772 for details
+                  CompletionServiceImpl.setCompletionPhase(CompletionPhase.NoCompletion)
+                  getInstance().handleKey(editor, key, context)
+                }
+                if (injector.messages.isError()) return@runnable
+              }
+            } finally {
+              keyStack.resetFirst()
             }
-            ProgressManager.getInstance().executeNonCancelableSection {
-              // Prevent autocompletion during macros.
-              // See https://github.com/JetBrains/ideavim/pull/772 for details
-              CompletionServiceImpl.setCompletionPhase(CompletionPhase.NoCompletion)
-              getInstance().handleKey(editor, key, context)
-            }
-            if (injector.messages.isError()) return@runnable
           }
-          keyStack.resetFirst()
+        } finally {
+          keyStack.removeFirst()
         }
-        keyStack.removeFirst()
       }
 
       if (isInternalMacro) {

diff --git a/src/test/java/org/jetbrains/plugins/ideavim/TestHelper.kt b/src/test/java/org/jetbrains/plugins/ideavim/TestHelper.kt
@@ -14,8 +14,13 @@ import com.intellij.openapi.editor.LogicalPosition
 import com.intellij.testFramework.EditorTestUtil
 import com.intellij.testFramework.fixtures.CodeInsightTestFixture
 import com.intellij.util.containers.toArray
+import com.maddyhome.idea.vim.api.ExecutionContext
+import com.maddyhome.idea.vim.api.VimEditor
 import com.maddyhome.idea.vim.api.injector
+import com.maddyhome.idea.vim.command.OperatorArguments
 import com.maddyhome.idea.vim.common.TextRange
+import com.maddyhome.idea.vim.extension.ExtensionHandler
+import com.maddyhome.idea.vim.key.MappingOwner
 import com.maddyhome.idea.vim.newapi.globalIjOptions
 import com.maddyhome.idea.vim.newapi.vim
 import com.maddyhome.idea.vim.state.mode.Mode
@@ -129,3 +134,15 @@ internal fun <T> product(vararg elements: List<T>): List<List<T>> {
   }
   return res
 }
+
+internal class ExceptionHandler : ExtensionHandler {
+  override fun execute(editor: VimEditor, context: ExecutionContext, operatorArguments: OperatorArguments) {
+    error(exceptionMessage)
+  }
+
+  companion object {
+    internal const val exceptionMessage = "Exception here"
+  }
+}
+
+internal val exceptionMappingOwner = MappingOwner.Plugin.get("Exception mapping owner")
diff --git a/src/test/java/org/jetbrains/plugins/ideavim/VimTestCase.kt b/src/test/java/org/jetbrains/plugins/ideavim/VimTestCase.kt
@@ -91,7 +91,6 @@ import java.util.*
 import javax.swing.KeyStroke
 import kotlin.math.roundToInt
 import kotlin.test.assertEquals
-import kotlin.test.assertFalse
 import kotlin.test.assertNotNull
 import kotlin.test.assertNull
 import kotlin.test.assertTrue
@@ -171,7 +170,7 @@ abstract class VimTestCase {
     injector.jumpService.resetJumps()
     VimPlugin.getChange().resetRepeat()
     VimPlugin.getKey().savedShortcutConflicts.clear()
-    assertFalse(KeyHandler.getInstance().keyStack.hasStroke())
+    assertTrue(KeyHandler.getInstance().keyStack.isEmpty())
 
     // Tear down neovim
     NeovimTesting.tearDown(testInfo)

diff --git a/src/test/java/org/jetbrains/plugins/ideavim/action/MacroActionTest.kt b/src/test/java/org/jetbrains/plugins/ideavim/action/MacroActionTest.kt
@@ -7,22 +7,40 @@
  */
 package org.jetbrains.plugins.ideavim.action
 
+import com.intellij.idea.TestFor
+import com.intellij.testFramework.LoggedErrorProcessor
+import com.maddyhome.idea.vim.KeyHandler
 import com.maddyhome.idea.vim.VimPlugin
 import com.maddyhome.idea.vim.api.injector
+import com.maddyhome.idea.vim.api.keys
+import com.maddyhome.idea.vim.command.MappingMode
 import com.maddyhome.idea.vim.helper.vimStateMachine
 import com.maddyhome.idea.vim.newapi.vim
+import org.jetbrains.plugins.ideavim.ExceptionHandler
+import org.jetbrains.plugins.ideavim.OnlyThrowLoggedErrorProcessor
 import org.jetbrains.plugins.ideavim.SkipNeovimReason
 import org.jetbrains.plugins.ideavim.TestWithoutNeovim
 import org.jetbrains.plugins.ideavim.VimTestCase
+import org.jetbrains.plugins.ideavim.exceptionMappingOwner
 import org.jetbrains.plugins.ideavim.rangeOf
 import org.jetbrains.plugins.ideavim.waitAndAssert
+import org.junit.jupiter.api.AfterEach
 import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.assertThrows
+import kotlin.test.assertEquals
 import kotlin.test.assertNotNull
+import kotlin.test.assertTrue
 
 /**
  * @author vlan
  */
 class MacroActionTest : VimTestCase() {
+
+  @AfterEach
+  fun tearDown() {
+    injector.keyGroup.removeKeyMapping(exceptionMappingOwner)
+  }
+
   // |q|
   @Test
   fun testRecordMacro() {
@@ -178,4 +196,33 @@ class MacroActionTest : VimTestCase() {
       """.trimIndent(),
     )
   }
+
+  @TestFor(issues = ["VIM-2929"])
+  @TestWithoutNeovim(reason = SkipNeovimReason.ACTION_COMMAND)
+  @Test
+  fun `macro to handler with exception`() {
+    configureByText(
+      """
+     Lorem Ipsum
+
+     Lorem ipsum dolor sit amet,
+     ${c}consectetur adipiscing elit
+     Sed in orci mauris.
+     Cras id tellus in ex imperdiet egestas. 
+    """.trimIndent()
+    )
+    injector.keyGroup.putKeyMapping(MappingMode.NXO, keys("abc"), exceptionMappingOwner, ExceptionHandler(), false)
+
+    injector.registerGroup.storeText('k', "abc")
+    injector.registerGroup.storeText('q', "x@ky")
+
+    val exception = assertThrows<Throwable> {
+      LoggedErrorProcessor.executeWith<Throwable>(OnlyThrowLoggedErrorProcessor) {
+        typeText("@q")
+      }
+    }
+    assertEquals(ExceptionHandler.exceptionMessage, exception.cause!!.cause!!.message)
+
+    assertTrue(KeyHandler.getInstance().keyStack.isEmpty())
+  }
 }
diff --git a/src/test/java/org/jetbrains/plugins/ideavim/ex/implementation/commands/MapCommandTest.kt b/src/test/java/org/jetbrains/plugins/ideavim/ex/implementation/commands/MapCommandTest.kt
@@ -11,27 +11,44 @@ import com.intellij.idea.TestFor
 import com.intellij.openapi.actionSystem.DataContext
 import com.intellij.openapi.editor.textarea.TextComponentEditorImpl
 import com.intellij.openapi.util.Disposer
+import com.intellij.testFramework.LoggedErrorProcessor
 import com.intellij.testFramework.TestLoggerFactory.TestLoggerAssertionError
+import com.maddyhome.idea.vim.KeyHandler
 import com.maddyhome.idea.vim.api.injector
+import com.maddyhome.idea.vim.api.keys
+import com.maddyhome.idea.vim.command.MappingMode
 import com.maddyhome.idea.vim.ex.ExException
 import com.maddyhome.idea.vim.helper.StringHelper.parseKeys
 import com.maddyhome.idea.vim.history.HistoryConstants
 import com.maddyhome.idea.vim.newapi.vim
 import com.maddyhome.idea.vim.state.mode.Mode
+import org.jetbrains.plugins.ideavim.ExceptionHandler
+import org.jetbrains.plugins.ideavim.OnlyThrowLoggedErrorProcessor
 import org.jetbrains.plugins.ideavim.SkipNeovimReason
 import org.jetbrains.plugins.ideavim.TestWithoutNeovim
 import org.jetbrains.plugins.ideavim.VimTestCase
 import org.jetbrains.plugins.ideavim.assertThrowsLogError
+import org.jetbrains.plugins.ideavim.exceptionMappingOwner
 import org.jetbrains.plugins.ideavim.waitAndAssert
+import org.junit.jupiter.api.AfterEach
 import org.junit.jupiter.api.Disabled
 import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.assertThrows
 import javax.swing.JTextArea
+import kotlin.test.assertEquals
 import kotlin.test.assertIs
+import kotlin.test.assertTrue
 
 /**
  * @author vlan
  */
 class MapCommandTest : VimTestCase() {
+
+  @AfterEach
+  fun tearDown() {
+    injector.keyGroup.removeKeyMapping(exceptionMappingOwner)
+  }
+
   @TestWithoutNeovim(reason = SkipNeovimReason.UNCLEAR)
   @Test
   fun testMapKtoJ() {
@@ -1099,4 +1116,32 @@ n  ,i            <Action>(Back)
      Cras id tellus in ex imperdiet egestas. 
     """.trimIndent())
   }
+
+  @TestFor(issues = ["VIM-2929"])
+  @TestWithoutNeovim(reason = SkipNeovimReason.ACTION_COMMAND)
+  @Test
+  fun `mapping to handler with exception`() {
+    configureByText(
+      """
+     Lorem Ipsum
+
+     Lorem ipsum dolor sit amet,
+     ${c}consectetur adipiscing elit
+     Sed in orci mauris.
+     Cras id tellus in ex imperdiet egestas. 
+    """.trimIndent()
+    )
+    injector.keyGroup.putKeyMapping(MappingMode.NXO, keys("abc"), exceptionMappingOwner, ExceptionHandler(), false)
+
+    typeText(commandToKeys("map k abcx"))
+
+    val exception = assertThrows<Throwable> {
+      LoggedErrorProcessor.executeWith<Throwable>(OnlyThrowLoggedErrorProcessor) {
+        typeText("k")
+      }
+    }
+    assertEquals(ExceptionHandler.exceptionMessage, exception.cause!!.cause!!.message)
+
+    assertTrue(KeyHandler.getInstance().keyStack.isEmpty())
+  }
 }
diff --git a/vim-engine/src/main/kotlin/com/maddyhome/idea/vim/key/KeyStack.kt b/vim-engine/src/main/kotlin/com/maddyhome/idea/vim/key/KeyStack.kt
@@ -28,6 +28,10 @@ public class KeyStack {
     return stack.isNotEmpty() && stack.first().hasStroke()
   }
 
+  public fun isEmpty(): Boolean {
+    return stack.none { it.hasStroke() }
+  }
+
   public fun feedSomeStroke(): KeyStroke? {
     stack.forEach {
       if (it.hasStroke()) {

diff --git a/vim-engine/src/main/kotlin/com/maddyhome/idea/vim/key/MappingInfo.kt b/vim-engine/src/main/kotlin/com/maddyhome/idea/vim/key/MappingInfo.kt
@@ -93,14 +93,17 @@ public class ToKeysMappingInfo(
     val keyHandler = KeyHandler.getInstance()
     LOG.trace { "Adding new keys to keyStack as toKeys of mapping. State before adding keys: ${keyHandler.keyStack.dump()}" }
     keyHandler.keyStack.addKeys(toKeys)
-    var first = true
-    while (keyHandler.keyStack.hasStroke()) {
-      val keyStroke = keyHandler.keyStack.feedStroke()
-      val recursive = isRecursive && !(first && fromIsPrefix)
-      keyHandler.handleKey(editor, keyStroke, editorDataContext, recursive, false)
-      first = false
+    try {
+      var first = true
+      while (keyHandler.keyStack.hasStroke()) {
+        val keyStroke = keyHandler.keyStack.feedStroke()
+        val recursive = isRecursive && !(first && fromIsPrefix)
+        keyHandler.handleKey(editor, keyStroke, editorDataContext, recursive, false)
+        first = false
+      }
+    } finally {
+      keyHandler.keyStack.removeFirst()
     }
-    keyHandler.keyStack.removeFirst()
   }
 
   public companion object {