vis: improve selection restore `gv`

2 files changed, 19 insertions, 2 deletions

diff --git a/main.c b/main.c
index 1db04fb..c4abe0b 100644
--- a/main.c
+++ b/main.c
@@ -1280,9 +1280,15 @@ static const char *selection_end(Vis *vis, const char *keys, const Arg *arg) {
 }
 
 static const char *selection_restore(Vis *vis, const char *keys, const Arg *arg) {
-	for (Cursor *c = view_cursors(vis_view(vis)); c; c = view_cursors_next(c))
+	Text *txt = vis_text(vis);
+	View *view = vis_view(vis);
+	for (Cursor *c = view_cursors(view); c; c = view_cursors_next(c))
 		view_cursors_selection_restore(c);
-	vis_mode_switch(vis, VIS_MODE_VISUAL);
+	Filerange sel = view_selection_get(view);
+	if (text_range_is_linewise(txt, &sel))
+		vis_mode_switch(vis, VIS_MODE_VISUAL_LINE);
+	else
+		vis_mode_switch(vis, VIS_MODE_VISUAL);
 	return keys;
 }
 
diff --git a/vis.c b/vis.c
index ba220c0..aceb0ed 100644
--- a/vis.c
+++ b/vis.c
@@ -519,7 +519,18 @@ void action_do(Vis *vis, Action *a) {
 			if (pos == EPOS) {
 				view_cursors_dispose(cursor);
 			} else if (pos <= text_size(txt)) {
+				/* moving the cursor will affect the selection.
+				 * because we want to be able to later restore
+				 * the old selection we update it again before
+				 * leaving visual mode.
+				 */
+				Filerange sel = view_cursors_selection_get(cursor);
 				view_cursors_to(cursor, pos);
+				if (vis->mode->visual) {
+					if (vis->mode == &vis_modes[VIS_MODE_VISUAL_LINE])
+						sel = text_range_linewise(txt, &sel);
+					view_cursors_selection_set(cursor, &sel);
+				}
 			}
 		}
 	}