Skip to content

Commit 01eff92

Browse files
authored
Merge pull request RustPython#3688 from youknowone/int
Int cleanup
2 parents c303127 + 251ce4e commit 01eff92

File tree

3 files changed

+130
-125
lines changed

3 files changed

+130
-125
lines changed

vm/src/anystr.rs

Lines changed: 33 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
use crate::{
2-
builtins::PyIntRef,
2+
builtins::{PyIntRef, PyTupleRef},
33
cformat::CFormatString,
4-
function::{single_or_tuple_any, OptionalOption},
4+
function::OptionalOption,
55
protocol::PyIterIter,
6-
AsObject, PyObjectRef, PyResult, TryFromObject, VirtualMachine,
6+
AsObject, PyObject, PyObjectRef, PyResult, TryFromObject, VirtualMachine,
77
};
88
use num_traits::{cast::ToPrimitive, sign::Signed};
99
use std::str::FromStr;
@@ -441,3 +441,33 @@ pub trait AnyStr<'s>: 's {
441441
.format(vm, values)
442442
}
443443
}
444+
445+
/// Tests that the predicate is True on a single value, or if the value is a tuple a tuple, then
446+
/// test that any of the values contained within the tuples satisfies the predicate. Type parameter
447+
/// T specifies the type that is expected, if the input value is not of that type or a tuple of
448+
/// values of that type, then a TypeError is raised.
449+
pub fn single_or_tuple_any<T, F, M>(
450+
obj: PyObjectRef,
451+
predicate: &F,
452+
message: &M,
453+
vm: &VirtualMachine,
454+
) -> PyResult<bool>
455+
where
456+
T: TryFromObject,
457+
F: Fn(&T) -> PyResult<bool>,
458+
M: Fn(&PyObject) -> String,
459+
{
460+
match T::try_from_object(vm, obj.clone()) {
461+
Ok(single) => (predicate)(&single),
462+
Err(_) => {
463+
let tuple = PyTupleRef::try_from_object(vm, obj.clone())
464+
.map_err(|_| vm.new_type_error((message)(&obj)))?;
465+
for obj in &tuple {
466+
if single_or_tuple_any(obj.clone(), predicate, message, vm)? {
467+
return Ok(true);
468+
}
469+
}
470+
Ok(false)
471+
}
472+
}
473+
}

vm/src/builtins/int.rs

