From ea32174522b7a5c143f3d12f0c6c7531295da25d Mon Sep 17 00:00:00 2001
From: Austen Adler <agadler@austenadler.com>
Date: Thu, 7 Jul 2022 00:14:23 -0400
Subject: [PATCH] Add ability to do set operations on current selection only

---
 kakplugin/src/lib.rs   |  5 +--
 kakplugin/src/types.rs |  6 ++++
 src/set.rs             | 76 ++++++++++++++++++++++++++++++++++--------
 3 files changed, 72 insertions(+), 15 deletions(-)

diff --git a/kakplugin/src/lib.rs b/kakplugin/src/lib.rs
index e505db9..4d64fe7 100644
--- a/kakplugin/src/lib.rs
+++ b/kakplugin/src/lib.rs
@@ -114,9 +114,10 @@ where
 /// # Errors
 ///
 /// Will return `Err` if command fifo could not be opened, read from, or written to
-pub fn set_selections_desc<'a, I>(selections: I) -> Result<(), KakError>
+pub fn set_selections_desc<'a, I, SD: 'a + std::fmt::Display>(selections: I) -> Result<(), KakError>
 where
-    I: IntoIterator<Item = &'a SelectionDesc>,
+    I: IntoIterator<Item = SD>,
+    SD: AsRef<SelectionDesc>,
 {
     let mut selections_iter = selections.into_iter().peekable();
     if selections_iter.peek().is_none() {
diff --git a/kakplugin/src/types.rs b/kakplugin/src/types.rs
index a8a2ad5..c9479b9 100644
--- a/kakplugin/src/types.rs
+++ b/kakplugin/src/types.rs
@@ -49,6 +49,12 @@ impl SelectionDesc {
     }
 }
 
+impl AsRef<SelectionDesc> for SelectionDesc {
+    fn as_ref(&self) -> &Self {
+        &self
+    }
+}
+
 impl fmt::Display for SelectionDesc {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "{},{}", self.left, self.right)
diff --git a/src/set.rs b/src/set.rs
index 649295a..6a74de4 100644
--- a/src/set.rs
+++ b/src/set.rs
@@ -118,7 +118,16 @@ pub fn set(options: &Options) -> Result<String, KakError> {
             &left_ordered_counts,
             &right_ordered_counts,
         )?,
-        Operation::Intersect | Operation::Subtract | Operation::Union => print_result(&result)?,
+        Operation::Union => print_result(&result)?,
+        Operation::Intersect | Operation::Subtract => {
+            if left_register == Register::Underscore {
+                // If the user asked for an intersection or subtraction from the current selection, we can update selection_descs only
+                reduce_selections(&options, &result)?
+            } else {
+                // The user asked for registers that *aren't* the current selection
+                print_result(&result)?
+            }
+        }
     }
 
     Ok(match &operation {
@@ -133,6 +142,36 @@ pub fn set(options: &Options) -> Result<String, KakError> {
     })
 }
 
+/// Reduces selections to those that are in the key_set_operation_result
+fn reduce_selections(
+    options: &Options,
+    key_set_operation_result: &LinkedHashSet<&Selection>,
+) -> Result<(), KakError> {
+    kakplugin::restore_register(&Register::Caret)?;
+    let selections_with_desc = {
+        let mut r = get_selections_with_desc()?;
+        r.sort_by_key(|s| s.desc.sort());
+        r
+    };
+
+    eprintln!("Key set operation result: {:?}", key_set_operation_result);
+
+    set_selections_desc(selections_with_desc.into_iter().filter_map(|swd| {
+        // Does not matter if the operation was - or &
+        // Since key_set_operation_result contains elements that should be in the set,
+        // we can just use contains here
+        let key = into_key(options, swd.content)?;
+        if key_set_operation_result.contains(&key) {
+            Some(swd.desc)
+        } else {
+            eprintln!("Key {key} not found");
+            None
+        }
+    }))?;
+
+    Ok(())
+}
+
 fn print_result(key_set_operation_result: &LinkedHashSet<&Selection>) -> Result<(), KakError> {
     // Manually set selections so we don't have to allocate a string
     let mut f = kakplugin::open_command_fifo()?;
@@ -225,23 +264,34 @@ fn to_ordered_counts(options: &Options, sels: Vec<Selection>) -> LinkedHashMap<S
     let mut ret = LinkedHashMap::new();
 
     for i in sels {
-        let key = if options.no_trim {
-            i
-        } else {
-            i.trim().to_string()
-        };
-
-        if key.is_empty() {
-            // We don't want to even pretend to look at empty keys
-            continue;
+        match into_key(options, i) {
+            Some(key) => {
+                let entry: &mut usize = ret.entry(key).or_insert(0);
+                *entry = entry.saturating_add(1);
+            }
+            None => {
+                // We don't want to even pretend to look at empty keys
+            }
         }
-
-        let entry: &mut usize = ret.entry(key).or_insert(0);
-        *entry = entry.saturating_add(1);
     }
     ret
 }
 
+fn into_key(options: &Options, sel: Selection) -> Option<Selection> {
+    let key = if options.no_trim {
+        sel
+    } else {
+        sel.trim().to_string()
+    };
+
+    if key.is_empty() {
+        // Never treat an empty string as a key
+        None
+    } else {
+        Some(key)
+    }
+}
+
 fn key_set_operation<'a>(
     operation: &Operation,
     left_keys: &LinkedHashSet<&'a Selection>,