Skip to content

Commit

Permalink
Add lshift() and rshift().
Browse files Browse the repository at this point in the history
  • Loading branch information
bgruening committed Sep 12, 2012
1 parent 214cf07 commit 0fd042e
Showing 1 changed file with 156 additions and 2 deletions.
158 changes: 156 additions & 2 deletions bitarray/_bitarray.c
Original file line number Diff line number Diff line change
Expand Up @@ -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)
{
Expand Down Expand Up @@ -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);
}

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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)
{
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit 0fd042e

Please sign in to comment.