Add support for count for visual and line motion surround

src/main/java/com/maddyhome/idea/vim/extension/surround/RepeatedCharSequence.kt

@@ -0,0 +1,30 @@
+package com.maddyhome.idea.vim.extension.surround
+
+import com.intellij.util.text.CharSequenceSubSequence
+
+internal data class RepeatedCharSequence(val text: CharSequence, val count: Int) : CharSequence {
+  override val length = text.length * count
+
+  override fun get(index: Int): Char {
+    if (index < 0 || index >= length) throw IndexOutOfBoundsException()
+    return text[index % text.length]
+  }
+
+  override fun subSequence(startIndex: Int, endIndex: Int): CharSequence {
+    return CharSequenceSubSequence(this, startIndex, endIndex)
+  }
+
+  override fun toString(): String {
+    return text.repeat(count)
+  }
+  
+  companion object {
+    fun of(text: CharSequence, count: Int): CharSequence {
+      return when (count) {
+        0 -> ""
+        1 -> text
+        else -> RepeatedCharSequence(text, count)
+      }
+    }
+  }
+}

src/main/java/com/maddyhome/idea/vim/extension/surround/VimSurroundExtension.kt

@@ -83,7 +83,7 @@ internal class VimSurroundExtension : VimExtension {
       putKeyMappingIfMissing(MappingMode.XO, injector.parser.parseKeys("S"), owner, injector.parser.parseKeys("<Plug>VSurround"), true)
     }
 
