From 80a69224f7c0aa11479381f352fb09e5f544f542 Mon Sep 17 00:00:00 2001
From: Reilly Wood <26268125+rgwood@users.noreply.github.com>
Date: Wed, 14 Dec 2022 11:31:54 -0800
Subject: [PATCH] Handle ctrl-c in `uniq` and `uniq-by` (#7478)

A partial fix for #7477. `uniq` can be slow sometimes, so we should
check `ctrl-c` when it's running.

Tested on [this
file](https://home.treasury.gov/system/files/276/yield-curve-rates-1990-2021.csv),
I ran `open yield-curve-rates-1990-2021.csv | uniq` and confirmed that I
can now cancel the operation.

Future work is needed to figure out why `uniq` is so slow.
---
 crates/nu-command/src/filters/uniq.rs  | 14 +++++++++-----
 crates/nu-command/src/input_handler.rs | 11 ++++++++++-
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/crates/nu-command/src/filters/uniq.rs b/crates/nu-command/src/filters/uniq.rs
index 1cdabf6a4b..aed95f24dd 100644
--- a/crates/nu-command/src/filters/uniq.rs
+++ b/crates/nu-command/src/filters/uniq.rs
@@ -1,3 +1,4 @@
+use crate::input_handler::ctrl_c_was_pressed;
 use nu_protocol::ast::Call;
 use nu_protocol::engine::{Command, EngineState, Stack};
 use nu_protocol::{
@@ -213,27 +214,30 @@ fn generate_results_with_count(head: Span, uniq_values: Vec<ValueCounter>) -> Ve
 }
 
 pub fn uniq(
-    _engine_state: &EngineState,
+    engine_state: &EngineState,
     _stack: &mut Stack,
     call: &Call,
     input: Vec<Value>,
     item_mapper: Box<dyn Fn(ItemMapperState) -> ValueCounter>,
     metadata: Option<PipelineMetadata>,
 ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
+    let ctrlc = engine_state.ctrlc.clone();
     let head = call.head;
     let flag_show_count = call.has_flag("count");
     let flag_show_repeated = call.has_flag("repeated");
     let flag_ignore_case = call.has_flag("ignore-case");
     let flag_only_uniques = call.has_flag("unique");
-    // let metadata = input.metadata();
 
     let mut uniq_values = input
         .into_iter()
-        .map(|item| {
-            item_mapper(ItemMapperState {
+        .map_while(|item| {
+            if ctrl_c_was_pressed(&ctrlc) {
+                return None;
+            }
+            Some(item_mapper(ItemMapperState {
                 item,
                 flag_ignore_case,
-            })
+            }))
         })
         .fold(Vec::<ValueCounter>::new(), |mut counter, item| {
             match counter
diff --git a/crates/nu-command/src/input_handler.rs b/crates/nu-command/src/input_handler.rs
index 342c00bf60..c9694bcfa9 100644
--- a/crates/nu-command/src/input_handler.rs
+++ b/crates/nu-command/src/input_handler.rs
@@ -1,6 +1,6 @@
 use nu_protocol::ast::CellPath;
 use nu_protocol::{PipelineData, ShellError, Span, Value};
-use std::sync::atomic::AtomicBool;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 
 pub trait CmdArgument {
@@ -71,3 +71,12 @@ where
         }
     }
 }
+
+// Helper method to avoid boilerplate every time we check ctrl+c
+pub fn ctrl_c_was_pressed(ctrlc: &Option<Arc<AtomicBool>>) -> bool {
+    if let Some(ctrlc) = ctrlc {
+        ctrlc.load(Ordering::SeqCst)
+    } else {
+        false
+    }
+}