jaq_json/
num.rs

1//! Integer / decimal numbers.
2use super::Rc;
3use alloc::string::{String, ToString};
4use core::cmp::Ordering;
5use core::fmt;
6use core::hash::{Hash, Hasher};
7use num_bigint::{BigInt, Sign};
8use num_traits::cast::ToPrimitive;
9
10/// Integer / decimal number.
11///
12/// The speciality of this type is that numbers are distinguished into
13/// integers and 64-bit floating-point numbers.
14/// This allows using integers to index arrays,
15/// while using floating-point numbers to do general math.
16///
17/// Operations on numbers follow a few principles:
18/// * The sum, difference, product, and remainder of two integers is integer.
19/// * Any other operation between two numbers yields a float.
20#[derive(Clone, Debug)]
21pub enum Num {
22    /// Machine-size integer
23    Int(isize),
24    /// Arbitrarily large integer
25    BigInt(Rc<BigInt>),
26    /// Floating-point number
27    Float(f64),
28    /// Decimal number
29    Dec(Rc<String>),
30}
31
32impl Num {
33    pub(crate) fn big_int(i: BigInt) -> Self {
34        Self::BigInt(i.into())
35    }
36
37    pub(crate) fn from_str(s: &str) -> Self {
38        Self::try_from_int_str(s, 10).unwrap_or_else(|| Self::Dec(Rc::new(s.to_string())))
39    }
40
41    pub(crate) fn from_integral<T: Copy + TryInto<isize> + Into<BigInt>>(x: T) -> Self {
42        x.try_into()
43            .map_or_else(|_| Num::big_int(x.into()), Num::Int)
44    }
45
46    pub(crate) fn try_from_int_str(i: &str, radix: u32) -> Option<Self> {
47        let big = || BigInt::parse_bytes(i.as_bytes(), radix).map(Self::big_int);
48        isize::from_str_radix(i, radix)
49            .ok()
50            .map(Num::Int)
51            .or_else(big)
52    }
53
54    /// Try to parse a decimal string to a [`Self::Float`], else return NaN.
55    pub(crate) fn from_dec_str(n: &str) -> Self {
56        // TODO: changed to NaN!
57        n.parse().map_or(Self::Float(f64::NAN), Self::Float)
58    }
59
60    pub(crate) fn is_int(&self) -> bool {
61        matches!(self, Self::Int(_) | Self::BigInt(_))
62    }
63
64    /// If the value is a machine-sized integer, return it, else fail.
65    pub(crate) fn as_isize(&self) -> Option<isize> {
66        match self {
67            Self::Int(i) => Some(*i),
68            Self::BigInt(i) => i.to_isize(),
69            _ => None,
70        }
71    }
72
73    pub(crate) fn as_pos_usize(&self) -> Option<PosUsize> {
74        match self {
75            Self::Int(i) => Some(PosUsize(*i >= 0, i.unsigned_abs())),
76            Self::BigInt(i) => i
77                .magnitude()
78                .to_usize()
79                .map(|u| PosUsize(i.sign() != Sign::Minus, u)),
80            _ => None,
81        }
82    }
83
84    /// If the value is or can be converted to float, return it, else fail.
85    pub(crate) fn as_f64(&self) -> f64 {
86        match self {
87            Self::Int(n) => *n as f64,
88            Self::BigInt(n) => n.to_f64().unwrap(),
89            Self::Float(n) => *n,
90            Self::Dec(n) => n.parse().unwrap(),
91        }
92    }
93
94    pub(crate) fn length(&self) -> Self {
95        match self {
96            Self::Int(i) => Self::Int(i.abs()),
97            Self::BigInt(i) => match i.sign() {
98                Sign::Plus | Sign::NoSign => Self::BigInt(i.clone()),
99                Sign::Minus => Self::BigInt(BigInt::from(i.magnitude().clone()).into()),
100            },
101            Self::Dec(n) => Self::from_dec_str(n).length(),
102            Self::Float(f) => Self::Float(f.abs()),
103        }
104    }
105}
106
107#[derive(Copy, Clone)]
108pub(crate) struct PosUsize(pub(crate) bool, pub(crate) usize);
109
110impl PosUsize {
111    pub fn wrap(&self, len: usize) -> Option<usize> {
112        self.0.then_some(self.1).or_else(|| len.checked_sub(self.1))
113    }
114}
115
116#[test]
117fn wrap_test() {
118    let len = 4;
119    let pos = |i| PosUsize(true, i);
120    let neg = |i| PosUsize(false, i);
121    assert_eq!(pos(0).wrap(len), Some(0));
122    assert_eq!(pos(8).wrap(len), Some(8));
123    assert_eq!(neg(1).wrap(len), Some(3));
124    assert_eq!(neg(4).wrap(len), Some(0));
125    assert_eq!(neg(8).wrap(len), None);
126}
127
128fn int_or_big<const N: usize>(
129    i: Option<isize>,
130    x: [isize; N],
131    f: fn([BigInt; N]) -> BigInt,
132) -> Num {
133    i.map_or_else(|| Num::big_int(f(x.map(BigInt::from))), Num::Int)
134}
135
136impl core::ops::Add for Num {
137    type Output = Num;
138    fn add(self, rhs: Self) -> Self::Output {
139        use num_bigint::BigInt;
140        use Num::*;
141        match (self, rhs) {
142            (Int(x), Int(y)) => int_or_big(x.checked_add(y), [x, y], |[x, y]| x + y),
143            (Int(i), BigInt(b)) | (BigInt(b), Int(i)) => Self::big_int(&BigInt::from(i) + &*b),
144            (Int(i), Float(f)) | (Float(f), Int(i)) => Float(f + i as f64),
145            (BigInt(x), BigInt(y)) => Self::big_int(&*x + &*y),
146            (BigInt(i), Float(f)) | (Float(f), BigInt(i)) => Float(f + i.to_f64().unwrap()),
147            (Float(x), Float(y)) => Float(x + y),
148            (Dec(n), r) => Self::from_dec_str(&n) + r,
149            (l, Dec(n)) => l + Self::from_dec_str(&n),
150        }
151    }
152}
153
154impl core::ops::Sub for Num {
155    type Output = Self;
156    fn sub(self, rhs: Self) -> Self::Output {
157        use num_bigint::BigInt;
158        use Num::*;
159        match (self, rhs) {
160            (Int(x), Int(y)) => int_or_big(x.checked_sub(y), [x, y], |[x, y]| x - y),
161            (Int(i), BigInt(b)) => Self::big_int(&BigInt::from(i) - &*b),
162            (BigInt(b), Int(i)) => Self::big_int(&*b - &BigInt::from(i)),
163            (BigInt(x), BigInt(y)) => Self::big_int(&*x - &*y),
164            (Float(f), Int(i)) => Float(f - i as f64),
165            (Int(i), Float(f)) => Float(i as f64 - f),
166            (Float(f), BigInt(i)) => Float(f - i.to_f64().unwrap()),
167            (BigInt(i), Float(f)) => Float(i.to_f64().unwrap() - f),
168            (Float(x), Float(y)) => Float(x - y),
169            (Dec(n), r) => Self::from_dec_str(&n) - r,
170            (l, Dec(n)) => l - Self::from_dec_str(&n),
171        }
172    }
173}
174
175impl core::ops::Mul for Num {
176    type Output = Self;
177    fn mul(self, rhs: Self) -> Self::Output {
178        use num_bigint::BigInt;
179        use Num::*;
180        match (self, rhs) {
181            (Int(x), Int(y)) => int_or_big(x.checked_mul(y), [x, y], |[x, y]| x * y),
182
183            (Int(i), BigInt(b)) | (BigInt(b), Int(i)) => Self::big_int(&BigInt::from(i) * &*b),
184            (BigInt(x), BigInt(y)) => Self::big_int(&*x * &*y),
185            (BigInt(i), Float(f)) | (Float(f), BigInt(i)) => Float(f * i.to_f64().unwrap()),
186            (Float(f), Int(i)) | (Int(i), Float(f)) => Float(f * i as f64),
187            (Float(x), Float(y)) => Float(x * y),
188            (Dec(n), r) => Self::from_dec_str(&n) * r,
189            (l, Dec(n)) => l * Self::from_dec_str(&n),
190        }
191    }
192}
193
194impl core::ops::Div for Num {
195    type Output = Self;
196    fn div(self, rhs: Self) -> Self::Output {
197        use Num::{BigInt, Dec, Float, Int};
198        match (self, rhs) {
199            (Int(l), r) => Float(l as f64) / r,
200            (l, Int(r)) => l / Float(r as f64),
201            (BigInt(l), r) => Float(l.to_f64().unwrap()) / r,
202            (l, BigInt(r)) => l / Float(r.to_f64().unwrap()),
203            (Float(x), Float(y)) => Float(x / y),
204            (Dec(n), r) => Self::from_dec_str(&n) / r,
205            (l, Dec(n)) => l / Self::from_dec_str(&n),
206        }
207    }
208}
209
210impl core::ops::Rem for Num {
211    type Output = Self;
212    fn rem(self, rhs: Self) -> Self::Output {
213        use num_bigint::BigInt;
214        use Num::*;
215        match (self, rhs) {
216            (Int(x), Int(y)) => Int(x % y),
217            (BigInt(x), BigInt(y)) => Num::big_int(&*x % &*y),
218            (Int(i), BigInt(b)) => Num::big_int(&BigInt::from(i) % &*b),
219            (BigInt(b), Int(i)) => Num::big_int(&*b % &BigInt::from(i)),
220            (Int(i), Float(f)) => Float(i as f64 % f),
221            (Float(f), Int(i)) => Float(f % i as f64),
222            (BigInt(i), Float(f)) => Float(i.to_f64().unwrap() % f),
223            (Float(f), BigInt(i)) => Float(f % i.to_f64().unwrap()),
224            (Float(x), Float(y)) => Float(x % y),
225            (Dec(n), r) => Self::from_dec_str(&n) % r,
226            (l, Dec(n)) => l % Self::from_dec_str(&n),
227        }
228    }
229}
230
231impl core::ops::Neg for Num {
232    type Output = Self;
233    fn neg(self) -> Self::Output {
234        match self {
235            Self::Int(x) => int_or_big(x.checked_neg(), [x], |[x]| -x),
236            Self::BigInt(x) => Self::big_int(-&*x),
237            Self::Float(x) => Self::Float(-x),
238            Self::Dec(n) => match n.strip_prefix('-') {
239                Some(pos) => Self::Dec(pos.to_string().into()),
240                None => Self::Dec(alloc::format!("-{n}").into()),
241            },
242        }
243    }
244}
245
246impl Hash for Num {
247    fn hash<H: Hasher>(&self, state: &mut H) {
248        match self {
249            // hash machine-sized integers like floating-point numbers,
250            // because they are also compared for equality that way
251            Self::Int(i) => Self::Float(*i as f64).hash(state),
252            // hash all non-finite floats, like NaN and infinity, to the same
253            // note that `Val::hash` assumes that `Num::hash` always starts
254            // with `state.write_u8(n)`, where `n < 2`
255            Self::Float(f) => {
256                state.write_u8(0);
257                if f.is_finite() {
258                    f.to_ne_bytes().hash(state);
259                }
260            }
261            Self::Dec(d) => Self::from_dec_str(d).hash(state),
262            Self::BigInt(i) => {
263                let f = i.to_f64().unwrap();
264                if f.is_finite() {
265                    Self::Float(f).hash(state)
266                } else {
267                    state.write_u8(1);
268                    i.hash(state)
269                }
270            }
271        }
272    }
273}
274
275#[test]
276fn hash_nums() {
277    use core::hash::BuildHasher;
278    let h = |n| foldhash::fast::FixedState::with_seed(42).hash_one(n);
279
280    assert_eq!(h(Num::Int(4096)), h(Num::big_int(4096.into())));
281    assert_eq!(h(Num::Float(0.0)), h(Num::Int(0)));
282    assert_eq!(h(Num::Float(3.0)), h(Num::big_int(3.into())));
283    assert_eq!(h(Num::Float(f64::NAN)), h(Num::Float(0.0 / 0.0)));
284
285    assert_ne!(h(Num::Float(0.2)), h(Num::Float(0.4)));
286    assert_ne!(h(Num::Float(3.1415)), h(Num::big_int(3.into())));
287    assert_ne!(h(Num::Float(0.2)), h(Num::Int(1)));
288    assert_ne!(h(Num::Int(1)), h(Num::Int(2)));
289}
290
291impl PartialEq for Num {
292    fn eq(&self, other: &Self) -> bool {
293        match (self, other) {
294            (Self::Int(x), Self::Int(y)) => x == y,
295            (Self::Int(i), Self::Float(f)) | (Self::Float(f), Self::Int(i)) => {
296                f.is_finite() && float_eq(*i as f64, *f)
297            }
298            (Self::BigInt(x), Self::BigInt(y)) => x == y,
299            (Self::Int(i), Self::BigInt(b)) | (Self::BigInt(b), Self::Int(i)) => {
300                **b == BigInt::from(*i)
301            }
302            (Self::BigInt(i), Self::Float(f)) | (Self::Float(f), Self::BigInt(i)) => {
303                f.is_finite() && float_eq(i.to_f64().unwrap(), *f)
304            }
305            (Self::Float(x), Self::Float(y)) => float_eq(*x, *y),
306            (Self::Dec(x), Self::Dec(y)) if Rc::ptr_eq(x, y) => true,
307            (Self::Dec(n), y) => &Self::from_dec_str(n) == y,
308            (x, Self::Dec(n)) => x == &Self::from_dec_str(n),
309        }
310    }
311}
312
313impl Eq for Num {}
314
315impl PartialOrd for Num {
316    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
317        Some(self.cmp(other))
318    }
319}
320
321impl Ord for Num {
322    fn cmp(&self, other: &Self) -> Ordering {
323        match (self, other) {
324            (Self::Int(x), Self::Int(y)) => x.cmp(y),
325            (Self::Int(x), Self::BigInt(y)) => BigInt::from(*x).cmp(y),
326            (Self::Int(i), Self::Float(f)) => float_cmp(*i as f64, *f),
327            (Self::BigInt(x), Self::Int(y)) => (**x).cmp(&BigInt::from(*y)),
328            (Self::BigInt(x), Self::BigInt(y)) => x.cmp(y),
329            // BigInt::to_f64 always yields Some, large values become f64::INFINITY
330            (Self::BigInt(x), Self::Float(y)) => float_cmp(x.to_f64().unwrap(), *y),
331            (Self::Float(f), Self::Int(i)) => float_cmp(*f, *i as f64),
332            (Self::Float(x), Self::BigInt(y)) => float_cmp(*x, y.to_f64().unwrap()),
333            (Self::Float(x), Self::Float(y)) => float_cmp(*x, *y),
334            (Self::Dec(x), Self::Dec(y)) if Rc::ptr_eq(x, y) => Ordering::Equal,
335            (Self::Dec(n), y) => Self::from_dec_str(n).cmp(y),
336            (x, Self::Dec(n)) => x.cmp(&Self::from_dec_str(n)),
337        }
338    }
339}
340
341fn float_eq(left: f64, right: f64) -> bool {
342    float_cmp(left, right) == Ordering::Equal
343}
344
345fn float_cmp(left: f64, right: f64) -> Ordering {
346    if left == 0. && right == 0. {
347        // consider negative and positive 0 as equal
348        Ordering::Equal
349    } else if left.is_nan() {
350        // there are more than 50 shades of NaN, and which of these
351        // you strike when you perform a calculation is not deterministic (!),
352        // therefore `total_cmp` may yield different results for the same calculation
353        // so we bite the bullet and handle this like in jq
354        Ordering::Less
355    } else if right.is_nan() {
356        Ordering::Greater
357    } else {
358        f64::total_cmp(&left, &right)
359    }
360}
361
362impl fmt::Display for Num {
363    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
364        match self {
365            Self::Int(i) => write!(f, "{i}"),
366            Self::BigInt(i) => write!(f, "{i}"),
367            Self::Float(x) if x.is_nan() => write!(f, "NaN"),
368            Self::Float(f64::INFINITY) => write!(f, "Infinity"),
369            Self::Float(f64::NEG_INFINITY) => write!(f, "-Infinity"),
370            Self::Float(x) => ryu::Buffer::new().format_finite(*x).fmt(f),
371            Self::Dec(n) => write!(f, "{n}"),
372        }
373    }
374}