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 type PreprocessOptions = (bool, bool);
233
234    pub use super::*;
235    pub struct AnsiInjector<I> {
236        pub injector: I,
237        parse: bool,
238        trim: bool,
239    }
240
241    impl<I: Injector<InputItem = Either<String, Text<'static>>>> Injector for AnsiInjector<I> {
242        type InputItem = String;
243        type Inner = I;
244        type Context = PreprocessOptions;
245
246        fn new(injector: Self::Inner, (parse, trim): Self::Context) -> Self {
247            Self {
248                injector,
249                parse,
250                trim,
251            }
252        }
253
254        fn wrap(
255            &self,
256            mut item: Self::InputItem,
257        ) -> Result<<Self::Inner as Injector>::InputItem, WorkerError> {
258            if self.trim {
259                item = item.trim().to_string();
260            }
261            let ret = if !self.parse {
262                Either::Left(item)
263            } else {
264                let mut parsed = item.as_bytes().into_text().unwrap_or(Text::from(item));
265                scrub_text_styles(&mut parsed);
266                Either::Right(parsed)
267            };
268            Ok(ret)
269        }
270
271        fn inner(&self) -> &Self::Inner {
272            &self.injector
273        }
274    }
275
276    impl SegmentableItem for Either<String, Text<'static>> {
277        fn slice(&self, range: Range<usize>) -> Text<'_> {
278            match self {
279                Either::Left(s) => {
280                    if range.start == range.end {
281                        Text::default()
282                    } else {
283                        Text::raw(&s[range.start..range.end])
284                    }
285                }
286                Either::Right(text) => slice_ratatui_text(text, range),
287            }
288        }
289    }
290}
291pub use ansi::*;
292
293// pub type SeenMap<T> = Arc<std::sync::Mutex<collections::HashSet<T>>>;
294// #[derive(Clone)]
295// pub struct UniqueInjector<T, I: Injector<InputItem = T>> {
296//     injector: I,
297//     seen: SeenMap<T>,
298// }
299// impl<T, I> Injector for UniqueInjector<T, I>
300// where
301//     T: Eq + std::hash::Hash + Clone,
302//     I: Injector<InputItem = T>,
303// {
304//     type InputItem = T;
305//     type Inner = I;
306//     type Context = SeenMap<T>;
307
308//     fn new(injector: Self::Inner, _ctx: Self::Context) -> Self {
309//         Self {
310//             injector,
311//             seen: _ctx,
312//         }
313//     }
314
315//     fn wrap(&self, item: Self::InputItem) -> Result<<Self::Inner as Injector>::InputItem, WorkerError> {
316//         let mut seen = self.seen.lock().unwrap();
317//         if seen.insert(item.clone()) {
318//             Ok(item)
319//         } else {
320//             Err(WorkerError::Custom("Duplicate"))
321//         }
322//     }
323
324//     fn inner(&self) -> &Self::Inner {
325//         &self.injector
326//     }
327// }
328
329// ----------- CLONE ----------------------------
330impl<T> Clone for WorkerInjector<T> {
331    fn clone(&self) -> Self {
332        Self {
333            inner: self.inner.clone(),
334            columns: Arc::clone(&self.columns),
335            version: self.version,
336            picker_version: Arc::clone(&self.picker_version),
337        }
338    }
339}
340
341impl<T: SegmentableItem, I: Injector<InputItem = Segmented<T>> + Clone> Clone
342    for SegmentedInjector<T, I>
343{
344    fn clone(&self) -> Self {
345        Self {
346            injector: self.injector.clone(),
347            splitter: Arc::clone(&self.splitter),
348        }
349    }
350}