1 #include "Python.h"
2 #include "pycore_bitutils.h" // _Py_popcount32()
3 #include "pycore_hamt.h"
4 #include "pycore_initconfig.h" // _PyStatus_OK()
5 #include "pycore_long.h" // _PyLong_Format()
6 #include "pycore_object.h" // _PyObject_GC_TRACK()
7
8 #include <stddef.h> // offsetof()
9
10 /*
11 This file provides an implementation of an immutable mapping using the
12 Hash Array Mapped Trie (or HAMT) datastructure.
13
14 This design allows to have:
15
16 1. Efficient copy: immutable mappings can be copied by reference,
17 making it an O(1) operation.
18
19 2. Efficient mutations: due to structural sharing, only a portion of
20 the trie needs to be copied when the collection is mutated. The
21 cost of set/delete operations is O(log N).
22
23 3. Efficient lookups: O(log N).
24
25 (where N is number of key/value items in the immutable mapping.)
26
27
28 HAMT
29 ====
30
31 The core idea of HAMT is that the shape of the trie is encoded into the
32 hashes of keys.
33
34 Say we want to store a K/V pair in our mapping. First, we calculate the
35 hash of K, let's say it's 19830128, or in binary:
36
37 0b1001011101001010101110000 = 19830128
38
39 Now let's partition this bit representation of the hash into blocks of
40 5 bits each:
41
42 0b00_00000_10010_11101_00101_01011_10000 = 19830128
43 (6) (5) (4) (3) (2) (1)
44
45 Each block of 5 bits represents a number between 0 and 31. So if we have
46 a tree that consists of nodes, each of which is an array of 32 pointers,
47 those 5-bit blocks will encode a position on a single tree level.
48
49 For example, storing the key K with hash 19830128, results in the following
50 tree structure:
51
52 (array of 32 pointers)
53 +---+ -- +----+----+----+ -- +----+
54 root node | 0 | .. | 15 | 16 | 17 | .. | 31 | 0b10000 = 16 (1)
55 (level 1) +---+ -- +----+----+----+ -- +----+
56 |
57 +---+ -- +----+----+----+ -- +----+
58 a 2nd level node | 0 | .. | 10 | 11 | 12 | .. | 31 | 0b01011 = 11 (2)
59 +---+ -- +----+----+----+ -- +----+
60 |
61 +---+ -- +----+----+----+ -- +----+
62 a 3rd level node | 0 | .. | 04 | 05 | 06 | .. | 31 | 0b00101 = 5 (3)
63 +---+ -- +----+----+----+ -- +----+
64 |
65 +---+ -- +----+----+----+----+
66 a 4th level node | 0 | .. | 04 | 29 | 30 | 31 | 0b11101 = 29 (4)
67 +---+ -- +----+----+----+----+
68 |
69 +---+ -- +----+----+----+ -- +----+
70 a 5th level node | 0 | .. | 17 | 18 | 19 | .. | 31 | 0b10010 = 18 (5)
71 +---+ -- +----+----+----+ -- +----+
72 |
73 +--------------+
74 |
75 +---+ -- +----+----+----+ -- +----+
76 a 6th level node | 0 | .. | 15 | 16 | 17 | .. | 31 | 0b00000 = 0 (6)
77 +---+ -- +----+----+----+ -- +----+
78 |
79 V -- our value (or collision)
80
81 To rehash: for a K/V pair, the hash of K encodes where in the tree V will
82 be stored.
83
84 To optimize memory footprint and handle hash collisions, our implementation
85 uses three different types of nodes:
86
87 * A Bitmap node;
88 * An Array node;
89 * A Collision node.
90
91 Because we implement an immutable dictionary, our nodes are also
92 immutable. Therefore, when we need to modify a node, we copy it, and
93 do that modification to the copy.
94
95
96 Array Nodes
97 -----------
98
99 These nodes are very simple. Essentially they are arrays of 32 pointers
100 we used to illustrate the high-level idea in the previous section.
101
102 We use Array nodes only when we need to store more than 16 pointers
103 in a single node.
104
105 Array nodes do not store key objects or value objects. They are used
106 only as an indirection level - their pointers point to other nodes in
107 the tree.
108
109
110 Bitmap Node
111 -----------
112
113 Allocating a new 32-pointers array for every node of our tree would be
114 very expensive. Unless we store millions of keys, most of tree nodes would
115 be very sparse.
116
117 When we have less than 16 elements in a node, we don't want to use the
118 Array node, that would mean that we waste a lot of memory. Instead,
119 we can use bitmap compression and can have just as many pointers
120 as we need!
121
122 Bitmap nodes consist of two fields:
123
124 1. An array of pointers. If a Bitmap node holds N elements, the
125 array will be of N pointers.
126
127 2. A 32bit integer -- a bitmap field. If an N-th bit is set in the
128 bitmap, it means that the node has an N-th element.
129
130 For example, say we need to store a 3 elements sparse array:
131
132 +---+ -- +---+ -- +----+ -- +----+
133 | 0 | .. | 4 | .. | 11 | .. | 17 |
134 +---+ -- +---+ -- +----+ -- +----+
135 | | |
136 o1 o2 o3
137
138 We allocate a three-pointer Bitmap node. Its bitmap field will be
139 then set to:
140
141 0b_00100_00010_00000_10000 == (1 << 17) | (1 << 11) | (1 << 4)
142
143 To check if our Bitmap node has an I-th element we can do:
144
145 bitmap & (1 << I)
146
147
148 And here's a formula to calculate a position in our pointer array
149 which would correspond to an I-th element:
150
151 popcount(bitmap & ((1 << I) - 1))
152
153
154 Let's break it down:
155
156 * `popcount` is a function that returns a number of bits set to 1;
157
158 * `((1 << I) - 1)` is a mask to filter the bitmask to contain bits
159 set to the *right* of our bit.
160
161
162 So for our 17, 11, and 4 indexes:
163
164 * bitmap & ((1 << 17) - 1) == 0b100000010000 => 2 bits are set => index is 2.
165
166 * bitmap & ((1 << 11) - 1) == 0b10000 => 1 bit is set => index is 1.
167
168 * bitmap & ((1 << 4) - 1) == 0b0 => 0 bits are set => index is 0.
169
170
171 To conclude: Bitmap nodes are just like Array nodes -- they can store
172 a number of pointers, but use bitmap compression to eliminate unused
173 pointers.
174
175
176 Bitmap nodes have two pointers for each item:
177
178 +----+----+----+----+ -- +----+----+
179 | k1 | v1 | k2 | v2 | .. | kN | vN |
180 +----+----+----+----+ -- +----+----+
181
182 When kI == NULL, vI points to another tree level.
183
184 When kI != NULL, the actual key object is stored in kI, and its
185 value is stored in vI.
186
187
188 Collision Nodes
189 ---------------
190
191 Collision nodes are simple arrays of pointers -- two pointers per
192 key/value. When there's a hash collision, say for k1/v1 and k2/v2
193 we have `hash(k1)==hash(k2)`. Then our collision node will be:
194
195 +----+----+----+----+
196 | k1 | v1 | k2 | v2 |
197 +----+----+----+----+
198
199
200 Tree Structure
201 --------------
202
203 All nodes are PyObjects.
204
205 The `PyHamtObject` object has a pointer to the root node (h_root),
206 and has a length field (h_count).
207
208 High-level functions accept a PyHamtObject object and dispatch to
209 lower-level functions depending on what kind of node h_root points to.
210
211
212 Operations
213 ==========
214
215 There are three fundamental operations on an immutable dictionary:
216
217 1. "o.assoc(k, v)" will return a new immutable dictionary, that will be
218 a copy of "o", but with the "k/v" item set.
219
220 Functions in this file:
221
222 hamt_node_assoc, hamt_node_bitmap_assoc,
223 hamt_node_array_assoc, hamt_node_collision_assoc
224
225 `hamt_node_assoc` function accepts a node object, and calls
226 other functions depending on its actual type.
227
228 2. "o.find(k)" will lookup key "k" in "o".
229
230 Functions:
231
232 hamt_node_find, hamt_node_bitmap_find,
233 hamt_node_array_find, hamt_node_collision_find
234
235 3. "o.without(k)" will return a new immutable dictionary, that will be
236 a copy of "o", buth without the "k" key.
237
238 Functions:
239
240 hamt_node_without, hamt_node_bitmap_without,
241 hamt_node_array_without, hamt_node_collision_without
242
243
244 Further Reading
245 ===============
246
247 1. http://blog.higher-order.net/2009/09/08/understanding-clojures-persistenthashmap-deftwice.html
248
249 2. http://blog.higher-order.net/2010/08/16/assoc-and-clojures-persistenthashmap-part-ii.html
250
251 3. Clojure's PersistentHashMap implementation:
252 https://github.com/clojure/clojure/blob/master/src/jvm/clojure/lang/PersistentHashMap.java
253
254
255 Debug
256 =====
257
258 The HAMT datatype is accessible for testing purposes under the
259 `_testcapi` module:
260
261 >>> from _testcapi import hamt
262 >>> h = hamt()
263 >>> h2 = h.set('a', 2)
264 >>> h3 = h2.set('b', 3)
265 >>> list(h3)
266 ['a', 'b']
267
268 When CPython is built in debug mode, a '__dump__()' method is available
269 to introspect the tree:
270
271 >>> print(h3.__dump__())
272 HAMT(len=2):
273 BitmapNode(size=4 count=2 bitmap=0b110 id=0x10eb9d9e8):
274 'a': 2
275 'b': 3
276 */
277
278
279 #define IS_ARRAY_NODE(node) Py_IS_TYPE(node, &_PyHamt_ArrayNode_Type)
280 #define IS_BITMAP_NODE(node) Py_IS_TYPE(node, &_PyHamt_BitmapNode_Type)
281 #define IS_COLLISION_NODE(node) Py_IS_TYPE(node, &_PyHamt_CollisionNode_Type)
282
283
284 /* Return type for 'find' (lookup a key) functions.
285
286 * F_ERROR - an error occurred;
287 * F_NOT_FOUND - the key was not found;
288 * F_FOUND - the key was found.
289 */
290 typedef enum {F_ERROR, F_NOT_FOUND, F_FOUND} hamt_find_t;
291
292
293 /* Return type for 'without' (delete a key) functions.
294
295 * W_ERROR - an error occurred;
296 * W_NOT_FOUND - the key was not found: there's nothing to delete;
297 * W_EMPTY - the key was found: the node/tree would be empty
298 if the key is deleted;
299 * W_NEWNODE - the key was found: a new node/tree is returned
300 without that key.
301 */
302 typedef enum {W_ERROR, W_NOT_FOUND, W_EMPTY, W_NEWNODE} hamt_without_t;
303
304
305 /* Low-level iterator protocol type.
306
307 * I_ITEM - a new item has been yielded;
308 * I_END - the whole tree was visited (similar to StopIteration).
309 */
310 typedef enum {I_ITEM, I_END} hamt_iter_t;
311
312
313 #define HAMT_ARRAY_NODE_SIZE 32
314
315
316 typedef struct {
317 PyObject_HEAD
318 PyHamtNode *a_array[HAMT_ARRAY_NODE_SIZE];
319 Py_ssize_t a_count;
320 } PyHamtNode_Array;
321
322
323 typedef struct {
324 PyObject_VAR_HEAD
325 int32_t c_hash;
326 PyObject *c_array[1];
327 } PyHamtNode_Collision;
328
329
330 static PyHamtObject *
331 hamt_alloc(void);
332
333 static PyHamtNode *
334 hamt_node_assoc(PyHamtNode *node,
335 uint32_t shift, int32_t hash,
336 PyObject *key, PyObject *val, int* added_leaf);
337
338 static hamt_without_t
339 hamt_node_without(PyHamtNode *node,
340 uint32_t shift, int32_t hash,
341 PyObject *key,
342 PyHamtNode **new_node);
343
344 static hamt_find_t
345 hamt_node_find(PyHamtNode *node,
346 uint32_t shift, int32_t hash,
347 PyObject *key, PyObject **val);
348
349 #ifdef Py_DEBUG
350 static int
351 hamt_node_dump(PyHamtNode *node,
352 _PyUnicodeWriter *writer, int level);
353 #endif
354
355 static PyHamtNode *
356 hamt_node_array_new(Py_ssize_t);
357
358 static PyHamtNode *
359 hamt_node_collision_new(int32_t hash, Py_ssize_t size);
360
361 static inline Py_ssize_t
362 hamt_node_collision_count(PyHamtNode_Collision *node);
363
364
365 #ifdef Py_DEBUG
366 static void
_hamt_node_array_validate(void * obj_raw)367 _hamt_node_array_validate(void *obj_raw)
368 {
369 PyObject *obj = _PyObject_CAST(obj_raw);
370 assert(IS_ARRAY_NODE(obj));
371 PyHamtNode_Array *node = (PyHamtNode_Array*)obj;
372 Py_ssize_t i = 0, count = 0;
373 for (; i < HAMT_ARRAY_NODE_SIZE; i++) {
374 if (node->a_array[i] != NULL) {
375 count++;
376 }
377 }
378 assert(count == node->a_count);
379 }
380
381 #define VALIDATE_ARRAY_NODE(NODE) \
382 do { _hamt_node_array_validate(NODE); } while (0);
383 #else
384 #define VALIDATE_ARRAY_NODE(NODE)
385 #endif
386
387
388 /* Returns -1 on error */
389 static inline int32_t
hamt_hash(PyObject * o)390 hamt_hash(PyObject *o)
391 {
392 Py_hash_t hash = PyObject_Hash(o);
393
394 #if SIZEOF_PY_HASH_T <= 4
395 return hash;
396 #else
397 if (hash == -1) {
398 /* exception */
399 return -1;
400 }
401
402 /* While it's somewhat suboptimal to reduce Python's 64 bit hash to
403 32 bits via XOR, it seems that the resulting hash function
404 is good enough (this is also how Long type is hashed in Java.)
405 Storing 10, 100, 1000 Python strings results in a relatively
406 shallow and uniform tree structure.
407
408 Also it's worth noting that it would be possible to adapt the tree
409 structure to 64 bit hashes, but that would increase memory pressure
410 and provide little to no performance benefits for collections with
411 fewer than billions of key/value pairs.
412
413 Important: do not change this hash reducing function. There are many
414 tests that need an exact tree shape to cover all code paths and
415 we do that by specifying concrete values for test data's `__hash__`.
416 If this function is changed most of the regression tests would
417 become useless.
418 */
419 int32_t xored = (int32_t)(hash & 0xffffffffl) ^ (int32_t)(hash >> 32);
420 return xored == -1 ? -2 : xored;
421 #endif
422 }
423
424 static inline uint32_t
hamt_mask(int32_t hash,uint32_t shift)425 hamt_mask(int32_t hash, uint32_t shift)
426 {
427 return (((uint32_t)hash >> shift) & 0x01f);
428 }
429
430 static inline uint32_t
hamt_bitpos(int32_t hash,uint32_t shift)431 hamt_bitpos(int32_t hash, uint32_t shift)
432 {
433 return (uint32_t)1 << hamt_mask(hash, shift);
434 }
435
436 static inline uint32_t
hamt_bitindex(uint32_t bitmap,uint32_t bit)437 hamt_bitindex(uint32_t bitmap, uint32_t bit)
438 {
439 return (uint32_t)_Py_popcount32(bitmap & (bit - 1));
440 }
441
442
443 /////////////////////////////////// Dump Helpers
444 #ifdef Py_DEBUG
445
446 static int
_hamt_dump_ident(_PyUnicodeWriter * writer,int level)447 _hamt_dump_ident(_PyUnicodeWriter *writer, int level)
448 {
449 /* Write `' ' * level` to the `writer` */
450 PyObject *str = NULL;
451 PyObject *num = NULL;
452 PyObject *res = NULL;
453 int ret = -1;
454
455 str = PyUnicode_FromString(" ");
456 if (str == NULL) {
457 goto error;
458 }
459
460 num = PyLong_FromLong((long)level);
461 if (num == NULL) {
462 goto error;
463 }
464
465 res = PyNumber_Multiply(str, num);
466 if (res == NULL) {
467 goto error;
468 }
469
470 ret = _PyUnicodeWriter_WriteStr(writer, res);
471
472 error:
473 Py_XDECREF(res);
474 Py_XDECREF(str);
475 Py_XDECREF(num);
476 return ret;
477 }
478
479 static int
_hamt_dump_format(_PyUnicodeWriter * writer,const char * format,...)480 _hamt_dump_format(_PyUnicodeWriter *writer, const char *format, ...)
481 {
482 /* A convenient helper combining _PyUnicodeWriter_WriteStr and
483 PyUnicode_FromFormatV.
484 */
485 PyObject* msg;
486 int ret;
487
488 va_list vargs;
489 va_start(vargs, format);
490 msg = PyUnicode_FromFormatV(format, vargs);
491 va_end(vargs);
492
493 if (msg == NULL) {
494 return -1;
495 }
496
497 ret = _PyUnicodeWriter_WriteStr(writer, msg);
498 Py_DECREF(msg);
499 return ret;
500 }
501
502 #endif /* Py_DEBUG */
503 /////////////////////////////////// Bitmap Node
504
505
506 static PyHamtNode *
hamt_node_bitmap_new(Py_ssize_t size)507 hamt_node_bitmap_new(Py_ssize_t size)
508 {
509 /* Create a new bitmap node of size 'size' */
510
511 PyHamtNode_Bitmap *node;
512 Py_ssize_t i;
513
514 if (size == 0) {
515 /* Since bitmap nodes are immutable, we can cache the instance
516 for size=0 and reuse it whenever we need an empty bitmap node.
517 */
518 return (PyHamtNode *)&_Py_SINGLETON(hamt_bitmap_node_empty);
519 }
520
521 assert(size >= 0);
522 assert(size % 2 == 0);
523
524 /* No freelist; allocate a new bitmap node */
525 node = PyObject_GC_NewVar(
526 PyHamtNode_Bitmap, &_PyHamt_BitmapNode_Type, size);
527 if (node == NULL) {
528 return NULL;
529 }
530
531 Py_SET_SIZE(node, size);
532
533 for (i = 0; i < size; i++) {
534 node->b_array[i] = NULL;
535 }
536
537 node->b_bitmap = 0;
538
539 _PyObject_GC_TRACK(node);
540
541 return (PyHamtNode *)node;
542 }
543
544 static inline Py_ssize_t
hamt_node_bitmap_count(PyHamtNode_Bitmap * node)545 hamt_node_bitmap_count(PyHamtNode_Bitmap *node)
546 {
547 return Py_SIZE(node) / 2;
548 }
549
550 static PyHamtNode_Bitmap *
hamt_node_bitmap_clone(PyHamtNode_Bitmap * node)551 hamt_node_bitmap_clone(PyHamtNode_Bitmap *node)
552 {
553 /* Clone a bitmap node; return a new one with the same child notes. */
554
555 PyHamtNode_Bitmap *clone;
556 Py_ssize_t i;
557
558 clone = (PyHamtNode_Bitmap *)hamt_node_bitmap_new(Py_SIZE(node));
559 if (clone == NULL) {
560 return NULL;
561 }
562
563 for (i = 0; i < Py_SIZE(node); i++) {
564 clone->b_array[i] = Py_XNewRef(node->b_array[i]);
565 }
566
567 clone->b_bitmap = node->b_bitmap;
568 return clone;
569 }
570
571 static PyHamtNode_Bitmap *
hamt_node_bitmap_clone_without(PyHamtNode_Bitmap * o,uint32_t bit)572 hamt_node_bitmap_clone_without(PyHamtNode_Bitmap *o, uint32_t bit)
573 {
574 assert(bit & o->b_bitmap);
575 assert(hamt_node_bitmap_count(o) > 1);
576
577 PyHamtNode_Bitmap *new = (PyHamtNode_Bitmap *)hamt_node_bitmap_new(
578 Py_SIZE(o) - 2);
579 if (new == NULL) {
580 return NULL;
581 }
582
583 uint32_t idx = hamt_bitindex(o->b_bitmap, bit);
584 uint32_t key_idx = 2 * idx;
585 uint32_t val_idx = key_idx + 1;
586 uint32_t i;
587
588 for (i = 0; i < key_idx; i++) {
589 new->b_array[i] = Py_XNewRef(o->b_array[i]);
590 }
591
592 assert(Py_SIZE(o) >= 0 && Py_SIZE(o) <= 32);
593 for (i = val_idx + 1; i < (uint32_t)Py_SIZE(o); i++) {
594 new->b_array[i - 2] = Py_XNewRef(o->b_array[i]);
595 }
596
597 new->b_bitmap = o->b_bitmap & ~bit;
598 return new;
599 }
600
601 static PyHamtNode *
hamt_node_new_bitmap_or_collision(uint32_t shift,PyObject * key1,PyObject * val1,int32_t key2_hash,PyObject * key2,PyObject * val2)602 hamt_node_new_bitmap_or_collision(uint32_t shift,
603 PyObject *key1, PyObject *val1,
604 int32_t key2_hash,
605 PyObject *key2, PyObject *val2)
606 {
607 /* Helper method. Creates a new node for key1/val and key2/val2
608 pairs.
609
610 If key1 hash is equal to the hash of key2, a Collision node
611 will be created. If they are not equal, a Bitmap node is
612 created.
613 */
614
615 int32_t key1_hash = hamt_hash(key1);
616 if (key1_hash == -1) {
617 return NULL;
618 }
619
620 if (key1_hash == key2_hash) {
621 PyHamtNode_Collision *n;
622 n = (PyHamtNode_Collision *)hamt_node_collision_new(key1_hash, 4);
623 if (n == NULL) {
624 return NULL;
625 }
626
627 n->c_array[0] = Py_NewRef(key1);
628 n->c_array[1] = Py_NewRef(val1);
629
630 n->c_array[2] = Py_NewRef(key2);
631 n->c_array[3] = Py_NewRef(val2);
632
633 return (PyHamtNode *)n;
634 }
635 else {
636 int added_leaf = 0;
637 PyHamtNode *n = hamt_node_bitmap_new(0);
638 if (n == NULL) {
639 return NULL;
640 }
641
642 PyHamtNode *n2 = hamt_node_assoc(
643 n, shift, key1_hash, key1, val1, &added_leaf);
644 Py_DECREF(n);
645 if (n2 == NULL) {
646 return NULL;
647 }
648
649 n = hamt_node_assoc(n2, shift, key2_hash, key2, val2, &added_leaf);
650 Py_DECREF(n2);
651 if (n == NULL) {
652 return NULL;
653 }
654
655 return n;
656 }
657 }
658
659 static PyHamtNode *
hamt_node_bitmap_assoc(PyHamtNode_Bitmap * self,uint32_t shift,int32_t hash,PyObject * key,PyObject * val,int * added_leaf)660 hamt_node_bitmap_assoc(PyHamtNode_Bitmap *self,
661 uint32_t shift, int32_t hash,
662 PyObject *key, PyObject *val, int* added_leaf)
663 {
664 /* assoc operation for bitmap nodes.
665
666 Return: a new node, or self if key/val already is in the
667 collection.
668
669 'added_leaf' is later used in '_PyHamt_Assoc' to determine if
670 `hamt.set(key, val)` increased the size of the collection.
671 */
672
673 uint32_t bit = hamt_bitpos(hash, shift);
674 uint32_t idx = hamt_bitindex(self->b_bitmap, bit);
675
676 /* Bitmap node layout:
677
678 +------+------+------+------+ --- +------+------+
679 | key1 | val1 | key2 | val2 | ... | keyN | valN |
680 +------+------+------+------+ --- +------+------+
681 where `N < Py_SIZE(node)`.
682
683 The `node->b_bitmap` field is a bitmap. For a given
684 `(shift, hash)` pair we can determine:
685
686 - If this node has the corresponding key/val slots.
687 - The index of key/val slots.
688 */
689
690 if (self->b_bitmap & bit) {
691 /* The key is set in this node */
692
693 uint32_t key_idx = 2 * idx;
694 uint32_t val_idx = key_idx + 1;
695
696 assert(val_idx < (size_t)Py_SIZE(self));
697
698 PyObject *key_or_null = self->b_array[key_idx];
699 PyObject *val_or_node = self->b_array[val_idx];
700
701 if (key_or_null == NULL) {
702 /* key is NULL. This means that we have a few keys
703 that have the same (hash, shift) pair. */
704
705 assert(val_or_node != NULL);
706
707 PyHamtNode *sub_node = hamt_node_assoc(
708 (PyHamtNode *)val_or_node,
709 shift + 5, hash, key, val, added_leaf);
710 if (sub_node == NULL) {
711 return NULL;
712 }
713
714 if (val_or_node == (PyObject *)sub_node) {
715 Py_DECREF(sub_node);
716 return (PyHamtNode *)Py_NewRef(self);
717 }
718
719 PyHamtNode_Bitmap *ret = hamt_node_bitmap_clone(self);
720 if (ret == NULL) {
721 return NULL;
722 }
723 Py_SETREF(ret->b_array[val_idx], (PyObject*)sub_node);
724 return (PyHamtNode *)ret;
725 }
726
727 assert(key != NULL);
728 /* key is not NULL. This means that we have only one other
729 key in this collection that matches our hash for this shift. */
730
731 int comp_err = PyObject_RichCompareBool(key, key_or_null, Py_EQ);
732 if (comp_err < 0) { /* exception in __eq__ */
733 return NULL;
734 }
735 if (comp_err == 1) { /* key == key_or_null */
736 if (val == val_or_node) {
737 /* we already have the same key/val pair; return self. */
738 return (PyHamtNode *)Py_NewRef(self);
739 }
740
741 /* We're setting a new value for the key we had before.
742 Make a new bitmap node with a replaced value, and return it. */
743 PyHamtNode_Bitmap *ret = hamt_node_bitmap_clone(self);
744 if (ret == NULL) {
745 return NULL;
746 }
747 Py_SETREF(ret->b_array[val_idx], Py_NewRef(val));
748 return (PyHamtNode *)ret;
749 }
750
751 /* It's a new key, and it has the same index as *one* another key.
752 We have a collision. We need to create a new node which will
753 combine the existing key and the key we're adding.
754
755 `hamt_node_new_bitmap_or_collision` will either create a new
756 Collision node if the keys have identical hashes, or
757 a new Bitmap node.
758 */
759 PyHamtNode *sub_node = hamt_node_new_bitmap_or_collision(
760 shift + 5,
761 key_or_null, val_or_node, /* existing key/val */
762 hash,
763 key, val /* new key/val */
764 );
765 if (sub_node == NULL) {
766 return NULL;
767 }
768
769 PyHamtNode_Bitmap *ret = hamt_node_bitmap_clone(self);
770 if (ret == NULL) {
771 Py_DECREF(sub_node);
772 return NULL;
773 }
774 Py_SETREF(ret->b_array[key_idx], NULL);
775 Py_SETREF(ret->b_array[val_idx], (PyObject *)sub_node);
776
777 *added_leaf = 1;
778 return (PyHamtNode *)ret;
779 }
780 else {
781 /* There was no key before with the same (shift,hash). */
782
783 uint32_t n = (uint32_t)_Py_popcount32(self->b_bitmap);
784
785 if (n >= 16) {
786 /* When we have a situation where we want to store more
787 than 16 nodes at one level of the tree, we no longer
788 want to use the Bitmap node with bitmap encoding.
789
790 Instead we start using an Array node, which has
791 simpler (faster) implementation at the expense of
792 having preallocated 32 pointers for its keys/values
793 pairs.
794
795 Small hamt objects (<30 keys) usually don't have any
796 Array nodes at all. Between ~30 and ~400 keys hamt
797 objects usually have one Array node, and usually it's
798 a root node.
799 */
800
801 uint32_t jdx = hamt_mask(hash, shift);
802 /* 'jdx' is the index of where the new key should be added
803 in the new Array node we're about to create. */
804
805 PyHamtNode *empty = NULL;
806 PyHamtNode_Array *new_node = NULL;
807 PyHamtNode *res = NULL;
808
809 /* Create a new Array node. */
810 new_node = (PyHamtNode_Array *)hamt_node_array_new(n + 1);
811 if (new_node == NULL) {
812 goto fin;
813 }
814
815 /* Create an empty bitmap node for the next
816 hamt_node_assoc call. */
817 empty = hamt_node_bitmap_new(0);
818 if (empty == NULL) {
819 goto fin;
820 }
821
822 /* Make a new bitmap node for the key/val we're adding.
823 Set that bitmap node to new-array-node[jdx]. */
824 new_node->a_array[jdx] = hamt_node_assoc(
825 empty, shift + 5, hash, key, val, added_leaf);
826 if (new_node->a_array[jdx] == NULL) {
827 goto fin;
828 }
829
830 /* Copy existing key/value pairs from the current Bitmap
831 node to the new Array node we've just created. */
832 Py_ssize_t i, j;
833 for (i = 0, j = 0; i < HAMT_ARRAY_NODE_SIZE; i++) {
834 if (((self->b_bitmap >> i) & 1) != 0) {
835 /* Ensure we don't accidentally override `jdx` element
836 we set few lines above.
837 */
838 assert(new_node->a_array[i] == NULL);
839
840 if (self->b_array[j] == NULL) {
841 new_node->a_array[i] =
842 (PyHamtNode *)Py_NewRef(self->b_array[j + 1]);
843 }
844 else {
845 int32_t rehash = hamt_hash(self->b_array[j]);
846 if (rehash == -1) {
847 goto fin;
848 }
849
850 new_node->a_array[i] = hamt_node_assoc(
851 empty, shift + 5,
852 rehash,
853 self->b_array[j],
854 self->b_array[j + 1],
855 added_leaf);
856
857 if (new_node->a_array[i] == NULL) {
858 goto fin;
859 }
860 }
861 j += 2;
862 }
863 }
864
865 VALIDATE_ARRAY_NODE(new_node)
866
867 /* That's it! */
868 res = (PyHamtNode *)new_node;
869
870 fin:
871 Py_XDECREF(empty);
872 if (res == NULL) {
873 Py_XDECREF(new_node);
874 }
875 return res;
876 }
877 else {
878 /* We have less than 16 keys at this level; let's just
879 create a new bitmap node out of this node with the
880 new key/val pair added. */
881
882 uint32_t key_idx = 2 * idx;
883 uint32_t val_idx = key_idx + 1;
884 uint32_t i;
885
886 *added_leaf = 1;
887
888 /* Allocate new Bitmap node which can have one more key/val
889 pair in addition to what we have already. */
890 PyHamtNode_Bitmap *new_node =
891 (PyHamtNode_Bitmap *)hamt_node_bitmap_new(2 * (n + 1));
892 if (new_node == NULL) {
893 return NULL;
894 }
895
896 /* Copy all keys/values that will be before the new key/value
897 we are adding. */
898 for (i = 0; i < key_idx; i++) {
899 new_node->b_array[i] = Py_XNewRef(self->b_array[i]);
900 }
901
902 /* Set the new key/value to the new Bitmap node. */
903 new_node->b_array[key_idx] = Py_NewRef(key);
904 new_node->b_array[val_idx] = Py_NewRef(val);
905
906 /* Copy all keys/values that will be after the new key/value
907 we are adding. */
908 assert(Py_SIZE(self) >= 0 && Py_SIZE(self) <= 32);
909 for (i = key_idx; i < (uint32_t)Py_SIZE(self); i++) {
910 new_node->b_array[i + 2] = Py_XNewRef(self->b_array[i]);
911 }
912
913 new_node->b_bitmap = self->b_bitmap | bit;
914 return (PyHamtNode *)new_node;
915 }
916 }
917 }
918
919 static hamt_without_t
hamt_node_bitmap_without(PyHamtNode_Bitmap * self,uint32_t shift,int32_t hash,PyObject * key,PyHamtNode ** new_node)920 hamt_node_bitmap_without(PyHamtNode_Bitmap *self,
921 uint32_t shift, int32_t hash,
922 PyObject *key,
923 PyHamtNode **new_node)
924 {
925 uint32_t bit = hamt_bitpos(hash, shift);
926 if ((self->b_bitmap & bit) == 0) {
927 return W_NOT_FOUND;
928 }
929
930 uint32_t idx = hamt_bitindex(self->b_bitmap, bit);
931
932 uint32_t key_idx = 2 * idx;
933 uint32_t val_idx = key_idx + 1;
934
935 PyObject *key_or_null = self->b_array[key_idx];
936 PyObject *val_or_node = self->b_array[val_idx];
937
938 if (key_or_null == NULL) {
939 /* key == NULL means that 'value' is another tree node. */
940
941 PyHamtNode *sub_node = NULL;
942
943 hamt_without_t res = hamt_node_without(
944 (PyHamtNode *)val_or_node,
945 shift + 5, hash, key, &sub_node);
946
947 switch (res) {
948 case W_EMPTY:
949 /* It's impossible for us to receive a W_EMPTY here:
950
951 - Array nodes are converted to Bitmap nodes when
952 we delete 16th item from them;
953
954 - Collision nodes are converted to Bitmap when
955 there is one item in them;
956
957 - Bitmap node's without() inlines single-item
958 sub-nodes.
959
960 So in no situation we can have a single-item
961 Bitmap child of another Bitmap node.
962 */
963 Py_UNREACHABLE();
964
965 case W_NEWNODE: {
966 assert(sub_node != NULL);
967
968 if (IS_BITMAP_NODE(sub_node)) {
969 PyHamtNode_Bitmap *sub_tree = (PyHamtNode_Bitmap *)sub_node;
970 if (hamt_node_bitmap_count(sub_tree) == 1 &&
971 sub_tree->b_array[0] != NULL)
972 {
973 /* A bitmap node with one key/value pair. Just
974 merge it into this node.
975
976 Note that we don't inline Bitmap nodes that
977 have a NULL key -- those nodes point to another
978 tree level, and we cannot simply move tree levels
979 up or down.
980 */
981
982 PyHamtNode_Bitmap *clone = hamt_node_bitmap_clone(self);
983 if (clone == NULL) {
984 Py_DECREF(sub_node);
985 return W_ERROR;
986 }
987
988 PyObject *key = sub_tree->b_array[0];
989 PyObject *val = sub_tree->b_array[1];
990
991 Py_XSETREF(clone->b_array[key_idx], Py_NewRef(key));
992 Py_SETREF(clone->b_array[val_idx], Py_NewRef(val));
993
994 Py_DECREF(sub_tree);
995
996 *new_node = (PyHamtNode *)clone;
997 return W_NEWNODE;
998 }
999 }
1000
1001 #ifdef Py_DEBUG
1002 /* Ensure that Collision.without implementation
1003 converts to Bitmap nodes itself.
1004 */
1005 if (IS_COLLISION_NODE(sub_node)) {
1006 assert(hamt_node_collision_count(
1007 (PyHamtNode_Collision*)sub_node) > 1);
1008 }
1009 #endif
1010
1011 PyHamtNode_Bitmap *clone = hamt_node_bitmap_clone(self);
1012 if (clone == NULL) {
1013 return W_ERROR;
1014 }
1015
1016 Py_SETREF(clone->b_array[val_idx],
1017 (PyObject *)sub_node); /* borrow */
1018
1019 *new_node = (PyHamtNode *)clone;
1020 return W_NEWNODE;
1021 }
1022
1023 case W_ERROR:
1024 case W_NOT_FOUND:
1025 assert(sub_node == NULL);
1026 return res;
1027
1028 default:
1029 Py_UNREACHABLE();
1030 }
1031 }
1032 else {
1033 /* We have a regular key/value pair */
1034
1035 int cmp = PyObject_RichCompareBool(key_or_null, key, Py_EQ);
1036 if (cmp < 0) {
1037 return W_ERROR;
1038 }
1039 if (cmp == 0) {
1040 return W_NOT_FOUND;
1041 }
1042
1043 if (hamt_node_bitmap_count(self) == 1) {
1044 return W_EMPTY;
1045 }
1046
1047 *new_node = (PyHamtNode *)
1048 hamt_node_bitmap_clone_without(self, bit);
1049 if (*new_node == NULL) {
1050 return W_ERROR;
1051 }
1052
1053 return W_NEWNODE;
1054 }
1055 }
1056
1057 static hamt_find_t
hamt_node_bitmap_find(PyHamtNode_Bitmap * self,uint32_t shift,int32_t hash,PyObject * key,PyObject ** val)1058 hamt_node_bitmap_find(PyHamtNode_Bitmap *self,
1059 uint32_t shift, int32_t hash,
1060 PyObject *key, PyObject **val)
1061 {
1062 /* Lookup a key in a Bitmap node. */
1063
1064 uint32_t bit = hamt_bitpos(hash, shift);
1065 uint32_t idx;
1066 uint32_t key_idx;
1067 uint32_t val_idx;
1068 PyObject *key_or_null;
1069 PyObject *val_or_node;
1070 int comp_err;
1071
1072 if ((self->b_bitmap & bit) == 0) {
1073 return F_NOT_FOUND;
1074 }
1075
1076 idx = hamt_bitindex(self->b_bitmap, bit);
1077 key_idx = idx * 2;
1078 val_idx = key_idx + 1;
1079
1080 assert(val_idx < (size_t)Py_SIZE(self));
1081
1082 key_or_null = self->b_array[key_idx];
1083 val_or_node = self->b_array[val_idx];
1084
1085 if (key_or_null == NULL) {
1086 /* There are a few keys that have the same hash at the current shift
1087 that match our key. Dispatch the lookup further down the tree. */
1088 assert(val_or_node != NULL);
1089 return hamt_node_find((PyHamtNode *)val_or_node,
1090 shift + 5, hash, key, val);
1091 }
1092
1093 /* We have only one key -- a potential match. Let's compare if the
1094 key we are looking at is equal to the key we are looking for. */
1095 assert(key != NULL);
1096 comp_err = PyObject_RichCompareBool(key, key_or_null, Py_EQ);
1097 if (comp_err < 0) { /* exception in __eq__ */
1098 return F_ERROR;
1099 }
1100 if (comp_err == 1) { /* key == key_or_null */
1101 *val = val_or_node;
1102 return F_FOUND;
1103 }
1104
1105 return F_NOT_FOUND;
1106 }
1107
1108 static int
hamt_node_bitmap_traverse(PyHamtNode_Bitmap * self,visitproc visit,void * arg)1109 hamt_node_bitmap_traverse(PyHamtNode_Bitmap *self, visitproc visit, void *arg)
1110 {
1111 /* Bitmap's tp_traverse */
1112
1113 Py_ssize_t i;
1114
1115 for (i = Py_SIZE(self); --i >= 0; ) {
1116 Py_VISIT(self->b_array[i]);
1117 }
1118
1119 return 0;
1120 }
1121
1122 static void
hamt_node_bitmap_dealloc(PyHamtNode_Bitmap * self)1123 hamt_node_bitmap_dealloc(PyHamtNode_Bitmap *self)
1124 {
1125 /* Bitmap's tp_dealloc */
1126
1127 Py_ssize_t len = Py_SIZE(self);
1128 Py_ssize_t i;
1129
1130 if (Py_SIZE(self) == 0) {
1131 /* The empty node is statically allocated. */
1132 assert(self == &_Py_SINGLETON(hamt_bitmap_node_empty));
1133 #ifdef Py_DEBUG
1134 _Py_FatalRefcountError("deallocating the empty hamt node bitmap singleton");
1135 #else
1136 return;
1137 #endif
1138 }
1139
1140 PyObject_GC_UnTrack(self);
1141 Py_TRASHCAN_BEGIN(self, hamt_node_bitmap_dealloc)
1142
1143 if (len > 0) {
1144 i = len;
1145 while (--i >= 0) {
1146 Py_XDECREF(self->b_array[i]);
1147 }
1148 }
1149
1150 Py_TYPE(self)->tp_free((PyObject *)self);
1151 Py_TRASHCAN_END
1152 }
1153
1154 #ifdef Py_DEBUG
1155 static int
hamt_node_bitmap_dump(PyHamtNode_Bitmap * node,_PyUnicodeWriter * writer,int level)1156 hamt_node_bitmap_dump(PyHamtNode_Bitmap *node,
1157 _PyUnicodeWriter *writer, int level)
1158 {
1159 /* Debug build: __dump__() method implementation for Bitmap nodes. */
1160
1161 Py_ssize_t i;
1162 PyObject *tmp1;
1163 PyObject *tmp2;
1164
1165 if (_hamt_dump_ident(writer, level + 1)) {
1166 goto error;
1167 }
1168
1169 if (_hamt_dump_format(writer, "BitmapNode(size=%zd count=%zd ",
1170 Py_SIZE(node), Py_SIZE(node) / 2))
1171 {
1172 goto error;
1173 }
1174
1175 tmp1 = PyLong_FromUnsignedLong(node->b_bitmap);
1176 if (tmp1 == NULL) {
1177 goto error;
1178 }
1179 tmp2 = _PyLong_Format(tmp1, 2);
1180 Py_DECREF(tmp1);
1181 if (tmp2 == NULL) {
1182 goto error;
1183 }
1184 if (_hamt_dump_format(writer, "bitmap=%S id=%p):\n", tmp2, node)) {
1185 Py_DECREF(tmp2);
1186 goto error;
1187 }
1188 Py_DECREF(tmp2);
1189
1190 for (i = 0; i < Py_SIZE(node); i += 2) {
1191 PyObject *key_or_null = node->b_array[i];
1192 PyObject *val_or_node = node->b_array[i + 1];
1193
1194 if (_hamt_dump_ident(writer, level + 2)) {
1195 goto error;
1196 }
1197
1198 if (key_or_null == NULL) {
1199 if (_hamt_dump_format(writer, "NULL:\n")) {
1200 goto error;
1201 }
1202
1203 if (hamt_node_dump((PyHamtNode *)val_or_node,
1204 writer, level + 2))
1205 {
1206 goto error;
1207 }
1208 }
1209 else {
1210 if (_hamt_dump_format(writer, "%R: %R", key_or_null,
1211 val_or_node))
1212 {
1213 goto error;
1214 }
1215 }
1216
1217 if (_hamt_dump_format(writer, "\n")) {
1218 goto error;
1219 }
1220 }
1221
1222 return 0;
1223 error:
1224 return -1;
1225 }
1226 #endif /* Py_DEBUG */
1227
1228
1229 /////////////////////////////////// Collision Node
1230
1231
1232 static PyHamtNode *
hamt_node_collision_new(int32_t hash,Py_ssize_t size)1233 hamt_node_collision_new(int32_t hash, Py_ssize_t size)
1234 {
1235 /* Create a new Collision node. */
1236
1237 PyHamtNode_Collision *node;
1238 Py_ssize_t i;
1239
1240 assert(size >= 4);
1241 assert(size % 2 == 0);
1242
1243 node = PyObject_GC_NewVar(
1244 PyHamtNode_Collision, &_PyHamt_CollisionNode_Type, size);
1245 if (node == NULL) {
1246 return NULL;
1247 }
1248
1249 for (i = 0; i < size; i++) {
1250 node->c_array[i] = NULL;
1251 }
1252
1253 Py_SET_SIZE(node, size);
1254 node->c_hash = hash;
1255
1256 _PyObject_GC_TRACK(node);
1257
1258 return (PyHamtNode *)node;
1259 }
1260
1261 static hamt_find_t
hamt_node_collision_find_index(PyHamtNode_Collision * self,PyObject * key,Py_ssize_t * idx)1262 hamt_node_collision_find_index(PyHamtNode_Collision *self, PyObject *key,
1263 Py_ssize_t *idx)
1264 {
1265 /* Lookup `key` in the Collision node `self`. Set the index of the
1266 found key to 'idx'. */
1267
1268 Py_ssize_t i;
1269 PyObject *el;
1270
1271 for (i = 0; i < Py_SIZE(self); i += 2) {
1272 el = self->c_array[i];
1273
1274 assert(el != NULL);
1275 int cmp = PyObject_RichCompareBool(key, el, Py_EQ);
1276 if (cmp < 0) {
1277 return F_ERROR;
1278 }
1279 if (cmp == 1) {
1280 *idx = i;
1281 return F_FOUND;
1282 }
1283 }
1284
1285 return F_NOT_FOUND;
1286 }
1287
1288 static PyHamtNode *
hamt_node_collision_assoc(PyHamtNode_Collision * self,uint32_t shift,int32_t hash,PyObject * key,PyObject * val,int * added_leaf)1289 hamt_node_collision_assoc(PyHamtNode_Collision *self,
1290 uint32_t shift, int32_t hash,
1291 PyObject *key, PyObject *val, int* added_leaf)
1292 {
1293 /* Set a new key to this level (currently a Collision node)
1294 of the tree. */
1295
1296 if (hash == self->c_hash) {
1297 /* The hash of the 'key' we are adding matches the hash of
1298 other keys in this Collision node. */
1299
1300 Py_ssize_t key_idx = -1;
1301 hamt_find_t found;
1302 PyHamtNode_Collision *new_node;
1303 Py_ssize_t i;
1304
1305 /* Let's try to lookup the new 'key', maybe we already have it. */
1306 found = hamt_node_collision_find_index(self, key, &key_idx);
1307 switch (found) {
1308 case F_ERROR:
1309 /* Exception. */
1310 return NULL;
1311
1312 case F_NOT_FOUND:
1313 /* This is a totally new key. Clone the current node,
1314 add a new key/value to the cloned node. */
1315
1316 new_node = (PyHamtNode_Collision *)hamt_node_collision_new(
1317 self->c_hash, Py_SIZE(self) + 2);
1318 if (new_node == NULL) {
1319 return NULL;
1320 }
1321
1322 for (i = 0; i < Py_SIZE(self); i++) {
1323 new_node->c_array[i] = Py_NewRef(self->c_array[i]);
1324 }
1325
1326 new_node->c_array[i] = Py_NewRef(key);
1327 new_node->c_array[i + 1] = Py_NewRef(val);
1328
1329 *added_leaf = 1;
1330 return (PyHamtNode *)new_node;
1331
1332 case F_FOUND:
1333 /* There's a key which is equal to the key we are adding. */
1334
1335 assert(key_idx >= 0);
1336 assert(key_idx < Py_SIZE(self));
1337 Py_ssize_t val_idx = key_idx + 1;
1338
1339 if (self->c_array[val_idx] == val) {
1340 /* We're setting a key/value pair that's already set. */
1341 return (PyHamtNode *)Py_NewRef(self);
1342 }
1343
1344 /* We need to replace old value for the key
1345 with a new value. Create a new Collision node.*/
1346 new_node = (PyHamtNode_Collision *)hamt_node_collision_new(
1347 self->c_hash, Py_SIZE(self));
1348 if (new_node == NULL) {
1349 return NULL;
1350 }
1351
1352 /* Copy all elements of the old node to the new one. */
1353 for (i = 0; i < Py_SIZE(self); i++) {
1354 new_node->c_array[i] = Py_NewRef(self->c_array[i]);
1355 }
1356
1357 /* Replace the old value with the new value for the our key. */
1358 Py_SETREF(new_node->c_array[val_idx], Py_NewRef(val));
1359
1360 return (PyHamtNode *)new_node;
1361
1362 default:
1363 Py_UNREACHABLE();
1364 }
1365 }
1366 else {
1367 /* The hash of the new key is different from the hash that
1368 all keys of this Collision node have.
1369
1370 Create a Bitmap node inplace with two children:
1371 key/value pair that we're adding, and the Collision node
1372 we're replacing on this tree level.
1373 */
1374
1375 PyHamtNode_Bitmap *new_node;
1376 PyHamtNode *assoc_res;
1377
1378 new_node = (PyHamtNode_Bitmap *)hamt_node_bitmap_new(2);
1379 if (new_node == NULL) {
1380 return NULL;
1381 }
1382 new_node->b_bitmap = hamt_bitpos(self->c_hash, shift);
1383 new_node->b_array[1] = Py_NewRef(self);
1384
1385 assoc_res = hamt_node_bitmap_assoc(
1386 new_node, shift, hash, key, val, added_leaf);
1387 Py_DECREF(new_node);
1388 return assoc_res;
1389 }
1390 }
1391
1392 static inline Py_ssize_t
hamt_node_collision_count(PyHamtNode_Collision * node)1393 hamt_node_collision_count(PyHamtNode_Collision *node)
1394 {
1395 return Py_SIZE(node) / 2;
1396 }
1397
1398 static hamt_without_t
hamt_node_collision_without(PyHamtNode_Collision * self,uint32_t shift,int32_t hash,PyObject * key,PyHamtNode ** new_node)1399 hamt_node_collision_without(PyHamtNode_Collision *self,
1400 uint32_t shift, int32_t hash,
1401 PyObject *key,
1402 PyHamtNode **new_node)
1403 {
1404 if (hash != self->c_hash) {
1405 return W_NOT_FOUND;
1406 }
1407
1408 Py_ssize_t key_idx = -1;
1409 hamt_find_t found = hamt_node_collision_find_index(self, key, &key_idx);
1410
1411 switch (found) {
1412 case F_ERROR:
1413 return W_ERROR;
1414
1415 case F_NOT_FOUND:
1416 return W_NOT_FOUND;
1417
1418 case F_FOUND:
1419 assert(key_idx >= 0);
1420 assert(key_idx < Py_SIZE(self));
1421
1422 Py_ssize_t new_count = hamt_node_collision_count(self) - 1;
1423
1424 if (new_count == 0) {
1425 /* The node has only one key/value pair and it's for the
1426 key we're trying to delete. So a new node will be empty
1427 after the removal.
1428 */
1429 return W_EMPTY;
1430 }
1431
1432 if (new_count == 1) {
1433 /* The node has two keys, and after deletion the
1434 new Collision node would have one. Collision nodes
1435 with one key shouldn't exist, so convert it to a
1436 Bitmap node.
1437 */
1438 PyHamtNode_Bitmap *node = (PyHamtNode_Bitmap *)
1439 hamt_node_bitmap_new(2);
1440 if (node == NULL) {
1441 return W_ERROR;
1442 }
1443
1444 if (key_idx == 0) {
1445 node->b_array[0] = Py_NewRef(self->c_array[2]);
1446 node->b_array[1] = Py_NewRef(self->c_array[3]);
1447 }
1448 else {
1449 assert(key_idx == 2);
1450 node->b_array[0] = Py_NewRef(self->c_array[0]);
1451 node->b_array[1] = Py_NewRef(self->c_array[1]);
1452 }
1453
1454 node->b_bitmap = hamt_bitpos(hash, shift);
1455
1456 *new_node = (PyHamtNode *)node;
1457 return W_NEWNODE;
1458 }
1459
1460 /* Allocate a new Collision node with capacity for one
1461 less key/value pair */
1462 PyHamtNode_Collision *new = (PyHamtNode_Collision *)
1463 hamt_node_collision_new(
1464 self->c_hash, Py_SIZE(self) - 2);
1465 if (new == NULL) {
1466 return W_ERROR;
1467 }
1468
1469 /* Copy all other keys from `self` to `new` */
1470 Py_ssize_t i;
1471 for (i = 0; i < key_idx; i++) {
1472 new->c_array[i] = Py_NewRef(self->c_array[i]);
1473 }
1474 for (i = key_idx + 2; i < Py_SIZE(self); i++) {
1475 new->c_array[i - 2] = Py_NewRef(self->c_array[i]);
1476 }
1477
1478 *new_node = (PyHamtNode*)new;
1479 return W_NEWNODE;
1480
1481 default:
1482 Py_UNREACHABLE();
1483 }
1484 }
1485
1486 static hamt_find_t
hamt_node_collision_find(PyHamtNode_Collision * self,uint32_t shift,int32_t hash,PyObject * key,PyObject ** val)1487 hamt_node_collision_find(PyHamtNode_Collision *self,
1488 uint32_t shift, int32_t hash,
1489 PyObject *key, PyObject **val)
1490 {
1491 /* Lookup `key` in the Collision node `self`. Set the value
1492 for the found key to 'val'. */
1493
1494 Py_ssize_t idx = -1;
1495 hamt_find_t res;
1496
1497 res = hamt_node_collision_find_index(self, key, &idx);
1498 if (res == F_ERROR || res == F_NOT_FOUND) {
1499 return res;
1500 }
1501
1502 assert(idx >= 0);
1503 assert(idx + 1 < Py_SIZE(self));
1504
1505 *val = self->c_array[idx + 1];
1506 assert(*val != NULL);
1507
1508 return F_FOUND;
1509 }
1510
1511
1512 static int
hamt_node_collision_traverse(PyHamtNode_Collision * self,visitproc visit,void * arg)1513 hamt_node_collision_traverse(PyHamtNode_Collision *self,
1514 visitproc visit, void *arg)
1515 {
1516 /* Collision's tp_traverse */
1517
1518 Py_ssize_t i;
1519
1520 for (i = Py_SIZE(self); --i >= 0; ) {
1521 Py_VISIT(self->c_array[i]);
1522 }
1523
1524 return 0;
1525 }
1526
1527 static void
hamt_node_collision_dealloc(PyHamtNode_Collision * self)1528 hamt_node_collision_dealloc(PyHamtNode_Collision *self)
1529 {
1530 /* Collision's tp_dealloc */
1531
1532 Py_ssize_t len = Py_SIZE(self);
1533
1534 PyObject_GC_UnTrack(self);
1535 Py_TRASHCAN_BEGIN(self, hamt_node_collision_dealloc)
1536
1537 if (len > 0) {
1538
1539 while (--len >= 0) {
1540 Py_XDECREF(self->c_array[len]);
1541 }
1542 }
1543
1544 Py_TYPE(self)->tp_free((PyObject *)self);
1545 Py_TRASHCAN_END
1546 }
1547
1548 #ifdef Py_DEBUG
1549 static int
hamt_node_collision_dump(PyHamtNode_Collision * node,_PyUnicodeWriter * writer,int level)1550 hamt_node_collision_dump(PyHamtNode_Collision *node,
1551 _PyUnicodeWriter *writer, int level)
1552 {
1553 /* Debug build: __dump__() method implementation for Collision nodes. */
1554
1555 Py_ssize_t i;
1556
1557 if (_hamt_dump_ident(writer, level + 1)) {
1558 goto error;
1559 }
1560
1561 if (_hamt_dump_format(writer, "CollisionNode(size=%zd id=%p):\n",
1562 Py_SIZE(node), node))
1563 {
1564 goto error;
1565 }
1566
1567 for (i = 0; i < Py_SIZE(node); i += 2) {
1568 PyObject *key = node->c_array[i];
1569 PyObject *val = node->c_array[i + 1];
1570
1571 if (_hamt_dump_ident(writer, level + 2)) {
1572 goto error;
1573 }
1574
1575 if (_hamt_dump_format(writer, "%R: %R\n", key, val)) {
1576 goto error;
1577 }
1578 }
1579
1580 return 0;
1581 error:
1582 return -1;
1583 }
1584 #endif /* Py_DEBUG */
1585
1586
1587 /////////////////////////////////// Array Node
1588
1589
1590 static PyHamtNode *
hamt_node_array_new(Py_ssize_t count)1591 hamt_node_array_new(Py_ssize_t count)
1592 {
1593 Py_ssize_t i;
1594
1595 PyHamtNode_Array *node = PyObject_GC_New(
1596 PyHamtNode_Array, &_PyHamt_ArrayNode_Type);
1597 if (node == NULL) {
1598 return NULL;
1599 }
1600
1601 for (i = 0; i < HAMT_ARRAY_NODE_SIZE; i++) {
1602 node->a_array[i] = NULL;
1603 }
1604
1605 node->a_count = count;
1606
1607 _PyObject_GC_TRACK(node);
1608 return (PyHamtNode *)node;
1609 }
1610
1611 static PyHamtNode_Array *
hamt_node_array_clone(PyHamtNode_Array * node)1612 hamt_node_array_clone(PyHamtNode_Array *node)
1613 {
1614 PyHamtNode_Array *clone;
1615 Py_ssize_t i;
1616
1617 VALIDATE_ARRAY_NODE(node)
1618
1619 /* Create a new Array node. */
1620 clone = (PyHamtNode_Array *)hamt_node_array_new(node->a_count);
1621 if (clone == NULL) {
1622 return NULL;
1623 }
1624
1625 /* Copy all elements from the current Array node to the new one. */
1626 for (i = 0; i < HAMT_ARRAY_NODE_SIZE; i++) {
1627 clone->a_array[i] = (PyHamtNode*)Py_XNewRef(node->a_array[i]);
1628 }
1629
1630 VALIDATE_ARRAY_NODE(clone)
1631 return clone;
1632 }
1633
1634 static PyHamtNode *
hamt_node_array_assoc(PyHamtNode_Array * self,uint32_t shift,int32_t hash,PyObject * key,PyObject * val,int * added_leaf)1635 hamt_node_array_assoc(PyHamtNode_Array *self,
1636 uint32_t shift, int32_t hash,
1637 PyObject *key, PyObject *val, int* added_leaf)
1638 {
1639 /* Set a new key to this level (currently a Collision node)
1640 of the tree.
1641
1642 Array nodes don't store values, they can only point to
1643 other nodes. They are simple arrays of 32 BaseNode pointers/
1644 */
1645
1646 uint32_t idx = hamt_mask(hash, shift);
1647 PyHamtNode *node = self->a_array[idx];
1648 PyHamtNode *child_node;
1649 PyHamtNode_Array *new_node;
1650 Py_ssize_t i;
1651
1652 if (node == NULL) {
1653 /* There's no child node for the given hash. Create a new
1654 Bitmap node for this key. */
1655
1656 PyHamtNode_Bitmap *empty = NULL;
1657
1658 /* Get an empty Bitmap node to work with. */
1659 empty = (PyHamtNode_Bitmap *)hamt_node_bitmap_new(0);
1660 if (empty == NULL) {
1661 return NULL;
1662 }
1663
1664 /* Set key/val to the newly created empty Bitmap, thus
1665 creating a new Bitmap node with our key/value pair. */
1666 child_node = hamt_node_bitmap_assoc(
1667 empty,
1668 shift + 5, hash, key, val, added_leaf);
1669 Py_DECREF(empty);
1670 if (child_node == NULL) {
1671 return NULL;
1672 }
1673
1674 /* Create a new Array node. */
1675 new_node = (PyHamtNode_Array *)hamt_node_array_new(self->a_count + 1);
1676 if (new_node == NULL) {
1677 Py_DECREF(child_node);
1678 return NULL;
1679 }
1680
1681 /* Copy all elements from the current Array node to the
1682 new one. */
1683 for (i = 0; i < HAMT_ARRAY_NODE_SIZE; i++) {
1684 new_node->a_array[i] = (PyHamtNode*)Py_XNewRef(self->a_array[i]);
1685 }
1686
1687 assert(new_node->a_array[idx] == NULL);
1688 new_node->a_array[idx] = child_node; /* borrow */
1689 VALIDATE_ARRAY_NODE(new_node)
1690 }
1691 else {
1692 /* There's a child node for the given hash.
1693 Set the key to it./ */
1694 child_node = hamt_node_assoc(
1695 node, shift + 5, hash, key, val, added_leaf);
1696 if (child_node == NULL) {
1697 return NULL;
1698 }
1699 else if (child_node == (PyHamtNode *)self) {
1700 Py_DECREF(child_node);
1701 return (PyHamtNode *)self;
1702 }
1703
1704 new_node = hamt_node_array_clone(self);
1705 if (new_node == NULL) {
1706 Py_DECREF(child_node);
1707 return NULL;
1708 }
1709
1710 Py_SETREF(new_node->a_array[idx], child_node); /* borrow */
1711 VALIDATE_ARRAY_NODE(new_node)
1712 }
1713
1714 return (PyHamtNode *)new_node;
1715 }
1716
1717 static hamt_without_t
hamt_node_array_without(PyHamtNode_Array * self,uint32_t shift,int32_t hash,PyObject * key,PyHamtNode ** new_node)1718 hamt_node_array_without(PyHamtNode_Array *self,
1719 uint32_t shift, int32_t hash,
1720 PyObject *key,
1721 PyHamtNode **new_node)
1722 {
1723 uint32_t idx = hamt_mask(hash, shift);
1724 PyHamtNode *node = self->a_array[idx];
1725
1726 if (node == NULL) {
1727 return W_NOT_FOUND;
1728 }
1729
1730 PyHamtNode *sub_node = NULL;
1731 hamt_without_t res = hamt_node_without(
1732 (PyHamtNode *)node,
1733 shift + 5, hash, key, &sub_node);
1734
1735 switch (res) {
1736 case W_NOT_FOUND:
1737 case W_ERROR:
1738 assert(sub_node == NULL);
1739 return res;
1740
1741 case W_NEWNODE: {
1742 /* We need to replace a node at the `idx` index.
1743 Clone this node and replace.
1744 */
1745 assert(sub_node != NULL);
1746
1747 PyHamtNode_Array *clone = hamt_node_array_clone(self);
1748 if (clone == NULL) {
1749 Py_DECREF(sub_node);
1750 return W_ERROR;
1751 }
1752
1753 Py_SETREF(clone->a_array[idx], sub_node); /* borrow */
1754 *new_node = (PyHamtNode*)clone; /* borrow */
1755 return W_NEWNODE;
1756 }
1757
1758 case W_EMPTY: {
1759 assert(sub_node == NULL);
1760 /* We need to remove a node at the `idx` index.
1761 Calculate the size of the replacement Array node.
1762 */
1763 Py_ssize_t new_count = self->a_count - 1;
1764
1765 if (new_count == 0) {
1766 return W_EMPTY;
1767 }
1768
1769 if (new_count >= 16) {
1770 /* We convert Bitmap nodes to Array nodes, when a
1771 Bitmap node needs to store more than 15 key/value
1772 pairs. So we will create a new Array node if we
1773 the number of key/values after deletion is still
1774 greater than 15.
1775 */
1776
1777 PyHamtNode_Array *new = hamt_node_array_clone(self);
1778 if (new == NULL) {
1779 return W_ERROR;
1780 }
1781 new->a_count = new_count;
1782 Py_CLEAR(new->a_array[idx]);
1783
1784 *new_node = (PyHamtNode*)new; /* borrow */
1785 return W_NEWNODE;
1786 }
1787
1788 /* New Array node would have less than 16 key/value
1789 pairs. We need to create a replacement Bitmap node. */
1790
1791 Py_ssize_t bitmap_size = new_count * 2;
1792 uint32_t bitmap = 0;
1793
1794 PyHamtNode_Bitmap *new = (PyHamtNode_Bitmap *)
1795 hamt_node_bitmap_new(bitmap_size);
1796 if (new == NULL) {
1797 return W_ERROR;
1798 }
1799
1800 Py_ssize_t new_i = 0;
1801 for (uint32_t i = 0; i < HAMT_ARRAY_NODE_SIZE; i++) {
1802 if (i == idx) {
1803 /* Skip the node we are deleting. */
1804 continue;
1805 }
1806
1807 PyHamtNode *node = self->a_array[i];
1808 if (node == NULL) {
1809 /* Skip any missing nodes. */
1810 continue;
1811 }
1812
1813 bitmap |= 1U << i;
1814
1815 if (IS_BITMAP_NODE(node)) {
1816 PyHamtNode_Bitmap *child = (PyHamtNode_Bitmap *)node;
1817
1818 if (hamt_node_bitmap_count(child) == 1 &&
1819 child->b_array[0] != NULL)
1820 {
1821 /* node is a Bitmap with one key/value pair, just
1822 merge it into the new Bitmap node we're building.
1823
1824 Note that we don't inline Bitmap nodes that
1825 have a NULL key -- those nodes point to another
1826 tree level, and we cannot simply move tree levels
1827 up or down.
1828 */
1829 PyObject *key = child->b_array[0];
1830 PyObject *val = child->b_array[1];
1831
1832 new->b_array[new_i] = Py_NewRef(key);
1833 new->b_array[new_i + 1] = Py_NewRef(val);
1834 }
1835 else {
1836 new->b_array[new_i] = NULL;
1837 new->b_array[new_i + 1] = Py_NewRef(node);
1838 }
1839 }
1840 else {
1841
1842 #ifdef Py_DEBUG
1843 if (IS_COLLISION_NODE(node)) {
1844 Py_ssize_t child_count = hamt_node_collision_count(
1845 (PyHamtNode_Collision*)node);
1846 assert(child_count > 1);
1847 }
1848 else if (IS_ARRAY_NODE(node)) {
1849 assert(((PyHamtNode_Array*)node)->a_count >= 16);
1850 }
1851 #endif
1852
1853 /* Just copy the node into our new Bitmap */
1854 new->b_array[new_i] = NULL;
1855 new->b_array[new_i + 1] = Py_NewRef(node);
1856 }
1857
1858 new_i += 2;
1859 }
1860
1861 new->b_bitmap = bitmap;
1862 *new_node = (PyHamtNode*)new; /* borrow */
1863 return W_NEWNODE;
1864 }
1865
1866 default:
1867 Py_UNREACHABLE();
1868 }
1869 }
1870
1871 static hamt_find_t
hamt_node_array_find(PyHamtNode_Array * self,uint32_t shift,int32_t hash,PyObject * key,PyObject ** val)1872 hamt_node_array_find(PyHamtNode_Array *self,
1873 uint32_t shift, int32_t hash,
1874 PyObject *key, PyObject **val)
1875 {
1876 /* Lookup `key` in the Array node `self`. Set the value
1877 for the found key to 'val'. */
1878
1879 uint32_t idx = hamt_mask(hash, shift);
1880 PyHamtNode *node;
1881
1882 node = self->a_array[idx];
1883 if (node == NULL) {
1884 return F_NOT_FOUND;
1885 }
1886
1887 /* Dispatch to the generic hamt_node_find */
1888 return hamt_node_find(node, shift + 5, hash, key, val);
1889 }
1890
1891 static int
hamt_node_array_traverse(PyHamtNode_Array * self,visitproc visit,void * arg)1892 hamt_node_array_traverse(PyHamtNode_Array *self,
1893 visitproc visit, void *arg)
1894 {
1895 /* Array's tp_traverse */
1896
1897 Py_ssize_t i;
1898
1899 for (i = 0; i < HAMT_ARRAY_NODE_SIZE; i++) {
1900 Py_VISIT(self->a_array[i]);
1901 }
1902
1903 return 0;
1904 }
1905
1906 static void
hamt_node_array_dealloc(PyHamtNode_Array * self)1907 hamt_node_array_dealloc(PyHamtNode_Array *self)
1908 {
1909 /* Array's tp_dealloc */
1910
1911 Py_ssize_t i;
1912
1913 PyObject_GC_UnTrack(self);
1914 Py_TRASHCAN_BEGIN(self, hamt_node_array_dealloc)
1915
1916 for (i = 0; i < HAMT_ARRAY_NODE_SIZE; i++) {
1917 Py_XDECREF(self->a_array[i]);
1918 }
1919
1920 Py_TYPE(self)->tp_free((PyObject *)self);
1921 Py_TRASHCAN_END
1922 }
1923
1924 #ifdef Py_DEBUG
1925 static int
hamt_node_array_dump(PyHamtNode_Array * node,_PyUnicodeWriter * writer,int level)1926 hamt_node_array_dump(PyHamtNode_Array *node,
1927 _PyUnicodeWriter *writer, int level)
1928 {
1929 /* Debug build: __dump__() method implementation for Array nodes. */
1930
1931 Py_ssize_t i;
1932
1933 if (_hamt_dump_ident(writer, level + 1)) {
1934 goto error;
1935 }
1936
1937 if (_hamt_dump_format(writer, "ArrayNode(id=%p):\n", node)) {
1938 goto error;
1939 }
1940
1941 for (i = 0; i < HAMT_ARRAY_NODE_SIZE; i++) {
1942 if (node->a_array[i] == NULL) {
1943 continue;
1944 }
1945
1946 if (_hamt_dump_ident(writer, level + 2)) {
1947 goto error;
1948 }
1949
1950 if (_hamt_dump_format(writer, "%zd::\n", i)) {
1951 goto error;
1952 }
1953
1954 if (hamt_node_dump(node->a_array[i], writer, level + 1)) {
1955 goto error;
1956 }
1957
1958 if (_hamt_dump_format(writer, "\n")) {
1959 goto error;
1960 }
1961 }
1962
1963 return 0;
1964 error:
1965 return -1;
1966 }
1967 #endif /* Py_DEBUG */
1968
1969
1970 /////////////////////////////////// Node Dispatch
1971
1972
1973 static PyHamtNode *
hamt_node_assoc(PyHamtNode * node,uint32_t shift,int32_t hash,PyObject * key,PyObject * val,int * added_leaf)1974 hamt_node_assoc(PyHamtNode *node,
1975 uint32_t shift, int32_t hash,
1976 PyObject *key, PyObject *val, int* added_leaf)
1977 {
1978 /* Set key/value to the 'node' starting with the given shift/hash.
1979 Return a new node, or the same node if key/value already
1980 set.
1981
1982 added_leaf will be set to 1 if key/value wasn't in the
1983 tree before.
1984
1985 This method automatically dispatches to the suitable
1986 hamt_node_{nodetype}_assoc method.
1987 */
1988
1989 if (IS_BITMAP_NODE(node)) {
1990 return hamt_node_bitmap_assoc(
1991 (PyHamtNode_Bitmap *)node,
1992 shift, hash, key, val, added_leaf);
1993 }
1994 else if (IS_ARRAY_NODE(node)) {
1995 return hamt_node_array_assoc(
1996 (PyHamtNode_Array *)node,
1997 shift, hash, key, val, added_leaf);
1998 }
1999 else {
2000 assert(IS_COLLISION_NODE(node));
2001 return hamt_node_collision_assoc(
2002 (PyHamtNode_Collision *)node,
2003 shift, hash, key, val, added_leaf);
2004 }
2005 }
2006
2007 static hamt_without_t
hamt_node_without(PyHamtNode * node,uint32_t shift,int32_t hash,PyObject * key,PyHamtNode ** new_node)2008 hamt_node_without(PyHamtNode *node,
2009 uint32_t shift, int32_t hash,
2010 PyObject *key,
2011 PyHamtNode **new_node)
2012 {
2013 if (IS_BITMAP_NODE(node)) {
2014 return hamt_node_bitmap_without(
2015 (PyHamtNode_Bitmap *)node,
2016 shift, hash, key,
2017 new_node);
2018 }
2019 else if (IS_ARRAY_NODE(node)) {
2020 return hamt_node_array_without(
2021 (PyHamtNode_Array *)node,
2022 shift, hash, key,
2023 new_node);
2024 }
2025 else {
2026 assert(IS_COLLISION_NODE(node));
2027 return hamt_node_collision_without(
2028 (PyHamtNode_Collision *)node,
2029 shift, hash, key,
2030 new_node);
2031 }
2032 }
2033
2034 static hamt_find_t
hamt_node_find(PyHamtNode * node,uint32_t shift,int32_t hash,PyObject * key,PyObject ** val)2035 hamt_node_find(PyHamtNode *node,
2036 uint32_t shift, int32_t hash,
2037 PyObject *key, PyObject **val)
2038 {
2039 /* Find the key in the node starting with the given shift/hash.
2040
2041 If a value is found, the result will be set to F_FOUND, and
2042 *val will point to the found value object.
2043
2044 If a value wasn't found, the result will be set to F_NOT_FOUND.
2045
2046 If an exception occurs during the call, the result will be F_ERROR.
2047
2048 This method automatically dispatches to the suitable
2049 hamt_node_{nodetype}_find method.
2050 */
2051
2052 if (IS_BITMAP_NODE(node)) {
2053 return hamt_node_bitmap_find(
2054 (PyHamtNode_Bitmap *)node,
2055 shift, hash, key, val);
2056
2057 }
2058 else if (IS_ARRAY_NODE(node)) {
2059 return hamt_node_array_find(
2060 (PyHamtNode_Array *)node,
2061 shift, hash, key, val);
2062 }
2063 else {
2064 assert(IS_COLLISION_NODE(node));
2065 return hamt_node_collision_find(
2066 (PyHamtNode_Collision *)node,
2067 shift, hash, key, val);
2068 }
2069 }
2070
2071 #ifdef Py_DEBUG
2072 static int
hamt_node_dump(PyHamtNode * node,_PyUnicodeWriter * writer,int level)2073 hamt_node_dump(PyHamtNode *node,
2074 _PyUnicodeWriter *writer, int level)
2075 {
2076 /* Debug build: __dump__() method implementation for a node.
2077
2078 This method automatically dispatches to the suitable
2079 hamt_node_{nodetype})_dump method.
2080 */
2081
2082 if (IS_BITMAP_NODE(node)) {
2083 return hamt_node_bitmap_dump(
2084 (PyHamtNode_Bitmap *)node, writer, level);
2085 }
2086 else if (IS_ARRAY_NODE(node)) {
2087 return hamt_node_array_dump(
2088 (PyHamtNode_Array *)node, writer, level);
2089 }
2090 else {
2091 assert(IS_COLLISION_NODE(node));
2092 return hamt_node_collision_dump(
2093 (PyHamtNode_Collision *)node, writer, level);
2094 }
2095 }
2096 #endif /* Py_DEBUG */
2097
2098
2099 /////////////////////////////////// Iterators: Machinery
2100
2101
2102 static hamt_iter_t
2103 hamt_iterator_next(PyHamtIteratorState *iter, PyObject **key, PyObject **val);
2104
2105
2106 static void
hamt_iterator_init(PyHamtIteratorState * iter,PyHamtNode * root)2107 hamt_iterator_init(PyHamtIteratorState *iter, PyHamtNode *root)
2108 {
2109 for (uint32_t i = 0; i < _Py_HAMT_MAX_TREE_DEPTH; i++) {
2110 iter->i_nodes[i] = NULL;
2111 iter->i_pos[i] = 0;
2112 }
2113
2114 iter->i_level = 0;
2115
2116 /* Note: we don't incref/decref nodes in i_nodes. */
2117 iter->i_nodes[0] = root;
2118 }
2119
2120 static hamt_iter_t
hamt_iterator_bitmap_next(PyHamtIteratorState * iter,PyObject ** key,PyObject ** val)2121 hamt_iterator_bitmap_next(PyHamtIteratorState *iter,
2122 PyObject **key, PyObject **val)
2123 {
2124 int8_t level = iter->i_level;
2125
2126 PyHamtNode_Bitmap *node = (PyHamtNode_Bitmap *)(iter->i_nodes[level]);
2127 Py_ssize_t pos = iter->i_pos[level];
2128
2129 if (pos + 1 >= Py_SIZE(node)) {
2130 #ifdef Py_DEBUG
2131 assert(iter->i_level >= 0);
2132 iter->i_nodes[iter->i_level] = NULL;
2133 #endif
2134 iter->i_level--;
2135 return hamt_iterator_next(iter, key, val);
2136 }
2137
2138 if (node->b_array[pos] == NULL) {
2139 iter->i_pos[level] = pos + 2;
2140
2141 int8_t next_level = level + 1;
2142 assert(next_level < _Py_HAMT_MAX_TREE_DEPTH);
2143 iter->i_level = next_level;
2144 iter->i_pos[next_level] = 0;
2145 iter->i_nodes[next_level] = (PyHamtNode *)
2146 node->b_array[pos + 1];
2147
2148 return hamt_iterator_next(iter, key, val);
2149 }
2150
2151 *key = node->b_array[pos];
2152 *val = node->b_array[pos + 1];
2153 iter->i_pos[level] = pos + 2;
2154 return I_ITEM;
2155 }
2156
2157 static hamt_iter_t
hamt_iterator_collision_next(PyHamtIteratorState * iter,PyObject ** key,PyObject ** val)2158 hamt_iterator_collision_next(PyHamtIteratorState *iter,
2159 PyObject **key, PyObject **val)
2160 {
2161 int8_t level = iter->i_level;
2162
2163 PyHamtNode_Collision *node = (PyHamtNode_Collision *)(iter->i_nodes[level]);
2164 Py_ssize_t pos = iter->i_pos[level];
2165
2166 if (pos + 1 >= Py_SIZE(node)) {
2167 #ifdef Py_DEBUG
2168 assert(iter->i_level >= 0);
2169 iter->i_nodes[iter->i_level] = NULL;
2170 #endif
2171 iter->i_level--;
2172 return hamt_iterator_next(iter, key, val);
2173 }
2174
2175 *key = node->c_array[pos];
2176 *val = node->c_array[pos + 1];
2177 iter->i_pos[level] = pos + 2;
2178 return I_ITEM;
2179 }
2180
2181 static hamt_iter_t
hamt_iterator_array_next(PyHamtIteratorState * iter,PyObject ** key,PyObject ** val)2182 hamt_iterator_array_next(PyHamtIteratorState *iter,
2183 PyObject **key, PyObject **val)
2184 {
2185 int8_t level = iter->i_level;
2186
2187 PyHamtNode_Array *node = (PyHamtNode_Array *)(iter->i_nodes[level]);
2188 Py_ssize_t pos = iter->i_pos[level];
2189
2190 if (pos >= HAMT_ARRAY_NODE_SIZE) {
2191 #ifdef Py_DEBUG
2192 assert(iter->i_level >= 0);
2193 iter->i_nodes[iter->i_level] = NULL;
2194 #endif
2195 iter->i_level--;
2196 return hamt_iterator_next(iter, key, val);
2197 }
2198
2199 for (Py_ssize_t i = pos; i < HAMT_ARRAY_NODE_SIZE; i++) {
2200 if (node->a_array[i] != NULL) {
2201 iter->i_pos[level] = i + 1;
2202
2203 int8_t next_level = level + 1;
2204 assert(next_level < _Py_HAMT_MAX_TREE_DEPTH);
2205 iter->i_pos[next_level] = 0;
2206 iter->i_nodes[next_level] = node->a_array[i];
2207 iter->i_level = next_level;
2208
2209 return hamt_iterator_next(iter, key, val);
2210 }
2211 }
2212
2213 #ifdef Py_DEBUG
2214 assert(iter->i_level >= 0);
2215 iter->i_nodes[iter->i_level] = NULL;
2216 #endif
2217
2218 iter->i_level--;
2219 return hamt_iterator_next(iter, key, val);
2220 }
2221
2222 static hamt_iter_t
hamt_iterator_next(PyHamtIteratorState * iter,PyObject ** key,PyObject ** val)2223 hamt_iterator_next(PyHamtIteratorState *iter, PyObject **key, PyObject **val)
2224 {
2225 if (iter->i_level < 0) {
2226 return I_END;
2227 }
2228
2229 assert(iter->i_level < _Py_HAMT_MAX_TREE_DEPTH);
2230
2231 PyHamtNode *current = iter->i_nodes[iter->i_level];
2232
2233 if (IS_BITMAP_NODE(current)) {
2234 return hamt_iterator_bitmap_next(iter, key, val);
2235 }
2236 else if (IS_ARRAY_NODE(current)) {
2237 return hamt_iterator_array_next(iter, key, val);
2238 }
2239 else {
2240 assert(IS_COLLISION_NODE(current));
2241 return hamt_iterator_collision_next(iter, key, val);
2242 }
2243 }
2244
2245
2246 /////////////////////////////////// HAMT high-level functions
2247
2248
2249 PyHamtObject *
_PyHamt_Assoc(PyHamtObject * o,PyObject * key,PyObject * val)2250 _PyHamt_Assoc(PyHamtObject *o, PyObject *key, PyObject *val)
2251 {
2252 int32_t key_hash;
2253 int added_leaf = 0;
2254 PyHamtNode *new_root;
2255 PyHamtObject *new_o;
2256
2257 key_hash = hamt_hash(key);
2258 if (key_hash == -1) {
2259 return NULL;
2260 }
2261
2262 new_root = hamt_node_assoc(
2263 (PyHamtNode *)(o->h_root),
2264 0, key_hash, key, val, &added_leaf);
2265 if (new_root == NULL) {
2266 return NULL;
2267 }
2268
2269 if (new_root == o->h_root) {
2270 Py_DECREF(new_root);
2271 return (PyHamtObject*)Py_NewRef(o);
2272 }
2273
2274 new_o = hamt_alloc();
2275 if (new_o == NULL) {
2276 Py_DECREF(new_root);
2277 return NULL;
2278 }
2279
2280 new_o->h_root = new_root; /* borrow */
2281 new_o->h_count = added_leaf ? o->h_count + 1 : o->h_count;
2282
2283 return new_o;
2284 }
2285
2286 PyHamtObject *
_PyHamt_Without(PyHamtObject * o,PyObject * key)2287 _PyHamt_Without(PyHamtObject *o, PyObject *key)
2288 {
2289 int32_t key_hash = hamt_hash(key);
2290 if (key_hash == -1) {
2291 return NULL;
2292 }
2293
2294 PyHamtNode *new_root = NULL;
2295
2296 hamt_without_t res = hamt_node_without(
2297 (PyHamtNode *)(o->h_root),
2298 0, key_hash, key,
2299 &new_root);
2300
2301 switch (res) {
2302 case W_ERROR:
2303 return NULL;
2304 case W_EMPTY:
2305 return _PyHamt_New();
2306 case W_NOT_FOUND:
2307 return (PyHamtObject*)Py_NewRef(o);
2308 case W_NEWNODE: {
2309 assert(new_root != NULL);
2310
2311 PyHamtObject *new_o = hamt_alloc();
2312 if (new_o == NULL) {
2313 Py_DECREF(new_root);
2314 return NULL;
2315 }
2316
2317 new_o->h_root = new_root; /* borrow */
2318 new_o->h_count = o->h_count - 1;
2319 assert(new_o->h_count >= 0);
2320 return new_o;
2321 }
2322 default:
2323 Py_UNREACHABLE();
2324 }
2325 }
2326
2327 static hamt_find_t
hamt_find(PyHamtObject * o,PyObject * key,PyObject ** val)2328 hamt_find(PyHamtObject *o, PyObject *key, PyObject **val)
2329 {
2330 if (o->h_count == 0) {
2331 return F_NOT_FOUND;
2332 }
2333
2334 int32_t key_hash = hamt_hash(key);
2335 if (key_hash == -1) {
2336 return F_ERROR;
2337 }
2338
2339 return hamt_node_find(o->h_root, 0, key_hash, key, val);
2340 }
2341
2342
2343 int
_PyHamt_Find(PyHamtObject * o,PyObject * key,PyObject ** val)2344 _PyHamt_Find(PyHamtObject *o, PyObject *key, PyObject **val)
2345 {
2346 hamt_find_t res = hamt_find(o, key, val);
2347 switch (res) {
2348 case F_ERROR:
2349 return -1;
2350 case F_NOT_FOUND:
2351 return 0;
2352 case F_FOUND:
2353 return 1;
2354 default:
2355 Py_UNREACHABLE();
2356 }
2357 }
2358
2359
2360 int
_PyHamt_Eq(PyHamtObject * v,PyHamtObject * w)2361 _PyHamt_Eq(PyHamtObject *v, PyHamtObject *w)
2362 {
2363 if (v == w) {
2364 return 1;
2365 }
2366
2367 if (v->h_count != w->h_count) {
2368 return 0;
2369 }
2370
2371 PyHamtIteratorState iter;
2372 hamt_iter_t iter_res;
2373 hamt_find_t find_res;
2374 PyObject *v_key;
2375 PyObject *v_val;
2376 PyObject *w_val;
2377
2378 hamt_iterator_init(&iter, v->h_root);
2379
2380 do {
2381 iter_res = hamt_iterator_next(&iter, &v_key, &v_val);
2382 if (iter_res == I_ITEM) {
2383 find_res = hamt_find(w, v_key, &w_val);
2384 switch (find_res) {
2385 case F_ERROR:
2386 return -1;
2387
2388 case F_NOT_FOUND:
2389 return 0;
2390
2391 case F_FOUND: {
2392 int cmp = PyObject_RichCompareBool(v_val, w_val, Py_EQ);
2393 if (cmp < 0) {
2394 return -1;
2395 }
2396 if (cmp == 0) {
2397 return 0;
2398 }
2399 }
2400 }
2401 }
2402 } while (iter_res != I_END);
2403
2404 return 1;
2405 }
2406
2407 Py_ssize_t
_PyHamt_Len(PyHamtObject * o)2408 _PyHamt_Len(PyHamtObject *o)
2409 {
2410 return o->h_count;
2411 }
2412
2413 static PyHamtObject *
hamt_alloc(void)2414 hamt_alloc(void)
2415 {
2416 PyHamtObject *o;
2417 o = PyObject_GC_New(PyHamtObject, &_PyHamt_Type);
2418 if (o == NULL) {
2419 return NULL;
2420 }
2421 o->h_count = 0;
2422 o->h_root = NULL;
2423 o->h_weakreflist = NULL;
2424 PyObject_GC_Track(o);
2425 return o;
2426 }
2427
2428 #define _empty_hamt \
2429 (&_Py_INTERP_SINGLETON(_PyInterpreterState_GET(), hamt_empty))
2430
2431 PyHamtObject *
_PyHamt_New(void)2432 _PyHamt_New(void)
2433 {
2434 /* HAMT is an immutable object so we can easily cache an
2435 empty instance. */
2436 return (PyHamtObject*)Py_NewRef(_empty_hamt);
2437 }
2438
2439 #ifdef Py_DEBUG
2440 static PyObject *
hamt_dump(PyHamtObject * self)2441 hamt_dump(PyHamtObject *self)
2442 {
2443 _PyUnicodeWriter writer;
2444
2445 _PyUnicodeWriter_Init(&writer);
2446
2447 if (_hamt_dump_format(&writer, "HAMT(len=%zd):\n", self->h_count)) {
2448 goto error;
2449 }
2450
2451 if (hamt_node_dump(self->h_root, &writer, 0)) {
2452 goto error;
2453 }
2454
2455 return _PyUnicodeWriter_Finish(&writer);
2456
2457 error:
2458 _PyUnicodeWriter_Dealloc(&writer);
2459 return NULL;
2460 }
2461 #endif /* Py_DEBUG */
2462
2463
2464 /////////////////////////////////// Iterators: Shared Iterator Implementation
2465
2466
2467 static int
hamt_baseiter_tp_clear(PyHamtIterator * it)2468 hamt_baseiter_tp_clear(PyHamtIterator *it)
2469 {
2470 Py_CLEAR(it->hi_obj);
2471 return 0;
2472 }
2473
2474 static void
hamt_baseiter_tp_dealloc(PyHamtIterator * it)2475 hamt_baseiter_tp_dealloc(PyHamtIterator *it)
2476 {
2477 PyObject_GC_UnTrack(it);
2478 (void)hamt_baseiter_tp_clear(it);
2479 PyObject_GC_Del(it);
2480 }
2481
2482 static int
hamt_baseiter_tp_traverse(PyHamtIterator * it,visitproc visit,void * arg)2483 hamt_baseiter_tp_traverse(PyHamtIterator *it, visitproc visit, void *arg)
2484 {
2485 Py_VISIT(it->hi_obj);
2486 return 0;
2487 }
2488
2489 static PyObject *
hamt_baseiter_tp_iternext(PyHamtIterator * it)2490 hamt_baseiter_tp_iternext(PyHamtIterator *it)
2491 {
2492 PyObject *key;
2493 PyObject *val;
2494 hamt_iter_t res = hamt_iterator_next(&it->hi_iter, &key, &val);
2495
2496 switch (res) {
2497 case I_END:
2498 PyErr_SetNone(PyExc_StopIteration);
2499 return NULL;
2500
2501 case I_ITEM: {
2502 return (*(it->hi_yield))(key, val);
2503 }
2504
2505 default: {
2506 Py_UNREACHABLE();
2507 }
2508 }
2509 }
2510
2511 static Py_ssize_t
hamt_baseiter_tp_len(PyHamtIterator * it)2512 hamt_baseiter_tp_len(PyHamtIterator *it)
2513 {
2514 return it->hi_obj->h_count;
2515 }
2516
2517 static PyMappingMethods PyHamtIterator_as_mapping = {
2518 (lenfunc)hamt_baseiter_tp_len,
2519 };
2520
2521 static PyObject *
hamt_baseiter_new(PyTypeObject * type,binaryfunc yield,PyHamtObject * o)2522 hamt_baseiter_new(PyTypeObject *type, binaryfunc yield, PyHamtObject *o)
2523 {
2524 PyHamtIterator *it = PyObject_GC_New(PyHamtIterator, type);
2525 if (it == NULL) {
2526 return NULL;
2527 }
2528
2529 it->hi_obj = (PyHamtObject*)Py_NewRef(o);
2530 it->hi_yield = yield;
2531
2532 hamt_iterator_init(&it->hi_iter, o->h_root);
2533
2534 return (PyObject*)it;
2535 }
2536
2537 #define ITERATOR_TYPE_SHARED_SLOTS \
2538 .tp_basicsize = sizeof(PyHamtIterator), \
2539 .tp_itemsize = 0, \
2540 .tp_as_mapping = &PyHamtIterator_as_mapping, \
2541 .tp_dealloc = (destructor)hamt_baseiter_tp_dealloc, \
2542 .tp_getattro = PyObject_GenericGetAttr, \
2543 .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, \
2544 .tp_traverse = (traverseproc)hamt_baseiter_tp_traverse, \
2545 .tp_clear = (inquiry)hamt_baseiter_tp_clear, \
2546 .tp_iter = PyObject_SelfIter, \
2547 .tp_iternext = (iternextfunc)hamt_baseiter_tp_iternext,
2548
2549
2550 /////////////////////////////////// _PyHamtItems_Type
2551
2552
2553 PyTypeObject _PyHamtItems_Type = {
2554 PyVarObject_HEAD_INIT(NULL, 0)
2555 "items",
2556 ITERATOR_TYPE_SHARED_SLOTS
2557 };
2558
2559 static PyObject *
hamt_iter_yield_items(PyObject * key,PyObject * val)2560 hamt_iter_yield_items(PyObject *key, PyObject *val)
2561 {
2562 return PyTuple_Pack(2, key, val);
2563 }
2564
2565 PyObject *
_PyHamt_NewIterItems(PyHamtObject * o)2566 _PyHamt_NewIterItems(PyHamtObject *o)
2567 {
2568 return hamt_baseiter_new(
2569 &_PyHamtItems_Type, hamt_iter_yield_items, o);
2570 }
2571
2572
2573 /////////////////////////////////// _PyHamtKeys_Type
2574
2575
2576 PyTypeObject _PyHamtKeys_Type = {
2577 PyVarObject_HEAD_INIT(NULL, 0)
2578 "keys",
2579 ITERATOR_TYPE_SHARED_SLOTS
2580 };
2581
2582 static PyObject *
hamt_iter_yield_keys(PyObject * key,PyObject * val)2583 hamt_iter_yield_keys(PyObject *key, PyObject *val)
2584 {
2585 return Py_NewRef(key);
2586 }
2587
2588 PyObject *
_PyHamt_NewIterKeys(PyHamtObject * o)2589 _PyHamt_NewIterKeys(PyHamtObject *o)
2590 {
2591 return hamt_baseiter_new(
2592 &_PyHamtKeys_Type, hamt_iter_yield_keys, o);
2593 }
2594
2595
2596 /////////////////////////////////// _PyHamtValues_Type
2597
2598
2599 PyTypeObject _PyHamtValues_Type = {
2600 PyVarObject_HEAD_INIT(NULL, 0)
2601 "values",
2602 ITERATOR_TYPE_SHARED_SLOTS
2603 };
2604
2605 static PyObject *
hamt_iter_yield_values(PyObject * key,PyObject * val)2606 hamt_iter_yield_values(PyObject *key, PyObject *val)
2607 {
2608 return Py_NewRef(val);
2609 }
2610
2611 PyObject *
_PyHamt_NewIterValues(PyHamtObject * o)2612 _PyHamt_NewIterValues(PyHamtObject *o)
2613 {
2614 return hamt_baseiter_new(
2615 &_PyHamtValues_Type, hamt_iter_yield_values, o);
2616 }
2617
2618
2619 /////////////////////////////////// _PyHamt_Type
2620
2621
2622 #ifdef Py_DEBUG
2623 static PyObject *
2624 hamt_dump(PyHamtObject *self);
2625 #endif
2626
2627
2628 static PyObject *
hamt_tp_new(PyTypeObject * type,PyObject * args,PyObject * kwds)2629 hamt_tp_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2630 {
2631 return (PyObject*)_PyHamt_New();
2632 }
2633
2634 static int
hamt_tp_clear(PyHamtObject * self)2635 hamt_tp_clear(PyHamtObject *self)
2636 {
2637 Py_CLEAR(self->h_root);
2638 return 0;
2639 }
2640
2641
2642 static int
hamt_tp_traverse(PyHamtObject * self,visitproc visit,void * arg)2643 hamt_tp_traverse(PyHamtObject *self, visitproc visit, void *arg)
2644 {
2645 Py_VISIT(self->h_root);
2646 return 0;
2647 }
2648
2649 static void
hamt_tp_dealloc(PyHamtObject * self)2650 hamt_tp_dealloc(PyHamtObject *self)
2651 {
2652 if (self == _empty_hamt) {
2653 /* The empty one is statically allocated. */
2654 #ifdef Py_DEBUG
2655 _Py_FatalRefcountError("deallocating the empty hamt singleton");
2656 #else
2657 return;
2658 #endif
2659 }
2660
2661 PyObject_GC_UnTrack(self);
2662 if (self->h_weakreflist != NULL) {
2663 PyObject_ClearWeakRefs((PyObject*)self);
2664 }
2665 (void)hamt_tp_clear(self);
2666 Py_TYPE(self)->tp_free(self);
2667 }
2668
2669
2670 static PyObject *
hamt_tp_richcompare(PyObject * v,PyObject * w,int op)2671 hamt_tp_richcompare(PyObject *v, PyObject *w, int op)
2672 {
2673 if (!PyHamt_Check(v) || !PyHamt_Check(w) || (op != Py_EQ && op != Py_NE)) {
2674 Py_RETURN_NOTIMPLEMENTED;
2675 }
2676
2677 int res = _PyHamt_Eq((PyHamtObject *)v, (PyHamtObject *)w);
2678 if (res < 0) {
2679 return NULL;
2680 }
2681
2682 if (op == Py_NE) {
2683 res = !res;
2684 }
2685
2686 if (res) {
2687 Py_RETURN_TRUE;
2688 }
2689 else {
2690 Py_RETURN_FALSE;
2691 }
2692 }
2693
2694 static int
hamt_tp_contains(PyHamtObject * self,PyObject * key)2695 hamt_tp_contains(PyHamtObject *self, PyObject *key)
2696 {
2697 PyObject *val;
2698 return _PyHamt_Find(self, key, &val);
2699 }
2700
2701 static PyObject *
hamt_tp_subscript(PyHamtObject * self,PyObject * key)2702 hamt_tp_subscript(PyHamtObject *self, PyObject *key)
2703 {
2704 PyObject *val;
2705 hamt_find_t res = hamt_find(self, key, &val);
2706 switch (res) {
2707 case F_ERROR:
2708 return NULL;
2709 case F_FOUND:
2710 return Py_NewRef(val);
2711 case F_NOT_FOUND:
2712 PyErr_SetObject(PyExc_KeyError, key);
2713 return NULL;
2714 default:
2715 Py_UNREACHABLE();
2716 }
2717 }
2718
2719 static Py_ssize_t
hamt_tp_len(PyHamtObject * self)2720 hamt_tp_len(PyHamtObject *self)
2721 {
2722 return _PyHamt_Len(self);
2723 }
2724
2725 static PyObject *
hamt_tp_iter(PyHamtObject * self)2726 hamt_tp_iter(PyHamtObject *self)
2727 {
2728 return _PyHamt_NewIterKeys(self);
2729 }
2730
2731 static PyObject *
hamt_py_set(PyHamtObject * self,PyObject * args)2732 hamt_py_set(PyHamtObject *self, PyObject *args)
2733 {
2734 PyObject *key;
2735 PyObject *val;
2736
2737 if (!PyArg_UnpackTuple(args, "set", 2, 2, &key, &val)) {
2738 return NULL;
2739 }
2740
2741 return (PyObject *)_PyHamt_Assoc(self, key, val);
2742 }
2743
2744 static PyObject *
hamt_py_get(PyHamtObject * self,PyObject * args)2745 hamt_py_get(PyHamtObject *self, PyObject *args)
2746 {
2747 PyObject *key;
2748 PyObject *def = NULL;
2749
2750 if (!PyArg_UnpackTuple(args, "get", 1, 2, &key, &def)) {
2751 return NULL;
2752 }
2753
2754 PyObject *val = NULL;
2755 hamt_find_t res = hamt_find(self, key, &val);
2756 switch (res) {
2757 case F_ERROR:
2758 return NULL;
2759 case F_FOUND:
2760 return Py_NewRef(val);
2761 case F_NOT_FOUND:
2762 if (def == NULL) {
2763 Py_RETURN_NONE;
2764 }
2765 return Py_NewRef(def);
2766 default:
2767 Py_UNREACHABLE();
2768 }
2769 }
2770
2771 static PyObject *
hamt_py_delete(PyHamtObject * self,PyObject * key)2772 hamt_py_delete(PyHamtObject *self, PyObject *key)
2773 {
2774 return (PyObject *)_PyHamt_Without(self, key);
2775 }
2776
2777 static PyObject *
hamt_py_items(PyHamtObject * self,PyObject * args)2778 hamt_py_items(PyHamtObject *self, PyObject *args)
2779 {
2780 return _PyHamt_NewIterItems(self);
2781 }
2782
2783 static PyObject *
hamt_py_values(PyHamtObject * self,PyObject * args)2784 hamt_py_values(PyHamtObject *self, PyObject *args)
2785 {
2786 return _PyHamt_NewIterValues(self);
2787 }
2788
2789 static PyObject *
hamt_py_keys(PyHamtObject * self,PyObject * Py_UNUSED (args))2790 hamt_py_keys(PyHamtObject *self, PyObject *Py_UNUSED(args))
2791 {
2792 return _PyHamt_NewIterKeys(self);
2793 }
2794
2795 #ifdef Py_DEBUG
2796 static PyObject *
hamt_py_dump(PyHamtObject * self,PyObject * Py_UNUSED (args))2797 hamt_py_dump(PyHamtObject *self, PyObject *Py_UNUSED(args))
2798 {
2799 return hamt_dump(self);
2800 }
2801 #endif
2802
2803
2804 static PyMethodDef PyHamt_methods[] = {
2805 {"set", _PyCFunction_CAST(hamt_py_set), METH_VARARGS, NULL},
2806 {"get", _PyCFunction_CAST(hamt_py_get), METH_VARARGS, NULL},
2807 {"delete", _PyCFunction_CAST(hamt_py_delete), METH_O, NULL},
2808 {"items", _PyCFunction_CAST(hamt_py_items), METH_NOARGS, NULL},
2809 {"keys", _PyCFunction_CAST(hamt_py_keys), METH_NOARGS, NULL},
2810 {"values", _PyCFunction_CAST(hamt_py_values), METH_NOARGS, NULL},
2811 #ifdef Py_DEBUG
2812 {"__dump__", _PyCFunction_CAST(hamt_py_dump), METH_NOARGS, NULL},
2813 #endif
2814 {NULL, NULL}
2815 };
2816
2817 static PySequenceMethods PyHamt_as_sequence = {
2818 0, /* sq_length */
2819 0, /* sq_concat */
2820 0, /* sq_repeat */
2821 0, /* sq_item */
2822 0, /* sq_slice */
2823 0, /* sq_ass_item */
2824 0, /* sq_ass_slice */
2825 (objobjproc)hamt_tp_contains, /* sq_contains */
2826 0, /* sq_inplace_concat */
2827 0, /* sq_inplace_repeat */
2828 };
2829
2830 static PyMappingMethods PyHamt_as_mapping = {
2831 (lenfunc)hamt_tp_len, /* mp_length */
2832 (binaryfunc)hamt_tp_subscript, /* mp_subscript */
2833 };
2834
2835 PyTypeObject _PyHamt_Type = {
2836 PyVarObject_HEAD_INIT(&PyType_Type, 0)
2837 "hamt",
2838 sizeof(PyHamtObject),
2839 .tp_methods = PyHamt_methods,
2840 .tp_as_mapping = &PyHamt_as_mapping,
2841 .tp_as_sequence = &PyHamt_as_sequence,
2842 .tp_iter = (getiterfunc)hamt_tp_iter,
2843 .tp_dealloc = (destructor)hamt_tp_dealloc,
2844 .tp_getattro = PyObject_GenericGetAttr,
2845 .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
2846 .tp_richcompare = hamt_tp_richcompare,
2847 .tp_traverse = (traverseproc)hamt_tp_traverse,
2848 .tp_clear = (inquiry)hamt_tp_clear,
2849 .tp_new = hamt_tp_new,
2850 .tp_weaklistoffset = offsetof(PyHamtObject, h_weakreflist),
2851 .tp_hash = PyObject_HashNotImplemented,
2852 };
2853
2854
2855 /////////////////////////////////// Tree Node Types
2856
2857
2858 PyTypeObject _PyHamt_ArrayNode_Type = {
2859 PyVarObject_HEAD_INIT(&PyType_Type, 0)
2860 "hamt_array_node",
2861 sizeof(PyHamtNode_Array),
2862 0,
2863 .tp_dealloc = (destructor)hamt_node_array_dealloc,
2864 .tp_getattro = PyObject_GenericGetAttr,
2865 .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
2866 .tp_traverse = (traverseproc)hamt_node_array_traverse,
2867 .tp_free = PyObject_GC_Del,
2868 .tp_hash = PyObject_HashNotImplemented,
2869 };
2870
2871 PyTypeObject _PyHamt_BitmapNode_Type = {
2872 PyVarObject_HEAD_INIT(&PyType_Type, 0)
2873 "hamt_bitmap_node",
2874 sizeof(PyHamtNode_Bitmap) - sizeof(PyObject *),
2875 sizeof(PyObject *),
2876 .tp_dealloc = (destructor)hamt_node_bitmap_dealloc,
2877 .tp_getattro = PyObject_GenericGetAttr,
2878 .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
2879 .tp_traverse = (traverseproc)hamt_node_bitmap_traverse,
2880 .tp_free = PyObject_GC_Del,
2881 .tp_hash = PyObject_HashNotImplemented,
2882 };
2883
2884 PyTypeObject _PyHamt_CollisionNode_Type = {
2885 PyVarObject_HEAD_INIT(&PyType_Type, 0)
2886 "hamt_collision_node",
2887 sizeof(PyHamtNode_Collision) - sizeof(PyObject *),
2888 sizeof(PyObject *),
2889 .tp_dealloc = (destructor)hamt_node_collision_dealloc,
2890 .tp_getattro = PyObject_GenericGetAttr,
2891 .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
2892 .tp_traverse = (traverseproc)hamt_node_collision_traverse,
2893 .tp_free = PyObject_GC_Del,
2894 .tp_hash = PyObject_HashNotImplemented,
2895 };
2896