Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
pola-rs
GitHub Repository: pola-rs/polars
Path: blob/main/crates/polars-utils/src/floor_divmod.rs
8420 views
1
use crate::float16::pf16;
2
3
pub 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
11
macro_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
24
macro_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
if other == 0 {
30
return (0, 0);
31
}
32
33
(self / other, self % other)
34
}
35
}
36
};
37
}
38
39
macro_rules! impl_signed_div_mod {
40
($T:ty) => {
41
impl FloorDivMod for $T {
42
#[inline]
43
fn wrapping_floor_div_mod(self, other: Self) -> (Self, Self) {
44
if other == 0 {
45
return (0, 0);
46
}
47
48
// Rust/C-style remainder is in the correct congruence
49
// class, but may not have the right sign. We want a
50
// remainder with the same sign as the RHS, which we
51
// can get by adding RHS to the remainder if the sign of
52
// the non-zero remainder differs from our RHS.
53
//
54
// Similarly, Rust/C-style division truncates instead of floors.
55
// If the remainder was non-zero and the signs were different
56
// (we'd have a negative result before truncating), we need to
57
// subtract 1 from the result.
58
let mut div = self.wrapping_div(other);
59
let mut mod_ = self.wrapping_rem(other);
60
if mod_ != 0 && (self < 0) != (other < 0) {
61
div -= 1;
62
mod_ += other;
63
}
64
(div, mod_)
65
}
66
}
67
};
68
}
69
70
impl_unsigned_div_mod!(u8);
71
impl_unsigned_div_mod!(u16);
72
impl_unsigned_div_mod!(u32);
73
impl_unsigned_div_mod!(u64);
74
impl_unsigned_div_mod!(u128);
75
impl_unsigned_div_mod!(usize);
76
impl_signed_div_mod!(i8);
77
impl_signed_div_mod!(i16);
78
impl_signed_div_mod!(i32);
79
impl_signed_div_mod!(i64);
80
impl_signed_div_mod!(i128);
81
impl_signed_div_mod!(isize);
82
impl_float_div_mod!(pf16);
83
impl_float_div_mod!(f32);
84
impl_float_div_mod!(f64);
85
86
#[cfg(test)]
87
mod test {
88
use super::*;
89
90
#[test]
91
fn test_signed_wrapping_div_mod() {
92
// Test for all i8, should transfer to other values.
93
for lhs in i8::MIN..=i8::MAX {
94
for rhs in i8::MIN..=i8::MAX {
95
let ans = if rhs != 0 {
96
let fdiv = (lhs as f64 / rhs as f64).floor();
97
let fmod = lhs as f64 - rhs as f64 * fdiv;
98
99
// float -> int conversion saturates, we want wrapping, double convert.
100
((fdiv as i32) as i8, (fmod as i32) as i8)
101
} else {
102
(0, 0)
103
};
104
105
assert_eq!(lhs.wrapping_floor_div_mod(rhs), ans);
106
}
107
}
108
}
109
}
110
111