From 0fd042ed49c91fff48735885a3eab1ae3862baa8 Mon Sep 17 00:00:00 2001 From: Bjoern Gruening Date: Wed, 12 Sep 2012 21:39:46 +0200 Subject: [PATCH] Add lshift() and rshift(). Based on https://bitbucket.org/quiark/bitarray. --- bitarray/_bitarray.c | 158 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 156 insertions(+), 2 deletions(-) diff --git a/bitarray/_bitarray.c b/bitarray/_bitarray.c index b1f55f93..ef97ca7d 100644 --- a/bitarray/_bitarray.c +++ b/bitarray/_bitarray.c @@ -94,6 +94,8 @@ static PyTypeObject Bitarraytype; #define GETBIT(self, i) \ ((self)->ob_item[(i) / 8] & BITMASK((self)->endian, i) ? 1 : 0) +#define ckb(i) assert(!(((i) < 0) || ((i) >= Py_SIZE(res)))) + static void setbit(bitarrayobject *self, idx_t i, int bit) { @@ -270,8 +272,10 @@ delete_n(bitarrayobject *self, idx_t start, idx_t n) assert(0 <= n && n <= self->nbits - start); if (n == 0) return 0; - - copy_n(self, start, self, start + n, self->nbits - start - n); + if (start != self->nbits + 1) { + /* default pop() wihtout index, in that case no copying is needed */ + copy_n(self, start, self, start + n, self->nbits - start - n); + } return resize(self, self->nbits - n); } @@ -371,6 +375,7 @@ bitwise(bitarrayobject *self, PyObject *arg, enum op_type oper) return 0; } + /* set the bits from start to stop (excluding) in self to val */ static void setrange(bitarrayobject *self, idx_t start, idx_t stop, int val) @@ -1780,6 +1785,7 @@ bitarray_pop(bitarrayobject *self, PyObject *args) return NULL; } vi = GETBIT(self, i); + if (delete_n(self, i, 1) < 0) return NULL; return PyBool_FromLong(vi); @@ -2376,6 +2382,150 @@ searchiter_dealloc(searchiterobject *it) PyObject_GC_Del(it); } +static PyObject * +bitarray_lshift(bitarrayobject *self, PyObject *other) +{ + if (self->endian != 1) { + PyErr_SetString(PyExc_ValueError, "Only implemented for bigendian."); + return NULL; + } + + if (!IS_INDEX(other)) { + PyErr_SetString(PyExc_TypeError, "integer argument to lshift expected"); + return NULL; + } + idx_t count = 0; + getIndex(other, &count); + + bitarrayobject *res = (bitarrayobject *)bitarray_copy(self); + + Py_ssize_t i; + + /* too much */ + if (count > res->nbits) { + memset(res->ob_item, 0, Py_SIZE(res)); + + } else if (count > 0) { + Py_ssize_t bytes = count / 8; + + /* move bytes around first */ + if (bytes > 0) { + for (i = 0; (i + bytes) < Py_SIZE(res); i++) { + ckb(i); + ckb(i + bytes); + res->ob_item[i] = res->ob_item[i + bytes]; + } + + /* clear bytes on the right */ + for (i = 0; i < bytes; i++) { + ckb(Py_SIZE(res) - 1 - i); + res->ob_item[Py_SIZE(res) - 1 - i] = 0; + } + + } + + /* perform shift smaller than 8 */ + idx_t small_shift = count % 8; + if (small_shift > 0) { + for (i = 0; i < Py_SIZE(res) - bytes; i++) { + ckb(i); + res->ob_item[i] <<= small_shift; + + if ((i + 1) < Py_SIZE(res)) { + ckb(i); + ckb(i + 1); + res->ob_item[i] |= (unsigned char)((unsigned char)(res->ob_item[i + 1]) >> (8 - small_shift)); + } + } + } + + /* clear right edge */ + idx_t border_bit = (res->nbits - count) % 8; + idx_t border_byte = (res->nbits - count) / 8; + ckb(border_byte); + res->ob_item[border_byte] &= (unsigned char)(0xff << (8 - border_bit)); + if (border_byte + 1 < Py_SIZE(res)) { + ckb(border_byte + 1); + res->ob_item[border_byte + 1] = 0; + } + + } + return (PyObject *)res; +} +PyDoc_STRVAR(lshift_doc, +"lshift([i]) -> item\n\ +\n\ +Shifts all bits i positions to the left and returns a new bitarray.\n\ +It works only on bigendian machines."); + + +static PyObject * +bitarray_rshift(bitarrayobject *self, PyObject *other) +{ + if (self->endian != 1) { + PyErr_SetString(PyExc_ValueError, "Only implemented for bigendian."); + return NULL; + } + if (!IS_INDEX(other)) { + PyErr_SetString(PyExc_TypeError, "integer argument to rshift expected"); + return NULL; + } + idx_t count = 0; + getIndex(other, &count); + + bitarrayobject *res = (bitarrayobject *)bitarray_copy(self); + + idx_t i; + + /* too much */ + if (count > res->nbits) { + memset(res->ob_item, 0, Py_SIZE(res)); + + } else if (count > 0) { + idx_t bytes = count / 8; + + /* move bytes around first */ + if (bytes > 0) { + for (i = Py_SIZE(res) - 1; i >= bytes; i--) { + ckb(i); + ckb(i - bytes); + res->ob_item[i] = res->ob_item[i - bytes]; + } + + /* clear bytes on the right */ + for (i = 0; i < bytes; i++) { + ckb(i); + res->ob_item[i] = 0; + } + + } + + /* perform shift smaller than 8 */ + idx_t small_shift = count % 8; + if (small_shift > 0) { + for (i = Py_SIZE(res) - 1; i >= bytes; i--) { + ckb(i); + res->ob_item[i] = ((unsigned char)res->ob_item[i]) >> small_shift; + + if (i >= 1) { + ckb(i - 1); + unsigned char from_left = ((unsigned char)(res->ob_item[i - 1]) << (8 - small_shift)); + ckb(i); + res->ob_item[i] |= from_left; + } + } + } + } + + return (PyObject*)res; +} +PyDoc_STRVAR(rshift_doc, +"rshift([i]) -> item\n\ +\n\ +Shifts all bits i positions to the right and returns a new bitarray.\n\ +It works only on bigendian machines."); + + static int searchiter_traverse(searchiterobject *it, visitproc visit, void *arg) { @@ -2466,6 +2616,10 @@ bitarray_methods[] = { pack_doc}, {"pop", (PyCFunction) bitarray_pop, METH_VARARGS, pop_doc}, + {"lshift", (PyCFunction) bitarray_lshift, METH_O, + lshift_doc}, + {"rshift", (PyCFunction) bitarray_rshift, METH_O, + rshift_doc}, {"remove", (PyCFunction) bitarray_remove, METH_O, remove_doc}, {"reverse", (PyCFunction) bitarray_reverse, METH_NOARGS,