N as extend with search (for now, N should be search_prev).

4 files changed, 70 insertions, 47 deletions

diff --git a/helix-core/src/selection.rs b/helix-core/src/selection.rs
index 064abec7..ca1a6d32 100644
--- a/helix-core/src/selection.rs
+++ b/helix-core/src/selection.rs
@@ -159,6 +159,12 @@ impl Selection {
         }
     }
 
+    pub fn add(mut self, range: Range) -> Self {
+        let index = self.ranges.len();
+        self.ranges.push(range);
+
+        Self::normalize(self.ranges, index)
+    }
     // add_range // push
     // replace_range
 
@@ -204,53 +210,53 @@ impl Selection {
         Self::single(pos, pos)
     }
 
-    // TODO: consume an iterator or a vec to reduce allocations?
-    #[must_use]
-    pub fn new(ranges: SmallVec<[Range; 1]>, primary_index: usize) -> Self {
-        assert!(!ranges.is_empty());
+    fn normalize(mut ranges: SmallVec<[Range; 1]>, mut primary_index: usize) -> Selection {
+        let primary = ranges[primary_index];
+        ranges.sort_unstable_by_key(Range::from);
+        primary_index = ranges.iter().position(|&range| range == primary).unwrap();
 
-        fn normalize(mut ranges: SmallVec<[Range; 1]>, mut primary_index: usize) -> Selection {
-            let primary = ranges[primary_index];
-            ranges.sort_unstable_by_key(Range::from);
-            primary_index = ranges.iter().position(|&range| range == primary).unwrap();
-
-            let mut result = SmallVec::with_capacity(ranges.len()); // approx
-
-            // TODO: we could do with one vec by removing elements as we mutate
-
-            for (i, range) in ranges.into_iter().enumerate() {
-                // if previous value exists
-                if let Some(prev) = result.last_mut() {
-                    // and we overlap it
-                    if range.overlaps(prev) {
-                        let from = prev.from();
-                        let to = std::cmp::max(range.to(), prev.to());
-
-                        if i <= primary_index {
-                            primary_index -= 1
-                        }
-
-                        // merge into previous
-                        if range.anchor > range.head {
-                            prev.anchor = to;
-                            prev.head = from;
-                        } else {
-                            prev.anchor = from;
-                            prev.head = to;
-                        }
-                        continue;
+        let mut result = SmallVec::with_capacity(ranges.len()); // approx
+
+        // TODO: we could do with one vec by removing elements as we mutate
+
+        for (i, range) in ranges.into_iter().enumerate() {
+            // if previous value exists
+            if let Some(prev) = result.last_mut() {
+                // and we overlap it
+                if range.overlaps(prev) {
+                    let from = prev.from();
+                    let to = std::cmp::max(range.to(), prev.to());
+
+                    if i <= primary_index {
+                        primary_index -= 1
                     }
-                }
 
-                result.push(range)
+                    // merge into previous
+                    if range.anchor > range.head {
+                        prev.anchor = to;
+                        prev.head = from;
+                    } else {
+                        prev.anchor = from;
+                        prev.head = to;
+                    }
+                    continue;
+                }
             }
 
-            Selection {
-                ranges: result,
-                primary_index,
-            }
+            result.push(range)
         }
 
+        Selection {
+            ranges: result,
+            primary_index,
+        }
+    }
+
+    // TODO: consume an iterator or a vec to reduce allocations?
+    #[must_use]
+    pub fn new(ranges: SmallVec<[Range; 1]>, primary_index: usize) -> Self {
+        assert!(!ranges.is_empty());
+
         // fast path for a single selection (cursor)
         if ranges.len() == 1 {
             return Self {
@@ -260,7 +266,7 @@ impl Selection {
         }
 
         // TODO: only normalize if needed (any ranges out of order)
-        normalize(ranges, primary_index)
+        Self::normalize(ranges, primary_index)
     }
 
     /// Takes a closure and maps each selection over the closure.
diff --git a/helix-term/src/commands.rs b/helix-term/src/commands.rs
index 53ec2fbf..98e8b20c 100644
--- a/helix-term/src/commands.rs
+++ b/helix-term/src/commands.rs
@@ -608,9 +608,10 @@ pub fn split_selection_on_newline(cx: &mut Context) {
 // I'd probably collect all the matches right now and store the current index. The cache needs
 // wiping if input happens.
 
-fn _search(doc: &mut Document, view_id: ViewId, contents: &str, regex: &Regex) {
+fn _search(doc: &mut Document, view_id: ViewId, contents: &str, regex: &Regex, extend: bool) {
     let text = doc.text();
-    let start = doc.selection(view_id).cursor();
+    let selection = doc.selection(view_id);
+    let start = selection.cursor();
 
     // use find_at to find the next match after the cursor, loop around the end
     let mat = regex
@@ -619,7 +620,13 @@ fn _search(doc: &mut Document, view_id: ViewId, contents: &str, regex: &Regex) {
     if let Some(mat) = mat {
         let start = text.byte_to_char(mat.start());
         let end = text.byte_to_char(mat.end());
-        let selection = Selection::single(start, end - 1);
+
+        let selection = if extend {
+            selection.clone().add(Range::new(start, end - 1))
+        } else {
+            Selection::single(start, end - 1)
+        };
+
         // TODO: (first_match, regex) stuff in register?
         doc.set_selection(view_id, selection);
     };
@@ -639,7 +646,7 @@ pub fn search(cx: &mut Context) {
     let prompt = ui::regex_prompt(cx, "search:".to_string(), move |doc, regex| {
         let text = doc.text();
         let start = doc.selection(view_id).cursor();
-        _search(doc, view_id, &contents, &regex);
+        _search(doc, view_id, &contents, &regex, false);
 
         // TODO: only store on enter (accept), not update
         register::set('\\', vec![regex.as_str().to_string()]);
@@ -648,17 +655,25 @@ pub fn search(cx: &mut Context) {
     cx.push_layer(Box::new(prompt));
 }
 
-pub fn search_next(cx: &mut Context) {
+pub fn _search_next(cx: &mut Context, extend: bool) {
     if let Some(query) = register::get('\\') {
         let query = query.first().unwrap();
         let view_id = cx.view_id;
         let doc = cx.doc();
         let contents = doc.text().slice(..).to_string();
         let regex = Regex::new(&query).unwrap();
-        _search(doc, view_id, &contents, &regex);
+        _search(doc, view_id, &contents, &regex, extend);
     }
 }
 
+pub fn search_next(cx: &mut Context) {
+    _search_next(cx, false);
+}
+
+pub fn extend_search_next(cx: &mut Context) {
+    _search_next(cx, true);
+}
+
 pub fn search_selection(cx: &mut Context) {
     let (view, doc) = cx.current();
     let contents = doc.text().slice(..);
diff --git a/helix-term/src/keymap.rs b/helix-term/src/keymap.rs
index 06b190aa..e20c8ebf 100644
--- a/helix-term/src/keymap.rs
+++ b/helix-term/src/keymap.rs
@@ -199,6 +199,7 @@ pub fn default() -> Keymaps {
         key!('/') => commands::search,
         // ? for search_reverse
         key!('n') => commands::search_next,
+        shift!('N') => commands::extend_search_next,
         // N for search_prev
         key!('*') => commands::search_selection,
 
diff --git a/helix-term/src/ui/editor.rs b/helix-term/src/ui/editor.rs
index d32f7864..35693d7b 100644
--- a/helix-term/src/ui/editor.rs
+++ b/helix-term/src/ui/editor.rs
@@ -61,6 +61,7 @@ impl EditorView {
             view.area.width - OFFSET,
             view.area.height.saturating_sub(1),
         ); // - 1 for statusline
+
         self.render_buffer(doc, view, area, surface, theme, is_focused);
 
         // if we're not at the edge of the screen, draw a right border