• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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