Skip to main content

rustpython_common/
int.rs

1use malachite_base::{num::conversion::traits::RoundingInto, rounding_modes::RoundingMode};
2use malachite_bigint::{BigInt, BigUint, Sign};
3use malachite_q::Rational;
4use num_traits::{One, ToPrimitive, Zero};
5
6#[must_use]
7pub fn true_div(numerator: &BigInt, denominator: &BigInt) -> f64 {
8    let rational = Rational::from_integers_ref(numerator.into(), denominator.into());
9    match rational.rounding_into(RoundingMode::Nearest) {
10        // returned value is $t::MAX but still less than the original
11        (val, core::cmp::Ordering::Less) if val == f64::MAX => f64::INFINITY,
12        // returned value is $t::MIN but still greater than the original
13        (val, core::cmp::Ordering::Greater) if val == f64::MIN => f64::NEG_INFINITY,
14        (val, _) => val,
15    }
16}
17
18#[must_use]
19pub fn float_to_ratio(value: f64) -> Option<(BigInt, BigInt)> {
20    let sign = match core::cmp::PartialOrd::partial_cmp(&value, &0.0)? {
21        core::cmp::Ordering::Less => Sign::Minus,
22        core::cmp::Ordering::Equal => return Some((BigInt::zero(), BigInt::one())),
23        core::cmp::Ordering::Greater => Sign::Plus,
24    };
25    Rational::try_from(value).ok().map(|x| {
26        let (numer, denom) = x.into_numerator_and_denominator();
27        (
28            BigInt::from_biguint(sign, numer.into()),
29            BigUint::from(denom).into(),
30        )
31    })
32}
33
34#[derive(Copy, Clone, Debug, Eq, PartialEq)]
35pub enum BytesToIntError {
36    InvalidLiteral { base: u32 },
37    InvalidBase,
38    DigitLimit { got: usize, limit: usize },
39}
40
41// https://github.com/python/cpython/blob/4e665351082c50018fb31d80db25b4693057393e/Objects/longobject.c#L2977
42// https://github.com/python/cpython/blob/4e665351082c50018fb31d80db25b4693057393e/Objects/longobject.c#L2884
43pub fn bytes_to_int(
44    buf: &[u8],
45    mut base: u32,
46    digit_limit: usize,
47) -> Result<BigInt, BytesToIntError> {
48    if base != 0 && !(2..=36).contains(&base) {
49        return Err(BytesToIntError::InvalidBase);
50    }
51    // A rejected literal names the base that was asked for, not the one a
52    // base of 0 turned out to mean.
53    let requested_base = base;
54
55    let mut buf = buf.trim_ascii();
56
57    // split sign
58    let sign = match buf.first() {
59        Some(b'+') => Some(Sign::Plus),
60        Some(b'-') => Some(Sign::Minus),
61        None => {
62            return Err(BytesToIntError::InvalidLiteral {
63                base: requested_base,
64            });
65        }
66        _ => None,
67    };
68
69    if sign.is_some() {
70        buf = &buf[1..];
71    }
72
73    let mut error_if_nonzero = false;
74    if base == 0 {
75        match (buf.first(), buf.get(1)) {
76            (Some(v), _) if *v != b'0' => base = 10,
77            // A sign on its own leaves nothing to read, which is not the
78            // old octal form either.
79            (None, _) => base = 10,
80            (_, Some(b'x' | b'X')) => base = 16,
81            (_, Some(b'o' | b'O')) => base = 8,
82            (_, Some(b'b' | b'B')) => base = 2,
83            (_, _) => {
84                // "old" (C-style) octal literal, now invalid. it might still be zero though
85                base = 10;
86                error_if_nonzero = true;
87            }
88        }
89    }
90
91    if error_if_nonzero {
92        if let [_first, others @ .., last] = buf {
93            let is_zero = *last == b'0' && others.iter().all(|&c| c == b'0' || c == b'_');
94            if !is_zero {
95                return Err(BytesToIntError::InvalidLiteral {
96                    base: requested_base,
97                });
98            }
99        }
100        return Ok(BigInt::zero());
101    }
102
103    if buf.first().is_some_and(|&v| v == b'0')
104        && buf.get(1).is_some_and(|&v| {
105            (base == 16 && (v == b'x' || v == b'X'))
106                || (base == 8 && (v == b'o' || v == b'O'))
107                || (base == 2 && (v == b'b' || v == b'B'))
108        })
109    {
110        buf = &buf[2..];
111
112        // One underscore allowed here
113        if buf.first().is_some_and(|&v| v == b'_') {
114            buf = &buf[1..];
115        }
116    }
117
118    // A pure digit string longer than the limit can be rejected from its
119    // length. The check is a wide digit test so it stays cheaper than parsing.
120    if digit_limit > 0
121        && !base.is_power_of_two()
122        && buf.len() > digit_limit
123        && let Some(digits) = plain_ascii_digit_len(buf)
124    {
125        return Err(BytesToIntError::DigitLimit {
126            got: digits,
127            limit: digit_limit,
128        });
129    }
130
131    // Reject empty strings
132    let mut prev = *buf.first().ok_or(BytesToIntError::InvalidLiteral {
133        base: requested_base,
134    })?;
135
136    // Leading underscore not allowed
137    if prev == b'_' || !prev.is_ascii_alphanumeric() {
138        return Err(BytesToIntError::InvalidLiteral {
139            base: requested_base,
140        });
141    }
142
143    // Verify all characters are digits and underscores
144    let mut digits = 1;
145    for &cur in buf.iter().skip(1) {
146        if cur == b'_' {
147            // Double underscore not allowed
148            if prev == b'_' {
149                return Err(BytesToIntError::InvalidLiteral {
150                    base: requested_base,
151                });
152            }
153        } else if cur.is_ascii_alphanumeric() {
154            digits += 1;
155        } else {
156            return Err(BytesToIntError::InvalidLiteral {
157                base: requested_base,
158            });
159        }
160
161        prev = cur;
162    }
163
164    // Trailing underscore not allowed
165    if prev == b'_' {
166        return Err(BytesToIntError::InvalidLiteral {
167            base: requested_base,
168        });
169    }
170
171    if digit_limit > 0 && !base.is_power_of_two() && digits > digit_limit {
172        return Err(BytesToIntError::DigitLimit {
173            got: digits,
174            limit: digit_limit,
175        });
176    }
177
178    let uint = BigUint::parse_bytes(buf, base).ok_or(BytesToIntError::InvalidLiteral {
179        base: requested_base,
180    })?;
181    Ok(BigInt::from_biguint(sign.unwrap_or(Sign::Plus), uint))
182}
183
184/// `Some(buf.len())` when every byte is an ASCII digit.
185fn plain_ascii_digit_len(buf: &[u8]) -> Option<usize> {
186    let mut rest = buf;
187    while rest.len() >= 8 {
188        let (head, tail) = rest.split_at(8);
189        let word = u64::from_le_bytes(head.try_into().unwrap());
190        if !word_all_ascii_digits(word) {
191            return None;
192        }
193        rest = tail;
194    }
195    rest.iter().all(u8::is_ascii_digit).then_some(buf.len())
196}
197
198/// True when each byte of `word` is in `b'0'..=b'9'`.
199fn word_all_ascii_digits(word: u64) -> bool {
200    let below = word.wrapping_sub(0x3030_3030_3030_3030);
201    let above = 0x3939_3939_3939_3939u64.wrapping_sub(word);
202    (below | above) & 0x8080_8080_8080_8080 == 0
203}
204
205// num-bigint now returns Some(inf) for to_f64() in some cases, so just keep that the same for now
206#[inline(always)]
207#[must_use]
208pub fn bigint_to_finite_float(int: &BigInt) -> Option<f64> {
209    int.to_f64().filter(|f| f.is_finite())
210}
211
212#[cfg(test)]
213mod tests {
214    use super::*;
215
216    const DIGIT_LIMIT: usize = 4300; // Default of Cpython
217
218    #[test]
219    fn bytes_to_int_valid() {
220        for ((buf, base), expected) in [
221            (("0b101", 2), BigInt::from(5)),
222            (("0x_10", 16), BigInt::from(16)),
223            (("0b", 16), BigInt::from(11)),
224            (("+0b101", 2), BigInt::from(5)),
225            (("0_0_0", 10), BigInt::from(0)),
226            (("000", 0), BigInt::from(0)),
227            (("0_100", 10), BigInt::from(100)),
228        ] {
229            assert_eq!(
230                bytes_to_int(buf.as_bytes(), base, DIGIT_LIMIT),
231                Ok(expected)
232            );
233        }
234    }
235
236    #[test]
237    fn bytes_to_int_invalid_literal() {
238        for ((buf, base), expected) in [
239            (("09_99", 0), BytesToIntError::InvalidLiteral { base: 0 }),
240            (("0_", 0), BytesToIntError::InvalidLiteral { base: 0 }),
241            (("0_", 2), BytesToIntError::InvalidLiteral { base: 2 }),
242            (("-", 0), BytesToIntError::InvalidLiteral { base: 0 }),
243            (("+", 0), BytesToIntError::InvalidLiteral { base: 0 }),
244        ] {
245            assert_eq!(
246                bytes_to_int(buf.as_bytes(), base, DIGIT_LIMIT),
247                Err(expected)
248            )
249        }
250    }
251
252    #[test]
253    fn bytes_to_int_invalid_base() {
254        for base in [1, 37] {
255            assert_eq!(
256                bytes_to_int("012345".as_bytes(), base, DIGIT_LIMIT),
257                Err(BytesToIntError::InvalidBase)
258            )
259        }
260    }
261
262    #[test]
263    fn bytes_to_int_digit_limit() {
264        assert_eq!(
265            bytes_to_int("012345".as_bytes(), 10, 5),
266            Err(BytesToIntError::DigitLimit { got: 6, limit: 5 })
267        );
268        let long = "8".repeat(20_000);
269        assert_eq!(
270            bytes_to_int(long.as_bytes(), 10, 100),
271            Err(BytesToIntError::DigitLimit {
272                got: 20_000,
273                limit: 100
274            })
275        );
276        // Underscores are not digits, so the plain-digit fast path must not count them.
277        assert_eq!(
278            bytes_to_int(b"1_1_1", 10, 2),
279            Err(BytesToIntError::DigitLimit { got: 3, limit: 2 })
280        );
281        let over = format!("{}{}", "1_".repeat(80), "2".repeat(40));
282        assert_eq!(
283            bytes_to_int(over.as_bytes(), 10, 100),
284            Err(BytesToIntError::DigitLimit {
285                got: 120,
286                limit: 100
287            })
288        );
289    }
290}