Introduce 'operatorfunc' option for custom operators

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

@@ -14,13 +14,15 @@ 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.api.globalOptions
 import com.maddyhome.idea.vim.api.injector
 import com.maddyhome.idea.vim.api.setChangeMarks
 import com.maddyhome.idea.vim.command.Argument
 import com.maddyhome.idea.vim.command.Command
 import com.maddyhome.idea.vim.command.OperatorArguments
 import com.maddyhome.idea.vim.common.TextRange
 import com.maddyhome.idea.vim.common.argumentCaptured
+import com.maddyhome.idea.vim.ex.ExException
 import com.maddyhome.idea.vim.group.MotionGroup
 import com.maddyhome.idea.vim.group.visual.VimSelection
 import com.maddyhome.idea.vim.handler.VimActionHandler
@@ -29,21 +31,67 @@ import com.maddyhome.idea.vim.helper.MessageHelper
 import com.maddyhome.idea.vim.helper.vimStateMachine
 import com.maddyhome.idea.vim.newapi.ij
 import com.maddyhome.idea.vim.state.mode.SelectionType
+import com.maddyhome.idea.vim.vimscript.model.CommandLineVimLContext
+import com.maddyhome.idea.vim.vimscript.model.datatypes.VimFuncref
+import com.maddyhome.idea.vim.vimscript.model.expressions.FunctionCallExpression
+import com.maddyhome.idea.vim.vimscript.model.expressions.SimpleExpression
 
 // todo make it multicaret
 private fun doOperatorAction(editor: VimEditor, context: ExecutionContext, textRange: TextRange, selectionType: SelectionType): Boolean {
-  val operatorFunction = injector.keyGroup.operatorFunction
-  if (operatorFunction == null) {
+  val func = injector.globalOptions().operatorfunc
+  if (func.isEmpty()) {
     VimPlugin.showMessage(MessageHelper.message("E774"))
     return false
   }
 
+  val scriptContext = CommandLineVimLContext
+
+  // The option value is either a function name, which should have a handler, or it might be a lambda expression, or a
+  // `function` or `funcref` call expression, all of which will return a funcref (with a handler)
+  var handler = injector.functionService.getFunctionHandlerOrNull(null, func, scriptContext)
+  if (handler == null) {
+    val expression = injector.vimscriptParser.parseExpression(func)
+    if (expression != null) {
+      try {
+        val value = expression.evaluate(editor, context, scriptContext)
+        if (value is VimFuncref) {
+          handler = value.handler
+        }
+      } catch (ex: ExException) {
+        // Get the argument for function('...') or funcref('...') for the error message
+        val functionName = if (expression is FunctionCallExpression && expression.arguments.size > 0) {
+          expression.arguments[0].evaluate(editor, context, scriptContext).toString()
+        }
+        else {
+          func
+        }
+
+        VimPlugin.showMessage("E117: Unknown function: $functionName")
+        return false
+      }
+    }
+  }
+
+  if (handler == null) {
+    VimPlugin.showMessage("E117: Unknown function: $func")
+    return false
+  }
+
+  val arg = when (selectionType) {
+    SelectionType.LINE_WISE -> "line"
+    SelectionType.CHARACTER_WISE -> "char"
+    SelectionType.BLOCK_WISE -> "block"
+  }
+
   val saveRepeatHandler = VimRepeater.repeatHandler
   injector.markService.setChangeMarks(editor.primaryCaret(), textRange)
   KeyHandler.getInstance().reset(editor)
-  val result = operatorFunction.apply(editor, context, selectionType)
+
+  val arguments = listOf(SimpleExpression(arg))
+  handler.executeFunction(arguments, editor, context, scriptContext)
+
   VimRepeater.repeatHandler = saveRepeatHandler
-  return result
+  return true
 }
 
 @CommandOrMotion(keys = ["g@"], modes = [Mode.NORMAL])

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

@@ -14,21 +14,34 @@ import com.intellij.openapi.editor.Editor
 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.api.ExecutionContext
 import com.maddyhome.idea.vim.api.ImmutableVimCaret
 import com.maddyhome.idea.vim.api.VimCaret
+import com.maddyhome.idea.vim.api.VimEditor
 import com.maddyhome.idea.vim.api.injector
 import com.maddyhome.idea.vim.command.MappingMode
 import com.maddyhome.idea.vim.common.CommandAlias
 import com.maddyhome.idea.vim.common.CommandAliasHandler
+import com.maddyhome.idea.vim.common.TextRange
 import com.maddyhome.idea.vim.helper.CommandLineHelper
 import com.maddyhome.idea.vim.helper.TestInputModel
+import com.maddyhome.idea.vim.helper.noneOfEnum
 import com.maddyhome.idea.vim.helper.vimStateMachine
 import com.maddyhome.idea.vim.key.MappingOwner
 import com.maddyhome.idea.vim.key.OperatorFunction
 import com.maddyhome.idea.vim.newapi.vim
 import com.maddyhome.idea.vim.state.mode.SelectionType
 import com.maddyhome.idea.vim.ui.ModalEntry
+import com.maddyhome.idea.vim.vimscript.model.Executable
+import com.maddyhome.idea.vim.vimscript.model.ExecutionResult
+import com.maddyhome.idea.vim.vimscript.model.VimLContext
+import com.maddyhome.idea.vim.vimscript.model.datatypes.VimDataType
+import com.maddyhome.idea.vim.vimscript.model.expressions.Expression
+import com.maddyhome.idea.vim.vimscript.model.expressions.Scope
+import com.maddyhome.idea.vim.vimscript.model.statements.FunctionDeclaration
+import com.maddyhome.idea.vim.vimscript.model.statements.FunctionFlag
 import java.awt.event.KeyEvent
+import java.util.*
 import javax.swing.KeyStroke
 
 /**
@@ -120,12 +133,6 @@ public object VimExtensionFacade {
       .setAlias(name, CommandAlias.Call(minimumNumberOfArguments, maximumNumberOfArguments, name, handler))
   }
 
-  /** Sets the value of 'operatorfunc' to be used as the operator function in 'g@'. */
-  @JvmStatic
-  public fun setOperatorFunction(function: OperatorFunction) {
-    VimPlugin.getKey().operatorFunction = function
-  }
-
   /**
    * Runs normal mode commands similar to ':normal! {commands}'.
    * Mappings doesn't work with this function
@@ -207,4 +214,65 @@ public object VimExtensionFacade {
   public fun setRegister(register: Char, keys: List<KeyStroke?>?, type: SelectionType) {
     VimPlugin.getRegister().setKeys(register, keys?.filterNotNull() ?: emptyList(), type)
   }
+
+  @JvmStatic
+  public fun exportScriptFunction(
+    scope: Scope?,
+    name: String,
+    args: List<String>,
+    defaultArgs: List<Pair<String, Expression>>,
+    hasOptionalArguments: Boolean,
+    flags: EnumSet<FunctionFlag>,
+    function: ScriptFunction
+  ) {
+    var functionDeclaration: FunctionDeclaration? = null
+    val body = listOf(object : Executable {
+      // This context is set to the function declaration during initialisation and then set to the function execution
+      // context during execution
+      override lateinit var vimContext: VimLContext
+      override var rangeInScript: TextRange = TextRange(0, 0)
+
+      override fun execute(editor: VimEditor, context: ExecutionContext): ExecutionResult {
+        return function.execute(editor, context, functionDeclaration!!.functionVariables)
+      }
+    })
+    functionDeclaration = FunctionDeclaration(
+      scope,
+      name,
+      args,
+      defaultArgs,
+      body,
+      replaceExisting = true,
+      flags,
+      hasOptionalArguments
+    )
+    functionDeclaration.rangeInScript = TextRange(0, 0)
+    body.forEach { it.vimContext = functionDeclaration }
+    injector.functionService.storeFunction(functionDeclaration)
+  }
 }
+
+public fun VimExtensionFacade.exportOperatorFunction(name: String, function: OperatorFunction) {
+  exportScriptFunction(null, name, listOf("type"), emptyList(), false, noneOfEnum()) {
+    editor, context, args ->
+
+    val type = args["type"]?.asString()
+    val selectionType = when (type) {
+      "line" -> SelectionType.LINE_WISE
+      "block" -> SelectionType.BLOCK_WISE
+      "char" -> SelectionType.CHARACTER_WISE
+      else -> return@exportScriptFunction ExecutionResult.Error
+    }
+
+    if (function.apply(editor, context, selectionType)) {
+      ExecutionResult.Success
+    }
+    else {
+      ExecutionResult.Error
+    }
+  }
+}
+
+public fun interface ScriptFunction {
+  public fun execute(editor: VimEditor, context: ExecutionContext, args: Map<String, VimDataType>): ExecutionResult
+}

src/main/java/com/maddyhome/idea/vim/extension/commentary/CommentaryExtension.kt

@@ -22,53 +22,57 @@ import com.maddyhome.idea.vim.api.ExecutionContext
 import com.maddyhome.idea.vim.api.ImmutableVimCaret
 import com.maddyhome.idea.vim.api.VimEditor
 import com.maddyhome.idea.vim.api.getLineEndOffset
+import com.maddyhome.idea.vim.api.globalOptions
 import com.maddyhome.idea.vim.api.injector
 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.MappingMode
-import com.maddyhome.idea.vim.state.mode.Mode
 import com.maddyhome.idea.vim.command.OperatorArguments
-import com.maddyhome.idea.vim.state.mode.SelectionType
 import com.maddyhome.idea.vim.command.TextObjectVisualType
 import com.maddyhome.idea.vim.common.CommandAliasHandler
 import com.maddyhome.idea.vim.common.TextRange
 import com.maddyhome.idea.vim.ex.ranges.Ranges
 import com.maddyhome.idea.vim.extension.ExtensionHandler
 import com.maddyhome.idea.vim.extension.VimExtension
+import com.maddyhome.idea.vim.extension.VimExtensionFacade
 import com.maddyhome.idea.vim.extension.VimExtensionFacade.addCommand
 import com.maddyhome.idea.vim.extension.VimExtensionFacade.executeNormalWithoutMapping
 import com.maddyhome.idea.vim.extension.VimExtensionFacade.putExtensionHandlerMapping
 import com.maddyhome.idea.vim.extension.VimExtensionFacade.putKeyMapping
 import com.maddyhome.idea.vim.extension.VimExtensionFacade.putKeyMappingIfMissing
-import com.maddyhome.idea.vim.extension.VimExtensionFacade.setOperatorFunction
+import com.maddyhome.idea.vim.extension.exportOperatorFunction
 import com.maddyhome.idea.vim.handler.TextObjectActionHandler
 import com.maddyhome.idea.vim.helper.PsiHelper
 import com.maddyhome.idea.vim.helper.vimStateMachine
 import com.maddyhome.idea.vim.key.OperatorFunction
 import com.maddyhome.idea.vim.newapi.IjVimEditor
 import com.maddyhome.idea.vim.newapi.ij
 import com.maddyhome.idea.vim.newapi.vim
+import com.maddyhome.idea.vim.state.mode.Mode
+import com.maddyhome.idea.vim.state.mode.SelectionType
 import java.util.*
 
 internal class CommentaryExtension : VimExtension {
 
   companion object {
+
+    private const val OPERATOR_FUNC = "CommentaryOperatorFunc"
+
     fun doCommentary(
       editor: VimEditor,
       context: ExecutionContext,
       range: TextRange,
       selectionType: SelectionType,
-      resetCaret: Boolean,
+      resetCaret: Boolean = true,
     ): Boolean {
       val mode = editor.vimStateMachine.mode
       if (mode !is Mode.VISUAL) {
         editor.ij.selectionModel.setSelection(range.startOffset, range.endOffset)
       }
 
       return runWriteAction {
-        // Treat block- and character-wise selections as block comments. Be ready to fall back to if the first action
-        // isn't available
+        // Treat block- and character-wise selections as block comments. Fall back if the first action isn't available
         val actions = if (selectionType === SelectionType.LINE_WISE) {
           listOf(IdeActions.ACTION_COMMENT_LINE, IdeActions.ACTION_COMMENT_BLOCK)
         } else {
@@ -113,6 +117,7 @@ internal class CommentaryExtension : VimExtension {
       // first non-whitespace character, then the caret is in the right place. If it's inserted at the first column,
       // then the caret is now in a bit of a weird place. We can't detect this scenario, so we just have to accept
       // the difference
+      // TODO: If we don't move the caret to the start offset, we should maintain the current logical position
       if (resetCaret) {
         editor.primaryCaret().moveToOffset(range.startOffset)
       }
@@ -145,6 +150,16 @@ internal class CommentaryExtension : VimExtension {
     putKeyMapping(MappingMode.N, injector.parser.parseKeys("<Plug>(CommentLine)"), owner, plugCommentaryLineKeys, true)
 
     addCommand("Commentary", CommentaryCommandAliasHandler())
+
+    VimExtensionFacade.exportOperatorFunction(OPERATOR_FUNC, CommentaryOperatorFunction())
+ }
+
+  private class CommentaryOperatorFunction : OperatorFunction {
+    // todo make it multicaret
+    override fun apply(editor: VimEditor, context: ExecutionContext, selectionType: SelectionType?): Boolean {
+      val range = injector.markService.getChangeMarks(editor.primaryCaret()) ?: return false
+      return doCommentary(editor, context, range, selectionType ?: SelectionType.CHARACTER_WISE, true)
+    }
   }
 
   /**
@@ -153,19 +168,13 @@ internal class CommentaryExtension : VimExtension {
    * E.g. handles the `gc` in `gc_`, by setting the operator function, then invoking `g@` to receive the `_` motion to
    * invoke the operator. This object is both the mapping handler and the operator function.
    */
-  private class CommentaryOperatorHandler : OperatorFunction, ExtensionHandler {
+  private class CommentaryOperatorHandler : ExtensionHandler {
     override val isRepeatable = true
 
     override fun execute(editor: VimEditor, context: ExecutionContext, operatorArguments: OperatorArguments) {
-      setOperatorFunction(this)
+      injector.globalOptions().operatorfunc = OPERATOR_FUNC
       executeNormalWithoutMapping(injector.parser.parseKeys("g@"), editor.ij)
     }
-
-    // todo make it multicaret
-    override fun apply(editor: VimEditor, context: ExecutionContext, selectionType: SelectionType?): Boolean {
-      val range = injector.markService.getChangeMarks(editor.primaryCaret()) ?: return false
-      return doCommentary(editor, context, range, selectionType ?: SelectionType.CHARACTER_WISE, true)
-    }
   }
 
   private class CommentaryMappingHandler : ExtensionHandler {