Fixed off-by-one bug in Vi visual block mode.

3 files changed, 56 insertions, 11 deletions

diff --git a/prompt_toolkit/document.py b/prompt_toolkit/document.py
index d057de77..36cba0c0 100644
--- a/prompt_toolkit/document.py
+++ b/prompt_toolkit/document.py
@@ -789,6 +789,8 @@ class Document(object):
         nothing was selected. The upper boundary is not included.
 
         This will yield several (from, to) tuples in case of a BLOCK selection.
+        This will return zero ranges, like (8,8) for empty lines in a block
+        selection.
         """
         if self.selection:
             from_, to = sorted([self.cursor_position, self.selection.original_cursor_position])
@@ -804,9 +806,10 @@ class Document(object):
 
                 for l in range(from_line, to_line + 1):
                     line_length = len(lines[l])
-                    if from_column < line_length:
+
+                    if from_column < line_length or (line_length == 0 and from_column == 0):
                         yield (self.translate_row_col_to_index(l, from_column),
-                               self.translate_row_col_to_index(l, min(line_length - 1, to_column)))
+                               self.translate_row_col_to_index(l, min(line_length, to_column)))
             else:
                 # In case of a LINES selection, go to the start/end of the lines.
                 if self.selection.type == SelectionType.LINES:
diff --git a/prompt_toolkit/layout/processors.py b/prompt_toolkit/layout/processors.py
index 05b982fb..ccc7b1ec 100644
--- a/prompt_toolkit/layout/processors.py
+++ b/prompt_toolkit/layout/processors.py
@@ -223,6 +223,8 @@ class HighlightSelectionProcessor(Processor):
                     if i < len(fragments):
                         old_fragment, old_text = fragments[i]
                         fragments[i] = (old_fragment + selected_fragment, old_text)
+                    elif i == len(fragments):
+                        fragments.append((selected_fragment, ' '))
 
         return Transformation(fragments)
 
@@ -329,7 +331,7 @@ class DisplayMultipleCursors(Processor):
         buff = buffer_control.buffer
 
         if vi_insert_multiple_mode():
-            positions = buff.multiple_cursor_positions
+            cursor_positions = buff.multiple_cursor_positions
             fragments = explode_text_fragments(fragments)
 
             # If any cursor appears on the current line, highlight that.
@@ -338,16 +340,19 @@ class DisplayMultipleCursors(Processor):
 
             fragment_suffix = ' class:multiple-cursors'
 
-            for p in positions:
-                if start_pos <= p < end_pos:
+            for p in cursor_positions:
+                if start_pos <= p <= end_pos:
                     column = source_to_display(p - start_pos)
 
                     # Replace fragment.
-                    style, text = fragments[column]
-                    style += fragment_suffix
-                    fragments[column] = (style, text)
-                elif p == end_pos:
-                    fragments.append((fragment_suffix, ' '))
+                    try:
+                        style, text = fragments[column]
+                    except IndexError:
+                        # Cursor needs to be displayed after the current text.
+                        fragments.append((fragment_suffix, ' '))
+                    else:
+                        style += fragment_suffix
+                        fragments[column] = (style, text)
 
             return Transformation(fragments)
         else:
@@ -587,7 +592,10 @@ class TabsProcessor(Processor):
 
         def source_to_display(from_position):
             " Maps original cursor position to the new one. "
-            return position_mappings[from_position]
+            try:
+                return position_mappings[from_position]
+            except KeyError:
+                return 0
 
         def display_to_source(display_pos):
             " Maps display cursor position to the original one. "
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 4ea180a1..3fc55fce 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -710,6 +710,40 @@ def test_vi_block_editing():
             '-line1\n-line***2\n-line***3\n-line***4\n-line5\n-line6')
 
 
+def test_vi_block_editing_empty_lines():
+    " Test block editing on empty lines. "
+    feed = partial(_feed_cli_with_input, editing_mode=EditingMode.VI,
+                   multiline=True)
+
+    operations = (
+        # Six empty lines.
+        '\r\r\r\r\r'
+        # Go to beginning of the document.
+        '\x1bgg'
+        # Enter Visual block mode.
+        '\x16'
+        # Go down two more lines.
+        'jj'
+        # Go 3 characters to the right.
+        'lll'
+        # Go to insert mode.
+        'insert'  # (Will be replaced.)
+        # Insert stars.
+        '***'
+        # Escape again.
+        '\x1b\r')
+
+    # Control-I
+    result, cli = feed(operations.replace('insert', 'I'))
+
+    assert result.text == '***\n***\n***\n\n\n'
+
+    # Control-A
+    result, cli = feed(operations.replace('insert', 'A'))
+
+    assert result.text == '***\n***\n***\n\n\n'
+
+
 def test_vi_visual_line_copy():
     feed = partial(_feed_cli_with_input, editing_mode=EditingMode.VI,
                    multiline=True)