polars_utils/
collection.rs

1use std::marker::PhantomData;
2use std::ops::{Deref, DerefMut, Index, IndexMut};
3
4pub trait Collection<T: ?Sized> {
5    fn is_empty(&self) -> bool {
6        self.len() == 0
7    }
8
9    fn len(&self) -> usize;
10    fn get(&self, idx: usize) -> Option<&T>;
11    fn get_mut(&mut self, idx: usize) -> Option<&mut T>;
12}
13
14/// Wrapper that implements indexing.
15pub struct CollectionWrap<T: ?Sized, C: Collection<T>> {
16    inner: C,
17    phantom: PhantomData<T>,
18}
19
20impl<T: ?Sized, C: Collection<T>> CollectionWrap<T, C> {
21    #[inline(always)]
22    pub fn new(inner: C) -> Self {
23        Self {
24            inner,
25            phantom: PhantomData,
26        }
27    }
28
29    pub fn iter(&self) -> CollectionIter<'_, T, C> {
30        CollectionIter {
31            idx: 0,
32            collection: &self.inner,
33            phantom: PhantomData,
34        }
35    }
36
37    pub fn for_each_mut<F>(&mut self, mut f: F)
38    where
39        F: for<'b> FnMut(&'b mut T),
40    {
41        (0..self.len()).for_each(move |i| f(self.get_mut(i).unwrap()))
42    }
43
44    pub fn map_mut<'a, B, F>(&'a mut self, mut f: F) -> impl Iterator<Item = B>
45    where
46        F: for<'b> FnMut(&'b mut T) -> B + 'a,
47    {
48        (0..self.len()).map(move |i| f(self.get_mut(i).unwrap()))
49    }
50
51    pub fn into_inner(self) -> C {
52        self.inner
53    }
54}
55
56impl<T: ?Sized, C: Collection<T>> Deref for CollectionWrap<T, C> {
57    type Target = C;
58
59    #[inline]
60    fn deref(&self) -> &Self::Target {
61        &self.inner
62    }
63}
64
65impl<T: ?Sized, C: Collection<T>> DerefMut for CollectionWrap<T, C> {
66    #[inline]
67    fn deref_mut(&mut self) -> &mut Self::Target {
68        &mut self.inner
69    }
70}
71
72impl<T: ?Sized, C: Collection<T>> Index<usize> for CollectionWrap<T, C> {
73    type Output = T;
74
75    fn index(&self, index: usize) -> &Self::Output {
76        self.get(index).unwrap()
77    }
78}
79
80impl<T: ?Sized, C: Collection<T>> IndexMut<usize> for CollectionWrap<T, C> {
81    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
82        self.get_mut(index).unwrap()
83    }
84}
85
86impl<T: Clone, C: Collection<T>, const N: usize> TryFrom<CollectionWrap<T, C>> for [T; N] {
87    type Error = ();
88
89    fn try_from(value: CollectionWrap<T, C>) -> Result<Self, Self::Error> {
90        if value.len() != N {
91            return Err(());
92        }
93
94        Ok(std::array::from_fn(|i| value.get(i).unwrap().clone()))
95    }
96}
97
98impl<T: ?Sized, C: Collection<T>> From<C> for CollectionWrap<T, C> {
99    #[inline]
100    fn from(value: C) -> Self {
101        Self {
102            inner: value,
103            phantom: PhantomData,
104        }
105    }
106}
107
108impl<T: ?Sized> Collection<T> for &mut dyn Collection<T> {
109    #[inline]
110    fn len(&self) -> usize {
111        (**self).len()
112    }
113
114    #[inline]
115    fn get(&self, idx: usize) -> Option<&T> {
116        (**self).get(idx)
117    }
118
119    #[inline]
120    fn get_mut(&mut self, idx: usize) -> Option<&mut T> {
121        (**self).get_mut(idx)
122    }
123}
124
125pub struct CollectionIter<'a, T: 'a + ?Sized, C: Collection<T>> {
126    idx: usize,
127    collection: &'a C,
128    phantom: PhantomData<T>,
129}
130
131impl<'a, T: 'a + ?Sized, C: Collection<T>> Iterator for CollectionIter<'a, T, C> {
132    type Item = &'a T;
133
134    fn next(&mut self) -> Option<Self::Item> {
135        let item = self.collection.get(self.idx);
136
137        if item.is_some() {
138            self.idx += 1;
139        }
140
141        item
142    }
143}
144
145impl<T> Collection<T> for [T] {
146    fn len(&self) -> usize {
147        <[T]>::len(self)
148    }
149
150    fn get(&self, idx: usize) -> Option<&T> {
151        <[T]>::get(self, idx)
152    }
153
154    fn get_mut(&mut self, idx: usize) -> Option<&mut T> {
155        <[T]>::get_mut(self, idx)
156    }
157}
158
159impl<T> Collection<T> for &mut [T] {
160    fn len(&self) -> usize {
161        <[T]>::len(self)
162    }
163
164    fn get(&self, idx: usize) -> Option<&T> {
165        <[T]>::get(self, idx)
166    }
167
168    fn get_mut(&mut self, idx: usize) -> Option<&mut T> {
169        <[T]>::get_mut(self, idx)
170    }
171}
172
173pub struct MappedCollection<'src, Src: ?Sized, T: ?Sized, U: ?Sized> {
174    src: &'src mut Src,
175    map: fn(&T) -> &U,
176    map_mut: fn(&mut T) -> &mut U,
177}
178
179impl<'src, Src: ?Sized, T: ?Sized, U: ?Sized> Collection<U> for MappedCollection<'src, Src, T, U>
180where
181    Src: Collection<T>,
182{
183    fn len(&self) -> usize {
184        self.src.len()
185    }
186
187    fn get(&self, idx: usize) -> Option<&U> {
188        self.src.get(idx).map(|t| (self.map)(t))
189    }
190
191    fn get_mut(&mut self, idx: usize) -> Option<&mut U> {
192        self.src.get_mut(idx).map(|t| (self.map_mut)(t))
193    }
194}
195
196impl<'src, Src: ?Sized, T: ?Sized, U: ?Sized> MappedCollection<'src, Src, T, U>
197where
198    Src: Collection<T>,
199{
200    pub fn new(src: &'src mut Src, map: fn(&T) -> &U, map_mut: fn(&mut T) -> &mut U) -> Self {
201        Self { src, map, map_mut }
202    }
203}
polars_utils/collection.rs

polars_utils/
collection.rs
