From 428c5c74fdf22c28b07a08062b44df5a859b005b Mon Sep 17 00:00:00 2001
From: Mathieu Agopian <mathieu@agopian.info>
Date: Sat, 18 Feb 2023 19:18:05 +0100
Subject: [PATCH] Expand selection to whole word when pressing `*`

---
 helix-term/src/commands.rs        | 35 +++++++++++++++++-----
 helix-term/tests/test/commands.rs | 50 +++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 8 deletions(-)

diff --git a/helix-term/src/commands.rs b/helix-term/src/commands.rs
index fb55ca2a8..41284aab1 100644
--- a/helix-term/src/commands.rs
+++ b/helix-term/src/commands.rs
@@ -1906,18 +1906,37 @@ fn extend_search_prev(cx: &mut Context) {
 }
 
 fn search_selection(cx: &mut Context) {
+    let count = cx.count();
     let (view, doc) = current!(cx.editor);
     let contents = doc.text().slice(..);
 
-    let regex = doc
-        .selection(view.id)
-        .iter()
-        .map(|selection| regex::escape(&selection.fragment(contents)))
-        .collect::<HashSet<_>>() // Collect into hashset to deduplicate identical regexes
-        .into_iter()
-        .collect::<Vec<_>>()
-        .join("|");
+    let regex; // The fragment to set in the search register
 
+    // Checks whether there is only one selection with a width of 1
+    let selections = doc.selection(view.id);
+    let primary = selections.primary();
+    if selections.len() == 1 && primary.len() == 1 {
+        let text = doc.text();
+        let text_slice = text.slice(..);
+        // In this case select the WORD under the cursor
+        let current_word = textobject::textobject_word(
+            text_slice,
+            primary,
+            textobject::TextObject::Inside,
+            count,
+            false,
+        );
+        let text_to_search = current_word.fragment(text_slice).to_string();
+        regex = regex::escape(&text_to_search);
+    } else {
+        regex = selections
+            .iter()
+            .map(|selection| regex::escape(&selection.fragment(contents)))
+            .collect::<HashSet<_>>() // Collect into hashset to deduplicate identical regexes
+            .into_iter()
+            .collect::<Vec<_>>()
+            .join("|");
+    }
     let msg = format!("register '{}' set to '{}'", '/', &regex);
     cx.editor.registers.push('/', regex);
     cx.editor.set_status(msg);
diff --git a/helix-term/tests/test/commands.rs b/helix-term/tests/test/commands.rs
index e8d16bfaf..c0276cd6f 100644
--- a/helix-term/tests/test/commands.rs
+++ b/helix-term/tests/test/commands.rs
@@ -258,6 +258,56 @@ async fn test_goto_file_impl() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn test_search_selection() -> anyhow::Result<()> {
+    // Single selection with a length of 1: search for the whole word
+    test_key_sequence(
+        &mut helpers::AppBuilder::new().build()?,
+        Some("ifoobar::baz<esc>3bl*"), // 3b places the cursor on the first letter of 'foobar', then move one to the right for good measure
+        Some(&|app| {
+            assert!(
+                r#"register '/' set to 'foobar'"# == app.editor.get_status().unwrap().0
+                    && Some(&"foobar".to_string()) == app.editor.registers.first('/')
+            );
+        }),
+        false,
+    )
+    .await?;
+
+    // Single selection with a length greather than 1: only search for the selection
+    test_key_sequence(
+        &mut helpers::AppBuilder::new().build()?,
+        Some("ifoobar::baz<esc>3blvll*"), // 3b places the cursor on the first letter of 'foobar', then move one to the right for good measure, then select two more chars for a total of three
+        Some(&|app| {
+            assert!(
+                r#"register '/' set to 'oob'"# == app.editor.get_status().unwrap().0
+                    && Some(&"oob".to_string()) == app.editor.registers.first('/')
+            );
+        }),
+        false,
+    )
+    .await?;
+
+    // Multiple selection of length 1 each : should still only search for the selection
+    test_key_sequence(
+        &mut helpers::AppBuilder::new().build()?,
+        Some("ifoobar::baz<ret>bar::crux<esc>k3blC*"), // k3b places the cursor on the first letter of 'foobar', then move one to the right for good measure, then adds a cursor on the line below
+        Some(&|app| {
+            assert!(
+                // The selections don't seem to be ordered, so we have to test for the two possible orders.
+                (r#"register '/' set to 'o|a'"# == app.editor.get_status().unwrap().0
+                    || r#"register '/' set to 'a|o'"# == app.editor.get_status().unwrap().0)
+                    && (Some(&"o|a".to_string()) == app.editor.registers.first('/')
+                        || Some(&"a|o".to_string()) == app.editor.registers.first('/'))
+            );
+        }),
+        false,
+    )
+    .await?;
+
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn test_multi_selection_paste() -> anyhow::Result<()> {
     test((