savvy/sexp/
numeric.rs

1use std::sync::OnceLock;
2
3use crate::{savvy_err, IntegerSexp, NotAvailableValue, RealSexp, Sexp};
4
5// --- Utils -------------------------
6
7const I32MAX: f64 = i32::MAX as f64;
8const I32MIN: f64 = i32::MIN as f64;
9
10// f64 can represent 2^53
11//
12// cf. https://en.wikipedia.org/wiki/Double-precision_floating-point_format,
13//     https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Number/MAX_SAFE_INTEGER
14#[cfg(target_pointer_width = "64")]
15const F64_MAX_CASTABLE_TO_USIZE: f64 = (2_u64.pow(53) - 1) as f64;
16
17// On 32-bit target, usize::MAX is less than 2^53
18#[cfg(target_pointer_width = "32")]
19const F64_MAX_CASTABLE_TO_USIZE: f64 = usize::MAX as f64;
20
21const TOLERANCE: f64 = 0.01; // This is super-tolerant than vctrs, but this should be sufficient.
22
23fn try_cast_f64_to_i32(f: f64) -> crate::Result<i32> {
24    if f.is_na() || f.is_nan() {
25        Ok(i32::na())
26    } else if f.is_infinite() || !(I32MIN..=I32MAX).contains(&f) {
27        Err(savvy_err!("{f:?} is out of range for integer"))
28    } else if (f - f.round()).abs() > TOLERANCE {
29        Err(savvy_err!("{f:?} is not integer-ish"))
30    } else {
31        Ok(f as i32)
32    }
33}
34
35fn cast_i32_to_f64(i: i32) -> f64 {
36    if i.is_na() {
37        f64::na()
38    } else {
39        i as f64
40    }
41}
42
43fn try_cast_i32_to_usize(i: i32) -> crate::error::Result<usize> {
44    if i.is_na() {
45        Err(savvy_err!("cannot convert NA to usize"))
46    } else {
47        Ok(<usize>::try_from(i)?)
48    }
49}
50
51fn try_cast_f64_to_usize(f: f64) -> crate::Result<usize> {
52    if f.is_na() || f.is_nan() {
53        Err(savvy_err!("cannot convert NA or NaN to usize"))
54    } else if f.is_infinite() || !(0f64..=F64_MAX_CASTABLE_TO_USIZE).contains(&f) {
55        Err(savvy_err!(
56            "{f:?} is out of range that can be safely converted to usize"
57        ))
58    } else if (f - f.round()).abs() > TOLERANCE {
59        Err(savvy_err!("{f:?} is not integer-ish"))
60    } else {
61        Ok(f as usize)
62    }
63}
64
65// --- Vector -------------------------
66
67/// A enum to hold both the original data and the converted version. Since it
68/// would be a bit confusing to expose the very implementational detail of
69/// `converted` field (this is needed to return a slice), this is private.
70enum PrivateNumericSexp {
71    Integer {
72        orig: IntegerSexp,
73        converted: OnceLock<Vec<f64>>,
74    },
75    Real {
76        orig: RealSexp,
77        converted: OnceLock<Vec<i32>>,
78    },
79}
80
81/// An enum to be used for `match`ing the content of `NumericSexp`.
82pub enum NumericTypedSexp {
83    Integer(IntegerSexp),
84    Real(RealSexp),
85}
86
87/// A struct that holds either an integer or a real vector.
88pub struct NumericSexp(PrivateNumericSexp);
89
90impl NumericSexp {
91    #[inline]
92    fn inner(&self) -> savvy_ffi::SEXP {
93        match &self.0 {
94            PrivateNumericSexp::Integer { orig, .. } => orig.0,
95            PrivateNumericSexp::Real { orig, .. } => orig.0,
96        }
97    }
98
99    /// Returns the reference to the raw SEXP. This is convenient when
100    /// the lifetime is needed (e.g. returning a slice).
101    #[inline]
102    pub(crate) fn inner_ref(&self) -> &savvy_ffi::SEXP {
103        match &self.0 {
104            PrivateNumericSexp::Integer { orig, .. } => &orig.0,
105            PrivateNumericSexp::Real { orig, .. } => &orig.0,
106        }
107    }
108
109    /// Returns the length of the SEXP.
110    pub fn len(&self) -> usize {
111        unsafe { savvy_ffi::Rf_xlength(self.inner()) as _ }
112    }
113
114    /// Returns `true` if the SEXP is of zero-length.
115    #[inline]
116    pub fn is_empty(&self) -> bool {
117        self.len() == 0
118    }
119
120    /// Returns the specified attribute.
121    pub fn get_attrib(&self, attr: &str) -> crate::error::Result<Option<Sexp>> {
122        crate::Sexp(self.inner()).get_attrib(attr)
123    }
124
125    /// Returns the names.
126    pub fn get_names(&self) -> Option<Vec<&'static str>> {
127        crate::Sexp(self.inner()).get_names()
128    }
129
130    /// Returns the S3 class.
131    pub fn get_class(&self) -> Option<Vec<&'static str>> {
132        crate::Sexp(self.inner()).get_class()
133    }
134
135    /// Returns the dimension.
136    pub fn get_dim(&self) -> Option<&[i32]> {
137        // In order to maintain the lifetime, this cannot rely on the
138        // Sexp's method. Otherwise, you'll see the "cannot return
139        // reference to temporary value" error.
140        unsafe { crate::sexp::get_dim_from_sexp(self.inner_ref()) }
141    }
142
143    /// Returns the typed SEXP.
144    pub fn into_typed(self) -> NumericTypedSexp {
145        match self.0 {
146            PrivateNumericSexp::Integer { orig, .. } => NumericTypedSexp::Integer(orig),
147            PrivateNumericSexp::Real { orig, .. } => NumericTypedSexp::Real(orig),
148        }
149    }
150
151    /// Extracts a slice containing the underlying data of the SEXP.
152    ///
153    /// If the data is real, allocates a new `Vec` and cache it. This fails when
154    /// the value is
155    ///
156    /// - infinite
157    /// - out of the range of `i32`
158    /// - not integer-ish (e.g. `1.1`)
159    ///
160    /// # Examples
161    ///
162    /// ```
163    /// # let int_sexp = savvy::OwnedRealSexp::try_from_slice([1.0, 2.0, 3.0])?.as_read_only();
164    /// # let num_sexp: savvy::NumericSexp = int_sexp.try_into()?;
165    /// // `num_sexp` is c(1, 2, 3)
166    /// assert_eq!(num_sexp.as_slice_i32().unwrap(), &[1, 2, 3]);
167    /// ```
168    pub fn as_slice_i32(&self) -> crate::error::Result<&[i32]> {
169        match &self.0 {
170            PrivateNumericSexp::Integer { orig, .. } => Ok(orig.as_slice()),
171            PrivateNumericSexp::Real { orig, converted } => {
172                if let Some(v) = converted.get() {
173                    return Ok(v);
174                }
175
176                // If `converted` is not created, convert the values.
177                let v_new = orig
178                    .iter()
179                    .map(|x| try_cast_f64_to_i32(*x))
180                    .collect::<crate::Result<Vec<i32>>>()?;
181
182                // Set v_new to converted. Otherwise, this is a temporary value and cannot be returned.
183                let v = converted.get_or_init(|| v_new);
184
185                Ok(v.as_slice())
186            }
187        }
188    }
189
190    /// Extracts a slice containing the underlying data of the SEXP.
191    ///
192    /// If the data is integer, allocates a new `Vec` and cache it.
193    ///
194    /// # Examples
195    ///
196    /// ```
197    /// # let int_sexp = savvy::OwnedIntegerSexp::try_from_slice([1, 2, 3])?.as_read_only();
198    /// # let num_sexp: savvy::NumericSexp = int_sexp.try_into()?;
199    /// // `num_sexp` is c(1L, 2L, 3L)
200    /// assert_eq!(num_sexp.as_slice_f64(), &[1.0, 2.0, 3.0]);
201    /// ```
202    pub fn as_slice_f64(&self) -> &[f64] {
203        match &self.0 {
204            PrivateNumericSexp::Real { orig, .. } => orig.as_slice(),
205            PrivateNumericSexp::Integer { orig, converted } => {
206                if let Some(v) = converted.get() {
207                    return v;
208                }
209
210                // If `converted` is not created, convert the values.
211                let v_new = orig.iter().map(|i| cast_i32_to_f64(*i)).collect();
212
213                // Set v_new to converted. Otherwise, this is a temporary value and cannot be returned.
214                let v = converted.get_or_init(|| v_new);
215
216                v.as_slice()
217            }
218        }
219    }
220
221    /// Returns an iterator over the underlying data of the SEXP.
222    ///
223    /// If the data is integer, allocates a new `Vec` and cache it. While this
224    /// method itself doesn't fail, the iterator might fail to return value in
225    /// case the conversion failed, i.e. when the value is
226    ///
227    /// - infinite
228    /// - out of the range of `i32`
229    /// - not integer-ish (e.g. `1.1`)
230    ///
231    /// # Examples
232    ///
233    /// ```
234    /// use savvy::NotAvailableValue;
235    ///
236    /// # let int_sexp = savvy::OwnedIntegerSexp::try_from_slice([1, i32::na()])?.as_read_only();
237    /// # let num_sexp: savvy::NumericSexp = int_sexp.try_into()?;
238    /// // `num_sexp` is c(1, NA)
239    /// let mut iter = num_sexp.iter_f64();
240    ///
241    /// assert_eq!(iter.next(), Some(1.0));
242    ///
243    /// // NA is propagated
244    /// let e1 = iter.next();
245    /// assert!(e1.is_some());
246    /// assert!(e1.unwrap().is_na());
247    /// ```
248    pub fn iter_i32<'a>(&'a self) -> NumericIteratorI32<'a> {
249        match &self.0 {
250            PrivateNumericSexp::Integer { orig, .. } => NumericIteratorI32 {
251                sexp: self,
252                raw: Some(orig.as_slice()),
253                i: 0,
254                len: self.len(),
255            },
256            PrivateNumericSexp::Real { converted, .. } => {
257                let raw = converted.get().map(|x| x.as_slice());
258                NumericIteratorI32 {
259                    sexp: self,
260                    raw,
261                    i: 0,
262                    len: self.len(),
263                }
264            }
265        }
266    }
267
268    /// Returns an iterator over the underlying data of the SEXP.
269    ///
270    /// If the data is integer, allocates a new `Vec` and cache it.
271    ///
272    ///
273    /// # Examples
274    ///
275    /// ```
276    /// use savvy::NotAvailableValue;
277    ///
278    /// # let int_sexp = savvy::OwnedRealSexp::try_from_slice([1.0, f64::na(), 1.1])?.as_read_only();
279    /// # let num_sexp: savvy::NumericSexp = int_sexp.try_into()?;
280    /// // `num_sexp` is c(1.0, NA, 1.1)
281    /// let mut iter = num_sexp.iter_i32();
282    ///
283    /// let e0 = iter.next();
284    /// assert!(e0.is_some());
285    /// assert_eq!(e0.unwrap()?, 1);
286    ///
287    /// // NA is propagated
288    /// let e1 = iter.next();
289    /// assert!(e1.is_some());
290    /// assert!(e1.unwrap()?.is_na());
291    ///
292    /// // 1.1 is not integer-ish, so the conversion fails.
293    /// let e2 = iter.next();
294    /// assert!(e2.is_some());
295    /// assert!(e2.unwrap().is_err());
296    /// ```
297    pub fn iter_f64<'a>(&'a self) -> NumericIteratorF64<'a> {
298        match &self.0 {
299            PrivateNumericSexp::Real { orig, .. } => NumericIteratorF64 {
300                sexp: self,
301                raw: Some(orig.as_slice()),
302                i: 0,
303                len: self.len(),
304            },
305            PrivateNumericSexp::Integer { converted, .. } => {
306                let raw = converted.get().map(|x| x.as_slice());
307                NumericIteratorF64 {
308                    sexp: self,
309                    raw,
310                    i: 0,
311                    len: self.len(),
312                }
313            }
314        }
315    }
316
317    /// Returns an iterator over the underlying data of the SEXP.
318    pub fn iter_usize<'a>(&'a self) -> NumericIteratorUsize<'a> {
319        NumericIteratorUsize {
320            sexp: self,
321            i: 0,
322            len: self.len(),
323        }
324    }
325
326    // Note: If the conversion is needed, to_vec_*() would copy the values twice
327    // because it creates a `Vec` from to_slice(). This is inefficient, but I'm
328    // not sure which is worse to always creates a `Vec` from scratch or use the
329    // cached one. So, I chose not to implement the method.
330}
331
332impl TryFrom<Sexp> for NumericSexp {
333    type Error = crate::error::Error;
334
335    fn try_from(value: Sexp) -> Result<Self, Self::Error> {
336        if !value.is_numeric() {
337            let expected = "numeric".to_string();
338            let actual = value.get_human_readable_type_name().to_string();
339            return Err(crate::error::Error::UnexpectedType { expected, actual });
340        }
341
342        match value.into_typed() {
343            crate::TypedSexp::Integer(i) => Ok(Self(PrivateNumericSexp::Integer {
344                orig: i,
345                converted: OnceLock::new(),
346            })),
347            crate::TypedSexp::Real(r) => Ok(Self(PrivateNumericSexp::Real {
348                orig: r,
349                converted: OnceLock::new(),
350            })),
351            _ => Err(crate::Error::GeneralError(
352                "Should not reach here!".to_string(),
353            )),
354        }
355    }
356}
357
358impl TryFrom<IntegerSexp> for NumericSexp {
359    type Error = crate::error::Error;
360
361    fn try_from(value: IntegerSexp) -> Result<Self, Self::Error> {
362        Ok(Self(PrivateNumericSexp::Integer {
363            orig: value,
364            converted: OnceLock::new(),
365        }))
366    }
367}
368
369impl TryFrom<RealSexp> for NumericSexp {
370    type Error = crate::error::Error;
371
372    fn try_from(value: RealSexp) -> Result<Self, Self::Error> {
373        Ok(Self(PrivateNumericSexp::Real {
374            orig: value,
375            converted: OnceLock::new(),
376        }))
377    }
378}
379
380// --- Iterator -----------------------
381
382/// An iterator that returns `i32` wrapped with `Result`.
383///
384/// - If the underlying data is integer, use the value as it is.
385/// - If the underlying data is real, but there's already the `i32` values
386///   converted from the real, use the values.
387/// - Otherwise, convert a real value to `i32` on the fly. This is fallible.
388pub struct NumericIteratorI32<'a> {
389    sexp: &'a NumericSexp,
390    raw: Option<&'a [i32]>,
391    i: usize,
392    len: usize,
393}
394
395impl Iterator for NumericIteratorI32<'_> {
396    type Item = crate::error::Result<i32>;
397
398    fn next(&mut self) -> Option<Self::Item> {
399        let i = self.i;
400        self.i += 1;
401
402        if i >= self.len {
403            return None;
404        }
405
406        match &self.raw {
407            Some(x) => Some(Ok(x[i])),
408            None => {
409                if let PrivateNumericSexp::Real { orig, .. } = &self.sexp.0 {
410                    Some(try_cast_f64_to_i32(orig.as_slice()[i]))
411                } else {
412                    unreachable!("Integer must have the raw slice.");
413                }
414            }
415        }
416    }
417}
418
419/// An iterator that returns `f64`.
420///
421/// - If the underlying data is real, use the value as it is.
422/// - If the underlying data is integer, but there's already the `f64` values
423///   converted from the integer, use the values.
424/// - Otherwise, convert an integer value to `f64` on the fly.
425pub struct NumericIteratorF64<'a> {
426    sexp: &'a NumericSexp,
427    raw: Option<&'a [f64]>,
428    i: usize,
429    len: usize,
430}
431
432impl Iterator for NumericIteratorF64<'_> {
433    type Item = f64;
434
435    fn next(&mut self) -> Option<Self::Item> {
436        let i = self.i;
437        self.i += 1;
438
439        if i >= self.len {
440            return None;
441        }
442
443        match &self.raw {
444            Some(x) => Some(x[i]),
445            None => {
446                if let PrivateNumericSexp::Integer { orig, .. } = &self.sexp.0 {
447                    Some(cast_i32_to_f64(orig.as_slice()[i]))
448                } else {
449                    unreachable!("Real must have the raw slice.");
450                }
451            }
452        }
453    }
454}
455
456/// An iterator that returns `usize` wrapped with `Result`.
457pub struct NumericIteratorUsize<'a> {
458    sexp: &'a NumericSexp,
459    i: usize,
460    len: usize,
461}
462
463impl Iterator for NumericIteratorUsize<'_> {
464    type Item = crate::error::Result<usize>;
465
466    fn next(&mut self) -> Option<Self::Item> {
467        let i = self.i;
468        self.i += 1;
469
470        if i >= self.len {
471            return None;
472        }
473
474        let elem = match &self.sexp.0 {
475            PrivateNumericSexp::Integer { orig, .. } => try_cast_i32_to_usize(orig.as_slice()[i]),
476            PrivateNumericSexp::Real { orig, .. } => try_cast_f64_to_usize(orig.as_slice()[i]),
477        };
478
479        Some(elem)
480    }
481}
482
483// --- Scalar -------------------------
484
485/// A struct that holds either an integer or a real scalar.
486pub enum NumericScalar {
487    Integer(i32),
488    Real(f64),
489}
490
491impl NumericScalar {
492    /// Extracts a slice containing the underlying data of the SEXP.
493    ///
494    /// If the data is real, allocates a new `Vec` and cache it. This fails when the value is
495    ///
496    /// - infinite
497    /// - out of the range of `i32`
498    /// - not integer-ish (e.g. `1.1`)
499    pub fn as_i32(&self) -> crate::error::Result<i32> {
500        match &self {
501            NumericScalar::Integer(i) => Ok(*i),
502            NumericScalar::Real(r) => try_cast_f64_to_i32(*r),
503        }
504    }
505
506    /// Extracts a slice containing the underlying data of the SEXP.
507    ///
508    /// If the data is integer, allocates a new `Vec` and cache it.
509    pub fn as_f64(&self) -> f64 {
510        match &self {
511            NumericScalar::Integer(i) => *i as f64,
512            NumericScalar::Real(r) => *r,
513        }
514    }
515
516    pub fn as_usize(&self) -> crate::error::Result<usize> {
517        match &self {
518            NumericScalar::Integer(i) => try_cast_i32_to_usize(*i),
519            NumericScalar::Real(r) => try_cast_f64_to_usize(*r),
520        }
521    }
522}
523
524impl TryFrom<Sexp> for NumericScalar {
525    type Error = crate::error::Error;
526
527    fn try_from(value: Sexp) -> Result<Self, Self::Error> {
528        if !value.is_numeric() {
529            let expected = "numeric".to_string();
530            let actual = value.get_human_readable_type_name().to_string();
531            return Err(crate::error::Error::UnexpectedType { expected, actual });
532        }
533
534        match value.into_typed() {
535            crate::TypedSexp::Integer(i) => {
536                if i.len() != 1 {
537                    return Err(crate::error::Error::NotScalar);
538                }
539
540                let i_scalar = *i.iter().next().unwrap();
541
542                if i_scalar.is_na() {
543                    return Err(crate::error::Error::NotScalar);
544                }
545
546                Ok(Self::Integer(i_scalar))
547            }
548            crate::TypedSexp::Real(r) => {
549                if r.len() != 1 {
550                    return Err(crate::error::Error::NotScalar);
551                }
552
553                let r_scalar = *r.iter().next().unwrap();
554
555                if r_scalar.is_na() {
556                    return Err(crate::error::Error::NotScalar);
557                }
558
559                Ok(Self::Real(r_scalar))
560            }
561
562            _ => Err(crate::Error::GeneralError(
563                "Should not reach here!".to_string(),
564            )),
565        }
566    }
567}