#define PY_SSIZE_T_CLEAN

#include "Python.h"

#include "pythoncapi_compat.h"

#include "bitarray.h"

#define BLOCKSIZE 65536

#ifdef IS_PY3K

#define Py_TPFLAGS_HAVE_WEAKREFS 0

#define BYTES_SIZE_FMT "y#"

#define BYTES_SIZE_FMT "s#"

#ifdef STDC_HEADERS

#include <stddef.h>

#else /* !STDC_HEADERS */

#ifdef HAVE_SYS_TYPES_H

#include <sys/types.h> /* For size_t */

#endif /* HAVE_SYS_TYPES_H */

#endif /* !STDC_HEADERS */

static int default_endian = ENDIAN_BIG;

static PyTypeObject Bitarray_Type;

#define bitarray_Check(obj) PyObject_TypeCheck((obj), &Bitarray_Type)

static int

resize(bitarrayobject *self, Py_ssize_t nbits)

{

const Py_ssize_t allocated = self->allocated, size = Py_SIZE(self);

Py_ssize_t newsize;

size_t new_allocated;

assert(allocated >= size && size == BYTES(self->nbits));

/* ob_item == NULL implies ob_size == allocated == 0 */

assert(self->ob_item != NULL || (size == 0 && allocated == 0));

/* allocated == 0 implies size == 0 */

assert(allocated != 0 || size == 0);

newsize = BYTES(nbits);

if (nbits < 0 || BITS(newsize) < 0) {

PyErr_Format(PyExc_OverflowError, "bitarray resize %zd", nbits);

return -1;

}

if (newsize == size) {

/* the buffer size hasn't changed - bypass everything */

self->nbits = nbits;

return 0;

}

if (self->ob_exports > 0) {

PyErr_SetString(PyExc_BufferError,

"cannot resize bitarray that is exporting buffers");

return -1;

}

/* Bypass reallocation when a allocation is large enough to accommodate

if (allocated >= newsize && newsize >= (allocated >> 1)) {

assert(self->ob_item != NULL || newsize == 0);

Py_SET_SIZE(self, newsize);

self->nbits = nbits;

return 0;

}

if (newsize == 0) {

PyMem_FREE(self->ob_item);

self->ob_item = NULL;

Py_SET_SIZE(self, 0);

self->allocated = 0;

self->nbits = 0;

return 0;

}

new_allocated = (size_t) newsize;

if (size == 0 && newsize <= 4)

/* When resizing an empty bitarray, we want at least 4 bytes. */

new_allocated = 4;

/* Over-allocate when the (previous) size is non-zero (as we often

else if (size != 0 && newsize > size)

/* This over-allocates proportional to the bitarray size, making

new_allocated += (newsize >> 4) + (newsize < 8 ? 3 : 7);

assert(new_allocated >= (size_t) newsize);

self->ob_item = PyMem_Realloc(self->ob_item, new_allocated);

if (self->ob_item == NULL) {

PyErr_NoMemory();

return -1;

}

Py_SET_SIZE(self, newsize);

self->allocated = new_allocated;

self->nbits = nbits;

return 0;

}

static PyObject *

newbitarrayobject(PyTypeObject *type, Py_ssize_t nbits, int endian)

{

const Py_ssize_t nbytes = BYTES(nbits);

bitarrayobject *obj;

assert(nbits >= 0);

obj = (bitarrayobject *) type->tp_alloc(type, 0);

if (obj == NULL)

return NULL;

Py_SET_SIZE(obj, nbytes);

if (nbytes == 0) {

obj->ob_item = NULL;

}

else {

obj->ob_item = (char *) PyMem_Malloc((size_t) nbytes);

if (obj->ob_item == NULL) {

PyObject_Del(obj);

return PyErr_NoMemory();

}

}

obj->allocated = nbytes;

obj->nbits = nbits;

obj->endian = endian;

obj->ob_exports = 0;

obj->weakreflist = NULL;

return (PyObject *) obj;

}

static void

bitarray_dealloc(bitarrayobject *self)

{

if (self->weakreflist != NULL)

PyObject_ClearWeakRefs((PyObject *) self);

if (self->ob_item != NULL)

PyMem_Free((void *) self->ob_item);

Py_TYPE(self)->tp_free((PyObject *) self);

}

static void

bytereverse(bitarrayobject *self, Py_ssize_t start, Py_ssize_t n)

{

static char trans[256];

static int setup = 0;

Py_ssize_t i, stop = start + n;

assert(n >= 0);

assert(0 <= start && start <= Py_SIZE(self));

assert(0 <= stop && stop <= Py_SIZE(self));

if (!setup) {

/* setup translation table, which maps each byte to it's reversed:

int j, k;

for (k = 0; k < 256; k++) {

trans[k] = 0x00;

for (j = 0; j < 8; j++)

if (1 << (7 - j) & k)

trans[k] |= 1 << j;

}

setup = 1;

}

setunused(self);

for (i = start; i < stop; i++)

self->ob_item[i] = trans[(unsigned char) self->ob_item[i]];

}

static void

copy_n(bitarrayobject *self, Py_ssize_t a,

bitarrayobject *other, Py_ssize_t b, Py_ssize_t n)

{

Py_ssize_t i;

assert(0 <= n && n <= self->nbits && n <= other->nbits);

assert(0 <= a && a <= self->nbits - n);

assert(0 <= b && b <= other->nbits - n);

if (n == 0)

return;

/* When the start positions are at byte positions, we can copy whole

if (a % 8 == 0 && b % 8 == 0 && n >= 8)

{

const size_t bytes = n / 8;

const Py_ssize_t bits = BITS(bytes);

assert(bytes > 0 && bits <= n && n < bits + 8);

if (a > b)

copy_n(self, bits + a, other, bits + b, n - bits);

memmove(self->ob_item + a / 8, other->ob_item + b / 8, bytes);

if (self->endian != other->endian)

bytereverse(self, a / 8, bytes);

if (a <= b)

copy_n(self, bits + a, other, bits + b, n - bits);

return;

}

/* The two different types of looping are only relevant when copying

if (a <= b) {

for (i = 0; i < n; i++) /* loop forward (delete) */

setbit(self, i + a, GETBIT(other, i + b));

}

