#include "Python.h"

#include "pycore_bitutils.h" // _Py_popcount32

#include "pycore_hamt.h"

#include "pycore_object.h" // _PyObject_GC_TRACK()

#include <stddef.h> // offsetof()

#define IS_ARRAY_NODE(node) Py_IS_TYPE(node, &_PyHamt_ArrayNode_Type)

#define IS_BITMAP_NODE(node) Py_IS_TYPE(node, &_PyHamt_BitmapNode_Type)

#define IS_COLLISION_NODE(node) Py_IS_TYPE(node, &_PyHamt_CollisionNode_Type)

typedef enum {F_ERROR, F_NOT_FOUND, F_FOUND} hamt_find_t;

typedef enum {W_ERROR, W_NOT_FOUND, W_EMPTY, W_NEWNODE} hamt_without_t;

typedef enum {I_ITEM, I_END} hamt_iter_t;

#define HAMT_ARRAY_NODE_SIZE 32

typedef struct {

PyObject_HEAD

PyHamtNode *a_array[HAMT_ARRAY_NODE_SIZE];

Py_ssize_t a_count;

} PyHamtNode_Array;

typedef struct {

PyObject_VAR_HEAD

uint32_t b_bitmap;

PyObject *b_array[1];

} PyHamtNode_Bitmap;

typedef struct {

PyObject_VAR_HEAD

int32_t c_hash;

PyObject *c_array[1];

} PyHamtNode_Collision;

static PyHamtNode_Bitmap *_empty_bitmap_node;

static PyHamtObject *_empty_hamt;

static PyHamtObject *

hamt_alloc(void);

static PyHamtNode *

hamt_node_assoc(PyHamtNode *node,

uint32_t shift, int32_t hash,

PyObject *key, PyObject *val, int* added_leaf);

static hamt_without_t

hamt_node_without(PyHamtNode *node,

uint32_t shift, int32_t hash,

PyObject *key,

PyHamtNode **new_node);

static hamt_find_t

hamt_node_find(PyHamtNode *node,

uint32_t shift, int32_t hash,

PyObject *key, PyObject **val);

#ifdef Py_DEBUG

static int

hamt_node_dump(PyHamtNode *node,

_PyUnicodeWriter *writer, int level);

static PyHamtNode *

hamt_node_array_new(Py_ssize_t);

static PyHamtNode *

hamt_node_collision_new(int32_t hash, Py_ssize_t size);

static inline Py_ssize_t

hamt_node_collision_count(PyHamtNode_Collision *node);

#ifdef Py_DEBUG

static void

_hamt_node_array_validate(void *obj_raw)

{

PyObject *obj = _PyObject_CAST(obj_raw);

assert(IS_ARRAY_NODE(obj));

PyHamtNode_Array *node = (PyHamtNode_Array*)obj;

Py_ssize_t i = 0, count = 0;

for (; i < HAMT_ARRAY_NODE_SIZE; i++) {

if (node->a_array[i] != NULL) {

count++;

}

}

assert(count == node->a_count);

}

#define VALIDATE_ARRAY_NODE(NODE) \

do { _hamt_node_array_validate(NODE); } while (0);

#define VALIDATE_ARRAY_NODE(NODE)

static inline int32_t

hamt_hash(PyObject *o)

{

Py_hash_t hash = PyObject_Hash(o);

#if SIZEOF_PY_HASH_T <= 4

return hash;

if (hash == -1) {

/* exception */

return -1;

}

/* While it's suboptimal to reduce Python's 64 bit hash to

int32_t xored = (int32_t)(hash & 0xffffffffl) ^ (int32_t)(hash >> 32);

return xored == -1 ? -2 : xored;

}

static inline uint32_t

hamt_mask(int32_t hash, uint32_t shift)

{

return (((uint32_t)hash >> shift) & 0x01f);

}

static inline uint32_t

hamt_bitpos(int32_t hash, uint32_t shift)

{

return (uint32_t)1 << hamt_mask(hash, shift);

}

static inline uint32_t

hamt_bitindex(uint32_t bitmap, uint32_t bit)

{

return (uint32_t)_Py_popcount32(bitmap & (bit - 1));

}

#ifdef Py_DEBUG

static int

_hamt_dump_ident(_PyUnicodeWriter *writer, int level)

{

/* Write `' ' * level` to the `writer` */

PyObject *str = NULL;

PyObject *num = NULL;

PyObject *res = NULL;

int ret = -1;

str = PyUnicode_FromString(" ");

if (str == NULL) {

goto error;

}

num = PyLong_FromLong((long)level);

if (num == NULL) {

goto error;

}

res = PyNumber_Multiply(str, num);

if (res == NULL) {

goto error;

}

ret = _PyUnicodeWriter_WriteStr(writer, res);

error:

Py_XDECREF(res);

Py_XDECREF(str);

Py_XDECREF(num);

return ret;

}

