fix(input): match a caret's word whole when selecting its occurrences

crates/base/src/input/base/occurrences.rs

@@ -1,5 +1,7 @@
 //! Selecting the occurrences of a selection: the next one added as another selection, or all
-//! of them at once. With only a caret, the first press selects the word under it.
+//! of them at once. With only a caret, the first press selects the word under it, and the
+//! presses after it match that word whole, as Sublime Text and VS Code do; a selection made by
+//! hand matches anywhere.
 
 use gpui::{Context, Window, actions};
 
@@ -39,10 +41,45 @@ impl<M: InputModeKind> InputBaseState<M> {
         active.end = word.end;
         active.reversed = false;
         active.column_anchor = None;
+        self.occurrence_word = Some((word.start, word.end, self.document_revision));
         cx.notify();
         true
     }
 
+    /// Whether the occurrences are of a word a caret selected, still selected and unedited.
+    fn occurrences_are_words(&mut self) -> bool {
+        let words = self.occurrence_word.is_some_and(|(start, end, revision)| {
+            revision == self.document_revision
+                && self
+                    .selections
+                    .iter()
+                    .any(|s| s.start == start && s.end == end)
+        });
+        if !words {
+            self.occurrence_word = None;
+        }
+        words
+    }
+
+    /// The starts of `needle` in `text`, only where it stands as a whole word when `words`.
+    fn occurrence_starts<'a>(
+        text: &'a str,
+        needle: &'a str,
+        words: bool,
+    ) -> impl Iterator<Item = usize> + 'a {
+        let is_word = |c: char| c.is_alphanumeric() || c == '_';
+        text.match_indices(needle)
+            .map(|(at, _)| at)
+            .filter(move |&at| {
+                !words
+                    || (!text[..at].chars().next_back().is_some_and(is_word)
+                        && !text[at + needle.len()..]
+                            .chars()
+                            .next()
+                            .is_some_and(is_word))
+            })
+    }
+
     /// The selected text the occurrences are of: the newest selection's.
     fn occurrence_needle(&self) -> Option<(String, usize)> {
         let newest = self.selections.iter().last()?;
@@ -83,19 +120,14 @@ impl<M: InputModeKind> InputBaseState<M> {
         let Some((needle, from)) = self.occurrence_needle() else {
             return;
         };
+        let words = self.occurrences_are_words();
         let text = self.text.to_string();
-        let after = text
-            .get(from..)
-            .unwrap_or_default()
-            .match_indices(&needle)
-            .map(|(at, _)| at + from);
-        let before = text
-            .get(..from)
-            .unwrap_or_default()
-            .match_indices(&needle)
-            .map(|(at, _)| at);
-        let next = after
-            .chain(before)
+        let starts: Vec<usize> = Self::occurrence_starts(&text, &needle, words).collect();
+        let split = starts.partition_point(|&at| at < from);
+        let next = starts[split..]
+            .iter()
+            .chain(&starts[..split])
+            .copied()
             .find(|&at| !self.overlaps_selection(at, at + needle.len()));
         if let Some(start) = next {
             self.add_occurrence(start, start + needle.len());
@@ -119,8 +151,9 @@ impl<M: InputModeKind> InputBaseState<M> {
         let Some((needle, _)) = self.occurrence_needle() else {
             return;
         };
+        let words = self.occurrences_are_words();
         let text = self.text.to_string();
-        let starts: Vec<usize> = text.match_indices(&needle).map(|(at, _)| at).collect();
+        let starts: Vec<usize> = Self::occurrence_starts(&text, &needle, words).collect();
         for start in starts {
             if !self.overlaps_selection(start, start + needle.len()) {
                 self.add_occurrence(start, start + needle.len());
@@ -229,6 +262,34 @@ mod tests {
         assert_eq!(selected(&state, cx), [(0, 2), (5, 7), (10, 12)]);
     }
 
+    #[gpui::test]
+    fn test_a_word_from_the_caret_matches_whole_words_and_a_selection_anywhere(
+        cx: &mut TestAppContext,
+    ) {
+        let (state, cx) = editor(cx, "n len n", 0);
+        cx.update(|window, cx| {
+            state.update(cx, |s, cx| {
+                s.select_all_occurrences(&SelectAllOccurrences, window, cx)
+            });
+        });
+        assert_eq!(
+            selected(&state, cx),
+            [(0, 1), (6, 7)],
+            "not the n inside `len`"
+        );
+        cx.update(|window, cx| {
+            state.update(cx, |s, cx| {
+                s.set_selected_range(6..7, cx);
+                s.select_all_occurrences(&SelectAllOccurrences, window, cx)
+            });
+        });
+        assert_eq!(
+            selected(&state, cx),
+            [(0, 1), (4, 5), (6, 7)],
+            "a selection made by hand matches inside words too"
+        );
+    }
+
     #[gpui::test]
     fn test_a_caret_off_any_word_selects_nothing(cx: &mut TestAppContext) {
         let (state, cx) = editor(cx, "a  b", 2);

crates/base/src/input/base/state.rs

@@ -392,6 +392,9 @@ pub struct InputBaseState<M: InputModeKind> {
     ///
     /// Always contains at least one selection where index 0 is the active cursor.
     pub(super) selections: Selections,
+    /// The word a caret's first "select next occurrence" selected, with the document revision it
+    /// was selected at: while that selection stands, later presses match whole words only.
+    pub(super) occurrence_word: Option<(usize, usize, u64)>,
     /// Range for save the selected word, use to keep word range when drag move.
     pub(super) selected_word_range: Option<CursorSelection>,
     /// The marked range is the temporary insert text on IME typing.
@@ -739,6 +742,7 @@ impl<M: InputModeKind> InputBaseState<M> {
             token_layout_cache: None,
             pressed_token: None,
             selections: Selections::default(),
+            occurrence_word: None,
             selected_word_range: None,
             ime_marked_range: None,
             input_bounds: Bounds::default(),

crates/base/src/input/mod.rs

@@ -113,10 +113,10 @@ pub use lsp::{
     DocumentRangeSemanticTokensProvider, HoverPopoverState, HoverProvider, InputOverlayKind, Lsp,
     ShowDocumentHandler,
 };
-pub use lsp_types::Position;
 /// The LSP types the providers speak ([`CompletionProvider`] and the others), at the version
 /// this crate builds with, so a provider needs no `lsp-types` of its own to match it.
 pub use lsp_types;
+pub use lsp_types::Position;
 pub use mask_pattern::MaskPattern;
 #[cfg(target_os = "macos")]
 #[doc(hidden)]