Optimize int parsing based on bytes rather than str

Lib/test/test_long.py

@@ -378,7 +378,6 @@ class LongTest(unittest.TestCase):
         # See bpo-34087
         self.assertRaises(ValueError, int, '\u3053\u3093\u306b\u3061\u306f')
 
-    @unittest.expectedFailure # TODO: RUSTPYTHON
     def test_long_a(self):
         self.assertRaises(ValueError, int, '-012395', 0)
 

vm/src/bytesinner.rs

@@ -124,7 +124,7 @@ impl ByteInnerNewOptions {
 
                         let mut data_bytes = vec![];
                         for elem in elements {
-                            let v = objint::to_int(vm, &elem, 10)?;
+                            let v = objint::to_int(vm, &elem)?;
                             if let Some(i) = v.to_u8() {
                                 data_bytes.push(i);
                             } else {

vm/src/obj/objint.rs

@@ -1,9 +1,10 @@
 use std::fmt;
 use std::mem::size_of;
 
-use num_bigint::{BigInt, Sign};
+use bstr::ByteSlice;
+use num_bigint::{BigInt, BigUint, Sign};
 use num_integer::Integer;
-use num_traits::{Num, One, Pow, Signed, ToPrimitive, Zero};
+use num_traits::{One, Pow, Signed, ToPrimitive, Zero};
 
 use super::objbool::IntoPyBool;
 use super::objbytearray::PyByteArray;
@@ -258,13 +259,13 @@ impl PyInt {
                     })?
                     .as_bigint()
                     .to_u32()
-                    .filter(|v| (2..=36).contains(v))
+                    .filter(|&v| v == 0 || (2..=36).contains(&v))
                     .ok_or_else(|| {
                         vm.new_value_error("int() base must be >= 2 and <= 36, or 0".to_owned())
                     })?;
-                to_int(vm, &val, base)
+                to_int_radix(vm, &val, base)
             } else {
-                to_int(vm, &val, 10)
+                to_int(vm, &val)
             }
         } else if let OptionalArg::Present(_) = options.base {
             Err(vm.new_type_error("int() missing string argument".to_owned()))
@@ -752,34 +753,30 @@ struct IntToByteArgs {
 }
 
 // Casting function:
-pub fn to_int(vm: &VirtualMachine, obj: &PyObjectRef, base: u32) -> PyResult<BigInt> {
-    debug_assert!((2..=36).contains(&base));
-    let bytes_to_int = |bytes: &[u8]| {
-        std::str::from_utf8(bytes)
-            .ok()
-            .and_then(|s| str_to_int(s, base))
-    };
-
+pub(crate) fn to_int(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<BigInt> {
+    let base = 10;
     let opt = match_class!(match obj.clone() {
         string @ PyString => {
             let s = string.as_str();
-            str_to_int(&s, base)
+            bytes_to_int(s.as_bytes(), base)
         }
         bytes @ PyBytes => {
             let bytes = bytes.get_value();
-            bytes_to_int(bytes)
+            bytes_to_int(bytes, base)
         }
         bytearray @ PyByteArray => {
             let inner = bytearray.borrow_value();
-            bytes_to_int(&inner.elements)
+            bytes_to_int(&inner.elements, base)
         }
         memoryview @ PyMemoryView => {
             // TODO: proper error handling instead of `unwrap()`
-            memoryview.try_bytes(|bytes| bytes_to_int(&bytes)).unwrap()
+            memoryview
+                .try_bytes(|bytes| bytes_to_int(&bytes, base))
+                .unwrap()
         }
         array @ PyArray => {
             let bytes = array.tobytes();
-            bytes_to_int(&bytes)
+            bytes_to_int(&bytes, base)
         }
         obj => {
             let method = vm.get_method_or_type_error(obj.clone(), "__int__", || {
@@ -808,79 +805,158 @@ pub fn to_int(vm: &VirtualMachine, obj: &PyObjectRef, base: u32) -> PyResult<Big
     }
 }
 
-fn str_to_int(literal: &str, mut base: u32) -> Option<BigInt> {
-    let mut buf = validate_literal(literal)?.to_owned();
-    let is_signed = buf.starts_with('+') || buf.starts_with('-');
-    let radix_range = if is_signed { 1..3 } else { 0..2 };
-    let radix_candidate = buf.get(radix_range.clone());
+fn to_int_radix(vm: &VirtualMachine, obj: &PyObjectRef, base: u32) -> PyResult<BigInt> {
+    debug_assert!(base == 0 || (2..=36).contains(&base));
 
-    // try to find base
-    if let Some(radix_candidate) = radix_candidate {
-        if let Some(matched_radix) = detect_base(&radix_candidate) {
-            if base == 0 || base == matched_radix {
-                /* If base is 0 or equal radix number, it means radix is validate
-                 * So change base to radix number and remove radix from literal
-                 */
-                base = matched_radix;
-                buf.drain(radix_range);
+    let opt = match_class!(match obj.clone() {
+        string @ PyString => {
+            let s = string.as_str();
+            bytes_to_int(s.as_bytes(), base)
+        }
+        bytes @ PyBytes => {
+            let bytes = bytes.get_value();
+            bytes_to_int(bytes, base)
+        }
+        bytearray @ PyByteArray => {
+            let inner = bytearray.borrow_value();
+            bytes_to_int(&inner.elements, base)
+        }
+        _ => {
+            unreachable!("this case must not enter this function");
+        }
+    });
+    match opt {
+        Some(int) => Ok(int),
+        None => Err(vm.new_value_error(format!(
+            "invalid literal for int() with base {}: {}",
+            base,
+            vm.to_repr(obj)?,
+        ))),
+    }
+}
 
-                /* first underscore with radix is validate
-                 * e.g : int(`0x_1`, base=0) = int('1', base=16)
-                 */
-                if buf.starts_with('_') {
-                    buf.remove(0);
+fn bytes_to_int(lit: &[u8], mut base: u32) -> Option<BigInt> {
+    // split sign
+    let mut lit = lit.trim();
+    let sign = match lit.first()? {
+        b'+' => Some(Sign::Plus),
+        b'-' => Some(Sign::Minus),
+        _ => None,
+    };
+    if sign.is_some() {
+        lit = unsafe {
+            // safe by sign check
+            lit.get_unchecked(1..)
+        };
+    }
+
+    // split radix
+    let first = *lit.first()?;
+    let has_radix = if first == b'0' {
+        match base {
+            0 => {
+                if let Some(parsed) = lit.get(1).and_then(detect_base) {
+                    base = parsed;
+                    true
+                } else {
+                    if let Some((last, body)) = lit[1..].split_last() {
+                        let is_zero = body.iter().all(|&c| c == b'0' || c == b'_') && *last == b'0';
+                        if !is_zero {
+                            return None;
+                        }
+                    }
+                    return Some(BigInt::zero());
                 }
-            } else if (matched_radix == 2 && base < 12)
-                || (matched_radix == 8 && base < 25)
-                || (matched_radix == 16 && base < 34)
-            {
-                return None;
             }
+            16 => lit.get(1).map_or(false, |&b| matches!(b, b'x' | b'X')),
+            2 => lit.get(1).map_or(false, |&b| matches!(b, b'b' | b'B')),
+            8 => lit.get(1).map_or(false, |&b| matches!(b, b'o' | b'O')),
+            _ => false,
+        }
+    } else {
+        if base == 0 {
+            base = 10;
+        }
+        false
+    };
+    if has_radix {
+        lit = unsafe {
+            // safe by radix check
+            lit.get_unchecked(2..)
+        };
+        if lit.first() == Some(&b'_') {
+            lit = unsafe {
+                // safe by b'_' check
+                lit.get_unchecked(1..)
+            };
         }
     }
 
-    // base still not found, try to use default
-    if base == 0 {
-        if buf.starts_with('0') {
-            if buf.chars().all(|c| matches!(c, '+' | '-' | '0' | '_')) {
-                return Some(BigInt::zero());
+    // remove zeroes
+    let mut last = *lit.first()?;
+    if last == b'0' {
+        let mut count = 0;
+        for pair in lit.windows(2) {
+            let c1 = pair[0];
+            let c2 = pair[1];
+            let skip = if c2 == b'_' {
+                if c1 == b'_' {
+                    return None;
+                }
+                true
+            } else {
+                c2 == b'0'
+            };
+            if skip {
+                count += 1;
+            } else {
+                break;
             }
+        }
+        let prefix_last = *unsafe {
+            // safe by upper loop: count < lit.len() - 1
+            lit.get_unchecked(count)
+        };
+        lit = unsafe {
+            // safe by upper loop: count < lit.len() - 1
+            lit.get_unchecked(count + 1..)
+        };
+        if lit.is_empty() && prefix_last == b'_' {
+            return None;
+        }
+    }
+
+    // validate
+    for c in lit.iter() {
+        let c = *c;
+        if !(c.is_ascii_alphanumeric() || c == b'_') {
             return None;
         }
 
-        base = 10;
+        if c == b'_' && last == b'_' {
+            return None;
+        }
+
+        last = c;
     }
-
-    BigInt::from_str_radix(&buf, base).ok()
-}
-
-fn validate_literal(literal: &str) -> Option<&str> {
-    let trimmed = literal.trim();
-    if trimmed.starts_with('_') || trimmed.ends_with('_') {
+    if last == b'_' {
         return None;
     }
 
-    let mut last_tok = None;
-    for c in trimmed.chars() {
-        if !(c.is_ascii_alphanumeric() || c == '_' || c == '+' || c == '-') {
-            return None;
-        }
-
-        if c == '_' && Some(c) == last_tok {
-            return None;
-        }
-
-        last_tok = Some(c);
-    }
-
-    Some(trimmed)
+    // parse
+    Some(if lit.is_empty() {
+        BigInt::zero()
+    } else {
+        let uint = BigUint::parse_bytes(&lit, base)?;
+        BigInt::from_biguint(sign.unwrap_or(Sign::Plus), uint)
+    })
 }
 
-fn detect_base(literal: &str) -> Option<u32> {
-    match literal {
-        "0x" | "0X" => Some(16),
-        "0o" | "0O" => Some(8),
-        "0b" | "0B" => Some(2),
+fn detect_base(c: &u8) -> Option<u32> {
+    match c {
+        b'x' | b'X' => Some(16),
+        b'b' | b'B' => Some(2),
+        b'o' | b'O' => Some(8),
         _ => None,
     }
 }
@@ -909,6 +985,19 @@ fn get_shift_amount(amount: &BigInt, vm: &VirtualMachine) -> PyResult<usize> {
     }
 }
 
-pub fn init(context: &PyContext) {
+pub(crate) fn init(context: &PyContext) {
     PyInt::extend_class(context, &context.types.int_type);
 }
+
+#[test]
+fn test_bytes_to_int() {
+    assert_eq!(bytes_to_int(&b"0b101"[..], 2).unwrap(), BigInt::from(5));
+    assert_eq!(bytes_to_int(&b"0x_10"[..], 16).unwrap(), BigInt::from(16));
+    assert_eq!(bytes_to_int(&b"0b"[..], 16).unwrap(), BigInt::from(11));
+    assert_eq!(bytes_to_int(&b"+0b101"[..], 2).unwrap(), BigInt::from(5));
+    assert_eq!(bytes_to_int(&b"0_0_0"[..], 10).unwrap(), BigInt::from(0));
+    assert_eq!(bytes_to_int(&b"09_99"[..], 0), None);
+    assert_eq!(bytes_to_int(&b"000"[..], 0).unwrap(), BigInt::from(0));
+    assert_eq!(bytes_to_int(&b"0_"[..], 0), None);
+    assert_eq!(bytes_to_int(&b"0_100"[..], 10).unwrap(), BigInt::from(100));
+}