else {

for (i = n - 1; i >= 0; i--) /* loop backwards (insert) */

setbit(self, i + a, GETBIT(other, i + b));

}

}

static int

delete_n(bitarrayobject *self, Py_ssize_t start, Py_ssize_t n)

{

assert(0 <= start && start <= self->nbits);

assert(0 <= n && n <= self->nbits - start);

if (n == 0)

return 0;

copy_n(self, start, self, start + n, self->nbits - start - n);

return resize(self, self->nbits - n);

}

static int

insert_n(bitarrayobject *self, Py_ssize_t start, Py_ssize_t n)

{

assert(0 <= start && start <= self->nbits);

assert(n >= 0);

if (n == 0)

return 0;

if (resize(self, self->nbits + n) < 0)

return -1;

copy_n(self, start + n, self, start, self->nbits - start - n);

return 0;

}

static void

invert(bitarrayobject *self)

{

const Py_ssize_t nbytes = Py_SIZE(self);

Py_ssize_t i;

for (i = 0; i < nbytes; i++)

self->ob_item[i] = ~self->ob_item[i];

}

static int

repeat(bitarrayobject *self, Py_ssize_t n)

{

const Py_ssize_t nbits = self->nbits;

Py_ssize_t i;

if (nbits == 0 || n == 1) /* nothing to do */

return 0;

if (n <= 0) /* clear */

return resize(self, 0);

assert(n > 1 && nbits > 0);

if (nbits > PY_SSIZE_T_MAX / n) {

PyErr_Format(PyExc_OverflowError,

"cannot repeat bitarray (of size %zd) %zd times",

nbits, n);

return -1;

}

if (resize(self, n * nbits) < 0)

return -1;

for (i = 1; i < n; i++)

copy_n(self, i * nbits, self, 0, nbits);

return 0;

}

static void

setrange(bitarrayobject *self, Py_ssize_t start, Py_ssize_t stop, int val)

