aorist_extendr_api/iter.rs
1use crate::*;
2
3use wrapper::symbol::levels_symbol;
4
5/// Generalised iterator of numbers and logical. See Int, Real and Logical.
6pub struct SliceIter<T> {
7 // Control lifetime of vector to make sure the memory is not freed.
8 #[allow(dead_code)]
9 vector: Robj,
10 i: usize,
11 len: usize,
12 ptr: *const T,
13}
14
15impl<T> Default for SliceIter<T> {
16 fn default() -> Self {
17 SliceIter::new()
18 }
19}
20
21impl<T> SliceIter<T> {
22 // A new, empty list iterator.
23 pub fn new() -> Self {
24 SliceIter {
25 vector: ().into(),
26 i: 0,
27 len: 0,
28 ptr: std::ptr::null(),
29 }
30 }
31
32 pub fn from_slice(vector: Robj, slice: &[T]) -> Self {
33 SliceIter {
34 vector,
35 len: slice.len(),
36 ptr: slice.as_ptr(),
37 ..Default::default()
38 }
39 }
40}
41
42/// Basis of Int, Real and Logical.
43impl<T: Copy> Iterator for SliceIter<T> {
44 type Item = T;
45
46 fn size_hint(&self) -> (usize, Option<usize>) {
47 (self.len, Some(self.len))
48 }
49
50 fn next(&mut self) -> Option<Self::Item> {
51 let i = self.i;
52 self.i += 1;
53 if i >= self.len {
54 self.i = self.len;
55 None
56 } else {
57 unsafe { Some(*self.ptr.add(i)) }
58 }
59 }
60
61 fn nth(&mut self, n: usize) -> Option<Self::Item> {
62 self.i += n;
63 self.next()
64 }
65}
66
67impl<T: Copy> ExactSizeIterator for SliceIter<T> {
68 fn len(&self) -> usize {
69 self.len - self.i
70 }
71}
72
73/// Iterator over name-value pairs in lists.
74pub type NamedListIter = std::iter::Zip<StrIter, ListIter>;
75
76/// Iterator over primitives in integer objects.
77/// ```
78/// use extendr_api::prelude::*;
79///
80/// fn add(a: Int, b: Int) -> Robj {
81/// a.zip(b).map(|(a, b)| a+b).collect_robj()
82/// }
83/// ```
84pub type Int = SliceIter<i32>;
85
86/// Iterator over primitives in real objects.
87/// ```
88/// use extendr_api::prelude::*;
89///
90/// fn add1(a: Real) -> Robj {
91/// a.map(|a| a + 1.0).collect_robj()
92/// }
93/// ```
94pub type Real = SliceIter<f64>;
95
96/// Iterator over primitives in logical objects.
97/// ```
98/// use extendr_api::prelude::*;
99///
100/// fn all_true(mut a: Logical) -> bool {
101/// a.all(|a| a.is_true())
102/// }
103/// ```
104pub type Logical = SliceIter<Bool>;
105
106/// Iterator over strings or string factors.
107///
108/// ```
109/// use extendr_api::prelude::*;
110/// test! {
111/// let robj = r!(["a", "b", "c"]);
112/// assert_eq!(robj.as_str_iter().unwrap().collect::<Vec<_>>(), vec!["a", "b", "c"]);
113///
114/// let factor = factor!(["abcd", "def", "fg", "fg"]);
115/// assert_eq!(factor.levels().unwrap().collect::<Vec<_>>(), vec!["abcd", "def", "fg"]);
116/// assert_eq!(factor.as_integer_vector().unwrap(), vec![1, 2, 3, 3]);
117/// assert_eq!(factor.as_str_iter().unwrap().collect::<Vec<_>>(), vec!["abcd", "def", "fg", "fg"]);
118/// assert_eq!(factor.as_str_iter().unwrap().collect::<Vec<_>>(), vec!["abcd", "def", "fg", "fg"]);
119/// }
120/// ```
121#[derive(Clone)]
122pub struct StrIter {
123 vector: Robj,
124 i: usize,
125 len: usize,
126 levels: SEXP,
127}
128
129impl Default for StrIter {
130 fn default() -> Self {
131 StrIter::new()
132 }
133}
134
135impl StrIter {
136 /// Make an empty str iterator.
137 pub fn new() -> Self {
138 unsafe {
139 Self {
140 vector: ().into(),
141 i: 0,
142 len: 0,
143 levels: R_NilValue,
144 }
145 }
146 }
147
148 pub fn na_iter(len: usize) -> StrIter {
149 Self {
150 len,
151 ..Default::default()
152 }
153 }
154}
155
156// Get a string reference from a CHARSXP
157fn str_from_strsxp<'a>(sexp: SEXP, index: isize) -> &'a str {
158 unsafe {
159 if index < 0 || index >= Rf_xlength(sexp) {
160 na_str()
161 } else {
162 let charsxp = STRING_ELT(sexp, index);
163 if charsxp == R_NaString {
164 na_str()
165 } else if TYPEOF(charsxp) == CHARSXP as i32 {
166 let ptr = R_CHAR(charsxp) as *const u8;
167 let slice = std::slice::from_raw_parts(ptr, Rf_xlength(charsxp) as usize);
168 std::str::from_utf8_unchecked(slice)
169 } else {
170 na_str()
171 }
172 }
173 }
174}
175
176impl Iterator for StrIter {
177 type Item = &'static str;
178
179 fn size_hint(&self) -> (usize, Option<usize>) {
180 (self.len, Some(self.len))
181 }
182
183 fn next(&mut self) -> Option<Self::Item> {
184 unsafe {
185 let i = self.i;
186 self.i += 1;
187 let vector = self.vector.get();
188 if i >= self.len {
189 None
190 } else if TYPEOF(vector) as u32 == STRSXP {
191 Some(str_from_strsxp(vector, i as isize))
192 } else if TYPEOF(vector) as u32 == INTSXP && TYPEOF(self.levels) as u32 == STRSXP {
193 let j = *(INTEGER(vector).add(i));
194 Some(str_from_strsxp(self.levels, j as isize - 1))
195 } else if TYPEOF(vector) as u32 == NILSXP {
196 Some(na_str())
197 } else {
198 None
199 }
200 }
201 }
202
203 fn nth(&mut self, n: usize) -> Option<Self::Item> {
204 self.i += n;
205 self.next()
206 }
207}
208
209impl ExactSizeIterator for StrIter {
210 fn len(&self) -> usize {
211 self.len - self.i
212 }
213}
214
215macro_rules! impl_iter_debug {
216 ($name: ty) => {
217 impl std::fmt::Debug for $name {
218 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
219 write!(f, "[")?;
220 let mut comma = "";
221 for s in self.clone() {
222 write!(f, "{}{:?}", comma, s)?;
223 comma = ", ";
224 }
225 write!(f, "]")
226 }
227 }
228 };
229}
230
231impl_iter_debug!(ListIter);
232impl_iter_debug!(PairlistIter);
233impl_iter_debug!(StrIter);
234impl_iter_debug!(EnvIter);
235
236impl Robj {
237 /// Get an iterator over a string vector.
238 /// Returns None if the object is not a string vector
239 /// but works for factors.
240 ///
241 /// ```
242 /// use extendr_api::prelude::*;
243 ///
244 /// test! {
245 /// let obj = Robj::from(vec!["a", "b", "c"]);
246 /// assert_eq!(obj.as_str_iter().unwrap().collect::<Vec<_>>(), vec!["a", "b", "c"]);
247 ///
248 /// let factor = factor!(vec!["abcd", "def", "fg", "fg"]);
249 /// assert_eq!(factor.levels().unwrap().collect::<Vec<_>>(), vec!["abcd", "def", "fg"]);
250 /// assert_eq!(factor.as_integer_vector().unwrap(), vec![1, 2, 3, 3]);
251 /// assert_eq!(factor.as_str_iter().unwrap().collect::<Vec<_>>(), vec!["abcd", "def", "fg", "fg"]);
252 /// assert_eq!(factor.as_str_iter().unwrap().collect::<Vec<_>>(), vec!["abcd", "def", "fg", "fg"]);
253 ///
254 /// let obj = Robj::from(vec![Some("a"), Some("b"), None]);
255 /// assert_eq!(obj.as_str_iter().unwrap().map(|s| s.is_na()).collect::<Vec<_>>(), vec![false, false, true]);
256 ///
257 /// let obj = Robj::from(vec!["a", "b", na_str()]);
258 /// assert_eq!(obj.as_str_iter().unwrap().map(|s| s.is_na()).collect::<Vec<_>>(), vec![false, false, true]);
259 ///
260 /// let obj = Robj::from(vec!["a", "b", "NA"]);
261 /// assert_eq!(obj.as_str_iter().unwrap().map(|s| s.is_na()).collect::<Vec<_>>(), vec![false, false, false]);
262 /// }
263 /// ```
264 pub fn as_str_iter(&self) -> Option<StrIter> {
265 let i = 0;
266 let len = self.len();
267 match self.sexptype() {
268 STRSXP => unsafe {
269 Some(StrIter {
270 vector: self.into(),
271 i,
272 len,
273 levels: R_NilValue,
274 })
275 },
276 INTSXP => unsafe {
277 if let Some(levels) = self.get_attrib(levels_symbol()) {
278 if self.is_factor() && levels.sexptype() == STRSXP {
279 Some(StrIter {
280 vector: self.into(),
281 i,
282 len,
283 levels: levels.get(),
284 })
285 } else {
286 None
287 }
288 } else {
289 None
290 }
291 },
292 _ => None,
293 }
294 }
295}