polars_utils/
slice_enum.rs

Help
1use std::num::TryFromIntError;
2use std::ops::Range;
3
4#[derive(Debug, Clone, PartialEq)]
5pub enum Slice {
6    /// Or zero
7    Positive {
8        offset: usize,
9        len: usize,
10    },
11    Negative {
12        offset_from_end: usize,
13        len: usize,
14    },
15}
16
17impl Slice {
18    #[allow(clippy::len_without_is_empty)]
19    pub fn len(&self) -> usize {
20        match self {
21            Slice::Positive { len, .. } => *len,
22            Slice::Negative { len, .. } => *len,
23        }
24    }
25
26    /// Returns the end position of the slice (offset + len).
27    ///
28    /// # Panics
29    /// Panics if self is negative.
30    pub fn end_position(&self) -> usize {
31        let Slice::Positive { offset, len } = self.clone() else {
32            panic!("cannot use end() on a negative slice");
33        };
34
35        offset.saturating_add(len)
36    }
37
38    /// Returns the equivalent slice to apply from an offsetted position.
39    ///
40    /// # Panics
41    /// Panics if self is negative.
42    pub fn offsetted(self, position: usize) -> Self {
43        let Slice::Positive { offset, len } = self else {
44            panic!("cannot use offsetted() on a negative slice");
45        };
46
47        let (offset, len) = if position <= offset {
48            (offset - position, len)
49        } else {
50            let n_past_offset = position - offset;
51            (0, len.saturating_sub(n_past_offset))
52        };
53
54        Slice::Positive { offset, len }
55    }
56
57    /// Restricts the bounds of the slice to within a number of rows. Negative slices will also
58    /// be translated to the positive equivalent.
59    pub fn restrict_to_bounds(self, n_rows: usize) -> Self {
60        match self {
61            Slice::Positive { offset, len } => {
62                let offset = offset.min(n_rows);
63                let len = len.min(n_rows - offset);
64                Slice::Positive { offset, len }
65            },
66            Slice::Negative {
67                offset_from_end,
68                len,
69            } => {
70                if n_rows >= offset_from_end {
71                    // Trim extra starting rows
72                    let offset = n_rows - offset_from_end;
73                    let len = len.min(n_rows - offset);
74                    Slice::Positive { offset, len }
75                } else {
76                    // Slice offset goes past start of data.
77                    let stop_at_n_from_end = offset_from_end.saturating_sub(len);
78                    let len = n_rows.saturating_sub(stop_at_n_from_end);
79
80                    Slice::Positive { offset: 0, len }
81                }
82            },
83        }
84    }
85}
86
87impl From<(usize, usize)> for Slice {
88    fn from((offset, len): (usize, usize)) -> Self {
89        Slice::Positive { offset, len }
90    }
91}
92
93impl From<(i64, usize)> for Slice {
94    fn from((offset, len): (i64, usize)) -> Self {
95        if offset >= 0 {
96            Slice::Positive {
97                offset: usize::try_from(offset).unwrap(),
98                len,
99            }
100        } else {
101            Slice::Negative {
102                offset_from_end: usize::try_from(-offset).unwrap(),
103                len,
104            }
105        }
106    }
107}
108
109impl TryFrom<Slice> for (i64, usize) {
110    type Error = TryFromIntError;
111
112    fn try_from(value: Slice) -> Result<Self, Self::Error> {
113        match value {
114            Slice::Positive { offset, len } => Ok((i64::try_from(offset)?, len)),
115            Slice::Negative {
116                offset_from_end,
117                len,
118            } => Ok((-i64::try_from(offset_from_end)?, len)),
119        }
120    }
121}
122
123impl From<Slice> for Range<usize> {
124    fn from(value: Slice) -> Self {
125        match value {
126            Slice::Positive { offset, len } => offset..offset.checked_add(len).unwrap(),
127            Slice::Negative { .. } => panic!("cannot convert negative slice into range"),
128        }
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::Slice;
135
136    #[test]
137    fn test_slice_offset() {
138        assert_eq!(
139            Slice::Positive { offset: 3, len: 10 }.offsetted(1),
140            Slice::Positive { offset: 2, len: 10 }
141        );
142        assert_eq!(
143            Slice::Positive { offset: 3, len: 10 }.offsetted(5),
144            Slice::Positive { offset: 0, len: 8 }
145        );
146    }
147
148    #[test]
149    fn test_slice_restrict_to_bounds() {
150        assert_eq!(
151            Slice::Positive { offset: 3, len: 10 }.restrict_to_bounds(7),
152            Slice::Positive { offset: 3, len: 4 },
153        );
154        assert_eq!(
155            Slice::Positive { offset: 3, len: 10 }.restrict_to_bounds(0),
156            Slice::Positive { offset: 0, len: 0 },
157        );
158        assert_eq!(
159            Slice::Positive { offset: 3, len: 10 }.restrict_to_bounds(1),
160            Slice::Positive { offset: 1, len: 0 },
161        );
162        assert_eq!(
163            Slice::Positive { offset: 2, len: 0 }.restrict_to_bounds(10),
164            Slice::Positive { offset: 2, len: 0 },
165        );
166        assert_eq!(
167            Slice::Negative {
168                offset_from_end: 3,
169                len: 1
170            }
171            .restrict_to_bounds(4),
172            Slice::Positive { offset: 1, len: 1 },
173        );
174        assert_eq!(
175            Slice::Negative {
176                offset_from_end: 3,
177                len: 1
178            }
179            .restrict_to_bounds(1),
180            Slice::Positive { offset: 0, len: 0 },
181        );
182    }
183}
polars_utils/slice_enum.rs

polars_utils/
slice_enum.rs