{

Py_ssize_t i;

assert(0 <= start && start <= self->nbits);

assert(0 <= stop && stop <= self->nbits);

if (self->nbits == 0 || start >= stop)

return;

if (stop >= start + 8) {

const Py_ssize_t byte_start = BYTES(start);

const Py_ssize_t byte_stop = stop / 8;

for (i = start; i < BITS(byte_start); i++)

setbit(self, i, val);

memset(self->ob_item + byte_start, val ? 0xff : 0x00,

(size_t) (byte_stop - byte_start));

for (i = BITS(byte_stop); i < stop; i++)

setbit(self, i, val);

}

else {

for (i = start; i < stop; i++)

setbit(self, i, val);

}

}

static Py_ssize_t

count(bitarrayobject *self, int vi, Py_ssize_t start, Py_ssize_t stop)

{

Py_ssize_t res = 0, i;

assert(0 <= start && start <= self->nbits);

assert(0 <= stop && stop <= self->nbits);

assert(0 <= vi && vi <= 1);

assert(BYTES(stop) <= Py_SIZE(self));

if (self->nbits == 0 || start >= stop)

return 0;

if (stop >= start + 8) {

const Py_ssize_t byte_start = BYTES(start);

const Py_ssize_t byte_stop = stop / 8;

Py_ssize_t j;

for (i = start; i < BITS(byte_start); i++)

res += GETBIT(self, i);

for (j = byte_start; j < byte_stop; j++)

res += bitcount_lookup[(unsigned char) self->ob_item[j]];

for (i = BITS(byte_stop); i < stop; i++)

res += GETBIT(self, i);

}

else {

for (i = start; i < stop; i++)

res += GETBIT(self, i);

}

return vi ? res : stop - start - res;

}

static Py_ssize_t

findfirst(bitarrayobject *self, int vi, Py_ssize_t start, Py_ssize_t stop)

{

Py_ssize_t i;

assert(0 <= start && start <= self->nbits);

assert(0 <= stop && stop <= self->nbits);

assert(0 <= vi && vi <= 1);

assert(BYTES(stop) <= Py_SIZE(self));

if (self->nbits == 0 || start >= stop)

return -1;

if (stop >= start + 8) {

/* seraching for 1 means: break when byte is not 0x00

const char c = vi ? 0x00 : 0xff;

/* skip ahead by checking whole bytes */

for (i = start / 8; i < BYTES(stop); i++) {

if (c ^ self->ob_item[i])

break;

}

if (start < BITS(i))

start = BITS(i);

}

/* fine grained search */

for (i = start; i < stop; i++) {

if (GETBIT(self, i) == vi)

return i;

}

return -1;

}

static Py_ssize_t

search(bitarrayobject *self, bitarrayobject *xa, Py_ssize_t p)

{

Py_ssize_t i;

assert(p >= 0);

if (xa->nbits == 1) /* faster for sparse bitarrays */

return findfirst(self, GETBIT(xa, 0), p, self->nbits);

while (p < self->nbits - xa->nbits + 1) {

for (i = 0; i < xa->nbits; i++)

if (GETBIT(self, p + i) != GETBIT(xa, i))

goto next;

return p;

next:

p++;

}

return -1;

}

static int

set_item(bitarrayobject *self, Py_ssize_t i, PyObject *v)

{

int vi;

assert(0 <= i && i < self->nbits);

vi = PyObject_IsTrue(v);

if (vi < 0)

return -1;

setbit(self, i, vi);

return 0;

}

static int

append_item(bitarrayobject *self, PyObject *item)

{

if (resize(self, self->nbits + 1) < 0)

return -1;

return set_item(self, self->nbits - 1, item);

}

static int

extend_bitarray(bitarrayobject *self, bitarrayobject *other)

