Skip to main content

extendr_api/scalar/
rcplx_full.rs

1use crate::scalar::macros::*;
2use crate::scalar::Rfloat;
3use crate::*;
4use extendr_ffi::{R_IsNA, R_NaReal, Rcomplex};
5use std::convert::TryFrom;
6use std::ops::{Add, Div, Mul, Neg, Sub};
7use std::ops::{AddAssign, DivAssign, MulAssign, SubAssign};
8
9#[allow(non_camel_case_types)]
10pub type c64 = num_complex::Complex<f64>;
11
12impl CanBeNA for c64 {
13    fn is_na(&self) -> bool {
14        unsafe { R_IsNA(self.re) != 0 }
15    }
16
17    fn na() -> c64 {
18        unsafe { c64::new(R_NaReal, R_NaReal) }
19    }
20}
21
22/// Rcplx is a wrapper for f64 in the context of an R's complex vector.
23///
24/// Rcplx has a special NA value, obtained from R headers via R_NaReal.
25///
26/// Rcplx has the same footprint as R's complex value allowing us to use it in zero copy slices.
27#[repr(transparent)]
28#[readonly::make]
29pub struct Rcplx(pub c64);
30
31impl Rcplx {
32    pub fn new(re: f64, im: f64) -> Self {
33        Self(c64::new(re, im))
34    }
35
36    pub fn is_nan(&self) -> bool {
37        self.0.is_nan()
38    }
39
40    pub fn is_infinite(&self) -> bool {
41        self.0.is_infinite()
42    }
43
44    pub fn re(&self) -> Rfloat {
45        Rfloat::from(self.0.re)
46    }
47
48    pub fn im(&self) -> Rfloat {
49        Rfloat::from(self.0.im)
50    }
51}
52
53impl From<f64> for Rcplx {
54    fn from(val: f64) -> Self {
55        Rcplx(c64::from(val))
56    }
57}
58
59impl From<(f64, f64)> for Rcplx {
60    fn from(val: (f64, f64)) -> Self {
61        Rcplx(c64::new(val.0, val.1))
62    }
63}
64
65impl From<(Rfloat, Rfloat)> for Rcplx {
66    fn from(val: (Rfloat, Rfloat)) -> Self {
67        Rcplx(c64::new(val.0 .0, val.1 .0))
68    }
69}
70
71impl From<Rfloat> for Rcplx {
72    fn from(val: Rfloat) -> Self {
73        Rcplx(c64::from(val.0))
74    }
75}
76
77impl From<Rcomplex> for Rcplx {
78    fn from(val: Rcomplex) -> Self {
79        Rcplx(c64::new(val.r, val.i))
80    }
81}
82
83impl From<Rcplx> for Option<c64> {
84    fn from(val: Rcplx) -> Self {
85        if val.is_na() {
86            None
87        } else {
88            Some(c64::new(val.re().0, val.im().0))
89        }
90    }
91}
92
93// `NA_real_` is a `NaN` with specific bit representation.
94// Check that underlying `f64` is `NA_real_`.
95gen_trait_impl!(Rcplx, c64, |x: &Rcplx| x.0.re.is_na(), c64::na());
96gen_from_primitive!(Rcplx, c64);
97// gen_from_scalar!(Rcplx, c64);
98gen_sum_iter!(Rcplx);
99
100// Generate binary ops for +, -, * and /
101gen_binop!(
102    Rcplx,
103    c64,
104    Add,
105    |lhs: c64, rhs: c64| Some(lhs + rhs),
106    "Add two Rcplx values or an option of c64."
107);
108gen_binop!(
109    Rcplx,
110    c64,
111    Sub,
112    |lhs: c64, rhs: c64| Some(lhs - rhs),
113    "Subtract two Rcplx values or an option of c64."
114);
115gen_binop!(
116    Rcplx,
117    c64,
118    Mul,
119    |lhs: c64, rhs: c64| Some(lhs * rhs),
120    "Multiply two Rcplx values or an option of c64."
121);
122gen_binop!(
123    Rcplx,
124    c64,
125    Div,
126    |lhs: c64, rhs: c64| Some(lhs / rhs),
127    "Divide two Rcplx values or an option of c64."
128);
129gen_binopassign!(
130    Rcplx,
131    c64,
132    AddAssign,
133    |lhs: c64, rhs: c64| Some(lhs + rhs),
134    "Add two Rcplx values or an option of c64, modifying the left-hand side in place. Overflows to NA."
135);
136gen_binopassign!(
137    Rcplx,
138    c64,
139    SubAssign,
140    |lhs: c64, rhs: c64| Some(lhs - rhs),
141    "Subtract two Rcplx values or an option of c64, modifying the left-hand side in place. Overflows to NA."
142);
143gen_binopassign!(
144    Rcplx,
145    c64,
146    MulAssign,
147    |lhs: c64, rhs: c64| Some(lhs * rhs),
148    "Multiply two Rcplx values or an option of c64, modifying the left-hand side in place. Overflows to NA."
149);
150gen_binopassign!(
151    Rcplx,
152    c64,
153    DivAssign,
154    |lhs: c64, rhs: c64| Some(lhs / rhs),
155    "Divide two Rcplx values or an option of c64, modifying the left-hand side in place. Overflows to NA."
156);
157
158// Generate unary ops for -, !
159gen_unop!(Rcplx, Neg, |lhs: c64| Some(-lhs), "Negate a Rcplx value.");
160
161impl PartialEq<f64> for Rcplx {
162    fn eq(&self, other: &f64) -> bool {
163        self.re().0 == *other && self.im() == 0.0
164    }
165}
166
167impl std::fmt::Debug for Rcplx {
168    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
169        if self.is_na() {
170            write!(f, "NA_COMPLEX")
171        } else {
172            write!(
173                f,
174                "{:?} {} {:?}i",
175                self.re(),
176                if self.im().is_sign_negative() {
177                    '-'
178                } else {
179                    '+'
180                },
181                self.im().abs()
182            )
183        }
184    }
185}