Add lshift() and rshift().

Based on https://bitbucket.org/quiark/bitarray.

Author: bgruening

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,