Fix(VIM-3166): Workaround to fix broken filtering of visual lines

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

@@ -19,10 +19,7 @@ import com.maddyhome.idea.vim.api.getLeadingCharacterOffset
 import com.maddyhome.idea.vim.api.injector
 import com.maddyhome.idea.vim.api.setChangeMarks
 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.state.mode.selectionType
 import com.maddyhome.idea.vim.common.TextRange
 import com.maddyhome.idea.vim.extension.ExtensionHandler
 import com.maddyhome.idea.vim.extension.VimExtension
@@ -34,7 +31,6 @@ import com.maddyhome.idea.vim.extension.VimExtensionFacade.putExtensionHandlerMa
 import com.maddyhome.idea.vim.extension.VimExtensionFacade.putKeyMappingIfMissing
 import com.maddyhome.idea.vim.extension.VimExtensionFacade.setOperatorFunction
 import com.maddyhome.idea.vim.extension.VimExtensionFacade.setRegisterForCaret
-import com.maddyhome.idea.vim.state.mode.mode
 import com.maddyhome.idea.vim.helper.runWithEveryCaretAndRestore
 import com.maddyhome.idea.vim.key.OperatorFunction
 import com.maddyhome.idea.vim.newapi.IjVimCaret
@@ -43,6 +39,10 @@ import com.maddyhome.idea.vim.newapi.ij
 import com.maddyhome.idea.vim.newapi.vim
 import com.maddyhome.idea.vim.options.helpers.ClipboardOptionHelper
 import com.maddyhome.idea.vim.put.PutData
+import com.maddyhome.idea.vim.state.mode.Mode
+import com.maddyhome.idea.vim.state.mode.SelectionType
+import com.maddyhome.idea.vim.state.mode.mode
+import com.maddyhome.idea.vim.state.mode.selectionType
 import org.jetbrains.annotations.NonNls
 import java.awt.event.KeyEvent
 import javax.swing.KeyStroke
@@ -83,7 +83,7 @@ internal class VimSurroundExtension : VimExtension {
     override val isRepeatable = true
 
     override fun execute(editor: VimEditor, context: ExecutionContext, operatorArguments: OperatorArguments) {
-      setOperatorFunction(Operator(supportsMultipleCursors = false)) // TODO
+      setOperatorFunction(Operator(supportsMultipleCursors = false, count = 1)) // TODO
       executeNormalWithoutMapping(injector.parser.parseKeys("g@"), editor.ij)
     }
   }
@@ -104,7 +104,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)
       }
@@ -125,7 +125,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 {
@@ -261,7 +261,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 editor = vimEditor.ij
       val c = getChar(editor)
@@ -273,11 +273,11 @@ internal class VimSurroundExtension : VimExtension {
         val change = VimPlugin.getChange()
         if (supportsMultipleCursors) {
           editor.runWithEveryCaretAndRestore {
-            applyOnce(editor, change, pair)
+            applyOnce(editor, change, pair, count)
           }
         }
         else {
-          applyOnce(editor, change, pair)
+          applyOnce(editor, change, pair, count)
           // Jump back to start
           executeNormalWithoutMapping(injector.parser.parseKeys("`["), editor)
         }
@@ -285,18 +285,15 @@ internal class VimSurroundExtension : VimExtension {
       return true
     }
     
-    private fun applyOnce(editor: Editor, change: VimChangeGroup, pair: Pair<String, String>) {
+    private fun applyOnce(editor: Editor, change: VimChangeGroup, pair: Pair<String, String>, 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)
-        change.insertText(
-          IjVimEditor(editor),
-          IjVimCaret(primaryCaret),
-          range.endOffset + pair.first.length,
-          pair.second
-        )
+        val start = RepeatedCharSequence.of(pair.first, count)
+        val end = RepeatedCharSequence.of(pair.second, count)
+        change.insertText(IjVimEditor(editor), IjVimCaret(primaryCaret), range.startOffset, start)
+        change.insertText(IjVimEditor(editor), IjVimCaret(primaryCaret), range.endOffset + start.length, end)
       }
     }
 
@@ -378,15 +375,15 @@ internal class VimSurroundExtension : VimExtension {
       }
     }
 
-    private fun performSurround(pair: Pair<String, String>, range: TextRange, caret: VimCaret, tagsOnNewLines: Boolean = false) {
+    private fun performSurround(pair: Pair<String, String>, 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 {
@@ -394,7 +391,7 @@ internal class VimSurroundExtension : VimExtension {
           }
         } 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/action/change/change/FilterVisualLinesAction.kt

@@ -27,6 +27,7 @@ public class FilterVisualLinesAction : VimActionHandler.SingleExecution() {
   override val flags: EnumSet<CommandFlags> = enumSetOf(CommandFlags.FLAG_MOT_LINEWISE)
 
   override fun execute(editor: VimEditor, context: ExecutionContext, cmd: Command, operatorArguments: OperatorArguments): Boolean {
+    injector.markService.setVisualSelectionMarks(editor)
     injector.processGroup.startFilterCommand(editor, context, cmd)
     editor.exitVisualMode()
     return true

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

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

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

@@ -202,7 +202,7 @@ public 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 {
     (editor as MutableVimEditor).insertText(Offset(offset), str)
     val newCaret = caret.moveToInlayAwareOffset(offset + str.length)