#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

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 || newsize < 0) {

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

return -1;

}

if (newsize == size) {

/* 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 stop)

{

static char trans[256];

static int setup = 0;

Py_ssize_t i;

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;

}

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

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

}

#ifdef PY_UINT64_T

#define UINT64_BUFFER(self) ((PY_UINT64_T *) (self)->ob_item)

#define UINT64_WORDS(bytes) ((bytes) >> 3)

#define UINT64_BUFFER(self) ((self)->ob_item)

#define UINT64_WORDS(bytes) 0

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 >> 3;

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 >> 3), other->ob_item + (b >> 3), bytes);

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

bytereverse(self, a >> 3, (a >> 3) + 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);

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 (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);

const Py_ssize_t nwords = UINT64_WORDS(nbytes);

Py_ssize_t i;

assert_nbits(self);

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

UINT64_BUFFER(self)[i] = ~UINT64_BUFFER(self)[i];

for (i = nwords << 3; i < nbytes; i++)

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

}

static int

repeat(bitarrayobject *self, Py_ssize_t m)

{

const Py_ssize_t nbits = self->nbits;

Py_ssize_t k;

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

return 0;

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

return resize(self, 0);

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

if (nbits > PY_SSIZE_T_MAX / m) {

PyErr_Format(PyExc_OverflowError,

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

nbits, m);

return -1;

}

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

return -1;

k = nbits; /* number of bits which have been copied */

while (k <= self->nbits / 2) { /* double copies */

copy_n(self, k, self, 0, k);

k *= 2;

}

assert(self->nbits - k >= 0);

copy_n(self, k, self, 0, self->nbits - k); /* copy remaining bits */

return 0;

}

static void

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

{

Py_ssize_t i;

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

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

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

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 >> 3;

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

setbit(self, i, vi);

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

(size_t) (byte_stop - byte_start));

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

setbit(self, i, vi);

}

else {

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

setbit(self, i, vi);

}

}

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 >> 3;

Py_ssize_t j;

assert(start + 8 > BITS(byte_start) && BITS(byte_stop) + 8 > stop);

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

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

{

Py_ssize_t i, j;

char c;

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;

/* search within lowest (partial) byte */

for (i = start; i < stop && i < BITS(BYTES(start)); i++)

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

return i;

if (i == stop) /* not found */

return -1;

assert(i % 8 == 0 && i < stop && i < start + 8);

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

c = vi ? 0x00 : 0xff;

/* skip ahead by checking whole bytes */

for (j = i >> 3; j < stop >> 3; j++) {

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

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

break;

}

/* search within found or highest (partial) byte */

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

if (getbit(self, i) == vi) {

assert(i < BITS(j + 1));

return i;

}

assert(j == stop / 8 && i == stop);

return -1;

}

static Py_ssize_t

find(bitarrayobject *self, bitarrayobject *xa,

Py_ssize_t start, Py_ssize_t stop)

{

Py_ssize_t i;

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

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

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

return find_bit(self, getbit(xa, 0), start, stop);

while (start <= stop - xa->nbits) {

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

if (getbit(self, start + i) != getbit(xa, i))

goto next;

return start;

next:

start++;

}

return -1;

}

static int

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

{

int vi;

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

if ((vi = pybit_as_int(v)) < 0)

return -1;

setbit(self, i, vi);

return 0;

}

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 (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)

{

const Py_ssize_t original_nbits = self->nbits;

PyObject *item;

assert(PyIter_Check(iter));

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

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

goto error;

if (set_item(self, self->nbits - 1, item) < 0)

goto error;

Py_DECREF(item);

}

if (PyErr_Occurred())

return -1;

return 0;

error:

Py_DECREF(item);

resize(self, original_nbits);

return -1;

}

static int

extend_sequence(bitarrayobject *self, PyObject *sequence)

