forked from RustPython/RustPython
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathbisect.rs
More file actions
150 lines (140 loc) · 4.31 KB
/
bisect.rs
File metadata and controls
150 lines (140 loc) · 4.31 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
pub(crate) use _bisect::make_module;
#[pymodule]
mod _bisect {
use crate::vm::{
function::{ArgIndex, OptionalArg},
types::PyComparisonOp,
PyObjectRef, PyResult, VirtualMachine,
};
#[derive(FromArgs)]
struct BisectArgs {
a: PyObjectRef,
x: PyObjectRef,
#[pyarg(any, optional)]
lo: OptionalArg<ArgIndex>,
#[pyarg(any, optional)]
hi: OptionalArg<ArgIndex>,
#[pyarg(named, default)]
key: Option<PyObjectRef>,
}
// Handles objects that implement __index__ and makes sure index fits in needed isize.
#[inline]
fn handle_default(arg: OptionalArg<ArgIndex>, vm: &VirtualMachine) -> PyResult<Option<isize>> {
arg.into_option()
.map(|v| v.try_to_primitive(vm))
.transpose()
}
// Handles defaults for lo, hi.
//
// - lo must be >= 0 with a default of 0.
// - hi, while it could be negative, defaults to `0` and the same effect is achieved
// (while loop isn't entered); this way we keep it a usize (and, if my understanding
// is correct, issue 13496 is handled). Its default value is set to the length of the
// input sequence.
#[inline]
fn as_usize(
lo: OptionalArg<ArgIndex>,
hi: OptionalArg<ArgIndex>,
seq_len: usize,
vm: &VirtualMachine,
) -> PyResult<(usize, usize)> {
// We only deal with positives for lo, try_from can't fail.
// Default is always a Some so we can safely unwrap.
let lo = handle_default(lo, vm)?
.map(|value| {
usize::try_from(value)
.map_err(|_| vm.new_value_error("lo must be non-negative".to_owned()))
})
.unwrap_or(Ok(0))?;
let hi = handle_default(hi, vm)?
.map(|value| usize::try_from(value).unwrap_or(0))
.unwrap_or(seq_len);
Ok((lo, hi))
}
#[inline]
#[pyfunction]
fn bisect_left(
BisectArgs { a, x, lo, hi, key }: BisectArgs,
vm: &VirtualMachine,
) -> PyResult<usize> {
let (mut lo, mut hi) = as_usize(lo, hi, a.length(vm)?, vm)?;
while lo < hi {
// Handles issue 13496.
let mid = (lo + hi) / 2;
let a_mid = a.get_item(&mid, vm)?;
let comp = if let Some(ref key) = key {
key.call((a_mid,), vm)?
} else {
a_mid
};
if comp.rich_compare_bool(&x, PyComparisonOp::Lt, vm)? {
lo = mid + 1;
} else {
hi = mid;
}
}
Ok(lo)
}
#[inline]
#[pyfunction]
fn bisect_right(
BisectArgs { a, x, lo, hi, key }: BisectArgs,
vm: &VirtualMachine,
) -> PyResult<usize> {
let (mut lo, mut hi) = as_usize(lo, hi, a.length(vm)?, vm)?;
while lo < hi {
// Handles issue 13496.
let mid = (lo + hi) / 2;
let a_mid = a.get_item(&mid, vm)?;
let comp = if let Some(ref key) = key {
key.call((a_mid,), vm)?
} else {
a_mid
};
if x.rich_compare_bool(&comp, PyComparisonOp::Lt, vm)? {
hi = mid;
} else {
lo = mid + 1;
}
}
Ok(lo)
}
#[pyfunction]
fn insort_left(BisectArgs { a, x, lo, hi, key }: BisectArgs, vm: &VirtualMachine) -> PyResult {
let x = if let Some(ref key) = key {
key.call((x,), vm)?
} else {
x
};
let index = bisect_left(
BisectArgs {
a: a.clone(),
x: x.clone(),
lo,
hi,
key,
},
vm,
)?;
vm.call_method(&a, "insert", (index, x))
}
#[pyfunction]
fn insort_right(BisectArgs { a, x, lo, hi, key }: BisectArgs, vm: &VirtualMachine) -> PyResult {
let x = if let Some(ref key) = key {
key.call((x,), vm)?
} else {
x
};
let index = bisect_right(
BisectArgs {
a: a.clone(),
x: x.clone(),
lo,
hi,
key,
},
vm,
)?;
vm.call_method(&a, "insert", (index, x))
}
}