-    VimExtensionFacade.exportOperatorFunction(OPERATOR_FUNC, Operator(supportsMultipleCursors = false)) // TODO
+    VimExtensionFacade.exportOperatorFunction(OPERATOR_FUNC, Operator(supportsMultipleCursors = false, count = 1)) // TODO
   }
 
   private class YSurroundHandler : ExtensionHandler {
@@ -111,7 +111,7 @@ internal class VimSurroundExtension : VimExtension {
         val lastNonWhiteSpaceOffset = getLastNonWhitespaceCharacterOffset(editor.text(), lineStartOffset, lineEndOffset)
         if (lastNonWhiteSpaceOffset != null) {
           val range = TextRange(lineStartOffset, lastNonWhiteSpaceOffset + 1)
-          performSurround(pair, range, it)
+          performSurround(pair, range, it, count = operatorArguments.count1)
         }
 //        it.moveToOffset(lineStartOffset)
       }
@@ -135,7 +135,7 @@ internal class VimSurroundExtension : VimExtension {
   private class VSurroundHandler : ExtensionHandler {
     override fun execute(editor: VimEditor, context: ExecutionContext, operatorArguments: OperatorArguments) {
       // NB: Operator ignores SelectionType anyway
-      if (!Operator(supportsMultipleCursors = true).apply(editor, context, editor.mode.selectionType)) {
+      if (!Operator(supportsMultipleCursors = true, count = operatorArguments.count1).apply(editor, context, editor.mode.selectionType)) {
         return
       }
       runWriteAction {
@@ -278,7 +278,7 @@ internal class VimSurroundExtension : VimExtension {
     }
   }
 
-  private class Operator(private val supportsMultipleCursors: Boolean) : OperatorFunction {
+  private class Operator(private val supportsMultipleCursors: Boolean, private val count: Int) : OperatorFunction {
     override fun apply(vimEditor: VimEditor, context: ExecutionContext, selectionType: SelectionType?): Boolean {
       val ijEditor = vimEditor.ij
       val c = getChar(ijEditor)
@@ -290,11 +290,11 @@ internal class VimSurroundExtension : VimExtension {
         val change = VimPlugin.getChange()
         if (supportsMultipleCursors) {
           ijEditor.runWithEveryCaretAndRestore {
-            applyOnce(ijEditor, change, pair)
+            applyOnce(ijEditor, change, pair, count)
           }
         }
         else {
-          applyOnce(ijEditor, change, pair)
+          applyOnce(ijEditor, change, pair, count)
           // Jump back to start
           executeNormalWithoutMapping(injector.parser.parseKeys("`["), ijEditor)
         }
@@ -302,18 +302,15 @@ internal class VimSurroundExtension : VimExtension {
       return true
     }
     
-    private fun applyOnce(editor: Editor, change: VimChangeGroup, pair: SurroundPair) {
+    private fun applyOnce(editor: Editor, change: VimChangeGroup, pair: SurroundPair, count: Int) {
       // XXX: Will it work with line-wise or block-wise selections?
       val primaryCaret = editor.caretModel.primaryCaret
       val range = getSurroundRange(primaryCaret.vim)
       if (range != null) {
-        change.insertText(IjVimEditor(editor), IjVimCaret(primaryCaret), range.startOffset, pair.first)
+        val start = RepeatedCharSequence.of(pair.first, count)
-        change.insertText(
+        val end = RepeatedCharSequence.of(pair.second, count)
-          IjVimEditor(editor),
+        change.insertText(IjVimEditor(editor), IjVimCaret(primaryCaret), range.startOffset, start)
-          IjVimCaret(primaryCaret),
+        change.insertText(IjVimEditor(editor), IjVimCaret(primaryCaret), range.endOffset + start.length, end)
-          range.endOffset + pair.first.length,
-          pair.second
-        )
       }
     }
 
@@ -416,15 +413,15 @@ private fun getChar(editor: Editor): Char {
   return res
 }
 
-private fun performSurround(pair: SurroundPair, range: TextRange, caret: VimCaret, tagsOnNewLines: Boolean = false) {
+private fun performSurround(pair: SurroundPair, range: TextRange, caret: VimCaret, count: Int, tagsOnNewLines: Boolean = false) {
   runWriteAction {
     val editor = caret.editor
     val change = VimPlugin.getChange()
-    val leftSurround = pair.first + if (tagsOnNewLines) "\n" else ""
+    val leftSurround = RepeatedCharSequence.of(pair.first + if (tagsOnNewLines) "\n" else "", count)
 
     val isEOF = range.endOffset == editor.text().length
     val hasNewLine = editor.endsWithNewLine()
-    val rightSurround = if (tagsOnNewLines) {
+    val rightSurround = (if (tagsOnNewLines) {
       if (isEOF && !hasNewLine) {
         "\n" + pair.second
       } else {
@@ -432,7 +429,7 @@ private fun performSurround(pair: SurroundPair, range: TextRange, caret: VimCare
       }
     } else {
       pair.second
-    }
+    }).let { RepeatedCharSequence.of(it, count) }
 
     change.insertText(editor, caret, range.startOffset, leftSurround)
     change.insertText(editor, caret, range.endOffset + leftSurround.length, rightSurround)

vim-engine/src/main/kotlin/com/maddyhome/idea/vim/api/VimChangeGroup.kt

@@ -271,7 +271,7 @@ interface VimChangeGroup {
     operatorArguments: OperatorArguments,
   )
 
-  fun insertText(editor: VimEditor, caret: VimCaret, offset: Int, str: String): VimCaret
+  fun insertText(editor: VimEditor, caret: VimCaret, offset: Int, str: CharSequence): VimCaret
 
   fun insertText(editor: VimEditor, caret: VimCaret, str: String): VimCaret
 

vim-engine/src/main/kotlin/com/maddyhome/idea/vim/api/VimChangeGroupBase.kt

@@ -216,7 +216,7 @@ abstract class VimChangeGroupBase : VimChangeGroup {
    * @param caret  The caret to start insertion in
    * @param str    The text to insert
    */
-  override fun insertText(editor: VimEditor, caret: VimCaret, offset: Int, str: String): VimCaret {
+  override fun insertText(editor: VimEditor, caret: VimCaret, offset: Int, str: CharSequence): VimCaret {
     injector.application.runWriteAction {
       (editor as MutableVimEditor).insertText(caret, offset, str)
     }