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

polars_utils/
collection.rs
