Improvements to Commentary extension

build.gradle.kts

@@ -147,7 +147,8 @@ intellij {
     downloadSources.set(downloadIdeaSources.toBoolean())
     instrumentCode.set(instrumentPluginCode.toBoolean())
     intellijRepository.set("https://www.jetbrains.com/intellij-repository")
-    plugins.set(listOf("java", "AceJump:3.8.4"))
+    // Yaml is only used for testing. It's part of the IdeaIC distribution, but needs to be included as a reference
+    plugins.set(listOf("java", "AceJump:3.8.4", "yaml"))
 }
 
 tasks {

src/main/java/com/maddyhome/idea/vim/action/change/OperatorAction.kt

@@ -20,56 +20,105 @@ package com.maddyhome.idea.vim.action.change
 import com.maddyhome.idea.vim.KeyHandler
 import com.maddyhome.idea.vim.VimPlugin
 import com.maddyhome.idea.vim.api.ExecutionContext
+import com.maddyhome.idea.vim.api.VimCaret
 import com.maddyhome.idea.vim.api.VimEditor
 import com.maddyhome.idea.vim.command.Argument
 import com.maddyhome.idea.vim.command.Command
+import com.maddyhome.idea.vim.command.CommandFlags
 import com.maddyhome.idea.vim.command.OperatorArguments
 import com.maddyhome.idea.vim.command.SelectionType
+import com.maddyhome.idea.vim.common.TextRange
 import com.maddyhome.idea.vim.common.argumentCaptured
 import com.maddyhome.idea.vim.group.MotionGroup
+import com.maddyhome.idea.vim.group.visual.VimSelection
 import com.maddyhome.idea.vim.handler.VimActionHandler
+import com.maddyhome.idea.vim.handler.VisualOperatorActionHandler
 import com.maddyhome.idea.vim.helper.MessageHelper
 import com.maddyhome.idea.vim.helper.commandState
+import com.maddyhome.idea.vim.helper.enumSetOf
 import com.maddyhome.idea.vim.newapi.ij
+import java.util.*
+
+private fun doOperatorAction(editor: VimEditor, context: ExecutionContext, textRange: TextRange, selectionType: SelectionType): Boolean {
+  val operatorFunction = VimPlugin.getKey().operatorFunction
+  if (operatorFunction == null) {
+    VimPlugin.showMessage(MessageHelper.message("E774"))
+    return false
+  }
+
+  val saveRepeatHandler = VimRepeater.repeatHandler
+  VimPlugin.getMark().setChangeMarks(editor, textRange)
+  KeyHandler.getInstance().reset(editor)
+  val result = operatorFunction.apply(editor.ij, context.ij, selectionType)
+  VimRepeater.repeatHandler = saveRepeatHandler
+  return result
+}
 
-/**
- * @author vlan
- */
 class OperatorAction : VimActionHandler.SingleExecution() {
   override val type: Command.Type = Command.Type.OTHER_SELF_SYNCHRONIZED
 
   override val argumentType: Argument.Type = Argument.Type.MOTION
 
   override fun execute(editor: VimEditor, context: ExecutionContext, cmd: Command, operatorArguments: OperatorArguments): Boolean {
-    val operatorFunction = VimPlugin.getKey().operatorFunction
-    if (operatorFunction != null) {
-      val argument = cmd.argument
-      if (argument != null) {
-        if (!editor.commandState.isDotRepeatInProgress) {
-          argumentCaptured = argument
-        }
-        val saveRepeatHandler = VimRepeater.repeatHandler
-        val motion = argument.motion
-        val range = MotionGroup
-          .getMotionRange(
-            editor.ij,
-            editor.ij.caretModel.primaryCaret,
-            context.ij,
-            argument,
-            operatorArguments
-          )
-        if (range != null) {
-          VimPlugin.getMark().setChangeMarks(editor, range)
-          val selectionType = if (motion.isLinewiseMotion()) SelectionType.LINE_WISE else SelectionType.CHARACTER_WISE
-          KeyHandler.getInstance().reset(editor)
-          val result = operatorFunction.apply(editor.ij, context.ij, selectionType)
-          VimRepeater.repeatHandler = saveRepeatHandler
-          return result
-        }
+    val argument = cmd.argument ?: return false
+    if (!editor.commandState.isDotRepeatInProgress) {
+      argumentCaptured = argument
+    }
+    val range = getMotionRange(editor, context, argument, operatorArguments)
+
+    if (range != null) {
+      val selectionType = if (argument.motion.isLinewiseMotion()) {
+        SelectionType.LINE_WISE
       }
-      return false
+      else {
+        SelectionType.CHARACTER_WISE
+      }
+      return doOperatorAction(editor, context, range, selectionType)
     }
-    VimPlugin.showMessage(MessageHelper.message("E774"))
     return false
   }
+
+  private fun getMotionRange(
+    editor: VimEditor,
+    context: ExecutionContext,
+    argument: Argument,
+    operatorArguments: OperatorArguments): TextRange? {
+
+    // Note that we're using getMotionRange2 in order to avoid normalising the linewise range into line start
+    // offsets that will be used to set the change marks. This affects things like the location of the caret in the
+    // Commentary extension
+    val ijEditor = editor.ij
+    return MotionGroup.getMotionRange2(
+      ijEditor,
+      ijEditor.caretModel.primaryCaret,
+      context.ij,
+      argument,
+      operatorArguments
+    )?.normalize()?.let {
+
+      // If we're linewise, make sure the end offset isn't just the EOL char
+      if (argument.motion.isLinewiseMotion() && it.endOffset < editor.fileSize()) {
+        TextRange(it.startOffset, it.endOffset + 1)
+      } else {
+        it
+      }
+    }
+  }
 }
+
+class VisualOperatorAction: VisualOperatorActionHandler.ForEachCaret() {
+  override val type: Command.Type = Command.Type.OTHER_SELF_SYNCHRONIZED
+
+  override val flags: EnumSet<CommandFlags> = enumSetOf(CommandFlags.FLAG_EXIT_VISUAL)
+
+  override fun executeAction(
+    editor: VimEditor,
+    caret: VimCaret,
+    context: ExecutionContext,
+    cmd: Command,
+    range: VimSelection,
+    operatorArguments: OperatorArguments,
+  ): Boolean {
+    return doOperatorAction(editor, context, range.toVimTextRange(), range.type)
+  }
+}

src/main/java/com/maddyhome/idea/vim/common/CommandAlias.kt

@@ -18,9 +18,10 @@
 
 package com.maddyhome.idea.vim.common
 
-import com.intellij.openapi.actionSystem.DataContext
-import com.intellij.openapi.editor.Editor
 import com.maddyhome.idea.vim.VimPlugin
+import com.maddyhome.idea.vim.api.ExecutionContext
+import com.maddyhome.idea.vim.api.VimEditor
+import com.maddyhome.idea.vim.ex.ranges.Ranges
 import com.maddyhome.idea.vim.helper.MessageHelper
 import org.jetbrains.annotations.NonNls
 
@@ -125,5 +126,5 @@ sealed class GoalCommand {
 }
 
 interface CommandAliasHandler {
-  fun execute(editor: Editor, context: DataContext)
+  fun execute(command: String, ranges: Ranges, editor: VimEditor, context: ExecutionContext)
 }

src/main/java/com/maddyhome/idea/vim/extension/VimExtensionFacade.kt

@@ -24,6 +24,8 @@ import com.maddyhome.idea.vim.KeyHandler
 import com.maddyhome.idea.vim.VimPlugin
 import com.maddyhome.idea.vim.action.change.Extension
 import com.maddyhome.idea.vim.command.SelectionType
+import com.maddyhome.idea.vim.common.CommandAlias
+import com.maddyhome.idea.vim.common.CommandAliasHandler
 import com.maddyhome.idea.vim.common.MappingMode
 import com.maddyhome.idea.vim.helper.CommandLineHelper
 import com.maddyhome.idea.vim.helper.EditorDataContext
@@ -81,6 +83,29 @@ object VimExtensionFacade {
     VimPlugin.getKey().putKeyMapping(filteredModes, fromKeys, pluginOwner, toKeys, recursive)
   }
 
+  /**
+   * Equivalent to calling 'command' to set up a user-defined command or alias
+   */
+  fun addCommand(
+    name: String,
+    handler: CommandAliasHandler
+  ) {
+    addCommand(name, 0, 0, handler)
+  }
+
+  /**
+   * Equivalent to calling 'command' to set up a user-defined command or alias
+   */
+  @JvmStatic
+  fun addCommand(
+    name: String,
+    minimumNumberOfArguments: Int,
+    maximumNumberOfArguments: Int,
+    handler: CommandAliasHandler
+  ) {
+    VimPlugin.getCommand().setAlias(name, CommandAlias.Call(minimumNumberOfArguments, maximumNumberOfArguments, name, handler))
+  }
+
   /** Sets the value of 'operatorfunc' to be used as the operator function in 'g@'. */
   @JvmStatic
   fun setOperatorFunction(function: OperatorFunction) {