static int

_hamt_dump_format(_PyUnicodeWriter *writer, const char *format, ...)

{

/* A convenient helper combining _PyUnicodeWriter_WriteStr and

PyObject* msg;

int ret;

va_list vargs;

#ifdef HAVE_STDARG_PROTOTYPES

va_start(vargs, format);

va_start(vargs);

msg = PyUnicode_FromFormatV(format, vargs);

va_end(vargs);

if (msg == NULL) {

return -1;

}

ret = _PyUnicodeWriter_WriteStr(writer, msg);

Py_DECREF(msg);

return ret;

}

#endif /* Py_DEBUG */

static PyHamtNode *

hamt_node_bitmap_new(Py_ssize_t size)

{

/* Create a new bitmap node of size 'size' */

PyHamtNode_Bitmap *node;

Py_ssize_t i;

assert(size >= 0);

assert(size % 2 == 0);

if (size == 0 && _empty_bitmap_node != NULL) {

Py_INCREF(_empty_bitmap_node);

return (PyHamtNode *)_empty_bitmap_node;

}

/* No freelist; allocate a new bitmap node */

node = PyObject_GC_NewVar(

PyHamtNode_Bitmap, &_PyHamt_BitmapNode_Type, size);

if (node == NULL) {

return NULL;

}

Py_SET_SIZE(node, size);

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

node->b_array[i] = NULL;

}

node->b_bitmap = 0;

_PyObject_GC_TRACK(node);

if (size == 0 && _empty_bitmap_node == NULL) {

/* Since bitmap nodes are immutable, we can cache the instance

_empty_bitmap_node = node;

Py_INCREF(_empty_bitmap_node);

}

return (PyHamtNode *)node;

}

static inline Py_ssize_t

hamt_node_bitmap_count(PyHamtNode_Bitmap *node)

{

return Py_SIZE(node) / 2;

}

static PyHamtNode_Bitmap *

hamt_node_bitmap_clone(PyHamtNode_Bitmap *node)

{

/* Clone a bitmap node; return a new one with the same child notes. */

PyHamtNode_Bitmap *clone;

Py_ssize_t i;

clone = (PyHamtNode_Bitmap *)hamt_node_bitmap_new(Py_SIZE(node));

if (clone == NULL) {

return NULL;

}

for (i = 0; i < Py_SIZE(node); i++) {

Py_XINCREF(node->b_array[i]);

clone->b_array[i] = node->b_array[i];

}

clone->b_bitmap = node->b_bitmap;

return clone;

}

static PyHamtNode_Bitmap *

hamt_node_bitmap_clone_without(PyHamtNode_Bitmap *o, uint32_t bit)

{

assert(bit & o->b_bitmap);

assert(hamt_node_bitmap_count(o) > 1);

PyHamtNode_Bitmap *new = (PyHamtNode_Bitmap *)hamt_node_bitmap_new(

Py_SIZE(o) - 2);

if (new == NULL) {

return NULL;

}

uint32_t idx = hamt_bitindex(o->b_bitmap, bit);

uint32_t key_idx = 2 * idx;

uint32_t val_idx = key_idx + 1;

uint32_t i;

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

Py_XINCREF(o->b_array[i]);

new->b_array[i] = o->b_array[i];

}

assert(Py_SIZE(o) >= 0 && Py_SIZE(o) <= 32);

for (i = val_idx + 1; i < (uint32_t)Py_SIZE(o); i++) {

Py_XINCREF(o->b_array[i]);

new->b_array[i - 2] = o->b_array[i];

}

new->b_bitmap = o->b_bitmap & ~bit;

return new;

}

static PyHamtNode *

hamt_node_new_bitmap_or_collision(uint32_t shift,

PyObject *key1, PyObject *val1,

int32_t key2_hash,

PyObject *key2, PyObject *val2)

