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
156 lines (146 loc) · 5.03 KB
/
bisect.rs
File metadata and controls
156 lines (146 loc) · 5.03 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
151
152
153
154
155
156
pub(crate) use _bisect::make_module;
#[pymodule]
mod _bisect {
use crate::vm::{
function::OptionalArg, slots::PyComparisonOp::Lt, ItemProtocol, PyObjectRef, PyResult,
VirtualMachine,
};
use std::convert::TryFrom;
#[derive(FromArgs)]
struct BisectArgs {
a: PyObjectRef,
x: PyObjectRef,
#[pyarg(any, optional)]
lo: OptionalArg<PyObjectRef>,
#[pyarg(any, optional)]
hi: OptionalArg<PyObjectRef>,
}
// Handles objects that implement __index__ and makes sure index fits in needed isize.
#[inline]
fn handle_default(
arg: OptionalArg<PyObjectRef>,
vm: &VirtualMachine,
) -> PyResult<Option<isize>> {
Ok(match arg {
OptionalArg::Present(v) => Some(vm.to_index(&v)?.try_to_primitive(vm)?),
OptionalArg::Missing => None,
})
}
// 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<PyObjectRef>,
hi: OptionalArg<PyObjectRef>,
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))
}
/// Return the index where to insert item x in list a, assuming a is sorted.
///
/// The return value i is such that all e in a[:i] have e < x, and all e in
/// a[i:] have e >= x. So if x already appears in the list, a.insert(x) will
/// insert just before the leftmost x already there.
///
/// Optional args lo (default 0) and hi (default len(a)) bound the
/// slice of a to be searched.
#[inline]
#[pyfunction]
fn bisect_left(
BisectArgs { a, x, lo, hi }: BisectArgs,
vm: &VirtualMachine,
) -> PyResult<usize> {
let (mut lo, mut hi) = as_usize(lo, hi, vm.obj_len(&a)?, vm)?;
while lo < hi {
// Handles issue 13496.
let mid = (lo + hi) / 2;
if vm.bool_cmp(&a.get_item(mid, vm)?, &x, Lt)? {
lo = mid + 1;
} else {
hi = mid;
}
}
Ok(lo)
}
/// Return the index where to insert item x in list a, assuming a is sorted.
///
/// The return value i is such that all e in a[:i] have e <= x, and all e in
/// a[i:] have e > x. So if x already appears in the list, a.insert(x) will
/// insert just after the rightmost x already there.
///
/// Optional args lo (default 0) and hi (default len(a)) bound the
/// slice of a to be searched.
#[inline]
#[pyfunction]
fn bisect_right(
BisectArgs { a, x, lo, hi }: BisectArgs,
vm: &VirtualMachine,
) -> PyResult<usize> {
let (mut lo, mut hi) = as_usize(lo, hi, vm.obj_len(&a)?, vm)?;
while lo < hi {
// Handles issue 13496.
let mid = (lo + hi) / 2;
if vm.bool_cmp(&x, &a.get_item(mid, vm)?, Lt)? {
hi = mid;
} else {
lo = mid + 1;
}
}
Ok(lo)
}
/// Insert item x in list a, and keep it sorted assuming a is sorted.
///
/// If x is already in a, insert it to the left of the leftmost x.
///
/// Optional args lo (default 0) and hi (default len(a)) bound the
/// slice of a to be searched.
#[pyfunction]
fn insort_left(BisectArgs { a, x, lo, hi }: BisectArgs, vm: &VirtualMachine) -> PyResult {
let index = bisect_left(
BisectArgs {
a: a.clone(),
x: x.clone(),
lo,
hi,
},
vm,
)?;
vm.call_method(&a, "insert", (index, x))
}
/// Insert item x in list a, and keep it sorted assuming a is sorted.
///
/// If x is already in a, insert it to the right of the rightmost x.
///
/// Optional args lo (default 0) and hi (default len(a)) bound the
/// slice of a to be searched
#[pyfunction]
fn insort_right(BisectArgs { a, x, lo, hi }: BisectArgs, vm: &VirtualMachine) -> PyResult {
let index = bisect_right(
BisectArgs {
a: a.clone(),
x: x.clone(),
lo,
hi,
},
vm,
)?;
vm.call_method(&a, "insert", (index, x))
}
}