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::ffi::{Rf_allocVector, Rf_install, Rf_protect, Rf_unprotect, SEXP, SEXPTYPE, SexpExt};
33use crate::from_r::{SexpError, TryFromSexp, charsxp_to_str};
34use crate::into_r::IntoR;
35
36static FACTOR_CLASS: OnceLock<SEXP> = OnceLock::new();
39
40pub(crate) fn factor_class_sexp() -> SEXP {
41 *FACTOR_CLASS.get_or_init(|| unsafe {
42 let class_sexp = Rf_allocVector(SEXPTYPE::STRSXP, 1);
43 crate::ffi::R_PreserveObject(class_sexp);
44 let sym = Rf_install(c"factor".as_ptr());
46 class_sexp.set_string_elt(0, sym.printname());
47 class_sexp
48 })
49}
50pub trait RFactor: crate::match_arg::MatchArg + Copy + 'static {
59 fn to_level_index(self) -> i32;
61
62 fn from_level_index(idx: i32) -> Option<Self>;
64}
65pub fn build_levels_sexp(levels: &[&str]) -> SEXP {
73 unsafe {
74 let sexp = Rf_allocVector(SEXPTYPE::STRSXP, levels.len() as isize);
75 for (i, level) in levels.iter().enumerate() {
76 let c_str = CString::new(*level).expect("level name contains null byte");
78 let sym = Rf_install(c_str.as_ptr());
79 sexp.set_string_elt(i as isize, sym.printname());
80 }
81 sexp
82 }
83}
84
85pub fn build_levels_sexp_cached(levels: &[&str]) -> SEXP {
87 unsafe {
88 let sexp = build_levels_sexp(levels);
89 crate::ffi::R_PreserveObject(sexp);
90 sexp
91 }
92}
93
94pub fn build_factor(indices: &[i32], levels: SEXP) -> SEXP {
96 unsafe {
97 let (sexp, dst) = crate::into_r::alloc_r_vector::<i32>(indices.len());
98 dst.copy_from_slice(indices);
99 sexp.set_levels(levels);
100 sexp.set_class(factor_class_sexp());
101 sexp
102 }
103}
104pub struct Factor<'a> {
127 indices: &'a [i32],
128 levels_sexp: SEXP,
129 _marker: PhantomData<&'a ()>,
130}
131
132impl<'a> Factor<'a> {
133 pub fn try_new(sexp: SEXP) -> Result<Self, SexpError> {
137 if !sexp.is_factor() {
138 return Err(SexpError::InvalidValue("expected a factor".into()));
139 }
140
141 let indices = unsafe { sexp.as_slice::<i32>() };
142 let levels_sexp = sexp.get_levels();
143
144 Ok(Self {
145 indices,
146 levels_sexp,
147 _marker: PhantomData,
148 })
149 }
150
151 #[inline]
153 pub fn len(&self) -> usize {
154 self.indices.len()
155 }
156
157 #[inline]
159 pub fn is_empty(&self) -> bool {
160 self.indices.is_empty()
161 }
162
163 #[inline]
165 pub fn levels_sexp(&self) -> SEXP {
166 self.levels_sexp
167 }
168
169 #[inline]
171 pub fn n_levels(&self) -> usize {
172 self.levels_sexp.len()
173 }
174
175 #[inline]
177 pub fn level(&self, idx: usize) -> &'a str {
178 assert!(
179 idx < self.n_levels(),
180 "level index {idx} out of bounds (n_levels = {})",
181 self.n_levels()
182 );
183 let charsxp = self.levels_sexp.string_elt(idx as isize);
184 unsafe { charsxp_to_str(charsxp) }
185 }
186}
187
188impl Deref for Factor<'_> {
189 type Target = [i32];
190
191 #[inline]
192 fn deref(&self) -> &Self::Target {
193 self.indices
194 }
195}
196
197impl<'a> TryFromSexp for Factor<'a> {
198 type Error = SexpError;
199
200 fn try_from_sexp(sexp: SEXP) -> Result<Self, Self::Error> {
201 Self::try_new(sexp)
202 }
203}
204pub struct FactorMut<'a> {
223 indices: &'a mut [i32],
224 levels_sexp: SEXP,
225 _marker: PhantomData<&'a mut ()>,
226}
227
228impl<'a> FactorMut<'a> {
229 pub fn try_new(sexp: SEXP) -> Result<Self, SexpError> {
233 if !sexp.is_factor() {
234 return Err(SexpError::InvalidValue("expected a factor".into()));
235 }
236
237 let indices = unsafe { sexp.as_mut_slice::<i32>() };
238 let levels_sexp = sexp.get_levels();
239
240 Ok(Self {
241 indices,
242 levels_sexp,
243 _marker: PhantomData,
244 })
245 }
246
247 #[inline]
249 pub fn len(&self) -> usize {
250 self.indices.len()
251 }
252
253 #[inline]
255 pub fn is_empty(&self) -> bool {
256 self.indices.is_empty()
257 }
258
259 #[inline]
261 pub fn levels_sexp(&self) -> SEXP {
262 self.levels_sexp
263 }
264
265 #[inline]
267 pub fn n_levels(&self) -> usize {
268 self.levels_sexp.len()
269 }
270
271 #[inline]
273 pub fn level(&self, idx: usize) -> &'a str {
274 assert!(
275 idx < self.n_levels(),
276 "level index {idx} out of bounds (n_levels = {})",
277 self.n_levels()
278 );
279 let charsxp = self.levels_sexp.string_elt(idx as isize);
280 unsafe { charsxp_to_str(charsxp) }
281 }
282}
283
284impl Deref for FactorMut<'_> {
285 type Target = [i32];
286
287 #[inline]
288 fn deref(&self) -> &Self::Target {
289 self.indices
290 }
291}
292
293impl std::ops::DerefMut for FactorMut<'_> {
294 #[inline]
295 fn deref_mut(&mut self) -> &mut Self::Target {
296 self.indices
297 }
298}
299pub(crate) fn validate_factor_levels(sexp: SEXP, expected: &[&str]) -> Result<(), SexpError> {
305 if !sexp.is_factor() {
306 return Err(SexpError::InvalidValue("expected a factor".into()));
307 }
308
309 let levels = sexp.get_levels();
310 if levels.type_of() != SEXPTYPE::STRSXP {
311 return Err(SexpError::InvalidValue("levels is not STRSXP".into()));
312 }
313
314 let n = levels.len();
315 if n != expected.len() {
316 return Err(SexpError::InvalidValue(format!(
317 "expected {} levels, got {}",
318 expected.len(),
319 n
320 )));
321 }
322
323 for (i, exp) in expected.iter().enumerate() {
324 let charsxp = levels.string_elt(i as isize);
325 let actual = unsafe { charsxp_to_str(charsxp) };
326 if actual != *exp {
327 return Err(SexpError::InvalidValue(format!(
328 "level {}: expected '{}', got '{}'",
329 i + 1,
330 exp,
331 actual
332 )));
333 }
334 }
335
336 Ok(())
337}
338#[inline]
344pub fn factor_from_sexp<T: RFactor>(sexp: SEXP) -> Result<T, SexpError> {
345 validate_factor_levels(sexp, T::CHOICES)?;
346
347 let len = sexp.xlength();
348 if len != 1 {
349 return Err(SexpError::InvalidValue(format!(
350 "expected length 1, got {}",
351 len
352 )));
353 }
354
355 let idx = sexp.integer_elt(0);
356 if idx == NA_INTEGER {
357 return Err(SexpError::InvalidValue("unexpected NA".into()));
358 }
359
360 T::from_level_index(idx).ok_or_else(|| SexpError::InvalidValue("index out of range".into()))
361}
362
363#[inline]
365pub(crate) fn factor_vec_from_sexp<T: RFactor>(sexp: SEXP) -> Result<Vec<T>, SexpError> {
366 validate_factor_levels(sexp, T::CHOICES)?;
367
368 let len = sexp.len();
369 let mut result = Vec::with_capacity(len);
370
371 for i in 0..len {
372 let idx = sexp.integer_elt(i as isize);
373 if idx == NA_INTEGER {
374 return Err(SexpError::InvalidValue(format!("NA at index {}", i)));
375 }
376 result.push(
377 T::from_level_index(idx)
378 .ok_or_else(|| SexpError::InvalidValue("index out of range".into()))?,
379 );
380 }
381
382 Ok(result)
383}
384
385#[inline]
387pub(crate) fn factor_option_vec_from_sexp<T: RFactor>(
388 sexp: SEXP,
389) -> Result<Vec<Option<T>>, SexpError> {
390 validate_factor_levels(sexp, T::CHOICES)?;
391
392 let len = sexp.len();
393 let mut result = Vec::with_capacity(len);
394
395 for i in 0..len {
396 let idx = sexp.integer_elt(i as isize);
397 if idx == NA_INTEGER {
398 result.push(None);
399 } else {
400 result.push(Some(T::from_level_index(idx).ok_or_else(|| {
401 SexpError::InvalidValue("index out of range".into())
402 })?));
403 }
404 }
405
406 Ok(result)
407}
408#[derive(Debug, Clone)]
414pub struct FactorVec<T>(pub Vec<T>);
415
416impl<T> FactorVec<T> {
417 pub fn new(vec: Vec<T>) -> Self {
419 Self(vec)
420 }
421
422 pub fn into_inner(self) -> Vec<T> {
424 self.0
425 }
426}
427
428impl<T> From<Vec<T>> for FactorVec<T> {
429 fn from(vec: Vec<T>) -> Self {
430 Self(vec)
431 }
432}
433
434impl<T> Deref for FactorVec<T> {
435 type Target = Vec<T>;
436 fn deref(&self) -> &Self::Target {
437 &self.0
438 }
439}
440
441impl<T> std::ops::DerefMut for FactorVec<T> {
442 fn deref_mut(&mut self) -> &mut Self::Target {
443 &mut self.0
444 }
445}
446
447impl<T: RFactor> IntoR for FactorVec<T> {
448 type Error = std::convert::Infallible;
449 fn try_into_sexp(self) -> Result<crate::ffi::SEXP, Self::Error> {
450 Ok(self.into_sexp())
451 }
452 unsafe fn try_into_sexp_unchecked(self) -> Result<crate::ffi::SEXP, Self::Error> {
453 self.try_into_sexp()
454 }
455 fn into_sexp(self) -> SEXP {
456 let indices: Vec<i32> = self.0.iter().map(|v| v.to_level_index()).collect();
457 unsafe {
461 let levels = Rf_protect(build_levels_sexp(T::CHOICES));
462 let result = build_factor(&indices, levels);
463 Rf_unprotect(1);
464 result
465 }
466 }
467}
468
469impl<T: RFactor> TryFromSexp for FactorVec<T> {
470 type Error = SexpError;
471 fn try_from_sexp(sexp: SEXP) -> Result<Self, Self::Error> {
472 factor_vec_from_sexp(sexp).map(FactorVec)
473 }
474}
475
476#[derive(Debug, Clone)]
478pub struct FactorOptionVec<T>(pub Vec<Option<T>>);
479
480impl<T> FactorOptionVec<T> {
481 pub fn new(vec: Vec<Option<T>>) -> Self {
483 Self(vec)
484 }
485
486 pub fn into_inner(self) -> Vec<Option<T>> {
488 self.0
489 }
490}
491
492impl<T> From<Vec<Option<T>>> for FactorOptionVec<T> {
493 fn from(vec: Vec<Option<T>>) -> Self {
494 Self(vec)
495 }
496}
497
498impl<T> Deref for FactorOptionVec<T> {
499 type Target = Vec<Option<T>>;
500 fn deref(&self) -> &Self::Target {
501 &self.0
502 }
503}
504
505impl<T> std::ops::DerefMut for FactorOptionVec<T> {
506 fn deref_mut(&mut self) -> &mut Self::Target {
507 &mut self.0
508 }
509}
510
511impl<T: RFactor + crate::match_arg::MatchArg> UnitEnumFactor for T {
513 const FACTOR_LEVELS: &'static [&'static str] = T::CHOICES;
514 fn to_factor_index(self) -> i32 {
515 self.to_level_index()
516 }
517}
518
519impl<T: RFactor> TryFromSexp for FactorOptionVec<T> {
520 type Error = SexpError;
521 fn try_from_sexp(sexp: SEXP) -> Result<Self, Self::Error> {
522 factor_option_vec_from_sexp(sexp).map(FactorOptionVec)
523 }
524}
525
526pub trait UnitEnumFactor {
545 const FACTOR_LEVELS: &'static [&'static str];
547
548 fn to_factor_index(self) -> i32;
550}
551
552impl<T: UnitEnumFactor> IntoR for FactorOptionVec<T> {
553 type Error = std::convert::Infallible;
554 fn try_into_sexp(self) -> Result<crate::ffi::SEXP, Self::Error> {
555 Ok(self.into_sexp())
556 }
557 unsafe fn try_into_sexp_unchecked(self) -> Result<crate::ffi::SEXP, Self::Error> {
558 self.try_into_sexp()
559 }
560 fn into_sexp(self) -> SEXP {
561 let indices: Vec<i32> = self
566 .0
567 .into_iter()
568 .map(|opt| match opt {
569 None => NA_INTEGER,
570 Some(v) => v.to_factor_index(),
571 })
572 .collect();
573 unsafe {
574 let levels = Rf_protect(build_levels_sexp(T::FACTOR_LEVELS));
575 let result = build_factor(&indices, levels);
576 Rf_unprotect(1);
577 result
578 }
579 }
580}
581#[cfg(test)]
586mod tests {
587 use super::*;
588 use crate::match_arg::MatchArg;
589
590 #[derive(Copy, Clone, Debug, PartialEq)]
591 enum TestColor {
592 Red,
593 Green,
594 Blue,
595 }
596
597 impl MatchArg for TestColor {
598 const CHOICES: &'static [&'static str] = &["Red", "Green", "Blue"];
599
600 fn from_choice(choice: &str) -> Option<Self> {
601 match choice {
602 "Red" => Some(TestColor::Red),
603 "Green" => Some(TestColor::Green),
604 "Blue" => Some(TestColor::Blue),
605 _ => None,
606 }
607 }
608
609 fn to_choice(self) -> &'static str {
610 match self {
611 TestColor::Red => "Red",
612 TestColor::Green => "Green",
613 TestColor::Blue => "Blue",
614 }
615 }
616 }
617
618 impl RFactor for TestColor {
619 fn to_level_index(self) -> i32 {
620 match self {
621 TestColor::Red => 1,
622 TestColor::Green => 2,
623 TestColor::Blue => 3,
624 }
625 }
626
627 fn from_level_index(idx: i32) -> Option<Self> {
628 match idx {
629 1 => Some(TestColor::Red),
630 2 => Some(TestColor::Green),
631 3 => Some(TestColor::Blue),
632 _ => None,
633 }
634 }
635 }
636
637 #[test]
638 fn test_level_index_roundtrip() {
639 assert_eq!(
640 TestColor::from_level_index(TestColor::Red.to_level_index()),
641 Some(TestColor::Red)
642 );
643 assert_eq!(
644 TestColor::from_level_index(TestColor::Green.to_level_index()),
645 Some(TestColor::Green)
646 );
647 assert_eq!(
648 TestColor::from_level_index(TestColor::Blue.to_level_index()),
649 Some(TestColor::Blue)
650 );
651 }
652
653 #[test]
654 fn test_invalid_index() {
655 assert_eq!(TestColor::from_level_index(0), None);
656 assert_eq!(TestColor::from_level_index(4), None);
657 assert_eq!(TestColor::from_level_index(-1), None);
658 }
659
660 #[test]
661 fn test_levels_array() {
662 assert_eq!(TestColor::CHOICES, &["Red", "Green", "Blue"]);
663 }
664
665 #[derive(Copy, Clone, Debug, PartialEq)]
667 enum Size {
668 Small,
669 Large,
670 }
671
672 impl MatchArg for Size {
673 const CHOICES: &'static [&'static str] = &["Small", "Large"];
674
675 fn from_choice(choice: &str) -> Option<Self> {
676 match choice {
677 "Small" => Some(Size::Small),
678 "Large" => Some(Size::Large),
679 _ => None,
680 }
681 }
682
683 fn to_choice(self) -> &'static str {
684 match self {
685 Size::Small => "Small",
686 Size::Large => "Large",
687 }
688 }
689 }
690
691 impl RFactor for Size {
692 fn to_level_index(self) -> i32 {
693 match self {
694 Size::Small => 1,
695 Size::Large => 2,
696 }
697 }
698
699 fn from_level_index(idx: i32) -> Option<Self> {
700 match idx {
701 1 => Some(Size::Small),
702 2 => Some(Size::Large),
703 _ => None,
704 }
705 }
706 }
707
708 #[derive(Copy, Clone, Debug, PartialEq)]
710 enum ColorSize {
711 Red(Size),
712 Green(Size),
713 Blue(Size),
714 }
715
716 impl MatchArg for ColorSize {
717 const CHOICES: &'static [&'static str] = &[
718 "Red.Small",
719 "Red.Large",
720 "Green.Small",
721 "Green.Large",
722 "Blue.Small",
723 "Blue.Large",
724 ];
725
726 fn from_choice(choice: &str) -> Option<Self> {
727 let idx_1 = Self::CHOICES
728 .iter()
729 .position(|&l| l == choice)
730 .map(|i| i as i32 + 1)?;
731 Self::from_level_index(idx_1)
732 }
733
734 fn to_choice(self) -> &'static str {
735 Self::CHOICES[(self.to_level_index() - 1) as usize]
736 }
737 }
738
739 impl RFactor for ColorSize {
740 fn to_level_index(self) -> i32 {
741 match self {
742 Self::Red(inner) => {
743 let inner_idx_0 = inner.to_level_index() - 1;
744 inner_idx_0 + 1
745 }
746 Self::Green(inner) => {
747 let inner_idx_0 = inner.to_level_index() - 1;
748 2 + inner_idx_0 + 1
749 }
750 Self::Blue(inner) => {
751 let inner_idx_0 = inner.to_level_index() - 1;
752 2 * 2 + inner_idx_0 + 1
753 }
754 }
755 }
756
757 fn from_level_index(idx: i32) -> Option<Self> {
758 match idx {
759 1..=2 => {
760 let inner_idx_1 = (idx - 1) % 2 + 1;
761 Size::from_level_index(inner_idx_1).map(Self::Red)
762 }
763 3..=4 => {
764 let inner_idx_1 = (idx - 1) % 2 + 1;
765 Size::from_level_index(inner_idx_1).map(Self::Green)
766 }
767 5..=6 => {
768 let inner_idx_1 = (idx - 1) % 2 + 1;
769 Size::from_level_index(inner_idx_1).map(Self::Blue)
770 }
771 _ => None,
772 }
773 }
774 }
775
776 #[test]
777 fn test_interaction_levels() {
778 assert_eq!(
779 ColorSize::CHOICES,
780 &[
781 "Red.Small",
782 "Red.Large",
783 "Green.Small",
784 "Green.Large",
785 "Blue.Small",
786 "Blue.Large"
787 ]
788 );
789 }
790
791 #[test]
792 fn test_interaction_to_index() {
793 assert_eq!(ColorSize::Red(Size::Small).to_level_index(), 1);
794 assert_eq!(ColorSize::Red(Size::Large).to_level_index(), 2);
795 assert_eq!(ColorSize::Green(Size::Small).to_level_index(), 3);
796 assert_eq!(ColorSize::Green(Size::Large).to_level_index(), 4);
797 assert_eq!(ColorSize::Blue(Size::Small).to_level_index(), 5);
798 assert_eq!(ColorSize::Blue(Size::Large).to_level_index(), 6);
799 }
800
801 #[test]
802 fn test_interaction_from_index() {
803 assert_eq!(
804 ColorSize::from_level_index(1),
805 Some(ColorSize::Red(Size::Small))
806 );
807 assert_eq!(
808 ColorSize::from_level_index(2),
809 Some(ColorSize::Red(Size::Large))
810 );
811 assert_eq!(
812 ColorSize::from_level_index(3),
813 Some(ColorSize::Green(Size::Small))
814 );
815 assert_eq!(
816 ColorSize::from_level_index(4),
817 Some(ColorSize::Green(Size::Large))
818 );
819 assert_eq!(
820 ColorSize::from_level_index(5),
821 Some(ColorSize::Blue(Size::Small))
822 );
823 assert_eq!(
824 ColorSize::from_level_index(6),
825 Some(ColorSize::Blue(Size::Large))
826 );
827 assert_eq!(ColorSize::from_level_index(0), None);
828 assert_eq!(ColorSize::from_level_index(7), None);
829 }
830
831 #[test]
832 fn test_interaction_roundtrip() {
833 for i in 1..=6 {
834 let color_size = ColorSize::from_level_index(i).unwrap();
835 assert_eq!(color_size.to_level_index(), i);
836 }
837 }
838}
839