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