aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorYury Selivanov <yury@magic.io>2018-11-19 12:35:02 -0500
committerYury Selivanov <yury@magic.io>2018-11-20 14:25:07 -0500
commit309f2991557673c67c3d8fae995c8b23cc0d4d7c (patch)
tree99d6f8516e3afcdcc297c42c51e7e9db9cb53403
parent5202b41186d7415ed3e236327c8ca782ce32533a (diff)
downloadimmutables-309f2991557673c67c3d8fae995c8b23cc0d4d7c.tar.gz
immutables-309f2991557673c67c3d8fae995c8b23cc0d4d7c.zip
Implement Map.mutate() method
-rw-r--r--immutables/_map.c628
-rw-r--r--immutables/_map.h28
-rw-r--r--immutables/map.py272
-rw-r--r--tests/test_map.py127
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):