{

const Py_ssize_t original_nbits = self->nbits;

PyObject *item;

Py_ssize_t n, i;

assert(PySequence_Check(sequence));

n = PySequence_Size(sequence);

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

return -1;

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

item = PySequence_GetItem(sequence, i);

if (item == NULL || set_item(self, self->nbits - n + i, item) < 0) {

Py_XDECREF(item);

resize(self, original_nbits);

return -1;

}

Py_DECREF(item);

}

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 (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 (PySequence_Check(obj)) /* sequence */

return extend_sequence(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 int

endian_from_string(const char* string)

{

assert(default_endian == ENDIAN_LITTLE || default_endian == ENDIAN_BIG);

if (string == NULL)

return default_endian;

if (strcmp(string, "little") == 0)

return ENDIAN_LITTLE;

if (strcmp(string, "big") == 0)

return ENDIAN_BIG;

PyErr_Format(PyExc_ValueError, "bit endianness must be either "

"'little' or 'big', got: '%s'", string);

return -1;

}

static PyObject *

bitarray_all(bitarrayobject *self)

{

return PyBool_FromLong(find_bit(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(find_bit(self, 1, 0, self->nbits) >= 0);

}

PyDoc_STRVAR(any_doc,

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

static PyObject *

bitarray_append(bitarrayobject *self, PyObject *value)

{

int vi;

if ((vi = pybit_as_int(value)) < 0)

return NULL;

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

return NULL;

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

Py_RETURN_NONE;

}

PyDoc_STRVAR(append_doc,

Append `item` to the end of the bitarray.");

static PyObject *

bitarray_bytereverse(bitarrayobject *self)

{

setunused(self);

bytereverse(self, 0, Py_SIZE(self));

Py_RETURN_NONE;

}

PyDoc_STRVAR(bytereverse_doc,

bitarray; it does *not* change the endianness of the bitarray object.");

static PyObject *

bitarray_buffer_info(bitarrayobject *self)

{

PyObject *res, *ptr;

Py_ssize_t size = Py_SIZE(self);

ptr = PyLong_FromVoidPtr(self->ob_item),

res = Py_BuildValue("Onsnn",

ptr,

size,

ENDIAN_STR(self->endian),

BITS(size) - self->nbits,

self->allocated);

Py_DECREF(ptr);

return res;

}

PyDoc_STRVAR(buffer_info_doc,

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

static PyObject *

bitarray_clear(bitarrayobject *self)

{

if (resize(self, 0) < 0)

return NULL;

Py_RETURN_NONE;

}

PyDoc_STRVAR(clear_doc,

Remove all items from the bitarray.");

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 *v = Py_True;

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

int vi;

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

return NULL;

if ((vi = pybit_as_int(v)) < 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 `value` in the bitarray.");

static PyObject *

bitarray_endian(bitarrayobject *self)

{

return Py_BuildValue("s", ENDIAN_STR(self->endian));

}

PyDoc_STRVAR(endian_doc,

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

static PyObject *

bitarray_extend(bitarrayobject *self, PyObject *obj)

{

if (extend_dispatch(self, obj) < 0)

return NULL;

Py_RETURN_NONE;

}

PyDoc_STRVAR(extend_doc,

bits (ignoring whitespace).");

static PyObject *

bitarray_fill(bitarrayobject *self)

{

long p;

p = setunused(self);

self->nbits += p;

assert(self->nbits % 8 == 0);

assert_nbits(self);

return PyLong_FromLong(p);

}

PyDoc_STRVAR(fill_doc,

will be a multiple of 8, and return the number of bits added (0..7).");

static PyObject *

bitarray_find(bitarrayobject *self, PyObject *args)

{

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

PyObject *x;

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

return NULL;

normalize_index(self->nbits, &start);

normalize_index(self->nbits, &stop);

if (PyIndex_Check(x)) {

int vi;

if ((vi = pybit_as_int(x)) < 0)

return NULL;

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

}

if (bitarray_Check(x))

return PyLong_FromSsize_t(

find(self, (bitarrayobject *) x, start, stop));

PyErr_SetString(PyExc_TypeError, "bitarray or int expected");

return NULL;

}

PyDoc_STRVAR(find_doc,

Return -1 when sub_bitarray is not found.");

static PyObject *

bitarray_index(bitarrayobject *self, PyObject *args)

{

PyObject *ret;

if ((ret = bitarray_find(self, args)) == NULL)

return NULL;

assert(PyLong_Check(ret));

if (PyLong_AsSsize_t(ret) < 0) {

Py_DECREF(ret);

#ifdef IS_PY3K

return PyErr_Format(PyExc_ValueError, "%A not in bitarray",

PyTuple_GET_ITEM(args, 0));

PyErr_SetString(PyExc_ValueError, "item not in bitarray");

return NULL;

}

return ret;

}

PyDoc_STRVAR(index_doc,

Raises `ValueError` when the sub_bitarray is not present.");

static PyObject *

bitarray_insert(bitarrayobject *self, PyObject *args)

{

Py_ssize_t i;

PyObject *v;

int vi;

if (!PyArg_ParseTuple(args, "nO:insert", &i, &v))

return NULL;

normalize_index(self->nbits, &i);

if ((vi = pybit_as_int(v)) < 0)

return NULL;

if (insert_n(self, i, 1) < 0)

return NULL;

setbit(self, i, vi);

Py_RETURN_NONE;

}

PyDoc_STRVAR(insert_doc,

Insert `value` into the bitarray before `index`.");

static PyObject *

bitarray_invert(bitarrayobject *self, PyObject *args)

{

Py_ssize_t i = PY_SSIZE_T_MAX;

if (!PyArg_ParseTuple(args, "|n:invert", &i))

return NULL;

if (i == PY_SSIZE_T_MAX) { /* default - invert all bits */

invert(self);

Py_RETURN_NONE;

}

if (i < 0)

i += self->nbits;

if (i < 0 || i >= self->nbits) {

PyErr_SetString(PyExc_IndexError, "index out of range");

return NULL;

}

setbit(self, i, 1 - getbit(self, i));

Py_RETURN_NONE;

}

PyDoc_STRVAR(invert_doc,

When the optional `index` is given, only invert the single bit at index.");

static PyObject *

bitarray_reduce(bitarrayobject *self)

{

const Py_ssize_t nbytes = Py_SIZE(self);

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