{

/* Helper method. Creates a new node for key1/val and key2/val2

int32_t key1_hash = hamt_hash(key1);

if (key1_hash == -1) {

return NULL;

}

if (key1_hash == key2_hash) {

PyHamtNode_Collision *n;

n = (PyHamtNode_Collision *)hamt_node_collision_new(key1_hash, 4);

if (n == NULL) {

return NULL;

}

Py_INCREF(key1);

n->c_array[0] = key1;

Py_INCREF(val1);

n->c_array[1] = val1;

Py_INCREF(key2);

n->c_array[2] = key2;

Py_INCREF(val2);

n->c_array[3] = val2;

return (PyHamtNode *)n;

}

else {

int added_leaf = 0;

PyHamtNode *n = hamt_node_bitmap_new(0);

if (n == NULL) {

return NULL;

}

PyHamtNode *n2 = hamt_node_assoc(

n, shift, key1_hash, key1, val1, &added_leaf);

Py_DECREF(n);

if (n2 == NULL) {

return NULL;

}

n = hamt_node_assoc(n2, shift, key2_hash, key2, val2, &added_leaf);

Py_DECREF(n2);

if (n == NULL) {

return NULL;

}

return n;

}

}

static PyHamtNode *

hamt_node_bitmap_assoc(PyHamtNode_Bitmap *self,

uint32_t shift, int32_t hash,

PyObject *key, PyObject *val, int* added_leaf)

{

/* assoc operation for bitmap nodes.

uint32_t bit = hamt_bitpos(hash, shift);

uint32_t idx = hamt_bitindex(self->b_bitmap, bit);

/* Bitmap node layout:

if (self->b_bitmap & bit) {

/* The key is set in this node */

uint32_t key_idx = 2 * idx;

uint32_t val_idx = key_idx + 1;

assert(val_idx < (size_t)Py_SIZE(self));

PyObject *key_or_null = self->b_array[key_idx];

PyObject *val_or_node = self->b_array[val_idx];

if (key_or_null == NULL) {

/* key is NULL. This means that we have a few keys

assert(val_or_node != NULL);

PyHamtNode *sub_node = hamt_node_assoc(

(PyHamtNode *)val_or_node,

shift + 5, hash, key, val, added_leaf);

if (sub_node == NULL) {

return NULL;

}

if (val_or_node == (PyObject *)sub_node) {

Py_DECREF(sub_node);

Py_INCREF(self);

return (PyHamtNode *)self;

}

PyHamtNode_Bitmap *ret = hamt_node_bitmap_clone(self);

if (ret == NULL) {

return NULL;

}

Py_SETREF(ret->b_array[val_idx], (PyObject*)sub_node);

return (PyHamtNode *)ret;

}

assert(key != NULL);

/* key is not NULL. This means that we have only one other

int comp_err = PyObject_RichCompareBool(key, key_or_null, Py_EQ);

if (comp_err < 0) { /* exception in __eq__ */

return NULL;

}

if (comp_err == 1) { /* key == key_or_null */

if (val == val_or_node) {

/* we already have the same key/val pair; return self. */

Py_INCREF(self);

return (PyHamtNode *)self;

}

/* We're setting a new value for the key we had before.

PyHamtNode_Bitmap *ret = hamt_node_bitmap_clone(self);

if (ret == NULL) {

return NULL;

}

Py_INCREF(val);

Py_SETREF(ret->b_array[val_idx], val);

return (PyHamtNode *)ret;

}

/* It's a new key, and it has the same index as *one* another key.

PyHamtNode *sub_node = hamt_node_new_bitmap_or_collision(

shift + 5,

key_or_null, val_or_node, /* existing key/val */

hash,

key, val /* new key/val */

);

if (sub_node == NULL) {

return NULL;

}

PyHamtNode_Bitmap *ret = hamt_node_bitmap_clone(self);

if (ret == NULL) {

Py_DECREF(sub_node);

return NULL;

}

Py_SETREF(ret->b_array[key_idx], NULL);

Py_SETREF(ret->b_array[val_idx], (PyObject *)sub_node);

*added_leaf = 1;

return (PyHamtNode *)ret;

}

