diff options
| -rw-r--r-- | immutables/_map.c | 628 | ||||
| -rw-r--r-- | immutables/_map.h | 28 | ||||
| -rw-r--r-- | immutables/map.py | 272 | ||||
| -rw-r--r-- | tests/test_map.py | 127 |
4 files changed, 838 insertions, 217 deletions
diff --git a/immutables/_map.c b/immutables/_map.c index 8eedd56..b276d8a 100644 --- a/immutables/_map.c +++ b/immutables/_map.c @@ -289,11 +289,13 @@ typedef struct { PyObject_HEAD MapNode *a_array[HAMT_ARRAY_NODE_SIZE]; Py_ssize_t a_count; + uint64_t a_mutid; } MapNode_Array; typedef struct { PyObject_VAR_HEAD + uint64_t b_mutid; uint32_t b_bitmap; PyObject *b_array[1]; } MapNode_Bitmap; @@ -301,11 +303,14 @@ typedef struct { typedef struct { PyObject_VAR_HEAD + uint64_t c_mutid; int32_t c_hash; PyObject *c_array[1]; } MapNode_Collision; +static volatile uint64_t mutid_counter = 1; + static MapNode_Bitmap *_empty_bitmap_node; static MapObject *_empty_map; @@ -331,14 +336,14 @@ map_without(MapObject *o, PyObject *key); - -1: An error occurred. */ static int -map_eq(MapObject *v, MapObject *w); +map_eq(BaseMapObject *v, BaseMapObject *w); static map_find_t -map_find(MapObject *o, PyObject *key, PyObject **val); +map_find(BaseMapObject *o, PyObject *key, PyObject **val); /* Return the size of "o"; equivalent of "len(o)". */ static Py_ssize_t -map_len(MapObject *o); +map_len(BaseMapObject *o); static MapObject * @@ -347,13 +352,15 @@ map_alloc(void); static MapNode * map_node_assoc(MapNode *node, uint32_t shift, int32_t hash, - PyObject *key, PyObject *val, int* added_leaf); + PyObject *key, PyObject *val, int* added_leaf, + uint64_t mutid); static map_without_t map_node_without(MapNode *node, uint32_t shift, int32_t hash, PyObject *key, - MapNode **new_node); + MapNode **new_node, + uint64_t mutid); static map_find_t map_node_find(MapNode *node, @@ -365,10 +372,10 @@ map_node_dump(MapNode *node, _PyUnicodeWriter *writer, int level); static MapNode * -map_node_array_new(Py_ssize_t); +map_node_array_new(Py_ssize_t, uint64_t mutid); static MapNode * -map_node_collision_new(int32_t hash, Py_ssize_t size); +map_node_collision_new(int32_t hash, Py_ssize_t size, uint64_t mutid); static inline Py_ssize_t map_node_collision_count(MapNode_Collision *node); @@ -529,7 +536,7 @@ _map_dump_format(_PyUnicodeWriter *writer, const char *format, ...) static MapNode * -map_node_bitmap_new(Py_ssize_t size) +map_node_bitmap_new(Py_ssize_t size, uint64_t mutid) { /* Create a new bitmap node of size 'size' */ @@ -539,7 +546,7 @@ map_node_bitmap_new(Py_ssize_t size) assert(size >= 0); assert(size % 2 == 0); - if (size == 0 && _empty_bitmap_node != NULL) { + if (size == 0 && _empty_bitmap_node != NULL && mutid == 0) { Py_INCREF(_empty_bitmap_node); return (MapNode *)_empty_bitmap_node; } @@ -558,10 +565,11 @@ map_node_bitmap_new(Py_ssize_t size) } node->b_bitmap = 0; + node->b_mutid = mutid; PyObject_GC_Track(node); - if (size == 0 && _empty_bitmap_node == NULL) { + if (size == 0 && _empty_bitmap_node == NULL && mutid == 0) { /* Since bitmap nodes are immutable, we can cache the instance for size=0 and reuse it whenever we need an empty bitmap node. */ @@ -579,14 +587,15 @@ map_node_bitmap_count(MapNode_Bitmap *node) } static MapNode_Bitmap * -map_node_bitmap_clone(MapNode_Bitmap *node) +map_node_bitmap_clone(MapNode_Bitmap *node, uint64_t mutid) { /* Clone a bitmap node; return a new one with the same child notes. */ MapNode_Bitmap *clone; Py_ssize_t i; - clone = (MapNode_Bitmap *)map_node_bitmap_new(Py_SIZE(node)); + clone = (MapNode_Bitmap *)map_node_bitmap_new( + Py_SIZE(node), mutid); if (clone == NULL) { return NULL; } @@ -601,13 +610,13 @@ map_node_bitmap_clone(MapNode_Bitmap *node) } static MapNode_Bitmap * -map_node_bitmap_clone_without(MapNode_Bitmap *o, uint32_t bit) +map_node_bitmap_clone_without(MapNode_Bitmap *o, uint32_t bit, uint64_t mutid) { assert(bit & o->b_bitmap); assert(map_node_bitmap_count(o) > 1); MapNode_Bitmap *new = (MapNode_Bitmap *)map_node_bitmap_new( - Py_SIZE(o) - 2); + Py_SIZE(o) - 2, mutid); if (new == NULL) { return NULL; } @@ -636,7 +645,8 @@ static MapNode * map_node_new_bitmap_or_collision(uint32_t shift, PyObject *key1, PyObject *val1, int32_t key2_hash, - PyObject *key2, PyObject *val2) + PyObject *key2, PyObject *val2, + uint64_t mutid) { /* Helper method. Creates a new node for key1/val and key2/val2 pairs. @@ -653,7 +663,7 @@ map_node_new_bitmap_or_collision(uint32_t shift, if (key1_hash == key2_hash) { MapNode_Collision *n; - n = (MapNode_Collision *)map_node_collision_new(key1_hash, 4); + n = (MapNode_Collision *)map_node_collision_new(key1_hash, 4, mutid); if (n == NULL) { return NULL; } @@ -672,19 +682,20 @@ map_node_new_bitmap_or_collision(uint32_t shift, } else { int added_leaf = 0; - MapNode *n = map_node_bitmap_new(0); + MapNode *n = map_node_bitmap_new(0, mutid); if (n == NULL) { return NULL; } MapNode *n2 = map_node_assoc( - n, shift, key1_hash, key1, val1, &added_leaf); + n, shift, key1_hash, key1, val1, &added_leaf, mutid); Py_DECREF(n); if (n2 == NULL) { return NULL; } - n = map_node_assoc(n2, shift, key2_hash, key2, val2, &added_leaf); + n = map_node_assoc( + n2, shift, key2_hash, key2, val2, &added_leaf, mutid); Py_DECREF(n2); if (n == NULL) { return NULL; @@ -697,7 +708,8 @@ map_node_new_bitmap_or_collision(uint32_t shift, static MapNode * map_node_bitmap_assoc(MapNode_Bitmap *self, uint32_t shift, int32_t hash, - PyObject *key, PyObject *val, int* added_leaf) + PyObject *key, PyObject *val, int* added_leaf, + uint64_t mutid) { /* assoc operation for bitmap nodes. @@ -744,7 +756,8 @@ map_node_bitmap_assoc(MapNode_Bitmap *self, MapNode *sub_node = map_node_assoc( (MapNode *)val_or_node, - shift + 5, hash, key, val, added_leaf); + shift + 5, hash, key, val, added_leaf, + mutid); if (sub_node == NULL) { return NULL; } @@ -755,12 +768,19 @@ map_node_bitmap_assoc(MapNode_Bitmap *self, return (MapNode *)self; } - MapNode_Bitmap *ret = map_node_bitmap_clone(self); - if (ret == NULL) { - return NULL; + if (mutid != 0 && self->b_mutid == mutid) { + Py_SETREF(self->b_array[val_idx], (PyObject*)sub_node); + Py_INCREF(self); + return (MapNode *)self; + } + else { + MapNode_Bitmap *ret = map_node_bitmap_clone(self, mutid); + if (ret == NULL) { + return NULL; + } + Py_SETREF(ret->b_array[val_idx], (PyObject*)sub_node); + return (MapNode *)ret; } - Py_SETREF(ret->b_array[val_idx], (PyObject*)sub_node); - return (MapNode *)ret; } assert(key != NULL); @@ -780,13 +800,21 @@ map_node_bitmap_assoc(MapNode_Bitmap *self, /* We're setting a new value for the key we had before. Make a new bitmap node with a replaced value, and return it. */ - MapNode_Bitmap *ret = map_node_bitmap_clone(self); - if (ret == NULL) { - return NULL; + if (mutid != 0 && self->b_mutid == mutid) { + Py_INCREF(val); + Py_SETREF(self->b_array[val_idx], val); + Py_INCREF(self); + return (MapNode *)self; + } + else { + MapNode_Bitmap *ret = map_node_bitmap_clone(self, mutid); + if (ret == NULL) { + return NULL; + } + Py_INCREF(val); + Py_SETREF(ret->b_array[val_idx], val); + return (MapNode *)ret; } - Py_INCREF(val); - Py_SETREF(ret->b_array[val_idx], val); - return (MapNode *)ret; } /* It's a new key, and it has the same index as *one* another key. @@ -801,22 +829,33 @@ map_node_bitmap_assoc(MapNode_Bitmap *self, shift + 5, key_or_null, val_or_node, /* existing key/val */ hash, - key, val /* new key/val */ + key, val, /* new key/val */ + self->b_mutid ); if (sub_node == NULL) { return NULL; } - MapNode_Bitmap *ret = map_node_bitmap_clone(self); - if (ret == NULL) { - Py_DECREF(sub_node); - return NULL; + if (mutid != 0 && self->b_mutid == mutid) { + Py_SETREF(self->b_array[key_idx], NULL); + Py_SETREF(self->b_array[val_idx], (PyObject *)sub_node); + Py_INCREF(self); + + *added_leaf = 1; + return (MapNode *)self; } - Py_SETREF(ret->b_array[key_idx], NULL); - Py_SETREF(ret->b_array[val_idx], (PyObject *)sub_node); + else { + MapNode_Bitmap *ret = map_node_bitmap_clone(self, mutid); + 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 (MapNode *)ret; + *added_leaf = 1; + return (MapNode *)ret; + } } else { /* There was no key before with the same (shift,hash). */ @@ -848,14 +887,14 @@ map_node_bitmap_assoc(MapNode_Bitmap *self, MapNode *res = NULL; /* Create a new Array node. */ - new_node = (MapNode_Array *)map_node_array_new(n + 1); + new_node = (MapNode_Array *)map_node_array_new(n + 1, mutid); if (new_node == NULL) { goto fin; } /* Create an empty bitmap node for the next map_node_assoc call. */ - empty = map_node_bitmap_new(0); + empty = map_node_bitmap_new(0, mutid); if (empty == NULL) { goto fin; } @@ -863,7 +902,7 @@ map_node_bitmap_assoc(MapNode_Bitmap *self, /* Make a new bitmap node for the key/val we're adding. Set that bitmap node to new-array-node[jdx]. */ new_node->a_array[jdx] = map_node_assoc( - empty, shift + 5, hash, key, val, added_leaf); + empty, shift + 5, hash, key, val, added_leaf, mutid); if (new_node->a_array[jdx] == NULL) { goto fin; } @@ -894,7 +933,8 @@ map_node_bitmap_assoc(MapNode_Bitmap *self, rehash, self->b_array[j], self->b_array[j + 1], - added_leaf); + added_leaf, + mutid); if (new_node->a_array[i] == NULL) { goto fin; @@ -930,7 +970,7 @@ map_node_bitmap_assoc(MapNode_Bitmap *self, /* Allocate new Bitmap node which can have one more key/val pair in addition to what we have already. */ MapNode_Bitmap *new_node = - (MapNode_Bitmap *)map_node_bitmap_new(2 * (n + 1)); + (MapNode_Bitmap *)map_node_bitmap_new(2 * (n + 1), mutid); if (new_node == NULL) { return NULL; } @@ -966,7 +1006,8 @@ static map_without_t map_node_bitmap_without(MapNode_Bitmap *self, uint32_t shift, int32_t hash, PyObject *key, - MapNode **new_node) + MapNode **new_node, + uint64_t mutid) { uint32_t bit = map_bitpos(hash, shift); if ((self->b_bitmap & bit) == 0) { @@ -985,10 +1026,12 @@ map_node_bitmap_without(MapNode_Bitmap *self, /* key == NULL means that 'value' is another tree node. */ MapNode *sub_node = NULL; + MapNode_Bitmap *target = NULL; map_without_t res = map_node_without( (MapNode *)val_or_node, - shift + 5, hash, key, &sub_node); + shift + 5, hash, key, &sub_node, + mutid); switch (res) { case W_EMPTY: @@ -1025,23 +1068,29 @@ map_node_bitmap_without(MapNode_Bitmap *self, up or down. */ - MapNode_Bitmap *clone = map_node_bitmap_clone(self); - if (clone == NULL) { - Py_DECREF(sub_node); - return W_ERROR; + if (mutid != 0 && self->b_mutid == mutid) { + target = self; + Py_INCREF(target); + } + else { + target = map_node_bitmap_clone(self, mutid); + if (target == NULL) { + Py_DECREF(sub_node); + return W_ERROR; + } } PyObject *key = sub_tree->b_array[0]; PyObject *val = sub_tree->b_array[1]; Py_INCREF(key); - Py_XSETREF(clone->b_array[key_idx], key); + Py_XSETREF(target->b_array[key_idx], key); Py_INCREF(val); - Py_SETREF(clone->b_array[val_idx], val); + Py_SETREF(target->b_array[val_idx], val); Py_DECREF(sub_tree); - *new_node = (MapNode *)clone; + *new_node = (MapNode *)target; return W_NEWNODE; } } @@ -1056,15 +1105,21 @@ map_node_bitmap_without(MapNode_Bitmap *self, } #endif - MapNode_Bitmap *clone = map_node_bitmap_clone(self); - if (clone == NULL) { - return W_ERROR; + if (mutid != 0 && self->b_mutid == mutid) { + target = self; + Py_INCREF(target); + } + else { + target = map_node_bitmap_clone(self, mutid); + if (target == NULL) { + return W_ERROR; + } } - Py_SETREF(clone->b_array[val_idx], + Py_SETREF(target->b_array[val_idx], (PyObject *)sub_node); /* borrow */ - *new_node = (MapNode *)clone; + *new_node = (MapNode *)target; return W_NEWNODE; } @@ -1093,7 +1148,7 @@ map_node_bitmap_without(MapNode_Bitmap *self, } *new_node = (MapNode *) - map_node_bitmap_clone_without(self, bit); + map_node_bitmap_clone_without(self, bit, mutid); if (*new_node == NULL) { return W_ERROR; } @@ -1266,7 +1321,7 @@ error: static MapNode * -map_node_collision_new(int32_t hash, Py_ssize_t size) +map_node_collision_new(int32_t hash, Py_ssize_t size, uint64_t mutid) { /* Create a new Collision node. */ @@ -1289,8 +1344,9 @@ map_node_collision_new(int32_t hash, Py_ssize_t size) Py_SIZE(node) = size; node->c_hash = hash; - PyObject_GC_Track(node); + node->c_mutid = mutid; + PyObject_GC_Track(node); return (MapNode *)node; } @@ -1324,7 +1380,8 @@ map_node_collision_find_index(MapNode_Collision *self, PyObject *key, static MapNode * map_node_collision_assoc(MapNode_Collision *self, uint32_t shift, int32_t hash, - PyObject *key, PyObject *val, int* added_leaf) + PyObject *key, PyObject *val, int* added_leaf, + uint64_t mutid) { /* Set a new key to this level (currently a Collision node) of the tree. */ @@ -1350,7 +1407,7 @@ map_node_collision_assoc(MapNode_Collision *self, add a new key/value to the cloned node. */ new_node = (MapNode_Collision *)map_node_collision_new( - self->c_hash, Py_SIZE(self) + 2); + self->c_hash, Py_SIZE(self) + 2, mutid); if (new_node == NULL) { return NULL; } @@ -1381,18 +1438,26 @@ map_node_collision_assoc(MapNode_Collision *self, return (MapNode *)self; } - /* We need to replace old value for the key - with a new value. Create a new Collision node.*/ - new_node = (MapNode_Collision *)map_node_collision_new( - self->c_hash, Py_SIZE(self)); - if (new_node == NULL) { - return NULL; + /* We need to replace old value for the key with + a new value. */ + + if (mutid != 0 && self->c_mutid == mutid) { + new_node = self; + Py_INCREF(self); } + else { + /* Create a new Collision node.*/ + new_node = (MapNode_Collision *)map_node_collision_new( + self->c_hash, Py_SIZE(self), mutid); + if (new_node == NULL) { + return NULL; + } - /* Copy all elements of the old node to the new one. */ - for (i = 0; i < Py_SIZE(self); i++) { - Py_INCREF(self->c_array[i]); - new_node->c_array[i] = self->c_array[i]; + /* Copy all elements of the old node to the new one. */ + for (i = 0; i < Py_SIZE(self); i++) { + Py_INCREF(self->c_array[i]); + new_node->c_array[i] = self->c_array[i]; + } } /* Replace the old value with the new value for the our key. */ @@ -1418,7 +1483,7 @@ map_node_collision_assoc(MapNode_Collision *self, MapNode_Bitmap *new_node; MapNode *assoc_res; - new_node = (MapNode_Bitmap *)map_node_bitmap_new(2); + new_node = (MapNode_Bitmap *)map_node_bitmap_new(2, mutid); if (new_node == NULL) { return NULL; } @@ -1427,7 +1492,7 @@ map_node_collision_assoc(MapNode_Collision *self, new_node->b_array[1] = (PyObject*) self; assoc_res = map_node_bitmap_assoc( - new_node, shift, hash, key, val, added_leaf); + new_node, shift, hash, key, val, added_leaf, mutid); Py_DECREF(new_node); return assoc_res; } @@ -1443,7 +1508,8 @@ static map_without_t map_node_collision_without(MapNode_Collision *self, uint32_t shift, int32_t hash, PyObject *key, - MapNode **new_node) + MapNode **new_node, + uint64_t mutid) { if (hash != self->c_hash) { return W_NOT_FOUND; @@ -1480,7 +1546,7 @@ map_node_collision_without(MapNode_Collision *self, Bitmap node. */ MapNode_Bitmap *node = (MapNode_Bitmap *) - map_node_bitmap_new(2); + map_node_bitmap_new(2, mutid); if (node == NULL) { return W_ERROR; } @@ -1509,7 +1575,7 @@ map_node_collision_without(MapNode_Collision *self, less key/value pair */ MapNode_Collision *new = (MapNode_Collision *) map_node_collision_new( - self->c_hash, Py_SIZE(self) - 2); + self->c_hash, Py_SIZE(self) - 2, mutid); if (new == NULL) { return W_ERROR; } @@ -1636,7 +1702,7 @@ error: static MapNode * -map_node_array_new(Py_ssize_t count) +map_node_array_new(Py_ssize_t count, uint64_t mutid) { Py_ssize_t i; @@ -1651,13 +1717,14 @@ map_node_array_new(Py_ssize_t count) } node->a_count = count; + node->a_mutid = mutid; PyObject_GC_Track(node); return (MapNode *)node; } static MapNode_Array * -map_node_array_clone(MapNode_Array *node) +map_node_array_clone(MapNode_Array *node, uint64_t mutid) { MapNode_Array *clone; Py_ssize_t i; @@ -1665,7 +1732,7 @@ map_node_array_clone(MapNode_Array *node) VALIDATE_ARRAY_NODE(node) /* Create a new Array node. */ - clone = (MapNode_Array *)map_node_array_new(node->a_count); + clone = (MapNode_Array *)map_node_array_new(node->a_count, mutid); if (clone == NULL) { return NULL; } @@ -1676,6 +1743,8 @@ map_node_array_clone(MapNode_Array *node) clone->a_array[i] = node->a_array[i]; } + clone->a_mutid = mutid; + VALIDATE_ARRAY_NODE(clone) return clone; } @@ -1683,7 +1752,8 @@ map_node_array_clone(MapNode_Array *node) static MapNode * map_node_array_assoc(MapNode_Array *self, uint32_t shift, int32_t hash, - PyObject *key, PyObject *val, int* added_leaf) + PyObject *key, PyObject *val, int* added_leaf, + uint64_t mutid) { /* Set a new key to this level (currently a Collision node) of the tree. @@ -1705,7 +1775,7 @@ map_node_array_assoc(MapNode_Array *self, MapNode_Bitmap *empty = NULL; /* Get an empty Bitmap node to work with. */ - empty = (MapNode_Bitmap *)map_node_bitmap_new(0); + empty = (MapNode_Bitmap *)map_node_bitmap_new(0, mutid); if (empty == NULL) { return NULL; } @@ -1714,24 +1784,31 @@ map_node_array_assoc(MapNode_Array *self, creating a new Bitmap node with our key/value pair. */ child_node = map_node_bitmap_assoc( empty, - shift + 5, hash, key, val, added_leaf); + shift + 5, hash, key, val, added_leaf, mutid); Py_DECREF(empty); if (child_node == NULL) { return NULL; } - /* Create a new Array node. */ - new_node = (MapNode_Array *)map_node_array_new(self->a_count + 1); - if (new_node == NULL) { - Py_DECREF(child_node); - return NULL; + if (mutid != 0 && self->a_mutid == mutid) { + new_node = self; + Py_INCREF(self); } + else { + /* Create a new Array node. */ + new_node = (MapNode_Array *)map_node_array_new( + self->a_count + 1, mutid); + if (new_node == NULL) { + Py_DECREF(child_node); + return NULL; + } - /* Copy all elements from the current Array node to the - new one. */ - for (i = 0; i < HAMT_ARRAY_NODE_SIZE; i++) { - Py_XINCREF(self->a_array[i]); - new_node->a_array[i] = self->a_array[i]; + /* Copy all elements from the current Array node to the + new one. */ + for (i = 0; i < HAMT_ARRAY_NODE_SIZE; i++) { + Py_XINCREF(self->a_array[i]); + new_node->a_array[i] = self->a_array[i]; + } } assert(new_node->a_array[idx] == NULL); @@ -1741,8 +1818,9 @@ map_node_array_assoc(MapNode_Array *self, else { /* There's a child node for the given hash. Set the key to it./ */ + child_node = map_node_assoc( - node, shift + 5, hash, key, val, added_leaf); + node, shift + 5, hash, key, val, added_leaf, mutid); if (child_node == NULL) { return NULL; } @@ -1751,7 +1829,14 @@ map_node_array_assoc(MapNode_Array *self, return (MapNode *)self; } - new_node = map_node_array_clone(self); + if (mutid != 0 && self->a_mutid == mutid) { + new_node = self; + Py_INCREF(self); + } + else { + new_node = map_node_array_clone(self, mutid); + } + if (new_node == NULL) { Py_DECREF(child_node); return NULL; @@ -1768,7 +1853,8 @@ static map_without_t map_node_array_without(MapNode_Array *self, uint32_t shift, int32_t hash, PyObject *key, - MapNode **new_node) + MapNode **new_node, + uint64_t mutid) { uint32_t idx = map_mask(hash, shift); MapNode *node = self->a_array[idx]; @@ -1778,9 +1864,10 @@ map_node_array_without(MapNode_Array *self, } MapNode *sub_node = NULL; + MapNode_Array *target = NULL; map_without_t res = map_node_without( (MapNode *)node, - shift + 5, hash, key, &sub_node); + shift + 5, hash, key, &sub_node, mutid); switch (res) { case W_NOT_FOUND: @@ -1794,14 +1881,20 @@ map_node_array_without(MapNode_Array *self, */ assert(sub_node != NULL); - MapNode_Array *clone = map_node_array_clone(self); - if (clone == NULL) { - Py_DECREF(sub_node); - return W_ERROR; + if (mutid != 0 && self->a_mutid == mutid) { + target = self; + Py_INCREF(self); + } + else { + target = map_node_array_clone(self, mutid); + if (target == NULL) { + Py_DECREF(sub_node); + return W_ERROR; + } } - Py_SETREF(clone->a_array[idx], sub_node); /* borrow */ - *new_node = (MapNode*)clone; /* borrow */ + Py_SETREF(target->a_array[idx], sub_node); /* borrow */ + *new_node = (MapNode*)target; /* borrow */ return W_NEWNODE; } @@ -1824,14 +1917,21 @@ map_node_array_without(MapNode_Array *self, greater than 15. */ - MapNode_Array *new = map_node_array_clone(self); - if (new == NULL) { - return W_ERROR; + if (mutid != 0 && self->a_mutid == mutid) { + target = self; + Py_INCREF(self); + } + else { + target = map_node_array_clone(self, mutid); + if (target == NULL) { + return W_ERROR; + } + target->a_count = new_count; } - new->a_count = new_count; - Py_CLEAR(new->a_array[idx]); - *new_node = (MapNode*)new; /* borrow */ + Py_CLEAR(target->a_array[idx]); + + *new_node = (MapNode*)target; /* borrow */ return W_NEWNODE; } @@ -1842,7 +1942,7 @@ map_node_array_without(MapNode_Array *self, uint32_t bitmap = 0; MapNode_Bitmap *new = (MapNode_Bitmap *) - map_node_bitmap_new(bitmap_size); + map_node_bitmap_new(bitmap_size, mutid); if (new == NULL) { return W_ERROR; } @@ -2025,7 +2125,8 @@ error: static MapNode * map_node_assoc(MapNode *node, uint32_t shift, int32_t hash, - PyObject *key, PyObject *val, int* added_leaf) + PyObject *key, PyObject *val, int* added_leaf, + uint64_t mutid) { /* Set key/value to the 'node' starting with the given shift/hash. Return a new node, or the same node if key/value already @@ -2041,18 +2142,18 @@ map_node_assoc(MapNode *node, if (IS_BITMAP_NODE(node)) { return map_node_bitmap_assoc( (MapNode_Bitmap *)node, - shift, hash, key, val, added_leaf); + shift, hash, key, val, added_leaf, mutid); } else if (IS_ARRAY_NODE(node)) { return map_node_array_assoc( (MapNode_Array *)node, - shift, hash, key, val, added_leaf); + shift, hash, key, val, added_leaf, mutid); } else { assert(IS_COLLISION_NODE(node)); return map_node_collision_assoc( (MapNode_Collision *)node, - shift, hash, key, val, added_leaf); + shift, hash, key, val, added_leaf, mutid); } } @@ -2060,26 +2161,30 @@ static map_without_t map_node_without(MapNode *node, uint32_t shift, int32_t hash, PyObject *key, - MapNode **new_node) + MapNode **new_node, + uint64_t mutid) { if (IS_BITMAP_NODE(node)) { return map_node_bitmap_without( (MapNode_Bitmap *)node, shift, hash, key, - new_node); + new_node, + mutid); } else if (IS_ARRAY_NODE(node)) { return map_node_array_without( (MapNode_Array *)node, shift, hash, key, - new_node); + new_node, + mutid); } else { assert(IS_COLLISION_NODE(node)); return map_node_collision_without( (MapNode_Collision *)node, shift, hash, key, - new_node); + new_node, + mutid); } } @@ -2311,7 +2416,8 @@ map_assoc(MapObject *o, PyObject *key, PyObject *val) new_root = map_node_assoc( (MapNode *)(o->h_root), - 0, key_hash, key, val, &added_leaf); + 0, key_hash, key, val, &added_leaf, + 0); if (new_root == NULL) { return NULL; } @@ -2347,7 +2453,8 @@ map_without(MapObject *o, PyObject *key) map_without_t res = map_node_without( (MapNode *)(o->h_root), 0, key_hash, key, - &new_root); + &new_root, + 0); switch (res) { case W_ERROR: @@ -2377,9 +2484,9 @@ map_without(MapObject *o, PyObject *key) } static map_find_t -map_find(MapObject *o, PyObject *key, PyObject **val) +map_find(BaseMapObject *o, PyObject *key, PyObject **val) { - if (o->h_count == 0) { + if (o->b_count == 0) { return F_NOT_FOUND; } @@ -2388,17 +2495,17 @@ map_find(MapObject *o, PyObject *key, PyObject **val) return F_ERROR; } - return map_node_find(o->h_root, 0, key_hash, key, val); + return map_node_find(o->b_root, 0, key_hash, key, val); } static int -map_eq(MapObject *v, MapObject *w) +map_eq(BaseMapObject *v, BaseMapObject *w) { if (v == w) { return 1; } - if (v->h_count != w->h_count) { + if (v->b_count != w->b_count) { return 0; } @@ -2409,7 +2516,7 @@ map_eq(MapObject *v, MapObject *w) PyObject *v_val; PyObject *w_val; - map_iterator_init(&iter, v->h_root); + map_iterator_init(&iter, v->b_root); do { iter_res = map_iterator_next(&iter, &v_key, &v_val); @@ -2439,9 +2546,9 @@ map_eq(MapObject *v, MapObject *w) } static Py_ssize_t -map_len(MapObject *o) +map_len(BaseMapObject *o) { - return o->h_count; + return o->b_count; } static MapObject * @@ -2475,7 +2582,7 @@ map_new(void) return NULL; } - o->h_root = map_node_bitmap_new(0); + o->h_root = map_node_bitmap_new(0, 0); if (o->h_root == NULL) { Py_DECREF(o); return NULL; @@ -2692,25 +2799,25 @@ map_tp_new(PyTypeObject *type, PyObject *args, PyObject *kwds) } static int -map_tp_clear(MapObject *self) +map_tp_clear(BaseMapObject *self) { - Py_CLEAR(self->h_root); + Py_CLEAR(self->b_root); return 0; } static int -map_tp_traverse(MapObject *self, visitproc visit, void *arg) +map_tp_traverse(BaseMapObject *self, visitproc visit, void *arg) { - Py_VISIT(self->h_root); + Py_VISIT(self->b_root); return 0; } static void -map_tp_dealloc(MapObject *self) +map_tp_dealloc(BaseMapObject *self) { PyObject_GC_UnTrack(self); - if (self->h_weakreflist != NULL) { + if (self->b_weakreflist != NULL) { PyObject_ClearWeakRefs((PyObject*)self); } (void)map_tp_clear(self); @@ -2725,7 +2832,7 @@ map_tp_richcompare(PyObject *v, PyObject *w, int op) Py_RETURN_NOTIMPLEMENTED; } - int res = map_eq((MapObject *)v, (MapObject *)w); + int res = map_eq((BaseMapObject *)v, (BaseMapObject *)w); if (res < 0) { return NULL; } @@ -2743,7 +2850,7 @@ map_tp_richcompare(PyObject *v, PyObject *w, int op) } static int -map_tp_contains(MapObject *self, PyObject *key) +map_tp_contains(BaseMapObject *self, PyObject *key) { PyObject *val; map_find_t res = map_find(self, key, &val); @@ -2760,7 +2867,7 @@ map_tp_contains(MapObject *self, PyObject *key) } static PyObject * -map_tp_subscript(MapObject *self, PyObject *key) +map_tp_subscript(BaseMapObject *self, PyObject *key) { PyObject *val; map_find_t res = map_find(self, key, &val); @@ -2779,7 +2886,7 @@ map_tp_subscript(MapObject *self, PyObject *key) } static Py_ssize_t -map_tp_len(MapObject *self) +map_tp_len(BaseMapObject *self) { return map_len(self); } @@ -2804,7 +2911,7 @@ map_py_set(MapObject *self, PyObject *args) } static PyObject * -map_py_get(MapObject *self, PyObject *args) +map_py_get(BaseMapObject *self, PyObject *args) { PyObject *key; PyObject *def = NULL; @@ -2839,6 +2946,27 @@ map_py_delete(MapObject *self, PyObject *key) } static PyObject * +map_py_mutate(MapObject *self, PyObject *args) +{ + + MapMutationObject *o; + o = PyObject_GC_New(MapMutationObject, &_MapMutation_Type); + if (o == NULL) { + return NULL; + } + o->m_weakreflist = NULL; + o->m_count = self->h_count; + + Py_INCREF(self->h_root); + o->m_root = self->h_root; + + o->m_mutid = mutid_counter++; + + PyObject_GC_Track(o); + return (PyObject *)o; +} + +static PyObject * map_py_items(MapObject *self, PyObject *args) { return map_new_iteritems(self); @@ -2864,28 +2992,37 @@ map_py_dump(MapObject *self, PyObject *args) static PyObject * -map_py_repr(MapObject *self) +map_py_repr(BaseMapObject *m) { Py_ssize_t i; _PyUnicodeWriter writer; - i = Py_ReprEnter((PyObject *)self); + i = Py_ReprEnter((PyObject *)m); if (i != 0) { return i > 0 ? PyUnicode_FromString("{...}") : NULL; } _PyUnicodeWriter_Init(&writer); - if (_PyUnicodeWriter_WriteASCIIString( - &writer, "<immutables.Map({", 17) < 0) - { - goto error; + if (MapMutation_Check(m)) { + if (_PyUnicodeWriter_WriteASCIIString( + &writer, "<immutables.MapMutation({", 25) < 0) + { + goto error; + } + } + else { + if (_PyUnicodeWriter_WriteASCIIString( + &writer, "<immutables.Map({", 17) < 0) + { + goto error; + } } MapIteratorState iter; map_iter_t iter_res; - map_iterator_init(&iter, self->h_root); + map_iterator_init(&iter, m->b_root); int second = 0; do { PyObject *v_key; @@ -2931,7 +3068,7 @@ map_py_repr(MapObject *self) goto error; } - PyObject *addr = PyUnicode_FromFormat(" at %p>", self); + PyObject *addr = PyUnicode_FromFormat(" at %p>", m); if (addr == NULL) { goto error; } @@ -2941,12 +3078,12 @@ map_py_repr(MapObject *self) } Py_DECREF(addr); - Py_ReprLeave((PyObject *)self); + Py_ReprLeave((PyObject *)m); return _PyUnicodeWriter_Finish(&writer); error: _PyUnicodeWriter_Dealloc(&writer); - Py_ReprLeave((PyObject *)self); + Py_ReprLeave((PyObject *)m); return NULL; } @@ -3014,6 +3151,7 @@ static PyMethodDef Map_methods[] = { {"set", (PyCFunction)map_py_set, METH_VARARGS, NULL}, {"get", (PyCFunction)map_py_get, METH_VARARGS, NULL}, {"delete", (PyCFunction)map_py_delete, METH_O, NULL}, + {"mutate", (PyCFunction)map_py_mutate, METH_NOARGS, NULL}, {"items", (PyCFunction)map_py_items, METH_NOARGS, NULL}, {"keys", (PyCFunction)map_py_keys, METH_NOARGS, NULL}, {"values", (PyCFunction)map_py_values, METH_NOARGS, NULL}, @@ -3060,6 +3198,187 @@ PyTypeObject _Map_Type = { }; +/////////////////////////////////// MapMutation + + +static PyObject * +mapmut_py_set(MapMutationObject *o, PyObject *args) +{ + PyObject *key; + PyObject *val; + + if (!PyArg_UnpackTuple(args, "set", 2, 2, &key, &val)) { + return NULL; + } + + int32_t key_hash; + int added_leaf = 0; + + key_hash = map_hash(key); + if (key_hash == -1) { + return NULL; + } + + MapNode *new_root = map_node_assoc( + (MapNode *)(o->m_root), + 0, key_hash, key, val, &added_leaf, + o->m_mutid); + if (new_root == NULL) { + return NULL; + } + + if (added_leaf) { + o->m_count++; + } + + if (new_root == o->m_root) { + Py_DECREF(new_root); + goto done; + } + + Py_SETREF(o->m_root, new_root); + +done: + Py_RETURN_NONE; +} + + +static PyObject * +mapmut_py_delete(MapMutationObject *o, PyObject *key) +{ + int32_t key_hash = map_hash(key); + if (key_hash == -1) { + return NULL; + } + + MapNode *new_root; + + map_without_t res = map_node_without( + (MapNode *)(o->m_root), + 0, key_hash, key, + &new_root, + o->m_mutid); + + switch (res) { + case W_ERROR: + return NULL; + case W_EMPTY: + new_root = map_node_bitmap_new(0, o->m_mutid); + if (new_root == NULL) { + return NULL; + } + Py_SETREF(o->m_root, new_root); + o->m_count = 0; + goto done; + + case W_NOT_FOUND: + PyErr_SetObject(PyExc_KeyError, key); + return NULL; + case W_NEWNODE: { + assert(new_root != NULL); + Py_SETREF(o->m_root, new_root); + o->m_count--; + goto done; + } + default: + abort(); + } + +done: + Py_RETURN_NONE; +} + + +static PyObject * +mapmut_tp_richcompare(PyObject *v, PyObject *w, int op) +{ + if (!MapMutation_Check(v) || !MapMutation_Check(w) || + (op != Py_EQ && op != Py_NE)) + { + Py_RETURN_NOTIMPLEMENTED; + } + + int res = map_eq((BaseMapObject *)v, (BaseMapObject *)w); + if (res < 0) { + return NULL; + } + + if (op == Py_NE) { + res = !res; + } + + if (res) { + Py_RETURN_TRUE; + } + else { + Py_RETURN_FALSE; + } +} + + +static PyObject * +mapmut_py_finalize(MapMutationObject *self, PyObject *args) +{ + self->m_mutid = 0; + + MapObject *o = map_alloc(); + if (o == NULL) { + return NULL; + } + + Py_INCREF(self->m_root); + o->h_root = self->m_root; + o->h_count = self->m_count; + + return (PyObject *)o; +} + + +static PyMethodDef MapMutation_methods[] = { + {"set", (PyCFunction)mapmut_py_set, METH_VARARGS, NULL}, + {"get", (PyCFunction)map_py_get, METH_VARARGS, NULL}, + {"delete", (PyCFunction)mapmut_py_delete, METH_O, NULL}, + {"finalize", (PyCFunction)mapmut_py_finalize, METH_NOARGS, NULL}, + {NULL, NULL} +}; + +static PySequenceMethods MapMutation_as_sequence = { + 0, /* sq_length */ + 0, /* sq_concat */ + 0, /* sq_repeat */ + 0, /* sq_item */ + 0, /* sq_slice */ + 0, /* sq_ass_item */ + 0, /* sq_ass_slice */ + (objobjproc)map_tp_contains, /* sq_contains */ + 0, /* sq_inplace_concat */ + 0, /* sq_inplace_repeat */ +}; + +static PyMappingMethods MapMutation_as_mapping = { + (lenfunc)map_tp_len, /* mp_length */ + (binaryfunc)map_tp_subscript, /* mp_subscript */ +}; + +PyTypeObject _MapMutation_Type = { + PyVarObject_HEAD_INIT(NULL, 0) + "immutables._map.MapMutation", + sizeof(MapMutationObject), + .tp_methods = MapMutation_methods, + .tp_as_mapping = &MapMutation_as_mapping, + .tp_as_sequence = &MapMutation_as_sequence, + .tp_dealloc = (destructor)map_tp_dealloc, + .tp_getattro = PyObject_GenericGetAttr, + .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, + .tp_traverse = (traverseproc)map_tp_traverse, + .tp_richcompare = mapmut_tp_richcompare, + .tp_clear = (inquiry)map_tp_clear, + .tp_weaklistoffset = offsetof(MapMutationObject, m_weakreflist), + .tp_repr = (reprfunc)map_py_repr, + .tp_hash = PyObject_HashNotImplemented, +}; + + /////////////////////////////////// Tree Node Types @@ -3113,7 +3432,7 @@ module_free(void *m) static struct PyModuleDef _mapmodule = { PyModuleDef_HEAD_INIT, /* m_base */ - "_map" , /* m_name */ + "_map", /* m_name */ NULL, /* m_doc */ -1, /* m_size */ NULL, /* m_methods */ @@ -3130,6 +3449,7 @@ PyInit__map(void) PyObject *m = PyModule_Create(&_mapmodule); if ((PyType_Ready(&_Map_Type) < 0) || + (PyType_Ready(&_MapMutation_Type) < 0) || (PyType_Ready(&_Map_ArrayNode_Type) < 0) || (PyType_Ready(&_Map_BitmapNode_Type) < 0) || (PyType_Ready(&_Map_CollisionNode_Type) < 0) || diff --git a/immutables/_map.h b/immutables/_map.h index 6e87c4b..483865c 100644 --- a/immutables/_map.h +++ b/immutables/_map.h @@ -8,6 +8,7 @@ #define Map_Check(o) (Py_TYPE(o) == &_Map_Type) +#define MapMutation_Check(o) (Py_TYPE(o) == &_MapMutation_Type) /* Abstract tree node. */ @@ -16,16 +17,34 @@ typedef struct { } MapNode; +#define _MapCommonFields(pref) \ + PyObject_HEAD \ + MapNode *pref##_root; \ + PyObject *pref##_weakreflist; \ + Py_ssize_t pref##_count; + + +/* Base mapping struct; used in methods shared between + MapObject and MapMutationObject types. */ +typedef struct { + _MapCommonFields(b) +} BaseMapObject; + + /* An HAMT immutable mapping collection. */ typedef struct { - PyObject_HEAD - MapNode *h_root; - PyObject *h_weakreflist; - Py_ssize_t h_count; + _MapCommonFields(h) Py_hash_t h_hash; } MapObject; +/* MapMutation object (returned from `map.mutate()`.) */ +typedef struct { + _MapCommonFields(m) + uint64_t m_mutid; +} MapMutationObject; + + /* A struct to hold the state of depth-first traverse of the tree. HAMT is an immutable collection. Iterators will hold a strong reference @@ -62,6 +81,7 @@ typedef struct { PyTypeObject _Map_Type; +PyTypeObject _MapMutation_Type; PyTypeObject _Map_ArrayNode_Type; PyTypeObject _Map_BitmapNode_Type; PyTypeObject _Map_CollisionNode_Type; diff --git a/immutables/map.py b/immutables/map.py index 81612bc..8c7d862 100644 --- a/immutables/map.py +++ b/immutables/map.py @@ -1,4 +1,5 @@ import collections.abc +import itertools import reprlib import sys @@ -6,6 +7,10 @@ import sys __all__ = ('Map',) +# Thread-safe counter. +_mut_id = itertools.count(1).__next__ + + # Python version of _map.c. The topmost comment there explains # all datastructures and algorithms. # The code here follows C code closely on purpose to make @@ -43,16 +48,17 @@ W_EMPTY, W_NEWNODE, W_NOT_FOUND = range(3) class BitmapNode: - def __init__(self, size, bitmap, array): + def __init__(self, size, bitmap, array, mutid): self.size = size self.bitmap = bitmap assert isinstance(array, list) and len(array) == size self.array = array + self.mutid = mutid - def clone(self): - return BitmapNode(self.size, self.bitmap, self.array.copy()) + def clone(self, mutid): + return BitmapNode(self.size, self.bitmap, self.array.copy(), mutid) - def assoc(self, shift, hash, key, val): + def assoc(self, shift, hash, key, val, mutid): bit = map_bitpos(hash, shift) idx = map_bitindex(self.bitmap, bit) @@ -65,38 +71,53 @@ class BitmapNode: if key_or_null is None: sub_node, added = val_or_node.assoc( - shift + 5, hash, key, val) + shift + 5, hash, key, val, mutid) if val_or_node is sub_node: - return self, False + return self, added - ret = self.clone() - ret.array[val_idx] = sub_node - return ret, added + if mutid and mutid == self.mutid: + self.array[val_idx] = sub_node + return self, added + else: + ret = self.clone(mutid) + ret.array[val_idx] = sub_node + return ret, added if key == key_or_null: if val is val_or_node: return self, False - ret = self.clone() - ret.array[val_idx] = val - return ret, False + if mutid and mutid == self.mutid: + self.array[val_idx] = val + return self, False + else: + ret = self.clone(mutid) + ret.array[val_idx] = val + return ret, False existing_key_hash = map_hash(key_or_null) if existing_key_hash == hash: sub_node = CollisionNode( - 4, hash, [key_or_null, val_or_node, key, val]) + 4, hash, [key_or_null, val_or_node, key, val], mutid) else: - sub_node = BitmapNode(0, 0, []) + sub_node = BitmapNode(0, 0, [], mutid) sub_node, _ = sub_node.assoc( shift + 5, existing_key_hash, - key_or_null, val_or_node) + key_or_null, val_or_node, + mutid) sub_node, _ = sub_node.assoc( - shift + 5, hash, key, val) + shift + 5, hash, key, val, + mutid) - ret = self.clone() - ret.array[key_idx] = None - ret.array[val_idx] = sub_node - return ret, True + if mutid and mutid == self.mutid: + self.array[key_idx] = None + self.array[val_idx] = sub_node + return self, True + else: + ret = self.clone(mutid) + ret.array[key_idx] = None + ret.array[val_idx] = sub_node + return ret, True else: key_idx = 2 * idx @@ -108,7 +129,15 @@ class BitmapNode: new_array.append(key) new_array.append(val) new_array.extend(self.array[key_idx:]) - return BitmapNode(2 * (n + 1), self.bitmap | bit, new_array), True + + if mutid and mutid == self.mutid: + self.size = 2 * (n + 1) + self.bitmap |= bit + self.array = new_array + return self, True + else: + return BitmapNode( + 2 * (n + 1), self.bitmap | bit, new_array, mutid), True def find(self, shift, hash, key): bit = map_bitpos(hash, shift) @@ -131,7 +160,7 @@ class BitmapNode: raise KeyError(key) - def without(self, shift, hash, key): + def without(self, shift, hash, key, mutid): bit = map_bitpos(hash, shift) if not (self.bitmap & bit): return W_NOT_FOUND, None @@ -144,7 +173,7 @@ class BitmapNode: val_or_node = self.array[val_idx] if key_or_null is None: - res, sub_node = val_or_node.without(shift + 5, hash, key) + res, sub_node = val_or_node.without(shift + 5, hash, key, mutid) if res is W_EMPTY: raise RuntimeError('unreachable code') # pragma: no cover @@ -153,14 +182,24 @@ class BitmapNode: if (type(sub_node) is BitmapNode and sub_node.size == 2 and sub_node.array[0] is not None): - clone = self.clone() - clone.array[key_idx] = sub_node.array[0] - clone.array[val_idx] = sub_node.array[1] - return W_NEWNODE, clone - clone = self.clone() - clone.array[val_idx] = sub_node - return W_NEWNODE, clone + if mutid and mutid == self.mutid: + self.array[key_idx] = sub_node.array[0] + self.array[val_idx] = sub_node.array[1] + return W_NEWNODE, self + else: + clone = self.clone(mutid) + clone.array[key_idx] = sub_node.array[0] + clone.array[val_idx] = sub_node.array[1] + return W_NEWNODE, clone + + if mutid and mutid == self.mutid: + self.array[val_idx] = sub_node + return W_NEWNODE, self + else: + clone = self.clone(mutid) + clone.array[val_idx] = sub_node + return W_NEWNODE, clone else: assert sub_node is None @@ -173,9 +212,16 @@ class BitmapNode: new_array = self.array[:key_idx] new_array.extend(self.array[val_idx + 1:]) - new_node = BitmapNode( - self.size - 2, self.bitmap & ~bit, new_array) - return W_NEWNODE, new_node + + if mutid and mutid == self.mutid: + self.size -= 2 + self.bitmap &= ~bit + self.array = new_array + return W_NEWNODE, self + else: + new_node = BitmapNode( + self.size - 2, self.bitmap & ~bit, new_array, mutid) + return W_NEWNODE, new_node else: return W_NOT_FOUND, None @@ -231,10 +277,11 @@ class BitmapNode: class CollisionNode: - def __init__(self, size, hash, array): + def __init__(self, size, hash, array, mutid): self.size = size self.hash = hash self.array = array + self.mutid = mutid def find_index(self, key): for i in range(0, self.size, 2): @@ -248,7 +295,7 @@ class CollisionNode: return self.array[i + 1] raise KeyError(key) - def assoc(self, shift, hash, key, val): + def assoc(self, shift, hash, key, val, mutid): if hash == self.hash: key_idx = self.find_index(key) @@ -256,23 +303,34 @@ class CollisionNode: new_array = self.array.copy() new_array.append(key) new_array.append(val) - new_node = CollisionNode(self.size + 2, hash, new_array) - return new_node, True + + if mutid and mutid == self.mutid: + self.size += 2 + self.array = new_array + return self, True + else: + new_node = CollisionNode( + self.size + 2, hash, new_array, mutid) + return new_node, True val_idx = key_idx + 1 if self.array[val_idx] is val: return self, False - new_array = self.array.copy() - new_array[val_idx] = val - return CollisionNode(self.size, hash, new_array), False + if mutid and mutid == self.mutid: + self.array[val_idx] = val + return self, False + else: + new_array = self.array.copy() + new_array[val_idx] = val + return CollisionNode(self.size, hash, new_array, mutid), False else: new_node = BitmapNode( - 2, map_bitpos(self.hash, shift), [None, self]) - return new_node.assoc(shift, hash, key, val) + 2, map_bitpos(self.hash, shift), [None, self], mutid) + return new_node.assoc(shift, hash, key, val, mutid) - def without(self, shift, hash, key): + def without(self, shift, hash, key, mutid): if hash != self.hash: return W_NOT_FOUND, None @@ -292,13 +350,20 @@ class CollisionNode: assert key_idx == 2 new_array = [self.array[0], self.array[1]] - new_node = BitmapNode(2, map_bitpos(hash, shift), new_array) + new_node = BitmapNode( + 2, map_bitpos(hash, shift), new_array, mutid) return W_NEWNODE, new_node new_array = self.array[:key_idx] new_array.extend(self.array[key_idx + 2:]) - new_node = CollisionNode(self.size - 2, self.hash, new_array) - return W_NEWNODE, new_node + if mutid and mutid == self.mutid: + self.array = new_array + self.size -= 2 + return W_NEWNODE, self + else: + new_node = CollisionNode( + self.size - 2, self.hash, new_array, mutid) + return W_NEWNODE, new_node def keys(self): for i in range(0, self.size, 2): @@ -346,9 +411,17 @@ class Map: def __init__(self): self.__count = 0 - self.__root = BitmapNode(0, 0, []) + self.__root = BitmapNode(0, 0, [], 0) self.__hash = -1 + @classmethod + def _new(cls, count, root): + m = Map.__new__(Map) + m.__count = count + m.__root = root + m.__hash = -1 + return m + def __len__(self): return self.__count @@ -370,9 +443,12 @@ class Map: return True + def mutate(self): + return MapMutation(self.__count, self.__root) + def set(self, key, val): new_count = self.__count - new_root, added = self.__root.assoc(0, map_hash(key), key, val) + new_root, added = self.__root.assoc(0, map_hash(key), key, val, 0) if new_root is self.__root: assert not added @@ -381,24 +457,16 @@ class Map: if added: new_count += 1 - m = Map.__new__(Map) - m.__count = new_count - m.__root = new_root - m.__hash = -1 - return m + return Map._new(new_count, new_root) def delete(self, key): - res, node = self.__root.without(0, map_hash(key), key) + res, node = self.__root.without(0, map_hash(key), key, 0) if res is W_EMPTY: return Map() elif res is W_NOT_FOUND: raise KeyError(key) else: - m = Map.__new__(Map) - m.__count = self.__count - 1 - m.__root = node - m.__hash = -1 - return m + return Map._new(self.__count - 1, node) def get(self, key, default=None): try: @@ -473,4 +541,90 @@ class Map: return '\n'.join(buf) +class MapMutation: + + def __init__(self, count, root): + self.__count = count + self.__root = root + self.__mutid = _mut_id() + + def set(self, key, val): + if self.__mutid == 0: + raise ValueError(f'mutation {self!r} has been finalized') + + self.__root, added = self.__root.assoc( + 0, map_hash(key), key, val, self.__mutid) + + if added: + self.__count += 1 + + def delete(self, key): + if self.__mutid == 0: + raise ValueError(f'mutation {self!r} has been finalized') + + res, new_root = self.__root.without( + 0, map_hash(key), key, self.__mutid) + if res is W_EMPTY: + self.__count = 0 + self.__root = BitmapNode(0, 0, [], self.__mutid) + elif res is W_NOT_FOUND: + raise KeyError(key) + else: + self.__root = new_root + self.__count -= 1 + + def get(self, key, default=None): + try: + return self.__root.find(0, map_hash(key), key) + except KeyError: + return default + + def __getitem__(self, key): + return self.__root.find(0, map_hash(key), key) + + def __contains__(self, key): + try: + self.__root.find(0, map_hash(key), key) + except KeyError: + return False + else: + return True + + def finalize(self): + self.__mutid = 0 + return Map._new(self.__count, self.__root) + + @reprlib.recursive_repr("{...}") + def __repr__(self): + items = [] + for key, val in self.__root.items(): + items.append("{!r}: {!r}".format(key, val)) + return '<immutables.MapMutation({{{}}}) at 0x{:0x}>'.format( + ', '.join(items), id(self)) + + def __len__(self): + return self.__count + + def __hash__(self): + raise TypeError(f'unhashable type: {type(self).__name__}') + + def __eq__(self, other): + if not isinstance(other, MapMutation): + return NotImplemented + + if len(self) != len(other): + return False + + for key, val in self.__root.items(): + try: + oval = other.__root.find(0, map_hash(key), key) + except KeyError: + return False + else: + if oval != val: + return False + + return True + + collections.abc.Mapping.register(Map) diff --git a/tests/test_map.py b/tests/test_map.py index fe2a131..e2ba2b4 100644 --- a/tests/test_map.py +++ b/tests/test_map.py @@ -906,6 +906,133 @@ class BaseMapTest: def test_abc_1(self): self.assertTrue(issubclass(self.Map, collections.abc.Mapping)) + def test_map_mut_1(self): + h = self.Map() + h = h.set('a', 1) + + hm1 = h.mutate() + hm2 = h.mutate() + + self.assertFalse(isinstance(hm1, self.Map)) + + self.assertIsNot(hm1, hm2) + self.assertEqual(hm1['a'], 1) + self.assertEqual(hm2['a'], 1) + + hm1.set('b', 2) + hm1.set('c', 3) + + hm2.set('x', 100) + hm2.set('a', 1000) + + self.assertEqual(hm1['a'], 1) + self.assertEqual(hm1.get('x', -1), -1) + + self.assertEqual(hm2['a'], 1000) + self.assertTrue('x' in hm2) + + h1 = hm1.finalize() + h2 = hm2.finalize() + + self.assertTrue(isinstance(h1, self.Map)) + + self.assertEqual(dict(h.items()), {'a': 1}) + self.assertEqual(dict(h1.items()), {'a': 1, 'b': 2, 'c': 3}) + self.assertEqual(dict(h2.items()), {'a': 1000, 'x': 100}) + + def test_map_mut_2(self): + h = self.Map() + h = h.set('a', 1) + + hm1 = h.mutate() + hm1.set('a', 2) + hm1.set('a', 3) + hm1.set('a', 4) + h2 = hm1.finalize() + + self.assertEqual(dict(h.items()), {'a': 1}) + self.assertEqual(dict(h2.items()), {'a': 4}) + + def test_map_mut_3(self): + h = self.Map() + h = h.set('a', 1) + hm1 = h.mutate() + + self.assertTrue(repr(hm1).startswith( + "<immutables.MapMutation({'a': 1})")) + + with self.assertRaisesRegex(TypeError, 'unhashable type'): + hash(hm1) + + def test_map_mut_4(self): + h = self.Map() + h = h.set('a', 1) + h = h.set('b', 2) + + hm1 = h.mutate() + hm2 = h.mutate() + + self.assertEqual(hm1, hm2) + + hm1.set('a', 10) + self.assertNotEqual(hm1, hm2) + + hm2.set('a', 10) + self.assertEqual(hm1, hm2) + + hm2.delete('a') + self.assertNotEqual(hm1, hm2) + + def test_map_mut_stress(self): + COLLECTION_SIZE = 7000 + TEST_ITERS_EVERY = 647 + RUN_XTIMES = 3 + + for _ in range(RUN_XTIMES): + h = self.Map() + d = dict() + + for i in range(COLLECTION_SIZE // TEST_ITERS_EVERY): + + hm = h.mutate() + for j in range(TEST_ITERS_EVERY): + key = random.randint(1, 100000) + key = HashKey(key % 271, str(key)) + + hm.set(key, key) + d[key] = key + + self.assertEqual(len(hm), len(d)) + + h2 = hm.finalize() + self.assertEqual(dict(h2.items()), d) + h = h2 + + self.assertEqual(dict(h.items()), d) + self.assertEqual(len(h), len(d)) + + it = iter(tuple(d.keys())) + for i in range(COLLECTION_SIZE // TEST_ITERS_EVERY): + + hm = h.mutate() + for j in range(TEST_ITERS_EVERY): + try: + key = next(it) + except StopIteration: + break + + del d[key] + hm.delete(key) + + self.assertEqual(len(hm), len(d)) + + h2 = hm.finalize() + self.assertEqual(dict(h2.items()), d) + h = h2 + + self.assertEqual(dict(h.items()), d) + self.assertEqual(len(h), len(d)) + class PyMapTest(BaseMapTest, unittest.TestCase): |