src: Select the data inserted by `!` and `<a-!>`

Closes #1468

4 files changed, 14 insertions, 6 deletions

diff --git a/doc/pages/keys.asciidoc b/doc/pages/keys.asciidoc
index 5f5c0e4a..e117b532 100644
--- a/doc/pages/keys.asciidoc
+++ b/doc/pages/keys.asciidoc
@@ -374,10 +374,10 @@ The default command comes from the *|* register by default (See <<registers#,`:d
     ignore its output.
 
 *!*::
-    insert command output before each selection.
+    insert and select command output before each selection.
 
 *<a-!>*::
-    append command output after each selection.
+    append and select command output after each selection.
 
 == Searching
 
diff --git a/src/normal.cc b/src/normal.cc
index ca7063fd..15f07211 100644
--- a/src/normal.cc
+++ b/src/normal.cc
@@ -765,6 +765,7 @@ void insert_output(Context& context, NormalParams params)
             auto& selections = context.selections();
             auto& buffer = context.buffer();
             const size_t old_main = selections.main_index();
+            Vector<BufferRange> ins_range;
 
             selections.for_each([&](size_t index, Selection& sel) {
                 selections.set_main_index(index);
@@ -772,10 +773,16 @@ void insert_output(Context& context, NormalParams params)
                     cmdline, context, content(context.buffer(), sel),
                     ShellManager::Flags::WaitForStdout);
 
-                insert(buffer, sel, paste_pos(buffer, sel, mode, false), out);
+                auto range = insert(buffer, sel, paste_pos(buffer, sel, mode, false), out);
+                ins_range.push_back(range);
             });
 
-            selections.set_main_index(old_main);
+            selections.set(ins_range | transform([&buffer](auto& range) {
+                if (range.empty())
+                    return Selection{range.begin, range.end};
+                return Selection{range.begin,
+                                 buffer.char_prev(range.end)};
+            }) | gather<Vector>(), old_main);
         });
 }
 
diff --git a/src/selection.cc b/src/selection.cc
index d408ddfd..cbf9bef2 100644
--- a/src/selection.cc
+++ b/src/selection.cc
@@ -399,11 +399,12 @@ void replace(Buffer& buffer, Selection& sel, StringView content)
     max = range.end > range.begin ? buffer.char_prev(range.end) : range.begin;
 }
 
-void insert(Buffer& buffer, Selection& sel, BufferCoord pos, StringView content)
+BufferRange insert(Buffer& buffer, Selection& sel, BufferCoord pos, StringView content)
 {
     auto range = buffer.insert(pos, content);
     sel.anchor() = buffer.clamp(update_insert(sel.anchor(), range.begin, range.end));
     sel.cursor() = buffer.clamp(update_insert(sel.cursor(), range.begin, range.end));
+    return range;
 }
 
 void SelectionList::replace(ConstArrayView<String> strings)
diff --git a/src/selection.hh b/src/selection.hh
index 1aebe41f..43ea412b 100644
--- a/src/selection.hh
+++ b/src/selection.hh
@@ -75,7 +75,7 @@ void merge_overlapping_selections(Vector<Selection>& selections, size_t& main);
 void clamp_selections(Vector<Selection>& sel, const Buffer& buffer);
 
 void replace(Buffer& buffer, Selection& sel, StringView content);
-void insert(Buffer& buffer, Selection& sel, BufferCoord pos, StringView content);
+BufferRange insert(Buffer& buffer, Selection& sel, BufferCoord pos, StringView content);
 
 struct SelectionList
 {