Skip to main content

pumpkin_checking/
int_ext.rs

1use std::cmp::Ordering;
2use std::fmt::Debug;
3use std::iter::Sum;
4use std::ops::Add;
5use std::ops::AddAssign;
6use std::ops::Mul;
7use std::ops::Neg;
8use std::ops::Sub;
9
10/// An integer or positive/negative infinity.
11///
12/// # Notes on arithmetic operations:
13/// - The result of the operation `infty + -infty` is undetermined, and if evaluated will cause a
14///   panic.
15/// - Multiplying [`IntExt::PositiveInf`] or [`IntExt::NegativeInf`] with `IntExt::I32(0)` will
16///   yield `IntExt::Int(0)`.
17#[derive(Clone, Copy, Debug, PartialEq, Eq)]
18pub enum IntExt<Int = i32> {
19    Int(Int),
20    NegativeInf,
21    PositiveInf,
22}
23
24impl<Int: Copy> IntExt<Int> {
25    pub fn as_int(&self) -> Option<Int> {
26        match self {
27            IntExt::Int(int) => Some(*int),
28            IntExt::NegativeInf | IntExt::PositiveInf => None,
29        }
30    }
31}
32
33/// Additional operations on integers.
34pub trait NumExt {
35    /// Division with rounding up.
36    fn div_ceil(self, other: Self) -> Self;
37
38    /// Division with rounding down.
39    ///
40    /// Note this is different from truncating, which is rounding toward zero.
41    fn div_floor(self, other: Self) -> Self;
42}
43
44macro_rules! impl_ops {
45    ($type:ty) => {
46        impl NumExt for $type {
47            fn div_ceil(self, other: Self) -> Self {
48                // TODO: The source is taken from the standard library nightly implementation of
49                // this function and div_floor. Once they are stabilized, these definitions
50                // can be removed. Tracking issue: https://github.com/rust-lang/rust/issues/88581
51                let d = self / other;
52                let r = self % other;
53                if (r > 0 && other > 0) || (r < 0 && other < 0) {
54                    d + 1
55                } else {
56                    d
57                }
58            }
59
60            fn div_floor(self, other: Self) -> Self {
61                // TODO: See todo in `div_ceil`.
62                let d = self / other;
63                let r = self % other;
64                if (r > 0 && other < 0) || (r < 0 && other > 0) {
65                    d - 1
66                } else {
67                    d
68                }
69            }
70        }
71
72        impl IntExt<$type> {
73            /// Division with rounding _up_, computed exactly with integer arithmetic.
74            ///
75            /// Returns `None` if both operands are infinite.
76            pub fn div_ceil(&self, other: IntExt<$type>) -> Option<IntExt<$type>> {
77                use IntExt::*;
78
79                match (*self, other) {
80                    (Int(n), Int(d)) => Some(Int(<$type as NumExt>::div_ceil(n, d))),
81
82                    // A finite value divided by an unboundedly large denominator approaches, but
83                    // for integers never exceeds, zero.
84                    (Int(_), NegativeInf | PositiveInf) => Some(Int(0)),
85
86                    (PositiveInf, Int(d)) => {
87                        if d > 0 {
88                            Some(PositiveInf)
89                        } else {
90                            Some(NegativeInf)
91                        }
92                    }
93
94                    (NegativeInf, Int(d)) => {
95                        if d > 0 {
96                            Some(NegativeInf)
97                        } else {
98                            Some(PositiveInf)
99                        }
100                    }
101
102                    (NegativeInf | PositiveInf, NegativeInf | PositiveInf) => None,
103                }
104            }
105
106            /// Division with rounding _down_, computed exactly with integer arithmetic.
107            ///
108            /// Returns `None` if both operands are infinite.
109            pub fn div_floor(&self, other: IntExt<$type>) -> Option<IntExt<$type>> {
110                use IntExt::*;
111
112                match (*self, other) {
113                    (Int(n), Int(d)) => Some(Int(<$type as NumExt>::div_floor(n, d))),
114
115                    (Int(_), NegativeInf | PositiveInf) => Some(Int(0)),
116
117                    (PositiveInf, Int(d)) => {
118                        if d > 0 {
119                            Some(PositiveInf)
120                        } else {
121                            Some(NegativeInf)
122                        }
123                    }
124
125                    (NegativeInf, Int(d)) => {
126                        if d > 0 {
127                            Some(NegativeInf)
128                        } else {
129                            Some(PositiveInf)
130                        }
131                    }
132
133                    (NegativeInf | PositiveInf, NegativeInf | PositiveInf) => None,
134                }
135            }
136        }
137    };
138}
139
140impl_ops!(i32);
141impl_ops!(i64);
142
143impl<Int: Into<f64>> From<IntExt<Int>> for f64 {
144    fn from(value: IntExt<Int>) -> Self {
145        match value {
146            IntExt::Int(inner) => inner.into(),
147            IntExt::NegativeInf => -f64::INFINITY,
148            IntExt::PositiveInf => f64::INFINITY,
149        }
150    }
151}
152
153impl From<i32> for IntExt {
154    fn from(value: i32) -> Self {
155        IntExt::Int(value)
156    }
157}
158
159impl From<IntExt<i32>> for IntExt<i64> {
160    fn from(value: IntExt<i32>) -> Self {
161        match value {
162            IntExt::Int(int) => IntExt::Int(int.into()),
163            IntExt::NegativeInf => IntExt::NegativeInf,
164            IntExt::PositiveInf => IntExt::PositiveInf,
165        }
166    }
167}
168
169// TODO: This is not a great pattern, but for now I do not want to touch this.
170impl TryInto<i32> for IntExt {
171    type Error = ();
172
173    fn try_into(self) -> Result<i32, Self::Error> {
174        match self {
175            IntExt::Int(inner) => Ok(inner),
176            IntExt::NegativeInf | IntExt::PositiveInf => Err(()),
177        }
178    }
179}
180
181impl<Int: PartialEq> PartialEq<Int> for IntExt<Int> {
182    fn eq(&self, other: &Int) -> bool {
183        match self {
184            IntExt::Int(v1) => v1 == other,
185            IntExt::NegativeInf | IntExt::PositiveInf => false,
186        }
187    }
188}
189
190impl PartialEq<IntExt> for i32 {
191    fn eq(&self, other: &IntExt) -> bool {
192        other.eq(self)
193    }
194}
195
196impl PartialOrd<IntExt> for i32 {
197    fn partial_cmp(&self, other: &IntExt) -> Option<Ordering> {
198        other.neg().partial_cmp(&self.neg())
199    }
200}
201
202impl<Int: Ord> PartialOrd for IntExt<Int> {
203    fn partial_cmp(&self, other: &IntExt<Int>) -> Option<Ordering> {
204        Some(self.cmp(other))
205    }
206}
207
208impl<Int: Ord> Ord for IntExt<Int> {
209    fn cmp(&self, other: &Self) -> Ordering {
210        match self {
211            IntExt::Int(v1) => match other {
212                IntExt::Int(v2) => v1.cmp(v2),
213                IntExt::NegativeInf => Ordering::Greater,
214                IntExt::PositiveInf => Ordering::Less,
215            },
216            IntExt::NegativeInf => match other {
217                IntExt::Int(_) => Ordering::Less,
218                IntExt::PositiveInf => Ordering::Less,
219                IntExt::NegativeInf => Ordering::Equal,
220            },
221            IntExt::PositiveInf => match other {
222                IntExt::Int(_) => Ordering::Greater,
223                IntExt::NegativeInf => Ordering::Greater,
224                IntExt::PositiveInf => Ordering::Greater,
225            },
226        }
227    }
228}
229
230impl PartialOrd<i32> for IntExt {
231    fn partial_cmp(&self, other: &i32) -> Option<Ordering> {
232        match self {
233            IntExt::Int(v1) => v1.partial_cmp(other),
234            IntExt::NegativeInf => Some(Ordering::Less),
235            IntExt::PositiveInf => Some(Ordering::Greater),
236        }
237    }
238}
239
240impl PartialOrd<i64> for IntExt<i64> {
241    fn partial_cmp(&self, other: &i64) -> Option<Ordering> {
242        match self {
243            IntExt::Int(v1) => v1.partial_cmp(other),
244            IntExt::NegativeInf => Some(Ordering::Less),
245            IntExt::PositiveInf => Some(Ordering::Greater),
246        }
247    }
248}
249
250impl Add<i32> for IntExt {
251    type Output = IntExt;
252
253    fn add(self, rhs: i32) -> Self::Output {
254        self + IntExt::Int(rhs)
255    }
256}
257
258impl<Int: Add<Output = Int> + Debug> Add for IntExt<Int> {
259    type Output = IntExt<Int>;
260
261    fn add(self, rhs: IntExt<Int>) -> Self::Output {
262        match (self, rhs) {
263            (IntExt::Int(lhs), IntExt::Int(rhs)) => IntExt::Int(lhs + rhs),
264
265            (IntExt::Int(_), Self::NegativeInf) => Self::NegativeInf,
266            (IntExt::Int(_), Self::PositiveInf) => Self::PositiveInf,
267            (Self::NegativeInf, IntExt::Int(_)) => Self::NegativeInf,
268            (Self::PositiveInf, IntExt::Int(_)) => Self::PositiveInf,
269
270            (IntExt::NegativeInf, IntExt::NegativeInf) => IntExt::NegativeInf,
271            (IntExt::PositiveInf, IntExt::PositiveInf) => IntExt::PositiveInf,
272
273            (lhs @ IntExt::NegativeInf, rhs @ IntExt::PositiveInf)
274            | (lhs @ IntExt::PositiveInf, rhs @ IntExt::NegativeInf) => {
275                panic!("the result of {lhs:?} + {rhs:?} is indeterminate")
276            }
277        }
278    }
279}
280
281impl Sub<IntExt<i64>> for i64 {
282    type Output = IntExt<i64>;
283
284    fn sub(self, rhs: IntExt<i64>) -> Self::Output {
285        IntExt::Int(self) - rhs
286    }
287}
288
289impl<Int: Sub<Output = Int> + Debug> Sub for IntExt<Int> {
290    type Output = IntExt<Int>;
291
292    fn sub(self, rhs: IntExt<Int>) -> Self::Output {
293        match (self, rhs) {
294            (IntExt::Int(lhs), IntExt::Int(rhs)) => IntExt::Int(lhs - rhs),
295
296            (IntExt::Int(_), Self::NegativeInf) => Self::PositiveInf,
297            (IntExt::Int(_), Self::PositiveInf) => Self::NegativeInf,
298            (Self::NegativeInf, IntExt::Int(_)) => Self::NegativeInf,
299            (Self::PositiveInf, IntExt::Int(_)) => Self::PositiveInf,
300
301            (lhs @ IntExt::NegativeInf, rhs @ IntExt::NegativeInf)
302            | (lhs @ IntExt::PositiveInf, rhs @ IntExt::PositiveInf)
303            | (lhs @ IntExt::NegativeInf, rhs @ IntExt::PositiveInf)
304            | (lhs @ IntExt::PositiveInf, rhs @ IntExt::NegativeInf) => {
305                panic!("the result of {lhs:?} - {rhs:?} is indeterminate")
306            }
307        }
308    }
309}
310
311impl<Int> AddAssign<Int> for IntExt<Int>
312where
313    Int: AddAssign<Int>,
314{
315    fn add_assign(&mut self, rhs: Int) {
316        match self {
317            IntExt::Int(value) => {
318                value.add_assign(rhs);
319            }
320
321            IntExt::NegativeInf | IntExt::PositiveInf => {}
322        }
323    }
324}
325
326impl Mul<i32> for IntExt {
327    type Output = IntExt;
328
329    fn mul(self, rhs: i32) -> Self::Output {
330        self * IntExt::Int(rhs)
331    }
332}
333
334impl Mul for IntExt {
335    type Output = Self;
336
337    fn mul(self, rhs: Self) -> Self::Output {
338        match (self, rhs) {
339            (IntExt::Int(lhs), IntExt::Int(rhs)) => IntExt::Int(lhs * rhs),
340
341            // Multiplication with 0 will always yield 0.
342            (IntExt::Int(0), Self::NegativeInf)
343            | (IntExt::Int(0), Self::PositiveInf)
344            | (Self::NegativeInf, IntExt::Int(0))
345            | (Self::PositiveInf, IntExt::Int(0)) => IntExt::Int(0),
346
347            (IntExt::Int(value), IntExt::NegativeInf)
348            | (IntExt::NegativeInf, IntExt::Int(value)) => {
349                if value >= 0 {
350                    IntExt::NegativeInf
351                } else {
352                    IntExt::PositiveInf
353                }
354            }
355
356            (IntExt::Int(value), IntExt::PositiveInf)
357            | (IntExt::PositiveInf, IntExt::Int(value)) => {
358                if value >= 0 {
359                    IntExt::PositiveInf
360                } else {
361                    IntExt::NegativeInf
362                }
363            }
364
365            (IntExt::NegativeInf, IntExt::NegativeInf)
366            | (IntExt::PositiveInf, IntExt::PositiveInf) => IntExt::PositiveInf,
367
368            (IntExt::NegativeInf, IntExt::PositiveInf)
369            | (IntExt::PositiveInf, IntExt::NegativeInf) => IntExt::NegativeInf,
370        }
371    }
372}
373
374impl Mul for IntExt<i64> {
375    type Output = IntExt<i64>;
376
377    fn mul(self, rhs: Self) -> Self::Output {
378        match (self, rhs) {
379            (IntExt::Int(lhs), IntExt::Int(rhs)) => IntExt::Int(lhs * rhs),
380
381            // Multiplication with 0 will always yield 0.
382            (IntExt::Int(0), Self::NegativeInf)
383            | (IntExt::Int(0), Self::PositiveInf)
384            | (Self::NegativeInf, IntExt::Int(0))
385            | (Self::PositiveInf, IntExt::Int(0)) => IntExt::Int(0),
386
387            (IntExt::Int(value), IntExt::NegativeInf)
388            | (IntExt::NegativeInf, IntExt::Int(value)) => {
389                if value >= 0 {
390                    IntExt::NegativeInf
391                } else {
392                    IntExt::PositiveInf
393                }
394            }
395
396            (IntExt::Int(value), IntExt::PositiveInf)
397            | (IntExt::PositiveInf, IntExt::Int(value)) => {
398                if value >= 0 {
399                    IntExt::PositiveInf
400                } else {
401                    IntExt::NegativeInf
402                }
403            }
404
405            (IntExt::NegativeInf, IntExt::NegativeInf)
406            | (IntExt::PositiveInf, IntExt::PositiveInf) => IntExt::PositiveInf,
407
408            (IntExt::NegativeInf, IntExt::PositiveInf)
409            | (IntExt::PositiveInf, IntExt::NegativeInf) => IntExt::NegativeInf,
410        }
411    }
412}
413
414impl Neg for IntExt {
415    type Output = Self;
416
417    fn neg(self) -> Self::Output {
418        match self {
419            IntExt::Int(value) => IntExt::Int(-value),
420            IntExt::NegativeInf => IntExt::PositiveInf,
421            IntExt::PositiveInf => Self::NegativeInf,
422        }
423    }
424}
425
426impl Sum for IntExt {
427    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
428        iter.fold(IntExt::Int(0), |acc, value| acc + value)
429    }
430}
431
432impl Sum for IntExt<i64> {
433    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
434        iter.fold(IntExt::Int(0), |acc, value| acc + value)
435    }
436}
437
438#[cfg(test)]
439mod tests {
440    use IntExt::*;
441
442    use super::*;
443
444    #[test]
445    fn ordering_of_i32_with_i32_ext() {
446        assert!(Int(2) < 3);
447        assert!(Int(-1) < 3);
448        assert!(Int(-10) < -1);
449    }
450
451    #[test]
452    fn ordering_of_i32_ext_with_i32() {
453        assert!(1 < Int(2));
454        assert!(-10 < Int(-1));
455        assert!(-11 < Int(-10));
456    }
457
458    #[test]
459    fn test_adding_i32s() {
460        assert_eq!(Int(3) + Int(4), Int(7));
461    }
462
463    #[test]
464    fn test_adding_negative_inf() {
465        assert_eq!(Int(3) + NegativeInf, NegativeInf);
466    }
467
468    #[test]
469    fn test_adding_positive_inf() {
470        assert_eq!(Int(3) + PositiveInf, PositiveInf);
471    }
472
473    #[test]
474    fn multiplying_i64s() {
475        let a: IntExt<i64> = Int(6);
476        let b: IntExt<i64> = Int(-2);
477        assert_eq!(a * b, Int(-12));
478    }
479
480    #[test]
481    fn multiplying_i64_zero_with_infinity_is_zero() {
482        let zero: IntExt<i64> = Int(0);
483        assert_eq!(zero * IntExt::<i64>::PositiveInf, Int(0));
484        assert_eq!(IntExt::<i64>::NegativeInf * zero, Int(0));
485    }
486
487    #[test]
488    fn multiplying_i64_large_products_do_not_overflow() {
489        let a: IntExt<i64> = Int(i32::MAX as i64);
490        let b: IntExt<i64> = Int(i32::MAX as i64);
491        assert_eq!(a * b, Int(i32::MAX as i64 * i32::MAX as i64));
492    }
493
494    #[test]
495    fn dividing_i64s_exactly() {
496        assert_eq!(Int(7_i64).div_ceil(Int(2)), Some(Int(4)));
497        assert_eq!(Int(7_i64).div_floor(Int(2)), Some(Int(3)));
498        assert_eq!(Int(-7_i64).div_ceil(Int(2)), Some(Int(-3)));
499        assert_eq!(Int(-7_i64).div_floor(Int(2)), Some(Int(-4)));
500    }
501
502    #[test]
503    fn dividing_i64s_exceeding_f64_precision() {
504        // `i32::MAX * i32::MAX` is well past `f64`'s exact-integer range (2^53), so a
505        // float-based division would round this incorrectly.
506        let huge = i32::MAX as i64 * i32::MAX as i64;
507        assert_eq!(
508            Int(huge).div_floor(Int(i32::MAX as i64)),
509            Some(Int(i32::MAX as i64))
510        );
511    }
512
513    #[test]
514    fn dividing_i64_finite_by_infinite_is_zero() {
515        assert_eq!(Int(5_i64).div_ceil(PositiveInf), Some(Int(0)));
516        assert_eq!(Int(-5_i64).div_floor(NegativeInf), Some(Int(0)));
517    }
518
519    #[test]
520    fn dividing_i64_infinite_by_finite_propagates_sign() {
521        assert_eq!(
522            IntExt::<i64>::PositiveInf.div_ceil(Int(2)),
523            Some(PositiveInf)
524        );
525        assert_eq!(
526            IntExt::<i64>::PositiveInf.div_ceil(Int(-2)),
527            Some(NegativeInf)
528        );
529    }
530
531    #[test]
532    fn dividing_i64_infinite_by_infinite_is_indeterminate() {
533        assert_eq!(IntExt::<i64>::PositiveInf.div_ceil(PositiveInf), None);
534        assert_eq!(IntExt::<i64>::NegativeInf.div_floor(PositiveInf), None);
535    }
536}