From 2f057785f7da4faf05af85dc667f8bba83ce027a Mon Sep 17 00:00:00 2001
From: Paul Eguisier <paul.eguisier@gmail.com>
Date: Thu, 6 Jun 2024 17:45:25 +0200
Subject: [PATCH] Maintain cursor to upper line in visual mode indent/outdent
 (#12582)

Release Notes:

- vim: Fix indent via `<` and `>` not being repeatable with `.`.
[#12351](https://github.com/zed-industries/zed/issues/12351)
---
 crates/vim/src/normal.rs | 43 ++++++++++++++++++++++++++++++++++++++--
 crates/vim/src/test.rs   |  7 ++++++-
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/crates/vim/src/normal.rs b/crates/vim/src/normal.rs
index db1bd8fd54..2f02fbb624 100644
--- a/crates/vim/src/normal.rs
+++ b/crates/vim/src/normal.rs
@@ -11,6 +11,7 @@ pub(crate) mod search;
 pub mod substitute;
 mod yank;
 
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use crate::{
@@ -21,8 +22,11 @@ use crate::{
     Vim,
 };
 use collections::BTreeSet;
+use editor::display_map::ToDisplayPoint;
 use editor::scroll::Autoscroll;
+use editor::Anchor;
 use editor::Bias;
+use editor::Editor;
 use gpui::{actions, ViewContext, WindowContext};
 use language::{Point, SelectionGoal};
 use log::error;
@@ -143,7 +147,11 @@ pub(crate) fn register(workspace: &mut Workspace, cx: &mut ViewContext<Workspace
         Vim::update(cx, |vim, cx| {
             vim.record_current_action(cx);
             vim.update_active_editor(cx, |_, editor, cx| {
-                editor.transact(cx, |editor, cx| editor.indent(&Default::default(), cx))
+                editor.transact(cx, |editor, cx| {
+                    let mut original_positions = save_selection_starts(editor, cx);
+                    editor.indent(&Default::default(), cx);
+                    restore_selection_cursors(editor, cx, &mut original_positions);
+                });
             });
             if vim.state().mode.is_visual() {
                 vim.switch_mode(Mode::Normal, false, cx)
@@ -155,7 +163,11 @@ pub(crate) fn register(workspace: &mut Workspace, cx: &mut ViewContext<Workspace
         Vim::update(cx, |vim, cx| {
             vim.record_current_action(cx);
             vim.update_active_editor(cx, |_, editor, cx| {
-                editor.transact(cx, |editor, cx| editor.outdent(&Default::default(), cx))
+                editor.transact(cx, |editor, cx| {
+                    let mut original_positions = save_selection_starts(editor, cx);
+                    editor.outdent(&Default::default(), cx);
+                    restore_selection_cursors(editor, cx, &mut original_positions);
+                });
             });
             if vim.state().mode.is_visual() {
                 vim.switch_mode(Mode::Normal, false, cx)
@@ -390,6 +402,33 @@ fn yank_line(_: &mut Workspace, _: &YankLine, cx: &mut ViewContext<Workspace>) {
     })
 }
 
+fn save_selection_starts(editor: &Editor, cx: &mut ViewContext<Editor>) -> HashMap<usize, Anchor> {
+    let (map, selections) = editor.selections.all_display(cx);
+    selections
+        .iter()
+        .map(|selection| {
+            (
+                selection.id,
+                map.display_point_to_anchor(selection.start, Bias::Right),
+            )
+        })
+        .collect::<HashMap<_, _>>()
+}
+
+fn restore_selection_cursors(
+    editor: &mut Editor,
+    cx: &mut ViewContext<Editor>,
+    positions: &mut HashMap<usize, Anchor>,
+) {
+    editor.change_selections(Some(Autoscroll::fit()), cx, |s| {
+        s.move_with(|map, selection| {
+            if let Some(anchor) = positions.remove(&selection.id) {
+                selection.collapse_to(anchor.to_display_point(map), SelectionGoal::None);
+            }
+        });
+    });
+}
+
 pub(crate) fn normal_replace(text: Arc<str>, cx: &mut WindowContext) {
     Vim::update(cx, |vim, cx| {
         vim.stop_recording();
diff --git a/crates/vim/src/test.rs b/crates/vim/src/test.rs
index ccfeb5b81e..b44622ec12 100644
--- a/crates/vim/src/test.rs
+++ b/crates/vim/src/test.rs
@@ -179,7 +179,7 @@ async fn test_indent_outdent(cx: &mut gpui::TestAppContext) {
 
     // works in visual mode
     cx.simulate_keystrokes("shift-v down >");
-    cx.assert_editor_state("aa\n    bb\n    cˇc");
+    cx.assert_editor_state("aa\n    bˇb\n    cc");
 
     // works as operator
     cx.set_state("aa\nbˇb\ncc\n", Mode::Normal);
@@ -202,11 +202,16 @@ async fn test_indent_outdent(cx: &mut gpui::TestAppContext) {
     cx.simulate_keystrokes("> 2 k");
     cx.assert_editor_state("    aa\n    bb\n    ˇcc\n");
 
+    // works with repeat
     cx.set_state("a\nb\nccˇc\n", Mode::Normal);
     cx.simulate_keystrokes("> 2 k");
     cx.assert_editor_state("    a\n    b\n    ccˇc\n");
     cx.simulate_keystrokes(".");
     cx.assert_editor_state("        a\n        b\n        ccˇc\n");
+    cx.simulate_keystrokes("v k <");
+    cx.assert_editor_state("        a\n    bˇ\n    ccc\n");
+    cx.simulate_keystrokes(".");
+    cx.assert_editor_state("        a\nbˇ\nccc\n");
 }
 
 #[gpui::test]