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