avoid excessive memory consumption in picker (#8127)

* avoid excessive memory consumption from file picker

* fix typos

Co-authored-by: Chris <75008413+cd-a@users.noreply.github.com>

---------

Co-authored-by: Chris <75008413+cd-a@users.noreply.github.com>

3 files changed, 40 insertions, 16 deletions

diff --git a/helix-term/src/compositor.rs b/helix-term/src/compositor.rs
index bcb3e449..dc962e1f 100644
--- a/helix-term/src/compositor.rs
+++ b/helix-term/src/compositor.rs
@@ -16,6 +16,7 @@ pub enum EventResult {
 }
 
 use crate::job::Jobs;
+use crate::ui::picker;
 use helix_view::Editor;
 
 pub use helix_view::input::Event;
@@ -100,6 +101,11 @@ impl Compositor {
 
     /// Add a layer to be rendered in front of all existing layers.
     pub fn push(&mut self, mut layer: Box<dyn Component>) {
+        // immediately clear last_picker field to avoid excessive memory
+        // consumption for picker with many items
+        if layer.id() == Some(picker::ID) {
+            self.last_picker = None;
+        }
         let size = self.size();
         // trigger required_size on init
         layer.required_size((size.width, size.height));
diff --git a/helix-term/src/ui/mod.rs b/helix-term/src/ui/mod.rs
index 3ca3d24d..b2160108 100644
--- a/helix-term/src/ui/mod.rs
+++ b/helix-term/src/ui/mod.rs
@@ -6,7 +6,7 @@ pub mod lsp;
 mod markdown;
 pub mod menu;
 pub mod overlay;
-mod picker;
+pub mod picker;
 pub mod popup;
 mod prompt;
 mod spinner;
diff --git a/helix-term/src/ui/picker.rs b/helix-term/src/ui/picker.rs
index 3073a697..1f94a72c 100644
--- a/helix-term/src/ui/picker.rs
+++ b/helix-term/src/ui/picker.rs
@@ -46,6 +46,7 @@ use helix_view::{
     Document, DocumentId, Editor,
 };
 
+pub const ID: &str = "picker";
 use super::{menu::Item, overlay::Overlay};
 
 pub const MIN_AREA_WIDTH_FOR_PREVIEW: u16 = 72;
@@ -802,11 +803,28 @@ impl<T: Item + 'static + Send + Sync> Component for Picker<T> {
             _ => return EventResult::Ignored(None),
         };
 
-        let close_fn =
-            EventResult::Consumed(Some(Box::new(|compositor: &mut Compositor, _ctx| {
-                // remove the layer
-                compositor.last_picker = compositor.pop();
-            })));
+        let close_fn = |picker: &mut Self| {
+            // if the picker is very large don't store it as last_picker to avoid
+            // excessive memory consumption
+            let callback: compositor::Callback = if picker.matcher.snapshot().item_count() > 100_000
+            {
+                Box::new(|compositor: &mut Compositor, _ctx| {
+                    // remove the layer
+                    compositor.pop();
+                })
+            } else {
+                // stop streaming in new items in the background, really we should
+                // be restarting the stream somehow once the picker gets
+                // reopened instead (like for an FS crawl) that would also remove the
+                // need for the special case above but that is pretty tricky
+                picker.shutdown.store(true, atomic::Ordering::Relaxed);
+                Box::new(|compositor: &mut Compositor, _ctx| {
+                    // remove the layer
+                    compositor.last_picker = compositor.pop();
+                })
+            };
+            EventResult::Consumed(Some(callback))
+        };
 
         // So that idle timeout retriggers
         ctx.editor.reset_idle_timer();
@@ -830,9 +848,7 @@ impl<T: Item + 'static + Send + Sync> Component for Picker<T> {
             key!(End) => {
                 self.to_end();
             }
-            key!(Esc) | ctrl!('c') => {
-                return close_fn;
-            }
+            key!(Esc) | ctrl!('c') => return close_fn(self),
             alt!(Enter) => {
                 if let Some(option) = self.selection() {
                     (self.callback_fn)(ctx, option, Action::Load);
@@ -842,19 +858,19 @@ impl<T: Item + 'static + Send + Sync> Component for Picker<T> {
                 if let Some(option) = self.selection() {
                     (self.callback_fn)(ctx, option, Action::Replace);
                 }
-                return close_fn;
+                return close_fn(self);
             }
             ctrl!('s') => {
                 if let Some(option) = self.selection() {
                     (self.callback_fn)(ctx, option, Action::HorizontalSplit);
                 }
-                return close_fn;
+                return close_fn(self);
             }
             ctrl!('v') => {
                 if let Some(option) = self.selection() {
                     (self.callback_fn)(ctx, option, Action::VerticalSplit);
                 }
-                return close_fn;
+                return close_fn(self);
             }
             ctrl!('t') => {
                 self.toggle_preview();
@@ -882,6 +898,10 @@ impl<T: Item + 'static + Send + Sync> Component for Picker<T> {
         self.completion_height = height.saturating_sub(4);
         Some((width, height))
     }
+
+    fn id(&self) -> Option<&'static str> {
+        Some(ID)
+    }
 }
 impl<T: Item> Drop for Picker<T> {
     fn drop(&mut self) {
@@ -906,8 +926,6 @@ pub struct DynamicPicker<T: ui::menu::Item + Send + Sync> {
 }
 
 impl<T: ui::menu::Item + Send + Sync> DynamicPicker<T> {
-    pub const ID: &'static str = "dynamic-picker";
-
     pub fn new(file_picker: Picker<T>, query_callback: DynQueryCallback<T>) -> Self {
         Self {
             file_picker,
@@ -939,7 +957,7 @@ impl<T: Item + Send + Sync + 'static> Component for DynamicPicker<T> {
             let callback = Callback::EditorCompositor(Box::new(move |editor, compositor| {
                 // Wrapping of pickers in overlay is done outside the picker code,
                 // so this is fragile and will break if wrapped in some other widget.
-                let picker = match compositor.find_id::<Overlay<DynamicPicker<T>>>(Self::ID) {
+                let picker = match compositor.find_id::<Overlay<DynamicPicker<T>>>(ID) {
                     Some(overlay) => &mut overlay.content.file_picker,
                     None => return,
                 };
@@ -960,6 +978,6 @@ impl<T: Item + Send + Sync + 'static> Component for DynamicPicker<T> {
     }
 
     fn id(&self) -> Option<&'static str> {
-        Some(Self::ID)
+        Some(ID)
     }
 }