1use std::ffi::CString;
27use std::marker::PhantomData;
28use std::ops::Deref;
29use std::sync::OnceLock;
30
31use crate::altrep_traits::NA_INTEGER;
32use crate::from_r::{SexpError, TryFromSexp, charsxp_to_str};
33use crate::gc_protect::OwnedProtect;
34use crate::into_r::IntoR;
35use crate::sys::{Rf_allocVector, Rf_install};
36use crate::{SEXP, SEXPTYPE, SexpExt};
37
38static FACTOR_CLASS: OnceLock<SEXP> = OnceLock::new();
41
42pub(crate) fn factor_class_sexp() -> SEXP {
43 *FACTOR_CLASS.get_or_init(|| unsafe {
44 let class_sexp = Rf_allocVector(SEXPTYPE::STRSXP, 1);
45 crate::sys::R_PreserveObject(class_sexp);
46 let sym = Rf_install(c"factor".as_ptr());
48 class_sexp.set_string_elt(0, sym.printname());
49 class_sexp
50 })
51}
52pub trait RFactor: crate::match_arg::MatchArg + Copy + 'static {
61 fn to_level_index(self) -> i32;
63
64 fn from_level_index(idx: i32) -> Option<Self>;
66}
67pub fn build_levels_sexp(levels: &[&str]) -> SEXP {
75 unsafe {
76 let sexp = Rf_allocVector(SEXPTYPE::STRSXP, levels.len() as isize);
77 for (i, level) in levels.iter().enumerate() {
78 let c_str = CString::new(*level).expect("level name contains null byte");
80 let sym = Rf_install(c_str.as_ptr());
81 sexp.set_string_elt(i as isize, sym.printname());
82 }
83 sexp
84 }
85}
86
87pub fn build_levels_sexp_cached(levels: &[&str]) -> SEXP {
89 unsafe {
90 let sexp = build_levels_sexp(levels);
91 crate::sys::R_PreserveObject(sexp);
92 sexp
93 }
94}
95
96pub fn build_factor(indices: &[i32], levels: SEXP) -> SEXP {
98 unsafe {
99 let (sexp, dst) = crate::into_r::alloc_r_vector::<i32>(indices.len());
100 dst.copy_from_slice(indices);
101 sexp.set_levels(levels);
102 sexp.set_class(factor_class_sexp());
103 sexp
104 }
105}
106
107pub fn build_factor_with_levels(indices: &[i32], level_names: &[&str]) -> SEXP {
124 unsafe {
125 let levels = OwnedProtect::new(build_levels_sexp(level_names));
126 build_factor(indices, levels.get())
127 }
128}
129pub struct Factor<'a> {
152 indices: &'a [i32],
153 levels_sexp: SEXP,
154 _marker: PhantomData<&'a ()>,
155}
156
157impl<'a> Factor<'a> {
158 pub fn try_new(sexp: SEXP) -> Result<Self, SexpError> {
162 if !sexp.is_factor() {
163 return Err(SexpError::InvalidValue("expected a factor".into()));
164 }
165
166 let indices = unsafe { sexp.as_slice::<i32>() };
167 let levels_sexp = sexp.get_levels();
168
169 Ok(Self {
170 indices,
171 levels_sexp,
172 _marker: PhantomData,
173 })
174 }
175
176 #[inline]
178 pub fn len(&self) -> usize {
179 self.indices.len()
180 }
181
182 #[inline]
184 pub fn is_empty(&self) -> bool {
185 self.indices.is_empty()
186 }
187
188 #[inline]
190 pub fn levels_sexp(&self) -> SEXP {
191 self.levels_sexp
192 }
193
194 #[inline]
196 pub fn n_levels(&self) -> usize {
197 self.levels_sexp.len()
198 }
199
200 #[inline]
202 pub fn level(&self, idx: usize) -> &'a str {
203 assert!(
204 idx < self.n_levels(),
205 "level index {idx} out of bounds (n_levels = {})",
206 self.n_levels()
207 );
208 let charsxp = self.levels_sexp.string_elt(idx as isize);
209 unsafe { charsxp_to_str(charsxp) }
210 }
211}
212
213impl Deref for Factor<'_> {
214 type Target = [i32];
215
216 #[inline]
217 fn deref(&self) -> &Self::Target {
218 self.indices
219 }
220}
221
222impl<'a> TryFromSexp for Factor<'a> {
223 type Error = SexpError;
224
225 fn try_from_sexp(sexp: SEXP) -> Result<Self, Self::Error> {
226 Self::try_new(sexp)
227 }
228}
229pub struct FactorMut<'a> {
248 indices: &'a mut [i32],
249 levels_sexp: SEXP,
250 _marker: PhantomData<&'a mut ()>,
251}
252
253impl<'a> FactorMut<'a> {
254 pub fn try_new(sexp: SEXP) -> Result<Self, SexpError> {
258 if !sexp.is_factor() {
259 return Err(SexpError::InvalidValue("expected a factor".into()));
260 }
261
262 let indices = unsafe { sexp.as_mut_slice::<i32>() };
263 let levels_sexp = sexp.get_levels();
264
265 Ok(Self {
266 indices,
267 levels_sexp,
268 _marker: PhantomData,
269 })
270 }
271
272 #[inline]
274 pub fn len(&self) -> usize {
275 self.indices.len()
276 }
277
278 #[inline]
280 pub fn is_empty(&self) -> bool {
281 self.indices.is_empty()
282 }
283
284 #[inline]
286 pub fn levels_sexp(&self) -> SEXP {
287 self.levels_sexp
288 }
289
290 #[inline]
292 pub fn n_levels(&self) -> usize {
293 self.levels_sexp.len()
294 }
295
296 #[inline]
298 pub fn level(&self, idx: usize) -> &'a str {
299 assert!(
300 idx < self.n_levels(),
301 "level index {idx} out of bounds (n_levels = {})",
302 self.n_levels()
303 );
304 let charsxp = self.levels_sexp.string_elt(idx as isize);
305 unsafe { charsxp_to_str(charsxp) }
306 }
307}
308
309impl Deref for FactorMut<'_> {
310 type Target = [i32];
311
312 #[inline]
313 fn deref(&self) -> &Self::Target {
314 self.indices
315 }
316}
317
318impl std::ops::DerefMut for FactorMut<'_> {
319 #[inline]
320 fn deref_mut(&mut self) -> &mut Self::Target {
321 self.indices
322 }
323}
324pub(crate) fn validate_factor_levels(sexp: SEXP, expected: &[&str]) -> Result<(), SexpError> {
330 if !sexp.is_factor() {
331 return Err(SexpError::InvalidValue("expected a factor".into()));
332 }
333
334 let levels = sexp.get_levels();
335 if levels.type_of() != SEXPTYPE::STRSXP {
336 return Err(SexpError::InvalidValue("levels is not STRSXP".into()));
337 }
338
339 let n = levels.len();
340 if n != expected.len() {
341 return Err(SexpError::InvalidValue(format!(
342 "expected {} levels, got {}",
343 expected.len(),
344 n
345 )));
346 }
347
348 for (i, exp) in expected.iter().enumerate() {
349 let charsxp = levels.string_elt(i as isize);
350 let actual = unsafe { charsxp_to_str(charsxp) };
351 if actual != *exp {
352 return Err(SexpError::InvalidValue(format!(
353 "level {}: expected '{}', got '{}'",
354 i + 1,
355 exp,
356 actual
357 )));
358 }
359 }
360
361 Ok(())
362}
363#[inline]
369pub fn factor_from_sexp<T: RFactor>(sexp: SEXP) -> Result<T, SexpError> {
370 validate_factor_levels(sexp, T::CHOICES)?;
371
372 let len = sexp.xlength();
373 if len != 1 {
374 return Err(SexpError::InvalidValue(format!(
375 "expected length 1, got {}",
376 len
377 )));
378 }
379
380 let idx = sexp.integer_elt(0);
381 if idx == NA_INTEGER {
382 return Err(SexpError::InvalidValue("unexpected NA".into()));
383 }
384
385 T::from_level_index(idx).ok_or_else(|| SexpError::InvalidValue("index out of range".into()))
386}
387
388#[inline]
390pub(crate) fn factor_vec_from_sexp<T: RFactor>(sexp: SEXP) -> Result<Vec<T>, SexpError> {
391 validate_factor_levels(sexp, T::CHOICES)?;
392
393 let len = sexp.len();
394 let mut result = Vec::with_capacity(len);
395
396 for i in 0..len {
397 let idx = sexp.integer_elt(i as isize);
398 if idx == NA_INTEGER {
399 return Err(SexpError::InvalidValue(format!("NA at index {}", i)));
400 }
401 result.push(
402 T::from_level_index(idx)
403 .ok_or_else(|| SexpError::InvalidValue("index out of range".into()))?,
404 );
405 }
406
407 Ok(result)
408}
409
410#[inline]
412pub(crate) fn factor_option_vec_from_sexp<T: RFactor>(
413 sexp: SEXP,
414) -> Result<Vec<Option<T>>, SexpError> {
415 validate_factor_levels(sexp, T::CHOICES)?;
416
417 let len = sexp.len();
418 let mut result = Vec::with_capacity(len);
419
420 for i in 0..len {
421 let idx = sexp.integer_elt(i as isize);
422 if idx == NA_INTEGER {
423 result.push(None);
424 } else {
425 result.push(Some(T::from_level_index(idx).ok_or_else(|| {
426 SexpError::InvalidValue("index out of range".into())
427 })?));
428 }
429 }
430
431 Ok(result)
432}
433
434#[inline]
441pub fn unit_factor_option_vec_from_sexp<T: UnitEnumFactor>(
442 sexp: SEXP,
443) -> Result<Vec<Option<T>>, SexpError> {
444 validate_factor_levels(sexp, T::FACTOR_LEVELS)?;
445
446 let len = sexp.len();
447 let mut result = Vec::with_capacity(len);
448
449 for i in 0..len {
450 let idx = sexp.integer_elt(i as isize);
451 if idx == NA_INTEGER {
452 result.push(None);
453 } else {
454 result.push(Some(T::from_factor_index(idx).ok_or_else(|| {
455 SexpError::InvalidValue("factor index out of range".into())
456 })?));
457 }
458 }
459
460 Ok(result)
461}
462#[derive(Debug, Clone)]
468pub struct FactorVec<T>(pub Vec<T>);
469
470impl<T> FactorVec<T> {
471 pub fn new(vec: Vec<T>) -> Self {
473 Self(vec)
474 }
475
476 pub fn into_inner(self) -> Vec<T> {
478 self.0
479 }
480}
481
482impl<T> From<Vec<T>> for FactorVec<T> {
483 fn from(vec: Vec<T>) -> Self {
484 Self(vec)
485 }
486}
487
488impl<T> Deref for FactorVec<T> {
489 type Target = Vec<T>;
490 fn deref(&self) -> &Self::Target {
491 &self.0
492 }
493}
494
495impl<T> std::ops::DerefMut for FactorVec<T> {
496 fn deref_mut(&mut self) -> &mut Self::Target {
497 &mut self.0
498 }
499}
500
501impl<T: RFactor> IntoR for FactorVec<T> {
502 type Error = std::convert::Infallible;
503 fn try_into_sexp(self) -> Result<crate::SEXP, Self::Error> {
504 Ok(self.into_sexp())
505 }
506 unsafe fn try_into_sexp_unchecked(self) -> Result<crate::SEXP, Self::Error> {
507 self.try_into_sexp()
508 }
509 fn into_sexp(self) -> SEXP {
510 let indices: Vec<i32> = self.0.iter().map(|v| v.to_level_index()).collect();
511 build_factor_with_levels(&indices, T::CHOICES)
514 }
515}
516
517impl<T: RFactor> TryFromSexp for FactorVec<T> {
518 type Error = SexpError;
519 fn try_from_sexp(sexp: SEXP) -> Result<Self, Self::Error> {
520 factor_vec_from_sexp(sexp).map(FactorVec)
521 }
522}
523
524#[derive(Debug, Clone)]
526pub struct FactorOptionVec<T>(pub Vec<Option<T>>);
527
528impl<T> FactorOptionVec<T> {
529 pub fn new(vec: Vec<Option<T>>) -> Self {
531 Self(vec)
532 }
533
534 pub fn into_inner(self) -> Vec<Option<T>> {
536 self.0
537 }
538}
539
540impl<T> From<Vec<Option<T>>> for FactorOptionVec<T> {
541 fn from(vec: Vec<Option<T>>) -> Self {
542 Self(vec)
543 }
544}
545
546impl<T> Deref for FactorOptionVec<T> {
547 type Target = Vec<Option<T>>;
548 fn deref(&self) -> &Self::Target {
549 &self.0
550 }
551}
552
553impl<T> std::ops::DerefMut for FactorOptionVec<T> {
554 fn deref_mut(&mut self) -> &mut Self::Target {
555 &mut self.0
556 }
557}
558
559impl<T: RFactor + crate::match_arg::MatchArg> UnitEnumFactor for T {
561 const FACTOR_LEVELS: &'static [&'static str] = T::CHOICES;
562 fn to_factor_index(self) -> i32 {
563 self.to_level_index()
564 }
565 fn from_factor_index(idx: i32) -> Option<Self> {
566 T::from_level_index(idx)
567 }
568}
569
570impl<T: RFactor> TryFromSexp for FactorOptionVec<T> {
571 type Error = SexpError;
572 fn try_from_sexp(sexp: SEXP) -> Result<Self, Self::Error> {
573 factor_option_vec_from_sexp(sexp).map(FactorOptionVec)
574 }
575}
576
577pub trait UnitEnumFactor: Sized {
596 const FACTOR_LEVELS: &'static [&'static str];
598
599 fn to_factor_index(self) -> i32;
601
602 fn from_factor_index(idx: i32) -> Option<Self>;
608}
609
610impl<T: UnitEnumFactor> IntoR for FactorOptionVec<T> {
611 type Error = std::convert::Infallible;
612 fn try_into_sexp(self) -> Result<crate::SEXP, Self::Error> {
613 Ok(self.into_sexp())
614 }
615 unsafe fn try_into_sexp_unchecked(self) -> Result<crate::SEXP, Self::Error> {
616 self.try_into_sexp()
617 }
618 fn into_sexp(self) -> SEXP {
619 let indices: Vec<i32> = self
623 .0
624 .into_iter()
625 .map(|opt| match opt {
626 None => NA_INTEGER,
627 Some(v) => v.to_factor_index(),
628 })
629 .collect();
630 build_factor_with_levels(&indices, T::FACTOR_LEVELS)
631 }
632}
633#[cfg(test)]
638mod tests {
639 use super::*;
640 use crate::match_arg::MatchArg;
641
642 #[derive(Copy, Clone, Debug, PartialEq)]
643 enum TestColor {
644 Red,
645 Green,
646 Blue,
647 }
648
649 impl MatchArg for TestColor {
650 const CHOICES: &'static [&'static str] = &["Red", "Green", "Blue"];
651
652 fn from_choice(choice: &str) -> Option<Self> {
653 match choice {
654 "Red" => Some(TestColor::Red),
655 "Green" => Some(TestColor::Green),
656 "Blue" => Some(TestColor::Blue),
657 _ => None,
658 }
659 }
660
661 fn to_choice(self) -> &'static str {
662 match self {
663 TestColor::Red => "Red",
664 TestColor::Green => "Green",
665 TestColor::Blue => "Blue",
666 }
667 }
668 }
669
670 impl RFactor for TestColor {
671 fn to_level_index(self) -> i32 {
672 match self {
673 TestColor::Red => 1,
674 TestColor::Green => 2,
675 TestColor::Blue => 3,
676 }
677 }
678
679 fn from_level_index(idx: i32) -> Option<Self> {
680 match idx {
681 1 => Some(TestColor::Red),
682 2 => Some(TestColor::Green),
683 3 => Some(TestColor::Blue),
684 _ => None,
685 }
686 }
687 }
688
689 #[test]
690 fn test_level_index_roundtrip() {
691 assert_eq!(
692 TestColor::from_level_index(TestColor::Red.to_level_index()),
693 Some(TestColor::Red)
694 );
695 assert_eq!(
696 TestColor::from_level_index(TestColor::Green.to_level_index()),
697 Some(TestColor::Green)
698 );
699 assert_eq!(
700 TestColor::from_level_index(TestColor::Blue.to_level_index()),
701 Some(TestColor::Blue)
702 );
703 }
704
705 #[test]
706 fn test_invalid_index() {
707 assert_eq!(TestColor::from_level_index(0), None);
708 assert_eq!(TestColor::from_level_index(4), None);
709 assert_eq!(TestColor::from_level_index(-1), None);
710 }
711
712 #[test]
713 fn test_levels_array() {
714 assert_eq!(TestColor::CHOICES, &["Red", "Green", "Blue"]);
715 }
716
717 #[derive(Copy, Clone, Debug, PartialEq)]
719 enum Size {
720 Small,
721 Large,
722 }
723
724 impl MatchArg for Size {
725 const CHOICES: &'static [&'static str] = &["Small", "Large"];
726
727 fn from_choice(choice: &str) -> Option<Self> {
728 match choice {
729 "Small" => Some(Size::Small),
730 "Large" => Some(Size::Large),
731 _ => None,
732 }
733 }
734
735 fn to_choice(self) -> &'static str {
736 match self {
737 Size::Small => "Small",
738 Size::Large => "Large",
739 }
740 }
741 }
742
743 impl RFactor for Size {
744 fn to_level_index(self) -> i32 {
745 match self {
746 Size::Small => 1,
747 Size::Large => 2,
748 }
749 }
750
751 fn from_level_index(idx: i32) -> Option<Self> {
752 match idx {
753 1 => Some(Size::Small),
754 2 => Some(Size::Large),
755 _ => None,
756 }
757 }
758 }
759
760 #[derive(Copy, Clone, Debug, PartialEq)]
762 enum ColorSize {
763 Red(Size),
764 Green(Size),
765 Blue(Size),
766 }
767
768 impl MatchArg for ColorSize {
769 const CHOICES: &'static [&'static str] = &[
770 "Red.Small",
771 "Red.Large",
772 "Green.Small",
773 "Green.Large",
774 "Blue.Small",
775 "Blue.Large",
776 ];
777
778 fn from_choice(choice: &str) -> Option<Self> {
779 let idx_1 = Self::CHOICES
780 .iter()
781 .position(|&l| l == choice)
782 .map(|i| i as i32 + 1)?;
783 Self::from_level_index(idx_1)
784 }
785
786 fn to_choice(self) -> &'static str {
787 Self::CHOICES[(self.to_level_index() - 1) as usize]
788 }
789 }
790
791 impl RFactor for ColorSize {
792 fn to_level_index(self) -> i32 {
793 match self {
794 Self::Red(inner) => {
795 let inner_idx_0 = inner.to_level_index() - 1;
796 inner_idx_0 + 1
797 }
798 Self::Green(inner) => {
799 let inner_idx_0 = inner.to_level_index() - 1;
800 2 + inner_idx_0 + 1
801 }
802 Self::Blue(inner) => {
803 let inner_idx_0 = inner.to_level_index() - 1;
804 2 * 2 + inner_idx_0 + 1
805 }
806 }
807 }
808
809 fn from_level_index(idx: i32) -> Option<Self> {
810 match idx {
811 1..=2 => {
812 let inner_idx_1 = (idx - 1) % 2 + 1;
813 Size::from_level_index(inner_idx_1).map(Self::Red)
814 }
815 3..=4 => {
816 let inner_idx_1 = (idx - 1) % 2 + 1;
817 Size::from_level_index(inner_idx_1).map(Self::Green)
818 }
819 5..=6 => {
820 let inner_idx_1 = (idx - 1) % 2 + 1;
821 Size::from_level_index(inner_idx_1).map(Self::Blue)
822 }
823 _ => None,
824 }
825 }
826 }
827
828 #[test]
829 fn test_interaction_levels() {
830 assert_eq!(
831 ColorSize::CHOICES,
832 &[
833 "Red.Small",
834 "Red.Large",
835 "Green.Small",
836 "Green.Large",
837 "Blue.Small",
838 "Blue.Large"
839 ]
840 );
841 }
842
843 #[test]
844 fn test_interaction_to_index() {
845 assert_eq!(ColorSize::Red(Size::Small).to_level_index(), 1);
846 assert_eq!(ColorSize::Red(Size::Large).to_level_index(), 2);
847 assert_eq!(ColorSize::Green(Size::Small).to_level_index(), 3);
848 assert_eq!(ColorSize::Green(Size::Large).to_level_index(), 4);
849 assert_eq!(ColorSize::Blue(Size::Small).to_level_index(), 5);
850 assert_eq!(ColorSize::Blue(Size::Large).to_level_index(), 6);
851 }
852
853 #[test]
854 fn test_interaction_from_index() {
855 assert_eq!(
856 ColorSize::from_level_index(1),
857 Some(ColorSize::Red(Size::Small))
858 );
859 assert_eq!(
860 ColorSize::from_level_index(2),
861 Some(ColorSize::Red(Size::Large))
862 );
863 assert_eq!(
864 ColorSize::from_level_index(3),
865 Some(ColorSize::Green(Size::Small))
866 );
867 assert_eq!(
868 ColorSize::from_level_index(4),
869 Some(ColorSize::Green(Size::Large))
870 );
871 assert_eq!(
872 ColorSize::from_level_index(5),
873 Some(ColorSize::Blue(Size::Small))
874 );
875 assert_eq!(
876 ColorSize::from_level_index(6),
877 Some(ColorSize::Blue(Size::Large))
878 );
879 assert_eq!(ColorSize::from_level_index(0), None);
880 assert_eq!(ColorSize::from_level_index(7), None);
881 }
882
883 #[test]
884 fn test_interaction_roundtrip() {
885 for i in 1..=6 {
886 let color_size = ColorSize::from_level_index(i).unwrap();
887 assert_eq!(color_size.to_level_index(), i);
888 }
889 }
890}
891