Lines changed: 79 additions & 93 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,10 @@ use crate::{
55
common::hash,
66
convert::{ToPyObject, ToPyResult},
77
format::FormatSpec,
8-
function::{ArgIntoBool, OptionalArg, OptionalOption, PyArithmeticValue, PyComparisonValue},
8+
function::{
9+
ArgByteOrder, ArgIntoBool, OptionalArg, OptionalOption, PyArithmeticValue,
10+
PyComparisonValue,
11+
},
912
types::{Comparable, Constructor, Hashable, PyComparisonOp},
1013
AsObject, Context, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromBorrowedObject,
1114
VirtualMachine,
@@ -425,62 +428,63 @@ impl PyInt {
425428
self.int_op(other, |a, b| a & b, vm)
426429
}
427430

431+
fn modpow(&self, other: PyObjectRef, modulus: PyObjectRef, vm: &VirtualMachine) -> PyResult {
432+
let modulus = match modulus.payload_if_subclass::<PyInt>(vm) {
433+
Some(val) => val.as_bigint(),
434+
None => return Ok(vm.ctx.not_implemented()),
435+
};
436+
if modulus.is_zero() {
437+
return Err(vm.new_value_error("pow() 3rd argument cannot be 0".to_owned()));
438+
}
439+
440+
self.general_op(
441+
other,
442+
|a, b| {
443+
let i = if b.is_negative() {
444+
// modular multiplicative inverse
445+
// based on rust-num/num-integer#10, should hopefully be published soon
446+
fn normalize(a: BigInt, n: &BigInt) -> BigInt {
447+
let a = a % n;
448+
if a.is_negative() {
449+
a + n
450+
} else {
451+
a
452+
}
453+
}
454+
fn inverse(a: BigInt, n: &BigInt) -> Option<BigInt> {
455+
use num_integer::*;
456+
let ExtendedGcd { gcd, x: c, .. } = a.extended_gcd(n);
457+
if gcd.is_one() {
458+
Some(normalize(c, n))
459+
} else {
460+
None
461+
}
462+
}
463+
let a = inverse(a % modulus, modulus).ok_or_else(|| {
464+
vm.new_value_error(
465+
"base is not invertible for the given modulus".to_owned(),
466+
)
467+
})?;
468+
let b = -b;
469+
a.modpow(&b, modulus)
470+
} else {
471+
a.modpow(b, modulus)
472+
};
473+
Ok(vm.ctx.new_int(i).into())
474+
},
475+
vm,
476+
)
477+
}
478+
428479
#[pymethod(magic)]
429480
fn pow(
430481
&self,
431482
other: PyObjectRef,
432-
mod_val: OptionalOption<PyObjectRef>,
483+
r#mod: OptionalOption<PyObjectRef>,
433484
vm: &VirtualMachine,
434485
) -> PyResult {
435-
match mod_val.flatten() {
436-
Some(int_ref) => {
437-
let int = match int_ref.payload_if_subclass::<PyInt>(vm) {
438-
Some(val) => val,
439-
None => return Ok(vm.ctx.not_implemented()),
440-
};
441-
442-
let modulus = int.as_bigint();
443-
if modulus.is_zero() {
444-
return Err(vm.new_value_error("pow() 3rd argument cannot be 0".to_owned()));
445-
}
446-
self.general_op(
447-
other,
448-
|a, b| {
449-
let i = if b.is_negative() {
450-
// modular multiplicative inverse
451-
// based on rust-num/num-integer#10, should hopefully be published soon
452-
fn normalize(a: BigInt, n: &BigInt) -> BigInt {
453-
let a = a % n;
454-
if a.is_negative() {
455-
a + n
456-
} else {
457-
a
458-
}
459-
}
460-
fn inverse(a: BigInt, n: &BigInt) -> Option<BigInt> {
461-
use num_integer::*;
462-
let ExtendedGcd { gcd, x: c, .. } = a.extended_gcd(n);
463-
if gcd.is_one() {
464-
Some(normalize(c, n))
465-
} else {
466-
None
467-
}
468-
}
469-
let a = inverse(a % modulus, modulus).ok_or_else(|| {
470-
vm.new_value_error(
471-
"base is not invertible for the given modulus".to_owned(),
472-
)
473-
})?;
474-
let b = -b;
475-
a.modpow(&b, modulus)
476-
} else {
477-
a.modpow(b, modulus)
478-
};
479-
Ok(vm.ctx.new_int(i).into())
480-
},
481-
vm,
482-
)
483-
}
486+
match r#mod.flatten() {
487+
Some(modulus) => self.modpow(other, modulus, vm),
484488
None => self.general_op(other, |a, b| inner_pow(a, b, vm), vm),
485489
}
486490
}
@@ -529,20 +533,13 @@ impl PyInt {
529533
match precision {
530534
OptionalArg::Missing => (),
531535
OptionalArg::Present(ref value) => {
532-
if !vm.is_none(value) {
533-
// Only accept int type ndigits
534-
let _ndigits = value.payload_if_subclass::<PyInt>(vm).ok_or_else(|| {
535-
vm.new_type_error(format!(
536-
"'{}' object cannot be interpreted as an integer",
537-
value.class().name()
538-
))
539-
})?;
540-
} else {
541-
return Err(vm.new_type_error(format!(
536+
// Only accept int type ndigits
537+
let _ndigits = value.payload_if_subclass::<PyInt>(vm).ok_or_else(|| {
538+
vm.new_type_error(format!(
542539
"'{}' object cannot be interpreted as an integer",
543540
value.class().name()
544-
)));
545-
}
541+
))
542+
})?;
546543
}
547544
}
548545
Ok(zelf)
@@ -595,12 +592,9 @@ impl PyInt {
595592

596593
#[pymethod(magic)]
597594
fn format(&self, spec: PyStrRef, vm: &VirtualMachine) -> PyResult<String> {
598-
match FormatSpec::parse(spec.as_str())
595+
FormatSpec::parse(spec.as_str())
599596
.and_then(|format_spec| format_spec.format_int(&self.value))
600-
{
601-
Ok(string) => Ok(string),
602-
Err(err) => Err(vm.new_value_error(err.to_string())),
603-
}
597+
.map_err(|msg| vm.new_value_error(msg.to_owned()))
604598
}
605599

606600
#[pymethod(magic)]
@@ -635,15 +629,12 @@ impl PyInt {
635629
vm: &VirtualMachine,
636630
) -> PyResult<PyRef<Self>> {
637631
let signed = args.signed.map_or(false, Into::into);
638-
let value = match (args.byteorder.as_str(), signed) {
639-
("big", true) => BigInt::from_signed_bytes_be(&args.bytes.elements),
640-
("big", false) => BigInt::from_bytes_be(Sign::Plus, &args.bytes.elements),
641-
("little", true) => BigInt::from_signed_bytes_le(&args.bytes.elements),
642-
("little", false) => BigInt::from_bytes_le(Sign::Plus, &args.bytes.elements),
643-
_ => {
644-
return Err(
645-
vm.new_value_error("byteorder must be either 'little' or 'big'".to_owned())
646-
)
632+
let value = match (args.byteorder, signed) {
633+
(ArgByteOrder::Big, true) => BigInt::from_signed_bytes_be(&args.bytes.elements),
634+
(ArgByteOrder::Big, false) => BigInt::from_bytes_be(Sign::Plus, &args.bytes.elements),
635+
(ArgByteOrder::Little, true) => BigInt::from_signed_bytes_le(&args.bytes.elements),
636+
(ArgByteOrder::Little, false) => {
637+
BigInt::from_bytes_le(Sign::Plus, &args.bytes.elements)
647638
}
648639
};
649640
Self::with_value(cls, value, vm)
@@ -665,16 +656,11 @@ impl PyInt {
665656
_ => {}
666657
}
667658

668-
let mut origin_bytes = match (args.byteorder.as_str(), signed) {
669-
("big", true) => value.to_signed_bytes_be(),
670-
("big", false) => value.to_bytes_be().1,
671-
("little", true) => value.to_signed_bytes_le(),
672-
("little", false) => value.to_bytes_le().1,
673-
_ => {
674-
return Err(
675-
vm.new_value_error("byteorder must be either 'little' or 'big'".to_owned())
676-
);
677-
}
659+
let mut origin_bytes = match (args.byteorder, signed) {
660+
(ArgByteOrder::Big, true) => value.to_signed_bytes_be(),
661+
(ArgByteOrder::Big, false) => value.to_bytes_be().1,
662+
(ArgByteOrder::Little, true) => value.to_signed_bytes_le(),
663+
(ArgByteOrder::Little, false) => value.to_bytes_le().1,
678664
};
679665

680666
let origin_len = origin_bytes.len();
@@ -687,21 +673,21 @@ impl PyInt {
687673
_ => vec![0u8; byte_len - origin_len],
688674
};
689675

690-
let bytes = match args.byteorder.as_str() {
691-
"big" => {
676+
let bytes = match args.byteorder {
677+
ArgByteOrder::Big => {
692678
let mut bytes = append_bytes;
693679
bytes.append(&mut origin_bytes);
694680
bytes
695681
}
696-
"little" => {
682+
ArgByteOrder::Little => {
697683
let mut bytes = origin_bytes;
698684
bytes.append(&mut append_bytes);
699685
bytes
700686
}
701-
_ => Vec::new(),
702687
};
703688
Ok(bytes.into())
704689
}
690+
705691
#[pyproperty]
706692
fn real(&self, vm: &VirtualMachine) -> PyRef<Self> {
707693
// subclasses must return int here
@@ -768,15 +754,15 @@ pub struct IntOptions {
768754
#[derive(FromArgs)]
769755
struct IntFromByteArgs {
770756
bytes: PyBytesInner,
771-
byteorder: PyStrRef,
757+
byteorder: ArgByteOrder,
772758
#[pyarg(named, optional)]
773759
signed: OptionalArg<ArgIntoBool>,
774760
}
775761

776762
#[derive(FromArgs)]
777763
struct IntToByteArgs {
778764
length: PyIntRef,
779-
byteorder: PyStrRef,
765+
byteorder: ArgByteOrder,
780766
#[pyarg(named, optional)]
781767
signed: OptionalArg<ArgIntoBool>,
782768
}

vm/src/function/mod.rs

Lines changed: 18 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -17,36 +17,25 @@ pub use either::Either;
1717
pub use number::{ArgIntoBool, ArgIntoComplex, ArgIntoFloat};
1818
pub use protocol::{ArgCallable, ArgIterable, ArgMapping, ArgSequence};
1919

20-
use crate::{
21-
builtins::PyTupleRef, convert::TryFromObject, PyObject, PyObjectRef, PyResult, VirtualMachine,
22-
};
20+
use crate::{builtins::PyStr, convert::TryFromBorrowedObject, PyObject, PyResult, VirtualMachine};
21+
22+
#[derive(Clone, Copy, PartialEq, Eq)]
23+
pub enum ArgByteOrder {
24+
Big,
25+
Little,
26+
}
2327

24-
/// Tests that the predicate is True on a single value, or if the value is a tuple a tuple, then
25-
/// test that any of the values contained within the tuples satisfies the predicate. Type parameter
26-
/// T specifies the type that is expected, if the input value is not of that type or a tuple of
27-
/// values of that type, then a TypeError is raised.
28-
pub fn single_or_tuple_any<T, F, M>(
29-
obj: PyObjectRef,
30-
predicate: &F,
31-
message: &M,
32-
vm: &VirtualMachine,
33-
) -> PyResult<bool>
34-
where
35-
T: TryFromObject,
36-
F: Fn(&T) -> PyResult<bool>,
37-
M: Fn(&PyObject) -> String,
38-
{
39-
match T::try_from_object(vm, obj.clone()) {
40-
Ok(single) => (predicate)(&single),
41-
Err(_) => {
42-
let tuple = PyTupleRef::try_from_object(vm, obj.clone())
43-
.map_err(|_| vm.new_type_error((message)(&obj)))?;
44-
for obj in &tuple {
45-
if single_or_tuple_any(obj.clone(), predicate, message, vm)? {
46-
return Ok(true);
28+
impl TryFromBorrowedObject for ArgByteOrder {
29+
fn try_from_borrowed_object(vm: &VirtualMachine, obj: &PyObject) -> PyResult<Self> {
30+
obj.try_value_with(
31+
|s: &PyStr| match s.as_str() {
32+
"big" => Ok(Self::Big),
33+
"little" => Ok(Self::Little),
34+
_ => {
35+
Err(vm.new_value_error("byteorder must be either 'little' or 'big'".to_owned()))
4736
}
48-
}
49-
Ok(false)
50-
}
37+
},
38+
vm,
39+
)
5140
}
5241
}

0 commit comments

Comments
 (0)