{

/* We have to store the sizes before we resize, and since

const Py_ssize_t self_nbits = self->nbits;

const Py_ssize_t other_nbits = other->nbits;

if (other_nbits == 0)

return 0;

if (resize(self, self_nbits + other_nbits) < 0)

return -1;

copy_n(self, self_nbits, other, 0, other_nbits);

return 0;

}

static int

extend_iter(bitarrayobject *self, PyObject *iter)

{

PyObject *item;

assert(PyIter_Check(iter));

while ((item = PyIter_Next(iter))) {

if (append_item(self, item) < 0) {

Py_DECREF(item);

return -1;

}

Py_DECREF(item);

}

if (PyErr_Occurred())

return -1;

return 0;

}

static int

extend_list(bitarrayobject *self, PyObject *list)

{

PyObject *item;

Py_ssize_t n, i;

assert(PyList_Check(list));

n = PyList_GET_SIZE(list);

if (n == 0)

return 0;

if (resize(self, self->nbits + n) < 0)

return -1;

for (i = 0; i < n; i++) {

item = PyList_GET_ITEM(list, i);

if (item == NULL)

return -1;

if (set_item(self, self->nbits - n + i, item) < 0)

return -1;

}

return 0;

}

static int

extend_bytes01(bitarrayobject *self, PyObject *bytes)

{

const Py_ssize_t original_nbits = self->nbits;

unsigned char c;

char *data;

int vi = 0; /* to avoid uninitialized warning for some compilers */

assert(PyBytes_Check(bytes));

data = PyBytes_AS_STRING(bytes);

while ((c = *data++)) {

switch (c) {

case '0': vi = 0; break;

case '1': vi = 1; break;

case ' ':

case '\n':

case '\r':

case '\t':

case '\v':

continue;

default:

PyErr_Format(PyExc_ValueError, "expected '0' or '1' "

"(or whitespace), got '%c' (0x%02x)", c, c);

resize(self, original_nbits); /* no bits added on error */

return -1;

}

if (resize(self, self->nbits + 1) < 0)

return -1;

setbit(self, self->nbits - 1, vi);

}

return 0;

}

static int

extend_unicode01(bitarrayobject *self, PyObject *unicode)

{

PyObject *bytes;

int res;

assert(PyUnicode_Check(unicode));

bytes = PyUnicode_AsASCIIString(unicode);

if (bytes == NULL)

return -1;

assert(PyBytes_Check(bytes));

res = extend_bytes01(self, bytes);

Py_DECREF(bytes); /* drop bytes */

return res;

}

static int

extend_dispatch(bitarrayobject *self, PyObject *obj)

{

PyObject *iter;

/* dispatch on type */

if (bitarray_Check(obj)) /* bitarray */

return extend_bitarray(self, (bitarrayobject *) obj);

if (PyList_Check(obj)) /* list */

return extend_list(self, obj);

if (PyBytes_Check(obj)) { /* bytes 01 */

#ifdef IS_PY3K

PyErr_SetString(PyExc_TypeError,

"cannot extend bitarray with 'bytes', "

"use .pack() or .frombytes() instead");

return -1;

return extend_bytes01(self, obj);

}

if (PyUnicode_Check(obj)) /* unicode 01 */

return extend_unicode01(self, obj);

if (PyIter_Check(obj)) /* iter */

return extend_iter(self, obj);

/* finally, try to get the iterator of the object */

iter = PyObject_GetIter(obj);

if (iter) {

int res;

res = extend_iter(self, iter);

Py_DECREF(iter);

return res;

}

PyErr_Format(PyExc_TypeError,

"'%s' object is not iterable", Py_TYPE(obj)->tp_name);

return -1;

}

static PyObject *

unpack(bitarrayobject *self, char zero, char one, const char *fmt)

{

PyObject *result;

Py_ssize_t i;

char *str;

str = (char *) PyMem_Malloc((size_t) self->nbits);

if (str == NULL)

return PyErr_NoMemory();

for (i = 0; i < self->nbits; i++)

str[i] = GETBIT(self, i) ? one : zero;

result = Py_BuildValue(fmt, str, self->nbits);

PyMem_Free((void *) str);

return result;

}

static void

normalize_index(Py_ssize_t n, Py_ssize_t *i)

{

if (*i < 0) {

*i += n;

if (*i < 0)

*i = 0;

}

if (*i > n)

*i = n;

}

static PyObject *

bitarray_length(bitarrayobject *self)

{

if (PyErr_WarnEx(PyExc_DeprecationWarning,

"self.length() has been deprecated since 1.5.1, "

"use len(self) instead", 1) < 0)

return NULL;

return PyLong_FromSsize_t(self->nbits);

}

