polars_utils/
floor_divmod.rs

1use crate::float16::pf16;
2
3pub trait FloorDivMod: Sized {
4    // Returns the flooring division and associated modulo of lhs / rhs.
5    // This is the same division / modulo combination as Python.
6    //
7    // Returns (0, 0) if other == 0.
8    fn wrapping_floor_div_mod(self, other: Self) -> (Self, Self);
9}
10
11macro_rules! impl_float_div_mod {
12    ($T:ty) => {
13        impl FloorDivMod for $T {
14            #[inline]
15            fn wrapping_floor_div_mod(self, other: Self) -> (Self, Self) {
16                let div = (self / other).floor();
17                let mod_ = self - other * div;
18                (div, mod_)
19            }
20        }
21    };
22}
23
24macro_rules! impl_unsigned_div_mod {
25    ($T:ty) => {
26        impl FloorDivMod for $T {
27            #[inline]
28            fn wrapping_floor_div_mod(self, other: Self) -> (Self, Self) {
29                (self / other, self % other)
30            }
31        }
32    };
33}
34
35macro_rules! impl_signed_div_mod {
36    ($T:ty) => {
37        impl FloorDivMod for $T {
38            #[inline]
39            fn wrapping_floor_div_mod(self, other: Self) -> (Self, Self) {
40                if other == 0 {
41                    return (0, 0);
42                }
43
44                // Rust/C-style remainder is in the correct congruence
45                // class, but may not have the right sign. We want a
46                // remainder with the same sign as the RHS, which we
47                // can get by adding RHS to the remainder if the sign of
48                // the non-zero remainder differs from our RHS.
49                //
50                // Similarly, Rust/C-style division truncates instead of floors.
51                // If the remainder was non-zero and the signs were different
52                // (we'd have a negative result before truncating), we need to
53                // subtract 1 from the result.
54                let mut div = self.wrapping_div(other);
55                let mut mod_ = self.wrapping_rem(other);
56                if mod_ != 0 && (self < 0) != (other < 0) {
57                    div -= 1;
58                    mod_ += other;
59                }
60                (div, mod_)
61            }
62        }
63    };
64}
65
66impl_unsigned_div_mod!(u8);
67impl_unsigned_div_mod!(u16);
68impl_unsigned_div_mod!(u32);
69impl_unsigned_div_mod!(u64);
70impl_unsigned_div_mod!(u128);
71impl_unsigned_div_mod!(usize);
72impl_signed_div_mod!(i8);
73impl_signed_div_mod!(i16);
74impl_signed_div_mod!(i32);
75impl_signed_div_mod!(i64);
76impl_signed_div_mod!(i128);
77impl_signed_div_mod!(isize);
78impl_float_div_mod!(pf16);
79impl_float_div_mod!(f32);
80impl_float_div_mod!(f64);
81
82#[cfg(test)]
83mod test {
84    use super::*;
85
86    #[test]
87    fn test_signed_wrapping_div_mod() {
88        // Test for all i8, should transfer to other values.
89        for lhs in i8::MIN..=i8::MAX {
90            for rhs in i8::MIN..=i8::MAX {
91                let ans = if rhs != 0 {
92                    let fdiv = (lhs as f64 / rhs as f64).floor();
93                    let fmod = lhs as f64 - rhs as f64 * fdiv;
94
95                    // float -> int conversion saturates, we want wrapping, double convert.
96                    ((fdiv as i32) as i8, (fmod as i32) as i8)
97                } else {
98                    (0, 0)
99                };
100
101                assert_eq!(lhs.wrapping_floor_div_mod(rhs), ans);
102            }
103        }
104    }
105}
polars_utils/floor_divmod.rs

polars_utils/
floor_divmod.rs
