1 /* Bisection algorithms. Drop in replacement for bisect.py
2
3 Converted to C by Dmitry Vasiliev (dima at hlabs.spb.ru).
4 */
5
6 #ifndef Py_BUILD_CORE_BUILTIN
7 # define Py_BUILD_CORE_MODULE 1
8 #endif
9
10 #include "Python.h"
11 #include "pycore_call.h" // _PyObject_CallMethod()
12
13 /*[clinic input]
14 module _bisect
15 [clinic start generated code]*/
16 /*[clinic end generated code: output=da39a3ee5e6b4b0d input=4d56a2b2033b462b]*/
17
18 #include "clinic/_bisectmodule.c.h"
19
20 typedef struct {
21 PyObject *str_insert;
22 } bisect_state;
23
24 static inline bisect_state*
get_bisect_state(PyObject * module)25 get_bisect_state(PyObject *module)
26 {
27 void *state = PyModule_GetState(module);
28 assert(state != NULL);
29 return (bisect_state *)state;
30 }
31
32 static ssizeargfunc
get_sq_item(PyObject * s)33 get_sq_item(PyObject *s)
34 {
35 // The parts of PySequence_GetItem that we only need to do once
36 PyTypeObject *tp = Py_TYPE(s);
37 PySequenceMethods *m = tp->tp_as_sequence;
38 if (m && m->sq_item) {
39 return m->sq_item;
40 }
41 const char *msg;
42 if (tp->tp_as_mapping && tp->tp_as_mapping->mp_subscript) {
43 msg = "%.200s is not a sequence";
44 }
45 else {
46 msg = "'%.200s' object does not support indexing";
47 }
48 PyErr_Format(PyExc_TypeError, msg, tp->tp_name);
49 return NULL;
50 }
51
52 static inline Py_ssize_t
internal_bisect_right(PyObject * list,PyObject * item,Py_ssize_t lo,Py_ssize_t hi,PyObject * key)53 internal_bisect_right(PyObject *list, PyObject *item, Py_ssize_t lo, Py_ssize_t hi,
54 PyObject* key)
55 {
56 PyObject *litem;
57 Py_ssize_t mid;
58 int res;
59
60 if (lo < 0) {
61 PyErr_SetString(PyExc_ValueError, "lo must be non-negative");
62 return -1;
63 }
64 if (hi == -1) {
65 hi = PySequence_Size(list);
66 if (hi < 0)
67 return -1;
68 }
69 ssizeargfunc sq_item = get_sq_item(list);
70 if (sq_item == NULL) {
71 return -1;
72 }
73 if (Py_EnterRecursiveCall(" in _bisect.bisect_right")) {
74 return -1;
75 }
76 PyTypeObject *tp = Py_TYPE(item);
77 richcmpfunc compare = tp->tp_richcompare;
78 while (lo < hi) {
79 /* The (size_t)cast ensures that the addition and subsequent division
80 are performed as unsigned operations, avoiding difficulties from
81 signed overflow. (See issue 13496.) */
82 mid = ((size_t)lo + hi) / 2;
83 assert(mid >= 0);
84 // PySequence_GetItem, but we already checked the types.
85 litem = sq_item(list, mid);
86 assert((PyErr_Occurred() == NULL) ^ (litem == NULL));
87 if (litem == NULL) {
88 goto error;
89 }
90 if (key != Py_None) {
91 PyObject *newitem = PyObject_CallOneArg(key, litem);
92 if (newitem == NULL) {
93 goto error;
94 }
95 Py_SETREF(litem, newitem);
96 }
97 /* if item < key(list[mid]):
98 * hi = mid
99 * else:
100 * lo = mid + 1
101 */
102 if (compare != NULL && Py_IS_TYPE(litem, tp)) {
103 // A fast path for comparing objects of the same type
104 PyObject *res_obj = compare(item, litem, Py_LT);
105 if (res_obj == Py_True) {
106 Py_DECREF(res_obj);
107 Py_DECREF(litem);
108 hi = mid;
109 continue;
110 }
111 if (res_obj == Py_False) {
112 Py_DECREF(res_obj);
113 Py_DECREF(litem);
114 lo = mid + 1;
115 continue;
116 }
117 if (res_obj == NULL) {
118 goto error;
119 }
120 if (res_obj == Py_NotImplemented) {
121 Py_DECREF(res_obj);
122 compare = NULL;
123 res = PyObject_RichCompareBool(item, litem, Py_LT);
124 }
125 else {
126 res = PyObject_IsTrue(res_obj);
127 Py_DECREF(res_obj);
128 }
129 }
130 else {
131 // A default path for comparing arbitrary objects
132 res = PyObject_RichCompareBool(item, litem, Py_LT);
133 }
134 if (res < 0) {
135 goto error;
136 }
137 Py_DECREF(litem);
138 if (res)
139 hi = mid;
140 else
141 lo = mid + 1;
142 }
143 Py_LeaveRecursiveCall();
144 return lo;
145 error:
146 Py_LeaveRecursiveCall();
147 Py_XDECREF(litem);
148 return -1;
149 }
150
151 /*[clinic input]
152 _bisect.bisect_right -> Py_ssize_t
153
154 a: object
155 x: object
156 lo: Py_ssize_t = 0
157 hi: Py_ssize_t(c_default='-1', accept={int, NoneType}) = None
158 *
159 key: object = None
160
161 Return the index where to insert item x in list a, assuming a is sorted.
162
163 The return value i is such that all e in a[:i] have e <= x, and all e in
164 a[i:] have e > x. So if x already appears in the list, a.insert(i, x) will
165 insert just after the rightmost x already there.
166
167 Optional args lo (default 0) and hi (default len(a)) bound the
168 slice of a to be searched.
169
170 A custom key function can be supplied to customize the sort order.
171 [clinic start generated code]*/
172
173 static Py_ssize_t
_bisect_bisect_right_impl(PyObject * module,PyObject * a,PyObject * x,Py_ssize_t lo,Py_ssize_t hi,PyObject * key)174 _bisect_bisect_right_impl(PyObject *module, PyObject *a, PyObject *x,
175 Py_ssize_t lo, Py_ssize_t hi, PyObject *key)
176 /*[clinic end generated code: output=3a4bc09cc7c8a73d input=43071869772dd53a]*/
177 {
178 return internal_bisect_right(a, x, lo, hi, key);
179 }
180
181 /*[clinic input]
182 _bisect.insort_right
183
184 a: object
185 x: object
186 lo: Py_ssize_t = 0
187 hi: Py_ssize_t(c_default='-1', accept={int, NoneType}) = None
188 *
189 key: object = None
190
191 Insert item x in list a, and keep it sorted assuming a is sorted.
192
193 If x is already in a, insert it to the right of the rightmost x.
194
195 Optional args lo (default 0) and hi (default len(a)) bound the
196 slice of a to be searched.
197
198 A custom key function can be supplied to customize the sort order.
199 [clinic start generated code]*/
200
201 static PyObject *
_bisect_insort_right_impl(PyObject * module,PyObject * a,PyObject * x,Py_ssize_t lo,Py_ssize_t hi,PyObject * key)202 _bisect_insort_right_impl(PyObject *module, PyObject *a, PyObject *x,
203 Py_ssize_t lo, Py_ssize_t hi, PyObject *key)
204 /*[clinic end generated code: output=ac3bf26d07aedda2 input=f60777d2b6ddb239]*/
205 {
206 PyObject *result, *key_x;
207 Py_ssize_t index;
208
209 if (key == Py_None) {
210 index = internal_bisect_right(a, x, lo, hi, key);
211 } else {
212 key_x = PyObject_CallOneArg(key, x);
213 if (key_x == NULL) {
214 return NULL;
215 }
216 index = internal_bisect_right(a, key_x, lo, hi, key);
217 Py_DECREF(key_x);
218 }
219 if (index < 0)
220 return NULL;
221 if (PyList_CheckExact(a)) {
222 if (PyList_Insert(a, index, x) < 0)
223 return NULL;
224 }
225 else {
226 bisect_state *state = get_bisect_state(module);
227 result = _PyObject_CallMethod(a, state->str_insert, "nO", index, x);
228 if (result == NULL)
229 return NULL;
230 Py_DECREF(result);
231 }
232
233 Py_RETURN_NONE;
234 }
235
236 static inline Py_ssize_t
internal_bisect_left(PyObject * list,PyObject * item,Py_ssize_t lo,Py_ssize_t hi,PyObject * key)237 internal_bisect_left(PyObject *list, PyObject *item, Py_ssize_t lo, Py_ssize_t hi,
238 PyObject *key)
239 {
240 PyObject *litem;
241 Py_ssize_t mid;
242 int res;
243
244 if (lo < 0) {
245 PyErr_SetString(PyExc_ValueError, "lo must be non-negative");
246 return -1;
247 }
248 if (hi == -1) {
249 hi = PySequence_Size(list);
250 if (hi < 0)
251 return -1;
252 }
253 ssizeargfunc sq_item = get_sq_item(list);
254 if (sq_item == NULL) {
255 return -1;
256 }
257 if (Py_EnterRecursiveCall(" in _bisect.bisect_left")) {
258 return -1;
259 }
260 PyTypeObject *tp = Py_TYPE(item);
261 richcmpfunc compare = tp->tp_richcompare;
262 while (lo < hi) {
263 /* The (size_t)cast ensures that the addition and subsequent division
264 are performed as unsigned operations, avoiding difficulties from
265 signed overflow. (See issue 13496.) */
266 mid = ((size_t)lo + hi) / 2;
267 assert(mid >= 0);
268 // PySequence_GetItem, but we already checked the types.
269 litem = sq_item(list, mid);
270 assert((PyErr_Occurred() == NULL) ^ (litem == NULL));
271 if (litem == NULL) {
272 goto error;
273 }
274 if (key != Py_None) {
275 PyObject *newitem = PyObject_CallOneArg(key, litem);
276 if (newitem == NULL) {
277 goto error;
278 }
279 Py_SETREF(litem, newitem);
280 }
281 /* if key(list[mid]) < item:
282 * lo = mid + 1
283 * else:
284 * hi = mid
285 */
286 if (compare != NULL && Py_IS_TYPE(litem, tp)) {
287 // A fast path for comparing objects of the same type
288 PyObject *res_obj = compare(litem, item, Py_LT);
289 if (res_obj == Py_True) {
290 Py_DECREF(res_obj);
291 Py_DECREF(litem);
292 lo = mid + 1;
293 continue;
294 }
295 if (res_obj == Py_False) {
296 Py_DECREF(res_obj);
297 Py_DECREF(litem);
298 hi = mid;
299 continue;
300 }
301 if (res_obj == NULL) {
302 goto error;
303 }
304 if (res_obj == Py_NotImplemented) {
305 Py_DECREF(res_obj);
306 compare = NULL;
307 res = PyObject_RichCompareBool(litem, item, Py_LT);
308 }
309 else {
310 res = PyObject_IsTrue(res_obj);
311 Py_DECREF(res_obj);
312 }
313 }
314 else {
315 // A default path for comparing arbitrary objects
316 res = PyObject_RichCompareBool(litem, item, Py_LT);
317 }
318 if (res < 0) {
319 goto error;
320 }
321 Py_DECREF(litem);
322 if (res)
323 lo = mid + 1;
324 else
325 hi = mid;
326 }
327 Py_LeaveRecursiveCall();
328 return lo;
329 error:
330 Py_LeaveRecursiveCall();
331 Py_XDECREF(litem);
332 return -1;
333 }
334
335
336 /*[clinic input]
337 _bisect.bisect_left -> Py_ssize_t
338
339 a: object
340 x: object
341 lo: Py_ssize_t = 0
342 hi: Py_ssize_t(c_default='-1', accept={int, NoneType}) = None
343 *
344 key: object = None
345
346 Return the index where to insert item x in list a, assuming a is sorted.
347
348 The return value i is such that all e in a[:i] have e < x, and all e in
349 a[i:] have e >= x. So if x already appears in the list, a.insert(i, x) will
350 insert just before the leftmost x already there.
351
352 Optional args lo (default 0) and hi (default len(a)) bound the
353 slice of a to be searched.
354
355 A custom key function can be supplied to customize the sort order.
356 [clinic start generated code]*/
357
358 static Py_ssize_t
_bisect_bisect_left_impl(PyObject * module,PyObject * a,PyObject * x,Py_ssize_t lo,Py_ssize_t hi,PyObject * key)359 _bisect_bisect_left_impl(PyObject *module, PyObject *a, PyObject *x,
360 Py_ssize_t lo, Py_ssize_t hi, PyObject *key)
361 /*[clinic end generated code: output=70749d6e5cae9284 input=f29c4fe7f9b797c7]*/
362 {
363 return internal_bisect_left(a, x, lo, hi, key);
364 }
365
366
367 /*[clinic input]
368 _bisect.insort_left
369
370 a: object
371 x: object
372 lo: Py_ssize_t = 0
373 hi: Py_ssize_t(c_default='-1', accept={int, NoneType}) = None
374 *
375 key: object = None
376
377 Insert item x in list a, and keep it sorted assuming a is sorted.
378
379 If x is already in a, insert it to the left of the leftmost x.
380
381 Optional args lo (default 0) and hi (default len(a)) bound the
382 slice of a to be searched.
383
384 A custom key function can be supplied to customize the sort order.
385 [clinic start generated code]*/
386
387 static PyObject *
_bisect_insort_left_impl(PyObject * module,PyObject * a,PyObject * x,Py_ssize_t lo,Py_ssize_t hi,PyObject * key)388 _bisect_insort_left_impl(PyObject *module, PyObject *a, PyObject *x,
389 Py_ssize_t lo, Py_ssize_t hi, PyObject *key)
390 /*[clinic end generated code: output=b1d33e5e7ffff11e input=0a700a82edbd472c]*/
391 {
392 PyObject *result, *key_x;
393 Py_ssize_t index;
394
395 if (key == Py_None) {
396 index = internal_bisect_left(a, x, lo, hi, key);
397 } else {
398 key_x = PyObject_CallOneArg(key, x);
399 if (key_x == NULL) {
400 return NULL;
401 }
402 index = internal_bisect_left(a, key_x, lo, hi, key);
403 Py_DECREF(key_x);
404 }
405 if (index < 0)
406 return NULL;
407 if (PyList_CheckExact(a)) {
408 if (PyList_Insert(a, index, x) < 0)
409 return NULL;
410 } else {
411 bisect_state *state = get_bisect_state(module);
412 result = _PyObject_CallMethod(a, state->str_insert, "nO", index, x);
413 if (result == NULL)
414 return NULL;
415 Py_DECREF(result);
416 }
417
418 Py_RETURN_NONE;
419 }
420
421 static PyMethodDef bisect_methods[] = {
422 _BISECT_BISECT_RIGHT_METHODDEF
423 _BISECT_INSORT_RIGHT_METHODDEF
424 _BISECT_BISECT_LEFT_METHODDEF
425 _BISECT_INSORT_LEFT_METHODDEF
426 {NULL, NULL} /* sentinel */
427 };
428
429 PyDoc_STRVAR(module_doc,
430 "Bisection algorithms.\n\
431 \n\
432 This module provides support for maintaining a list in sorted order without\n\
433 having to sort the list after each insertion. For long lists of items with\n\
434 expensive comparison operations, this can be an improvement over the more\n\
435 common approach.\n");
436
437 static int
bisect_clear(PyObject * module)438 bisect_clear(PyObject *module)
439 {
440 bisect_state *state = get_bisect_state(module);
441 Py_CLEAR(state->str_insert);
442 return 0;
443 }
444
445 static void
bisect_free(void * module)446 bisect_free(void *module)
447 {
448 bisect_clear((PyObject *)module);
449 }
450
451 static int
bisect_modexec(PyObject * m)452 bisect_modexec(PyObject *m)
453 {
454 bisect_state *state = get_bisect_state(m);
455 state->str_insert = PyUnicode_InternFromString("insert");
456 if (state->str_insert == NULL) {
457 return -1;
458 }
459 return 0;
460 }
461
462 static PyModuleDef_Slot bisect_slots[] = {
463 {Py_mod_exec, bisect_modexec},
464 {Py_mod_multiple_interpreters, Py_MOD_PER_INTERPRETER_GIL_SUPPORTED},
465 {Py_mod_gil, Py_MOD_GIL_NOT_USED},
466 {0, NULL}
467 };
468
469 static struct PyModuleDef _bisectmodule = {
470 PyModuleDef_HEAD_INIT,
471 .m_name = "_bisect",
472 .m_size = sizeof(bisect_state),
473 .m_doc = module_doc,
474 .m_methods = bisect_methods,
475 .m_slots = bisect_slots,
476 .m_clear = bisect_clear,
477 .m_free = bisect_free,
478 };
479
480 PyMODINIT_FUNC
PyInit__bisect(void)481 PyInit__bisect(void)
482 {
483 return PyModuleDef_Init(&_bisectmodule);
484 }
485