PyDoc_STRVAR(length_doc,

Deprecated since 1.5.1, use len().");

static PyObject *

bitarray_copy(bitarrayobject *self)

{

PyObject *res;

res = newbitarrayobject(Py_TYPE(self), self->nbits, self->endian);

if (res == NULL)

return NULL;

memcpy(((bitarrayobject *) res)->ob_item, self->ob_item,

(size_t) Py_SIZE(self));

return res;

}

PyDoc_STRVAR(copy_doc,

Return a copy of the bitarray.");

static PyObject *

bitarray_count(bitarrayobject *self, PyObject *args)

{

PyObject *x = Py_True;

Py_ssize_t start = 0, stop = self->nbits;

int vi;

if (!PyArg_ParseTuple(args, "|Onn:count", &x, &start, &stop))

return NULL;

vi = PyObject_IsTrue(x);

if (vi < 0)

return NULL;

normalize_index(self->nbits, &start);

normalize_index(self->nbits, &stop);

return PyLong_FromSsize_t(count(self, vi, start, stop));

}

PyDoc_STRVAR(count_doc,

Count the number of occurrences of bool(value) in the bitarray.");

static PyObject *

bitarray_index(bitarrayobject *self, PyObject *args)

{

PyObject *x;

Py_ssize_t start = 0, stop = self->nbits, i;

int vi;

if (!PyArg_ParseTuple(args, "O|nn:index", &x, &start, &stop))

return NULL;

vi = PyObject_IsTrue(x);

if (vi < 0)

return NULL;

normalize_index(self->nbits, &start);

normalize_index(self->nbits, &stop);

i = findfirst(self, vi, start, stop);

if (i < 0)

return PyErr_Format(PyExc_ValueError, "%d is not in bitarray", vi);

return PyLong_FromSsize_t(i);

}

PyDoc_STRVAR(index_doc,

Raises `ValueError` if the value is not present.");

static PyObject *

bitarray_extend(bitarrayobject *self, PyObject *obj)

{

if (extend_dispatch(self, obj) < 0)

return NULL;

Py_RETURN_NONE;

}

PyDoc_STRVAR(extend_doc,

as bits (whitespace is ignored).");

static PyObject *

bitarray_search(bitarrayobject *self, PyObject *args)

{

PyObject *list = NULL; /* list of matching positions to be returned */

PyObject *item = NULL, *x;

Py_ssize_t limit = -1;

bitarrayobject *xa;

Py_ssize_t p = 0;

if (!PyArg_ParseTuple(args, "O|n:search", &x, &limit))

return NULL;

if (!bitarray_Check(x)) {

PyErr_SetString(PyExc_TypeError, "bitarray expected for search");

return NULL;

}

xa = (bitarrayobject *) x;

if (xa->nbits == 0) {

PyErr_SetString(PyExc_ValueError, "can't search for empty bitarray");

return NULL;

}

list = PyList_New(0);

if (list == NULL)

return NULL;

if (xa->nbits > self->nbits || limit == 0)

return list;

while ((p = search(self, xa, p)) >= 0) {

item = PyLong_FromSsize_t(p++);

if (item == NULL || PyList_Append(list, item) < 0) {

Py_XDECREF(item);

Py_XDECREF(list);

return NULL;

}

Py_DECREF(item);

if (limit > 0 && PyList_Size(list) >= limit)

break;

}

return list;

}

PyDoc_STRVAR(search_doc,

specified. By default, all search results are returned.");

static PyObject *

bitarray_buffer_info(bitarrayobject *self)

{

PyObject *res, *ptr;

ptr = PyLong_FromVoidPtr(self->ob_item),

res = Py_BuildValue("Onsin",

ptr,

Py_SIZE(self),

ENDIAN_STR(self),

(int) (BITS(Py_SIZE(self)) - self->nbits),

self->allocated);

Py_DECREF(ptr);

return res;

}

PyDoc_STRVAR(buffer_info_doc,

byte, and the allocated memory for the buffer (in bytes).");

static PyObject *

bitarray_endian(bitarrayobject *self)

{

return Py_BuildValue("s", ENDIAN_STR(self));

}

PyDoc_STRVAR(endian_doc,

Return the bit endianness of the bitarray as a string (`little` or `big`).");

static PyObject *

bitarray_append(bitarrayobject *self, PyObject *v)

{

if (append_item(self, v) < 0)

return NULL;

Py_RETURN_NONE;

}

PyDoc_STRVAR(append_doc,

Append the truth value `bool(item)` to the end of the bitarray.");

static PyObject *

bitarray_all(bitarrayobject *self)

{

return PyBool_FromLong(findfirst(self, 0, 0, self->nbits) == -1);

}

PyDoc_STRVAR(all_doc,

Note that `a.all()` is faster than `all(a)`.");

static PyObject *

bitarray_any(bitarrayobject *self)

{

return PyBool_FromLong(findfirst(self, 1, 0, self->nbits) >= 0);

}

PyDoc_STRVAR(any_doc,

Note that `a.any()` is faster than `any(a)`.");

static PyObject *

bitarray_reduce(bitarrayobject *self)

{

const Py_ssize_t nbytes = Py_SIZE(self);

PyObject *dict, *repr = NULL, *result = NULL;

char *data;

dict = PyObject_GetAttrString((PyObject *) self, "__dict__");

if (dict == NULL) {

PyErr_Clear();

dict = Py_None;

Py_INCREF(dict);

}

data = (char *) PyMem_Malloc(nbytes + 1);

if (data == NULL) {

PyErr_NoMemory();

goto error;

}

/* first byte contains the number of unused bits */

*data = (char) setunused(self);

/* remaining bytes contain buffer */

memcpy(data + 1, self->ob_item, (size_t) nbytes);

repr = PyBytes_FromStringAndSize(data, nbytes + 1);

if (repr == NULL)

goto error;

PyMem_Free((void *) data);

result = Py_BuildValue("O(Os)O", Py_TYPE(self),

repr, ENDIAN_STR(self), dict);

error:

Py_DECREF(dict);

Py_XDECREF(repr);

return result;

}

PyDoc_STRVAR(reduce_doc, "state information for pickling");

static PyObject *

unpickle(PyTypeObject *type, PyObject *bytes, int endian)

{

PyObject *res;

Py_ssize_t nbytes;

unsigned char head;

char *data;

assert(PyBytes_Check(bytes));

nbytes = PyBytes_GET_SIZE(bytes);

assert(nbytes > 0);

data = PyBytes_AS_STRING(bytes);

head = *data;

if (nbytes == 1 && head % 8)

return PyErr_Format(PyExc_ValueError,

"invalid header byte 0x%02x", head);

res = newbitarrayobject(type,

BITS(nbytes - 1) - ((Py_ssize_t) (head % 8)),

endian);

if (res == NULL)

return NULL;

memcpy(((bitarrayobject *) res)->ob_item, data + 1, (size_t) nbytes - 1);

return res;

}

static PyObject *

bitarray_reverse(bitarrayobject *self)

{

const Py_ssize_t m = self->nbits - 1; /* index of last item */

PyObject *t; /* temp bitarray to store lower half of self */

Py_ssize_t i;

if (self->nbits < 2) /* nothing needs to be done */

Py_RETURN_NONE;

t = newbitarrayobject(Py_TYPE(self), self->nbits / 2, self->endian);

if (t == NULL)

return NULL;

#define tt ((bitarrayobject *) t)

/* copy lower half of array into temporary array */

memcpy(tt->ob_item, self->ob_item, (size_t) Py_SIZE(tt));

/* reverse upper half onto the lower half. */

for (i = 0; i < tt->nbits; i++)

setbit(self, i, GETBIT(self, m - i));

/* reverse the stored away lower half onto the upper half of self. */

for (i = 0; i < tt->nbits; i++)

setbit(self, m - i, GETBIT(tt, i));

#undef tt

Py_DECREF(t);

Py_RETURN_NONE;

}

PyDoc_STRVAR(reverse_doc,

Reverse the order of bits in the array (in-place).");

static PyObject *

bitarray_fill(bitarrayobject *self)

{

long p;

p = setunused(self);

self->nbits += p;

return PyLong_FromLong(p);