From 2a0e84828e14fa4ce77b6265758a91b9683ded58 Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Fri, 7 Aug 2020 01:56:29 +0900 Subject: [PATCH] Optimize int parsing based on bytes rather than str --- Lib/test/test_long.py | 1 - vm/src/bytesinner.rs | 2 +- vm/src/obj/objint.rs | 239 +++++++++++++++++++++++++++++------------- 3 files changed, 165 insertions(+), 77 deletions(-) diff --git a/Lib/test/test_long.py b/Lib/test/test_long.py index 6615b33a4..eebb87754 100644 --- a/Lib/test/test_long.py +++ b/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) diff --git a/vm/src/bytesinner.rs b/vm/src/bytesinner.rs index 980c21cfb..de71b5318 100644 --- a/vm/src/bytesinner.rs +++ b/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 { diff --git a/vm/src/obj/objint.rs b/vm/src/obj/objint.rs index fcfd662f3..7690e3017 100644 --- a/vm/src/obj/objint.rs +++ b/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 { - 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 { + 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 Option { - 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 { + 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 { + // 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 { - match literal { - "0x" | "0X" => Some(16), - "0o" | "0O" => Some(8), - "0b" | "0B" => Some(2), +fn detect_base(c: &u8) -> Option { + 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 { } } -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)); +}