Simplify and speedup the internals of the heapq module.

This commit is contained in:
Raymond Hettinger 2014-05-03 18:36:48 -07:00
parent 4ce5f3f203
commit 871620d951
1 changed files with 36 additions and 71 deletions

View File

@ -11,7 +11,7 @@ annotated by François Pinard, and converted to C by Raymond Hettinger.
static int static int
_siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos) _siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos)
{ {
PyObject *newitem, *parent, *olditem; PyObject *newitem, *parent;
int cmp; int cmp;
Py_ssize_t parentpos; Py_ssize_t parentpos;
Py_ssize_t size; Py_ssize_t size;
@ -23,39 +23,28 @@ _siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos)
return -1; return -1;
} }
newitem = PyList_GET_ITEM(heap, pos);
Py_INCREF(newitem);
/* Follow the path to the root, moving parents down until finding /* Follow the path to the root, moving parents down until finding
a place newitem fits. */ a place newitem fits. */
newitem = PyList_GET_ITEM(heap, pos);
while (pos > startpos){ while (pos > startpos){
parentpos = (pos - 1) >> 1; parentpos = (pos - 1) >> 1;
parent = PyList_GET_ITEM(heap, parentpos); parent = PyList_GET_ITEM(heap, parentpos);
cmp = PyObject_RichCompareBool(newitem, parent, Py_LT); cmp = PyObject_RichCompareBool(newitem, parent, Py_LT);
if (cmp == -1) { if (cmp == -1)
Py_DECREF(newitem);
return -1; return -1;
}
if (size != PyList_GET_SIZE(heap)) { if (size != PyList_GET_SIZE(heap)) {
Py_DECREF(newitem);
PyErr_SetString(PyExc_RuntimeError, PyErr_SetString(PyExc_RuntimeError,
"list changed size during iteration"); "list changed size during iteration");
return -1; return -1;
} }
if (cmp == 0) if (cmp == 0)
break; break;
Py_INCREF(parent); parent = PyList_GET_ITEM(heap, parentpos);
olditem = PyList_GET_ITEM(heap, pos); newitem = PyList_GET_ITEM(heap, pos);
PyList_SET_ITEM(heap, parentpos, newitem);
PyList_SET_ITEM(heap, pos, parent); PyList_SET_ITEM(heap, pos, parent);
Py_DECREF(olditem);
pos = parentpos; pos = parentpos;
if (size != PyList_GET_SIZE(heap)) {
PyErr_SetString(PyExc_RuntimeError,
"list changed size during iteration");
return -1;
}
} }
Py_DECREF(PyList_GET_ITEM(heap, pos));
PyList_SET_ITEM(heap, pos, newitem);
return 0; return 0;
} }
@ -63,20 +52,16 @@ static int
_siftup(PyListObject *heap, Py_ssize_t pos) _siftup(PyListObject *heap, Py_ssize_t pos)
{ {
Py_ssize_t startpos, endpos, childpos, rightpos, limit; Py_ssize_t startpos, endpos, childpos, rightpos, limit;
PyObject *tmp1, *tmp2;
int cmp; int cmp;
PyObject *newitem, *tmp, *olditem;
Py_ssize_t size;
assert(PyList_Check(heap)); assert(PyList_Check(heap));
size = PyList_GET_SIZE(heap); endpos = PyList_GET_SIZE(heap);
endpos = size;
startpos = pos; startpos = pos;
if (pos >= endpos) { if (pos >= endpos) {
PyErr_SetString(PyExc_IndexError, "index out of range"); PyErr_SetString(PyExc_IndexError, "index out of range");
return -1; return -1;
} }
newitem = PyList_GET_ITEM(heap, pos);
Py_INCREF(newitem);
/* Bubble up the smaller child until hitting a leaf. */ /* Bubble up the smaller child until hitting a leaf. */
limit = endpos / 2; /* smallest pos that has no child */ limit = endpos / 2; /* smallest pos that has no child */
@ -89,37 +74,24 @@ _siftup(PyListObject *heap, Py_ssize_t pos)
PyList_GET_ITEM(heap, childpos), PyList_GET_ITEM(heap, childpos),
PyList_GET_ITEM(heap, rightpos), PyList_GET_ITEM(heap, rightpos),
Py_LT); Py_LT);
if (cmp == -1) { if (cmp == -1)
Py_DECREF(newitem);
return -1; return -1;
}
if (cmp == 0) if (cmp == 0)
childpos = rightpos; childpos = rightpos;
} if (endpos != PyList_GET_SIZE(heap)) {
if (size != PyList_GET_SIZE(heap)) { PyErr_SetString(PyExc_RuntimeError,
Py_DECREF(newitem); "list changed size during iteration");
PyErr_SetString(PyExc_RuntimeError, return -1;
"list changed size during iteration"); }
return -1;
} }
/* Move the smaller child up. */ /* Move the smaller child up. */
tmp = PyList_GET_ITEM(heap, childpos); tmp1 = PyList_GET_ITEM(heap, childpos);
Py_INCREF(tmp); tmp2 = PyList_GET_ITEM(heap, pos);
olditem = PyList_GET_ITEM(heap, pos); PyList_SET_ITEM(heap, childpos, tmp2);
PyList_SET_ITEM(heap, pos, tmp); PyList_SET_ITEM(heap, pos, tmp1);
Py_DECREF(olditem);
pos = childpos; pos = childpos;
if (size != PyList_GET_SIZE(heap)) {
PyErr_SetString(PyExc_RuntimeError,
"list changed size during iteration");
return -1;
}
} }
/* Bubble it up to its final resting place (by sifting its parents down). */
/* The leaf at pos is empty now. Put newitem there, and bubble
it up to its final resting place (by sifting its parents down). */
Py_DECREF(PyList_GET_ITEM(heap, pos));
PyList_SET_ITEM(heap, pos, newitem);
return _siftdown(heap, startpos, pos); return _siftdown(heap, startpos, pos);
} }
@ -392,27 +364,23 @@ _siftdownmax(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos)
return -1; return -1;
} }
newitem = PyList_GET_ITEM(heap, pos);
Py_INCREF(newitem);
/* Follow the path to the root, moving parents down until finding /* Follow the path to the root, moving parents down until finding
a place newitem fits. */ a place newitem fits. */
newitem = PyList_GET_ITEM(heap, pos);
while (pos > startpos){ while (pos > startpos){
parentpos = (pos - 1) >> 1; parentpos = (pos - 1) >> 1;
parent = PyList_GET_ITEM(heap, parentpos); parent = PyList_GET_ITEM(heap, parentpos);
cmp = PyObject_RichCompareBool(parent, newitem, Py_LT); cmp = PyObject_RichCompareBool(parent, newitem, Py_LT);
if (cmp == -1) { if (cmp == -1)
Py_DECREF(newitem);
return -1; return -1;
}
if (cmp == 0) if (cmp == 0)
break; break;
Py_INCREF(parent); parent = PyList_GET_ITEM(heap, parentpos);
Py_DECREF(PyList_GET_ITEM(heap, pos)); newitem = PyList_GET_ITEM(heap, pos);
PyList_SET_ITEM(heap, parentpos, newitem);
PyList_SET_ITEM(heap, pos, parent); PyList_SET_ITEM(heap, pos, parent);
pos = parentpos; pos = parentpos;
} }
Py_DECREF(PyList_GET_ITEM(heap, pos));
PyList_SET_ITEM(heap, pos, newitem);
return 0; return 0;
} }
@ -420,8 +388,8 @@ static int
_siftupmax(PyListObject *heap, Py_ssize_t pos) _siftupmax(PyListObject *heap, Py_ssize_t pos)
{ {
Py_ssize_t startpos, endpos, childpos, rightpos, limit; Py_ssize_t startpos, endpos, childpos, rightpos, limit;
PyObject *tmp1, *tmp2;
int cmp; int cmp;
PyObject *newitem, *tmp;
assert(PyList_Check(heap)); assert(PyList_Check(heap));
endpos = PyList_GET_SIZE(heap); endpos = PyList_GET_SIZE(heap);
@ -430,8 +398,6 @@ _siftupmax(PyListObject *heap, Py_ssize_t pos)
PyErr_SetString(PyExc_IndexError, "index out of range"); PyErr_SetString(PyExc_IndexError, "index out of range");
return -1; return -1;
} }
newitem = PyList_GET_ITEM(heap, pos);
Py_INCREF(newitem);
/* Bubble up the smaller child until hitting a leaf. */ /* Bubble up the smaller child until hitting a leaf. */
limit = endpos / 2; /* smallest pos that has no child */ limit = endpos / 2; /* smallest pos that has no child */
@ -444,25 +410,24 @@ _siftupmax(PyListObject *heap, Py_ssize_t pos)
PyList_GET_ITEM(heap, rightpos), PyList_GET_ITEM(heap, rightpos),
PyList_GET_ITEM(heap, childpos), PyList_GET_ITEM(heap, childpos),
Py_LT); Py_LT);
if (cmp == -1) { if (cmp == -1)
Py_DECREF(newitem);
return -1; return -1;
}
if (cmp == 0) if (cmp == 0)
childpos = rightpos; childpos = rightpos;
if (endpos != PyList_GET_SIZE(heap)) {
PyErr_SetString(PyExc_RuntimeError,
"list changed size during iteration");
return -1;
}
} }
/* Move the smaller child up. */ /* Move the smaller child up. */
tmp = PyList_GET_ITEM(heap, childpos); tmp1 = PyList_GET_ITEM(heap, childpos);
Py_INCREF(tmp); tmp2 = PyList_GET_ITEM(heap, pos);
Py_DECREF(PyList_GET_ITEM(heap, pos)); PyList_SET_ITEM(heap, childpos, tmp2);
PyList_SET_ITEM(heap, pos, tmp); PyList_SET_ITEM(heap, pos, tmp1);
pos = childpos; pos = childpos;
} }
/* Bubble it up to its final resting place (by sifting its parents down). */
/* The leaf at pos is empty now. Put newitem there, and bubble
it up to its final resting place (by sifting its parents down). */
Py_DECREF(PyList_GET_ITEM(heap, pos));
PyList_SET_ITEM(heap, pos, newitem);
return _siftdownmax(heap, startpos, pos); return _siftdownmax(heap, startpos, pos);
} }