polars_utils/
slice_enum.rs

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

polars_utils/
slice_enum.rs
