Skip to main content

matchmaker/nucleo/
injector.rs

1// Original code from https://github.com/helix-editor/helix (MPL 2.0)
2// Modified by Squirreljetpack, 2025
3
4use std::{
5    marker::PhantomData,
6    sync::{
7        Arc,
8        atomic::{AtomicU32, Ordering},
9    },
10};
11
12use super::worker::{Column, Worker, WorkerError};
13use super::{Indexed, Segmented};
14use crate::{SSS, nucleo::SegmentableItem};
15
16pub trait Injector {
17    type InputItem;
18    type Inner: Injector;
19    type Context;
20
21    fn new(injector: Self::Inner, data: Self::Context) -> Self;
22    fn inner(&self) -> &Self::Inner;
23    fn wrap(
24        &self,
25        item: Self::InputItem,
26    ) -> Result<<Self::Inner as Injector>::InputItem, WorkerError>;
27
28    fn push(&self, item: Self::InputItem) -> Result<(), WorkerError> {
29        let item = self.wrap(item)?;
30        self.inner().push(item)
31    }
32
33    #[cfg(feature = "experimental")]
34    fn extend(
35        &self,
36        items: impl IntoIterator<Item = Self::InputItem> + ExactSizeIterator,
37    ) -> Result<(), WorkerError> {
38        let items =
39        items.into_iter().map(|item| self.wrap(item)).collect::<Result<Vec<<<Self as Injector>::Inner as Injector>::InputItem>, WorkerError>>()?;
40        self.inner().extend(items.into_iter())
41    }
42}
43
44impl Injector for () {
45    fn inner(&self) -> &Self::Inner {
46        unreachable!()
47    }
48    fn new(_: Self::Inner, _: Self::Context) -> Self {
49        unreachable!()
50    }
51    fn wrap(
52        &self,
53        _: Self::InputItem,
54    ) -> Result<<Self::Inner as Injector>::InputItem, WorkerError> {
55        unreachable!()
56    }
57
58    type Context = ();
59    type Inner = ();
60    type InputItem = ();
61}
62
63pub struct WorkerInjector<T> {
64    pub(super) inner: nucleo::Injector<T>,
65    pub(super) columns: Arc<[Column<T>]>,
66    pub(super) version: u32,
67    pub(super) picker_version: Arc<AtomicU32>,
68}
69
70impl<T: SSS> Injector for WorkerInjector<T> {
71    type InputItem = T;
72    type Inner = ();
73    type Context = Worker<T>;
74
75    fn new(_: Self::Inner, data: Self::Context) -> Self {
76        data.injector()
77    }
78
79    fn inner(&self) -> &Self::Inner {
80        &()
81    }
82
83    fn wrap(
84        &self,
85        _: Self::InputItem,
86    ) -> Result<<Self::Inner as Injector>::InputItem, WorkerError> {
87        Ok(())
88    }
89
90    fn push(&self, item: T) -> Result<(), WorkerError> {
91        if self.version != self.picker_version.load(Ordering::Relaxed) {
92            return Err(WorkerError::InjectorShutdown);
93        }
94        push_impl(&self.inner, &self.columns, item);
95        Ok(())
96    }
97
98    #[cfg(feature = "experimental")]
99    fn extend(
100        &self,
101        items: impl IntoIterator<Item = T> + ExactSizeIterator,
102    ) -> Result<(), WorkerError> {
103        if self.version != self.picker_version.load(Ordering::Relaxed) {
104            return Err(WorkerError::InjectorShutdown);
105        }
106        extend_impl(&self.inner, &self.columns, items);
107        Ok(())
108    }
109}
110
111pub(super) fn push_impl<T>(injector: &nucleo::Injector<T>, columns: &[Column<T>], item: T) {
112    injector.push(item, |item, dst| {
113        for (column, text) in columns.iter().filter(|column| column.filter).zip(dst) {
114            *text = column.format_text(item).into()
115        }
116    });
117}
118
119#[cfg(feature = "experimental")]
120pub(super) fn extend_impl<T, I>(injector: &nucleo::Injector<T>, columns: &[Column<T>], items: I)
121where
122    I: IntoIterator<Item = T> + ExactSizeIterator,
123{
124    injector.extend(items, |item, dst| {
125        for (column, text) in columns.iter().filter(|column| column.filter).zip(dst) {
126            *text = column.format_text(item).into()
127        }
128    });
129}
130
131// ----- Injectors
132
133/// Wraps the injected item with an atomic index which is incremented on push.
134#[derive(Clone)]
135pub struct IndexedInjector<T, I: Injector<InputItem = Indexed<T>>> {
136    injector: I,
137    counter: &'static AtomicU32,
138    input_type: PhantomData<T>,
139}
140
141// note that invalidation can be handled
142impl<T, I: Injector<InputItem = Indexed<T>>> Injector for IndexedInjector<T, I> {
143    type InputItem = T;
144    type Inner = I;
145    type Context = &'static AtomicU32;
146
147    fn new(injector: Self::Inner, counter: Self::Context) -> Self {
148        Self {
149            injector,
150            counter,
151            input_type: PhantomData,
152        }
153    }
154
155    fn wrap(
156        &self,
157        item: Self::InputItem,
158    ) -> Result<<Self::Inner as Injector>::InputItem, WorkerError> {
159        let index = self.counter.fetch_add(1, Ordering::SeqCst);
160        Ok(Indexed { index, inner: item })
161    }
162
163    fn inner(&self) -> &Self::Inner {
164        &self.injector
165    }
166}
167
168static GLOBAL_COUNTER: AtomicU32 = AtomicU32::new(0);
169
170impl<T, I> IndexedInjector<T, I>
171where
172    I: Injector<InputItem = Indexed<T>>,
173{
174    pub fn new_globally_indexed(injector: <Self as Injector>::Inner) -> Self {
175        Self::global_reset();
176        Self::new(injector, &GLOBAL_COUNTER)
177    }
178
179    pub fn global_reset() {
180        GLOBAL_COUNTER.store(0, Ordering::SeqCst);
181    }
182}
183
184// ------------------------------------------------------------------------------------------------
185pub type SplitterFn<T, const MAX_SPLITS: usize = { crate::MAX_SPLITS }> = std::sync::Arc<
186    dyn for<'a> Fn(&'a T) -> arrayvec::ArrayVec<(usize, usize), MAX_SPLITS> + Send + Sync,
187>;
188
189pub struct SegmentedInjector<T, I: Injector<InputItem = Segmented<T>>> {
190    injector: I,
191    splitter: SplitterFn<T>,
192}
193
194impl<T, I: Injector<InputItem = Segmented<T>>> Injector for SegmentedInjector<T, I> {
195    type InputItem = T;
196    type Inner = I;
197    type Context = SplitterFn<T>;
198
199    fn new(injector: Self::Inner, data: Self::Context) -> Self {
200        Self {
201            injector,
202            splitter: data,
203        }
204    }
205
206    fn wrap(
207        &self,
208        item: Self::InputItem,
209    ) -> Result<<Self::Inner as Injector>::InputItem, WorkerError> {
210        let ranges = (self.splitter)(&item);
211        Ok(Segmented {
212            inner: item,
213            ranges,
214        })
215    }
216
217    fn inner(&self) -> &Self::Inner {
218        &self.injector
219    }
220}
221
222mod ansi {
223    use std::ops::Range;
224
225    pub use crate::utils::Either;
226    use crate::{
227        nucleo::Text,
228        utils::text::{scrub_text_styles, slice_ratatui_text},
229    };
230    use ansi_to_tui::IntoText;
231
232    pub use super::*;
233    pub struct AnsiInjector<I> {
234        pub injector: I,
235        parse: bool,
236    }
237
238    impl<I: Injector<InputItem = Either<String, Text<'static>>>> Injector for AnsiInjector<I> {
239        type InputItem = String;
240        type Inner = I;
241        type Context = bool;
242
243        fn new(injector: Self::Inner, wrap: Self::Context) -> Self {
244            Self {
245                injector,
246                parse: wrap,
247            }
248        }
249
250        fn wrap(
251            &self,
252            item: Self::InputItem,
253        ) -> Result<<Self::Inner as Injector>::InputItem, WorkerError> {
254            let ret = if !self.parse {
255                Either::Left(item)
256            } else {
257                let mut parsed = item.as_bytes().into_text().unwrap_or(Text::from(item));
258                scrub_text_styles(&mut parsed);
259                Either::Right(parsed)
260            };
261            Ok(ret)
262        }
263
264        fn inner(&self) -> &Self::Inner {
265            &self.injector
266        }
267    }
268
269    impl SegmentableItem for Either<String, Text<'static>> {
270        fn slice(&self, range: Range<usize>) -> Text<'_> {
271            match self {
272                Either::Left(s) => {
273                    if range.start == range.end {
274                        Text::default()
275                    } else {
276                        Text::raw(&s[range.start..range.end])
277                    }
278                }
279                Either::Right(text) => slice_ratatui_text(text, range),
280            }
281        }
282    }
283}
284pub use ansi::*;
285
286// pub type SeenMap<T> = Arc<std::sync::Mutex<collections::HashSet<T>>>;
287// #[derive(Clone)]
288// pub struct UniqueInjector<T, I: Injector<InputItem = T>> {
289//     injector: I,
290//     seen: SeenMap<T>,
291// }
292// impl<T, I> Injector for UniqueInjector<T, I>
293// where
294//     T: Eq + std::hash::Hash + Clone,
295//     I: Injector<InputItem = T>,
296// {
297//     type InputItem = T;
298//     type Inner = I;
299//     type Context = SeenMap<T>;
300
301//     fn new(injector: Self::Inner, _ctx: Self::Context) -> Self {
302//         Self {
303//             injector,
304//             seen: _ctx,
305//         }
306//     }
307
308//     fn wrap(&self, item: Self::InputItem) -> Result<<Self::Inner as Injector>::InputItem, WorkerError> {
309//         let mut seen = self.seen.lock().unwrap();
310//         if seen.insert(item.clone()) {
311//             Ok(item)
312//         } else {
313//             Err(WorkerError::Custom("Duplicate"))
314//         }
315//     }
316
317//     fn inner(&self) -> &Self::Inner {
318//         &self.injector
319//     }
320// }
321
322// ----------- CLONE ----------------------------
323impl<T> Clone for WorkerInjector<T> {
324    fn clone(&self) -> Self {
325        Self {
326            inner: self.inner.clone(),
327            columns: Arc::clone(&self.columns),
328            version: self.version,
329            picker_version: Arc::clone(&self.picker_version),
330        }
331    }
332}
333
334impl<T: SegmentableItem, I: Injector<InputItem = Segmented<T>> + Clone> Clone
335    for SegmentedInjector<T, I>
336{
337    fn clone(&self) -> Self {
338        Self {
339            injector: self.injector.clone(),
340            splitter: Arc::clone(&self.splitter),
341        }
342    }
343}