else {

/* There was no key before with the same (shift,hash). */

uint32_t n = (uint32_t)_Py_popcount32(self->b_bitmap);

if (n >= 16) {

/* When we have a situation where we want to store more

uint32_t jdx = hamt_mask(hash, shift);

/* 'jdx' is the index of where the new key should be added

PyHamtNode *empty = NULL;

PyHamtNode_Array *new_node = NULL;

PyHamtNode *res = NULL;

/* Create a new Array node. */

new_node = (PyHamtNode_Array *)hamt_node_array_new(n + 1);

if (new_node == NULL) {

goto fin;

}

/* Create an empty bitmap node for the next

empty = hamt_node_bitmap_new(0);

if (empty == NULL) {

goto fin;

}

/* Make a new bitmap node for the key/val we're adding.

new_node->a_array[jdx] = hamt_node_assoc(

empty, shift + 5, hash, key, val, added_leaf);

if (new_node->a_array[jdx] == NULL) {

goto fin;

}

/* Copy existing key/value pairs from the current Bitmap

Py_ssize_t i, j;

for (i = 0, j = 0; i < HAMT_ARRAY_NODE_SIZE; i++) {

if (((self->b_bitmap >> i) & 1) != 0) {

/* Ensure we don't accidentally override `jdx` element

assert(new_node->a_array[i] == NULL);

if (self->b_array[j] == NULL) {

new_node->a_array[i] =

(PyHamtNode *)self->b_array[j + 1];

Py_INCREF(new_node->a_array[i]);

}

else {

int32_t rehash = hamt_hash(self->b_array[j]);

if (rehash == -1) {

goto fin;

}

new_node->a_array[i] = hamt_node_assoc(

empty, shift + 5,

rehash,

self->b_array[j],

self->b_array[j + 1],

added_leaf);

if (new_node->a_array[i] == NULL) {

goto fin;

}

}

j += 2;

}

}

VALIDATE_ARRAY_NODE(new_node)

/* That's it! */

res = (PyHamtNode *)new_node;

fin:

Py_XDECREF(empty);

if (res == NULL) {

Py_XDECREF(new_node);

}

return res;

}

else {

/* We have less than 16 keys at this level; let's just

uint32_t key_idx = 2 * idx;

uint32_t val_idx = key_idx + 1;

uint32_t i;

*added_leaf = 1;

/* Allocate new Bitmap node which can have one more key/val

PyHamtNode_Bitmap *new_node =

(PyHamtNode_Bitmap *)hamt_node_bitmap_new(2 * (n + 1));

if (new_node == NULL) {

return NULL;

}

/* Copy all keys/values that will be before the new key/value

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

Py_XINCREF(self->b_array[i]);

new_node->b_array[i] = self->b_array[i];

}

/* Set the new key/value to the new Bitmap node. */

Py_INCREF(key);

new_node->b_array[key_idx] = key;

Py_INCREF(val);

new_node->b_array[val_idx] = val;

/* Copy all keys/values that will be after the new key/value

assert(Py_SIZE(self) >= 0 && Py_SIZE(self) <= 32);

for (i = key_idx; i < (uint32_t)Py_SIZE(self); i++) {

Py_XINCREF(self->b_array[i]);

new_node->b_array[i + 2] = self->b_array[i];

}

new_node->b_bitmap = self->b_bitmap | bit;

return (PyHamtNode *)new_node;

}

}

}

static hamt_without_t

hamt_node_bitmap_without(PyHamtNode_Bitmap *self,

uint32_t shift, int32_t hash,

PyObject *key,

PyHamtNode **new_node)

{

uint32_t bit = hamt_bitpos(hash, shift);

if ((self->b_bitmap & bit) == 0) {

return W_NOT_FOUND;

}

uint32_t idx = hamt_bitindex(self->b_bitmap, bit);

uint32_t key_idx = 2 * idx;

uint32_t val_idx = key_idx + 1;

PyObject *key_or_null = self->b_array[key_idx];

PyObject *val_or_node = self->b_array[val_idx];

if (key_or_null == NULL) {

/* key == NULL means that 'value' is another tree node. */

PyHamtNode *sub_node = NULL;

hamt_without_t res = hamt_node_without(

(PyHamtNode *)val_or_node,

shift + 5, hash, key, &sub_node);

switch (res) {

case W_EMPTY:

/* It's impossible for us to receive a W_EMPTY here:

Py_UNREACHABLE();

case W_NEWNODE: {

assert(sub_node != NULL);

if (IS_BITMAP_NODE(sub_node)) {

PyHamtNode_Bitmap *sub_tree = (PyHamtNode_Bitmap *)sub_node;

if (hamt_node_bitmap_count(sub_tree) == 1 &&

sub_tree->b_array[0] != NULL)

{

/* A bitmap node with one key/value pair. Just