From 0642b909c34416b8d8006749dd40684c7351f567 Mon Sep 17 00:00:00 2001
From: Austen Adler <agadler@austenadler.com>
Date: Sun, 2 Oct 2022 16:40:56 -0400
Subject: [PATCH] Use cow for keys

---
 src/set.rs     |  7 +++++--
 src/utils.rs   | 50 ++++++++++++++++++++++++++++++++++----------------
 src/xlookup.rs |  6 +++---
 3 files changed, 42 insertions(+), 21 deletions(-)

diff --git a/src/set.rs b/src/set.rs
index adf0ad3..36b1ec7 100644
--- a/src/set.rs
+++ b/src/set.rs
@@ -168,7 +168,8 @@ fn reduce_selections(
             options.ignore_case,
         );
 
-        if key_set_operation_result.contains(&key) {
+        // TODO: Do not allocate
+        if key_set_operation_result.contains(&key.into_owned()) {
             Some(swd.desc)
         } else {
             None
@@ -268,6 +269,7 @@ fn to_ordered_counts(options: &Options, sels: Vec<Selection>) -> LinkedHashMap<S
 
     for i in sels {
         let key = crate::utils::get_key(
+            // TODO: Do not allocate
             &i,
             options.skip_whitespace,
             options.regex.as_ref(),
@@ -279,7 +281,8 @@ fn to_ordered_counts(options: &Options, sels: Vec<Selection>) -> LinkedHashMap<S
             continue;
         }
 
-        let entry: &mut usize = ret.entry(key).or_insert(0);
+        // TODO: Do not allocate
+        let entry: &mut usize = ret.entry(key.into_owned()).or_insert(0);
         *entry = entry.saturating_add(1);
     }
     ret
diff --git a/src/utils.rs b/src/utils.rs
index f17a66b..17f13ae 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -1,56 +1,74 @@
-use kakplugin::Selection;
+// use kakplugin::Selection;
 use regex::Regex;
 use std::{
+    borrow::{Borrow, Cow},
     collections::hash_map::DefaultHasher,
     hash::{Hash, Hasher},
 };
 
-pub fn get_key(
-    // TODO: Use Cow
-    selection: &Selection,
-    skip_whitespace: bool,
+/// Gets a key out of a selection
+///
+/// # Examples
+///
+/// ```
+/// assert_eq!(get_key("  asdf\n", false, None, false), "asdf\n");
+/// assert_eq!(get_key("  asdf\n", true, None, false), "  asdf\n");
+/// assert_eq!(get_key("  as1f\n", false, Some("\w+"), false), "as");
+/// assert_eq!(get_key("  aS1F\n", false, Some("\w+"), true), "as1f");
+/// ```
+pub fn get_key<'sel>(
+    selection: &'sel str,
+    preserve_whitespace: bool,
     regex: Option<&Regex>,
     ignore_case: bool,
-) -> String {
+) -> Cow<'sel, str> {
     // Strip whitespace if requested
-    let mut key = if skip_whitespace {
-        selection.as_str()
+    let mut key = if preserve_whitespace {
+        // TODO: Does this need to be swapped?
+        selection
     } else {
         selection.trim()
     };
 
     // If they requested a regex match, set the key to the string slice of that match
     if let Some(regex_match) = (|| {
-        let captures = regex.as_ref()?.captures(key)?;
+        // let captures = regex.as_ref()?.captures(&key)?;
+        let captures = regex.as_ref()?.captures(key.borrow())?;
         captures
             .get(1)
             .or_else(|| captures.get(0))
             .map(|m| m.as_str())
     })() {
         key = regex_match;
+        // Cow::Borrowed(regex_match)
     }
 
     // Ignore case if requested
-    // Lowercase at the end to not mangle regex
     if ignore_case {
-        key.to_lowercase()
+        // Lowercase at the end to not mangle regex
+        // TODO: Do not allocate if it is already lowercased
+        // Need to_lowercase(&self) -> Cow<str>
+        if !key.as_bytes().iter().any(u8::is_ascii_uppercase) {
+            Cow::Borrowed(key)
+        } else {
+            Cow::Owned(key.to_ascii_lowercase())
+        }
     } else {
-        // TODO: Do not perform an allocation here
-        key.to_string()
+        Cow::Borrowed(key)
     }
 }
 
 /// Get a key out of a selection based on options
 pub fn get_hash(
     // TODO: Accept any Into<AsRef<Selection>>
-    selection: &Selection,
-    skip_whitespace: bool,
+    selection: &str,
+    preserve_whitespace: bool,
     regex: Option<&Regex>,
     ignore_case: bool,
 ) -> u64 {
     let mut hasher = DefaultHasher::new();
 
-    get_key(selection, skip_whitespace, regex, ignore_case).hash(&mut hasher);
+    get_key(&selection, preserve_whitespace, regex, ignore_case).hash(&mut hasher);
 
     hasher.finish()
 }
diff --git a/src/xlookup.rs b/src/xlookup.rs
index 41475ec..a1c579f 100644
--- a/src/xlookup.rs
+++ b/src/xlookup.rs
@@ -22,7 +22,7 @@ pub fn xlookup(options: &Options) -> Result<String, KakError> {
 
     set_selections(selections.iter().map(|key| {
         lookup_table
-            .get(&get_hash(key, false, None, false))
+            .get(&get_hash(&key, false, None, false))
             .map_or_else(
                 || {
                     eprintln!("Key '{key}' not found",);
@@ -48,7 +48,7 @@ pub fn xlookup(options: &Options) -> Result<String, KakError> {
 fn build_lookuptable(mut selections: Vec<Selection>) -> Result<BTreeMap<u64, Selection>, KakError> {
     let mut iter = selections.array_chunks_mut();
     let ret = iter.try_fold(BTreeMap::new(), |mut acc, [key, value]| {
-        match acc.entry(get_hash(key, false, None, false)) {
+        match acc.entry(get_hash(&key, false, None, false)) {
             Occupied(_) => Err(KakError::Custom(format!("Duplicate key '{key}'"))),
             Vacant(v) => {
                 v.insert(value.clone());
@@ -76,7 +76,7 @@ mod tests {
     }
     macro_rules! hsh {
         ($expr:expr) => {
-            get_hash(&$expr.to_string(), false, None, false)
+            get_hash($expr, false, None, false)
         };
     }
     #[test]