Branch data 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 : : _hamt_node_array_validate(void *obj_raw)
378 : : {
379 : : PyObject *obj = _PyObject_CAST(obj_raw);
380 : : assert(IS_ARRAY_NODE(obj));
381 : : PyHamtNode_Array *node = (PyHamtNode_Array*)obj;
382 : : Py_ssize_t i = 0, count = 0;
383 : : for (; i < HAMT_ARRAY_NODE_SIZE; i++) {
384 : : if (node->a_array[i] != NULL) {
385 : : count++;
386 : : }
387 : : }
388 : : assert(count == node->a_count);
389 : : }
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 : 170291 : hamt_hash(PyObject *o)
401 : : {
402 : 170291 : Py_hash_t hash = PyObject_Hash(o);
403 : :
404 : : #if SIZEOF_PY_HASH_T <= 4
405 : : return hash;
406 : : #else
407 [ + + ]: 170291 : 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 : 169851 : int32_t xored = (int32_t)(hash & 0xffffffffl) ^ (int32_t)(hash >> 32);
430 [ + - ]: 169851 : return xored == -1 ? -2 : xored;
431 : : #endif
432 : : }
433 : :
434 : : static inline uint32_t
435 : 448828 : hamt_mask(int32_t hash, uint32_t shift)
436 : : {
437 : 448828 : return (((uint32_t)hash >> shift) & 0x01f);
438 : : }
439 : :
440 : : static inline uint32_t
441 : 179000 : hamt_bitpos(int32_t hash, uint32_t shift)
442 : : {
443 : 179000 : return (uint32_t)1 << hamt_mask(hash, shift);
444 : : }
445 : :
446 : : static inline uint32_t
447 : 162977 : hamt_bitindex(uint32_t bitmap, uint32_t bit)
448 : : {
449 : 162977 : return (uint32_t)_Py_popcount32(bitmap & (bit - 1));
450 : : }
451 : :
452 : :
453 : : /////////////////////////////////// Dump Helpers
454 : : #ifdef Py_DEBUG
455 : :
456 : : static int
457 : : _hamt_dump_ident(_PyUnicodeWriter *writer, int level)
458 : : {
459 : : /* Write `' ' * level` to the `writer` */
460 : : PyObject *str = NULL;
461 : : PyObject *num = NULL;
462 : : PyObject *res = NULL;
463 : : int ret = -1;
464 : :
465 : : str = PyUnicode_FromString(" ");
466 : : if (str == NULL) {
467 : : goto error;
468 : : }
469 : :
470 : : num = PyLong_FromLong((long)level);
471 : : if (num == NULL) {
472 : : goto error;
473 : : }
474 : :
475 : : res = PyNumber_Multiply(str, num);
476 : : if (res == NULL) {
477 : : goto error;
478 : : }
479 : :
480 : : ret = _PyUnicodeWriter_WriteStr(writer, res);
481 : :
482 : : error:
483 : : Py_XDECREF(res);
484 : : Py_XDECREF(str);
485 : : Py_XDECREF(num);
486 : : return ret;
487 : : }
488 : :
489 : : static int
490 : : _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 : : va_start(vargs, format);
500 : : msg = PyUnicode_FromFormatV(format, vargs);
501 : : va_end(vargs);
502 : :
503 : : if (msg == NULL) {
504 : : return -1;
505 : : }
506 : :
507 : : ret = _PyUnicodeWriter_WriteStr(writer, msg);
508 : : Py_DECREF(msg);
509 : : return ret;
510 : : }
511 : :
512 : : #endif /* Py_DEBUG */
513 : : /////////////////////////////////// Bitmap Node
514 : :
515 : :
516 : : static PyHamtNode *
517 : 72623 : 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 : : assert(size >= 0);
525 : : assert(size % 2 == 0);
526 : :
527 [ + + + + ]: 72623 : if (size == 0 && _empty_bitmap_node != NULL) {
528 : 4107 : Py_INCREF(_empty_bitmap_node);
529 : 4107 : return (PyHamtNode *)_empty_bitmap_node;
530 : : }
531 : :
532 : : /* No freelist; allocate a new bitmap node */
533 : 68516 : node = PyObject_GC_NewVar(
534 : : PyHamtNode_Bitmap, &_PyHamt_BitmapNode_Type, size);
535 [ - + ]: 68516 : if (node == NULL) {
536 : 0 : return NULL;
537 : : }
538 : :
539 : 68516 : Py_SET_SIZE(node, size);
540 : :
541 [ + + ]: 574190 : for (i = 0; i < size; i++) {
542 : 505674 : node->b_array[i] = NULL;
543 : : }
544 : :
545 : 68516 : node->b_bitmap = 0;
546 : :
547 : 68516 : _PyObject_GC_TRACK(node);
548 : :
549 [ + + + - ]: 68516 : 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 : 68516 : return (PyHamtNode *)node;
558 : : }
559 : :
560 : : static inline Py_ssize_t
561 : 38282 : hamt_node_bitmap_count(PyHamtNode_Bitmap *node)
562 : : {
563 : 38282 : return Py_SIZE(node) / 2;
564 : : }
565 : :
566 : : static PyHamtNode_Bitmap *
567 : 14410 : 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 : 14410 : clone = (PyHamtNode_Bitmap *)hamt_node_bitmap_new(Py_SIZE(node));
575 [ - + ]: 14410 : if (clone == NULL) {
576 : 0 : return NULL;
577 : : }
578 : :
579 [ + + ]: 128848 : for (i = 0; i < Py_SIZE(node); i++) {
580 : 114438 : Py_XINCREF(node->b_array[i]);
581 : 114438 : clone->b_array[i] = node->b_array[i];
582 : : }
583 : :
584 : 14410 : clone->b_bitmap = node->b_bitmap;
585 : 14410 : return clone;
586 : : }
587 : :
588 : : static PyHamtNode_Bitmap *
589 : 28259 : hamt_node_bitmap_clone_without(PyHamtNode_Bitmap *o, uint32_t bit)
590 : : {
591 : : assert(bit & o->b_bitmap);
592 : : assert(hamt_node_bitmap_count(o) > 1);
593 : :
594 : 28259 : PyHamtNode_Bitmap *new = (PyHamtNode_Bitmap *)hamt_node_bitmap_new(
595 : 28259 : Py_SIZE(o) - 2);
596 [ - + ]: 28259 : if (new == NULL) {
597 : 0 : return NULL;
598 : : }
599 : :
600 : 28259 : uint32_t idx = hamt_bitindex(o->b_bitmap, bit);
601 : 28259 : uint32_t key_idx = 2 * idx;
602 : 28259 : uint32_t val_idx = key_idx + 1;
603 : : uint32_t i;
604 : :
605 [ + + ]: 124243 : for (i = 0; i < key_idx; i++) {
606 : 95984 : Py_XINCREF(o->b_array[i]);
607 : 95984 : new->b_array[i] = o->b_array[i];
608 : : }
609 : :
610 : : assert(Py_SIZE(o) >= 0 && Py_SIZE(o) <= 32);
611 [ + + ]: 123437 : for (i = val_idx + 1; i < (uint32_t)Py_SIZE(o); i++) {
612 : 95178 : Py_XINCREF(o->b_array[i]);
613 : 95178 : new->b_array[i - 2] = o->b_array[i];
614 : : }
615 : :
616 : 28259 : new->b_bitmap = o->b_bitmap & ~bit;
617 : 28259 : return new;
618 : : }
619 : :
620 : : static PyHamtNode *
621 : 2534 : 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 : 2534 : int32_t key1_hash = hamt_hash(key1);
635 [ - + ]: 2534 : if (key1_hash == -1) {
636 : 0 : return NULL;
637 : : }
638 : :
639 [ + + ]: 2534 : 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 : 2525 : int added_leaf = 0;
660 : 2525 : PyHamtNode *n = hamt_node_bitmap_new(0);
661 [ - + ]: 2525 : if (n == NULL) {
662 : 0 : return NULL;
663 : : }
664 : :
665 : 2525 : PyHamtNode *n2 = hamt_node_assoc(
666 : : n, shift, key1_hash, key1, val1, &added_leaf);
667 : 2525 : Py_DECREF(n);
668 [ - + ]: 2525 : if (n2 == NULL) {
669 : 0 : return NULL;
670 : : }
671 : :
672 : 2525 : n = hamt_node_assoc(n2, shift, key2_hash, key2, val2, &added_leaf);
673 : 2525 : Py_DECREF(n2);
674 [ - + ]: 2525 : if (n == NULL) {
675 : 0 : return NULL;
676 : : }
677 : :
678 : 2525 : return n;
679 : : }
680 : : }
681 : :
682 : : static PyHamtNode *
683 : 36348 : 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 : 36348 : uint32_t bit = hamt_bitpos(hash, shift);
697 : 36348 : 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 [ + + ]: 36348 : if (self->b_bitmap & bit) {
714 : : /* The key is set in this node */
715 : :
716 : 10644 : uint32_t key_idx = 2 * idx;
717 : 10644 : uint32_t val_idx = key_idx + 1;
718 : :
719 : : assert(val_idx < (size_t)Py_SIZE(self));
720 : :
721 : 10644 : PyObject *key_or_null = self->b_array[key_idx];
722 : 10644 : PyObject *val_or_node = self->b_array[val_idx];
723 : :
724 [ + + ]: 10644 : 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 : : assert(val_or_node != NULL);
729 : :
730 : 296 : PyHamtNode *sub_node = hamt_node_assoc(
731 : : (PyHamtNode *)val_or_node,
732 : : shift + 5, hash, key, val, added_leaf);
733 [ - + ]: 296 : if (sub_node == NULL) {
734 : 0 : return NULL;
735 : : }
736 : :
737 [ - + ]: 296 : 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 : 296 : PyHamtNode_Bitmap *ret = hamt_node_bitmap_clone(self);
744 [ - + ]: 296 : if (ret == NULL) {
745 : 0 : return NULL;
746 : : }
747 : 296 : Py_SETREF(ret->b_array[val_idx], (PyObject*)sub_node);
748 : 296 : return (PyHamtNode *)ret;
749 : : }
750 : :
751 : : 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 : 10348 : int comp_err = PyObject_RichCompareBool(key, key_or_null, Py_EQ);
756 [ - + ]: 10348 : if (comp_err < 0) { /* exception in __eq__ */
757 : 0 : return NULL;
758 : : }
759 [ + + ]: 10348 : 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 : 2534 : 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 [ - + ]: 2534 : if (sub_node == NULL) {
792 : 0 : return NULL;
793 : : }
794 : :
795 : 2534 : PyHamtNode_Bitmap *ret = hamt_node_bitmap_clone(self);
796 [ - + ]: 2534 : if (ret == NULL) {
797 : 0 : Py_DECREF(sub_node);
798 : 0 : return NULL;
799 : : }
800 : 2534 : Py_SETREF(ret->b_array[key_idx], NULL);
801 : 2534 : Py_SETREF(ret->b_array[val_idx], (PyObject *)sub_node);
802 : :
803 : 2534 : *added_leaf = 1;
804 : 2534 : return (PyHamtNode *)ret;
805 : : }
806 : : else {
807 : : /* There was no key before with the same (shift,hash). */
808 : :
809 : 25704 : uint32_t n = (uint32_t)_Py_popcount32(self->b_bitmap);
810 : :
811 [ + + ]: 25704 : 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 : : assert(new_node->a_array[i] == NULL);
865 : :
866 [ + + ]: 1600 : if (self->b_array[j] == NULL) {
867 : 498 : new_node->a_array[i] =
868 : 498 : (PyHamtNode *)self->b_array[j + 1];
869 : 498 : Py_INCREF(new_node->a_array[i]);
870 : : }
871 : : else {
872 : 1102 : int32_t rehash = hamt_hash(self->b_array[j]);
873 [ - + ]: 1102 : if (rehash == -1) {
874 : 0 : goto fin;
875 : : }
876 : :
877 : 2204 : new_node->a_array[i] = hamt_node_assoc(
878 : : empty, shift + 5,
879 : : rehash,
880 : : self->b_array[j],
881 : 1102 : self->b_array[j + 1],
882 : : added_leaf);
883 : :
884 [ - + ]: 1102 : if (new_node->a_array[i] == NULL) {
885 : 0 : goto fin;
886 : : }
887 : : }
888 : 1600 : j += 2;
889 : : }
890 : : }
891 : :
892 : : 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 : 25604 : uint32_t key_idx = 2 * idx;
910 : 25604 : uint32_t val_idx = key_idx + 1;
911 : : uint32_t i;
912 : :
913 : 25604 : *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 : 25604 : (PyHamtNode_Bitmap *)hamt_node_bitmap_new(2 * (n + 1));
919 [ - + ]: 25604 : 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 [ + + ]: 96510 : for (i = 0; i < key_idx; i++) {
926 : 70906 : Py_XINCREF(self->b_array[i]);
927 : 70906 : new_node->b_array[i] = self->b_array[i];
928 : : }
929 : :
930 : : /* Set the new key/value to the new Bitmap node. */
931 : 25604 : Py_INCREF(key);
932 : 25604 : new_node->b_array[key_idx] = key;
933 : 25604 : Py_INCREF(val);
934 : 25604 : 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 : : assert(Py_SIZE(self) >= 0 && Py_SIZE(self) <= 32);
939 [ + + ]: 97576 : for (i = key_idx; i < (uint32_t)Py_SIZE(self); i++) {
940 : 71972 : Py_XINCREF(self->b_array[i]);
941 : 71972 : new_node->b_array[i + 2] = self->b_array[i];
942 : : }
943 : :
944 : 25604 : new_node->b_bitmap = self->b_bitmap | bit;
945 : 25604 : return (PyHamtNode *)new_node;
946 : : }
947 : : }
948 : : }
949 : :
950 : : static hamt_without_t
951 : 47645 : hamt_node_bitmap_without(PyHamtNode_Bitmap *self,
952 : : uint32_t shift, int32_t hash,
953 : : PyObject *key,
954 : : PyHamtNode **new_node)
955 : : {
956 : 47645 : uint32_t bit = hamt_bitpos(hash, shift);
957 [ + + ]: 47645 : if ((self->b_bitmap & bit) == 0) {
958 : 9132 : return W_NOT_FOUND;
959 : : }
960 : :
961 : 38513 : uint32_t idx = hamt_bitindex(self->b_bitmap, bit);
962 : :
963 : 38513 : uint32_t key_idx = 2 * idx;
964 : 38513 : uint32_t val_idx = key_idx + 1;
965 : :
966 : 38513 : PyObject *key_or_null = self->b_array[key_idx];
967 : 38513 : PyObject *val_or_node = self->b_array[val_idx];
968 : :
969 [ + + ]: 38513 : if (key_or_null == NULL) {
970 : : /* key == NULL means that 'value' is another tree node. */
971 : :
972 : 4058 : PyHamtNode *sub_node = NULL;
973 : :
974 : 4058 : hamt_without_t res = hamt_node_without(
975 : : (PyHamtNode *)val_or_node,
976 : : shift + 5, hash, key, &sub_node);
977 : :
978 [ + + - ]: 4058 : switch (res) {
979 : : 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 : : Py_UNREACHABLE();
995 : :
996 : 3768 : case W_NEWNODE: {
997 : : assert(sub_node != NULL);
998 : :
999 [ + + ]: 3768 : if (IS_BITMAP_NODE(sub_node)) {
1000 : 3767 : PyHamtNode_Bitmap *sub_tree = (PyHamtNode_Bitmap *)sub_node;
1001 [ + + ]: 3767 : if (hamt_node_bitmap_count(sub_tree) == 1 &&
1002 [ + + ]: 3421 : 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 : 3415 : PyHamtNode_Bitmap *clone = hamt_node_bitmap_clone(self);
1014 [ - + ]: 3415 : if (clone == NULL) {
1015 : 0 : Py_DECREF(sub_node);
1016 : 0 : return W_ERROR;
1017 : : }
1018 : :
1019 : 3415 : PyObject *key = sub_tree->b_array[0];
1020 : 3415 : PyObject *val = sub_tree->b_array[1];
1021 : :
1022 : 3415 : Py_INCREF(key);
1023 : 3415 : Py_XSETREF(clone->b_array[key_idx], key);
1024 : 3415 : Py_INCREF(val);
1025 : 3415 : Py_SETREF(clone->b_array[val_idx], val);
1026 : :
1027 : 3415 : Py_DECREF(sub_tree);
1028 : :
1029 : 3415 : *new_node = (PyHamtNode *)clone;
1030 : 3415 : return W_NEWNODE;
1031 : : }
1032 : : }
1033 : :
1034 : : #ifdef Py_DEBUG
1035 : : /* Ensure that Collision.without implementation
1036 : : converts to Bitmap nodes itself.
1037 : : */
1038 : : if (IS_COLLISION_NODE(sub_node)) {
1039 : : assert(hamt_node_collision_count(
1040 : : (PyHamtNode_Collision*)sub_node) > 1);
1041 : : }
1042 : : #endif
1043 : :
1044 : 353 : PyHamtNode_Bitmap *clone = hamt_node_bitmap_clone(self);
1045 [ - + ]: 353 : if (clone == NULL) {
1046 : 0 : return W_ERROR;
1047 : : }
1048 : :
1049 : 353 : Py_SETREF(clone->b_array[val_idx],
1050 : : (PyObject *)sub_node); /* borrow */
1051 : :
1052 : 353 : *new_node = (PyHamtNode *)clone;
1053 : 353 : return W_NEWNODE;
1054 : : }
1055 : :
1056 : 290 : case W_ERROR:
1057 : : case W_NOT_FOUND:
1058 : : assert(sub_node == NULL);
1059 : 290 : return res;
1060 : :
1061 : 0 : default:
1062 : 0 : Py_UNREACHABLE();
1063 : : }
1064 : : }
1065 : : else {
1066 : : /* We have a regular key/value pair */
1067 : :
1068 : 34455 : int cmp = PyObject_RichCompareBool(key_or_null, key, Py_EQ);
1069 [ + + ]: 34455 : if (cmp < 0) {
1070 : 1913 : return W_ERROR;
1071 : : }
1072 [ + + ]: 32542 : if (cmp == 0) {
1073 : 1012 : return W_NOT_FOUND;
1074 : : }
1075 : :
1076 [ + + ]: 31530 : if (hamt_node_bitmap_count(self) == 1) {
1077 : 3271 : return W_EMPTY;
1078 : : }
1079 : :
1080 : 28259 : *new_node = (PyHamtNode *)
1081 : 28259 : hamt_node_bitmap_clone_without(self, bit);
1082 [ - + ]: 28259 : if (*new_node == NULL) {
1083 : 0 : return W_ERROR;
1084 : : }
1085 : :
1086 : 28259 : return W_NEWNODE;
1087 : : }
1088 : : }
1089 : :
1090 : : static hamt_find_t
1091 : 94998 : 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 : 94998 : 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 [ + + ]: 94998 : if ((self->b_bitmap & bit) == 0) {
1106 : 35141 : return F_NOT_FOUND;
1107 : : }
1108 : :
1109 : 59857 : idx = hamt_bitindex(self->b_bitmap, bit);
1110 : 59857 : key_idx = idx * 2;
1111 : 59857 : val_idx = key_idx + 1;
1112 : :
1113 : : assert(val_idx < (size_t)Py_SIZE(self));
1114 : :
1115 : 59857 : key_or_null = self->b_array[key_idx];
1116 : 59857 : val_or_node = self->b_array[val_idx];
1117 : :
1118 [ + + ]: 59857 : 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 : : assert(val_or_node != NULL);
1122 : 5993 : 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 : : assert(key != NULL);
1129 : 53864 : comp_err = PyObject_RichCompareBool(key, key_or_null, Py_EQ);
1130 [ + + ]: 53864 : if (comp_err < 0) { /* exception in __eq__ */
1131 : 1915 : return F_ERROR;
1132 : : }
1133 [ + + ]: 51949 : if (comp_err == 1) { /* key == key_or_null */
1134 : 47589 : *val = val_or_node;
1135 : 47589 : return F_FOUND;
1136 : : }
1137 : :
1138 : 4360 : return F_NOT_FOUND;
1139 : : }
1140 : :
1141 : : static int
1142 : 83000 : 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 [ + + ]: 637852 : for (i = Py_SIZE(self); --i >= 0; ) {
1149 [ + + - + ]: 554852 : Py_VISIT(self->b_array[i]);
1150 : : }
1151 : :
1152 : 83000 : return 0;
1153 : : }
1154 : :
1155 : : static void
1156 : 68516 : hamt_node_bitmap_dealloc(PyHamtNode_Bitmap *self)
1157 : : {
1158 : : /* Bitmap's tp_dealloc */
1159 : :
1160 : 68516 : Py_ssize_t len = Py_SIZE(self);
1161 : : Py_ssize_t i;
1162 : :
1163 : 68516 : PyObject_GC_UnTrack(self);
1164 [ + - - + ]: 68516 : Py_TRASHCAN_BEGIN(self, hamt_node_bitmap_dealloc)
1165 : :
1166 [ + + ]: 68516 : if (len > 0) {
1167 : 68481 : i = len;
1168 [ + + ]: 574155 : while (--i >= 0) {
1169 : 505674 : Py_XDECREF(self->b_array[i]);
1170 : : }
1171 : : }
1172 : :
1173 : 68516 : Py_TYPE(self)->tp_free((PyObject *)self);
1174 [ + - ]: 68516 : Py_TRASHCAN_END
1175 : 68516 : }
1176 : :
1177 : : #ifdef Py_DEBUG
1178 : : static int
1179 : : 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 : : if (_hamt_dump_ident(writer, level + 1)) {
1189 : : goto error;
1190 : : }
1191 : :
1192 : : if (_hamt_dump_format(writer, "BitmapNode(size=%zd count=%zd ",
1193 : : Py_SIZE(node), Py_SIZE(node) / 2))
1194 : : {
1195 : : goto error;
1196 : : }
1197 : :
1198 : : tmp1 = PyLong_FromUnsignedLong(node->b_bitmap);
1199 : : if (tmp1 == NULL) {
1200 : : goto error;
1201 : : }
1202 : : tmp2 = _PyLong_Format(tmp1, 2);
1203 : : Py_DECREF(tmp1);
1204 : : if (tmp2 == NULL) {
1205 : : goto error;
1206 : : }
1207 : : if (_hamt_dump_format(writer, "bitmap=%S id=%p):\n", tmp2, node)) {
1208 : : Py_DECREF(tmp2);
1209 : : goto error;
1210 : : }
1211 : : Py_DECREF(tmp2);
1212 : :
1213 : : for (i = 0; i < Py_SIZE(node); i += 2) {
1214 : : PyObject *key_or_null = node->b_array[i];
1215 : : PyObject *val_or_node = node->b_array[i + 1];
1216 : :
1217 : : if (_hamt_dump_ident(writer, level + 2)) {
1218 : : goto error;
1219 : : }
1220 : :
1221 : : if (key_or_null == NULL) {
1222 : : if (_hamt_dump_format(writer, "NULL:\n")) {
1223 : : goto error;
1224 : : }
1225 : :
1226 : : if (hamt_node_dump((PyHamtNode *)val_or_node,
1227 : : writer, level + 2))
1228 : : {
1229 : : goto error;
1230 : : }
1231 : : }
1232 : : else {
1233 : : if (_hamt_dump_format(writer, "%R: %R", key_or_null,
1234 : : val_or_node))
1235 : : {
1236 : : goto error;
1237 : : }
1238 : : }
1239 : :
1240 : : if (_hamt_dump_format(writer, "\n")) {
1241 : : goto error;
1242 : : }
1243 : : }
1244 : :
1245 : : return 0;
1246 : : error:
1247 : : 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 : : assert(size >= 4);
1264 : : 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 : : 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 : : assert(key_idx >= 0);
1362 : : 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 : 4 : hamt_node_collision_count(PyHamtNode_Collision *node)
1425 : : {
1426 : 4 : 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 : : assert(key_idx >= 0);
1451 : : 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 : : 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 : : assert(idx >= 0);
1540 : : assert(idx + 1 < Py_SIZE(self));
1541 : :
1542 : 19 : *val = self->c_array[idx + 1];
1543 : : 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 : : 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 : : if (_hamt_dump_ident(writer, level + 1)) {
1595 : : goto error;
1596 : : }
1597 : :
1598 : : if (_hamt_dump_format(writer, "CollisionNode(size=%zd id=%p):\n",
1599 : : Py_SIZE(node), node))
1600 : : {
1601 : : goto error;
1602 : : }
1603 : :
1604 : : for (i = 0; i < Py_SIZE(node); i += 2) {
1605 : : PyObject *key = node->c_array[i];
1606 : : PyObject *val = node->c_array[i + 1];
1607 : :
1608 : : if (_hamt_dump_ident(writer, level + 2)) {
1609 : : goto error;
1610 : : }
1611 : :
1612 : : if (_hamt_dump_format(writer, "%R: %R\n", key, val)) {
1613 : : goto error;
1614 : : }
1615 : : }
1616 : :
1617 : : return 0;
1618 : : error:
1619 : : return -1;
1620 : : }
1621 : : #endif /* Py_DEBUG */
1622 : :
1623 : :
1624 : : /////////////////////////////////// Array Node
1625 : :
1626 : :
1627 : : static PyHamtNode *
1628 : 98575 : hamt_node_array_new(Py_ssize_t count)
1629 : : {
1630 : : Py_ssize_t i;
1631 : :
1632 : 98575 : PyHamtNode_Array *node = PyObject_GC_New(
1633 : : PyHamtNode_Array, &_PyHamt_ArrayNode_Type);
1634 [ - + ]: 98575 : if (node == NULL) {
1635 : 0 : return NULL;
1636 : : }
1637 : :
1638 [ + + ]: 3252975 : for (i = 0; i < HAMT_ARRAY_NODE_SIZE; i++) {
1639 : 3154400 : node->a_array[i] = NULL;
1640 : : }
1641 : :
1642 : 98575 : node->a_count = count;
1643 : :
1644 : 98575 : _PyObject_GC_TRACK(node);
1645 : 98575 : return (PyHamtNode *)node;
1646 : : }
1647 : :
1648 : : static PyHamtNode_Array *
1649 : 96993 : hamt_node_array_clone(PyHamtNode_Array *node)
1650 : : {
1651 : : PyHamtNode_Array *clone;
1652 : : Py_ssize_t i;
1653 : :
1654 : : VALIDATE_ARRAY_NODE(node)
1655 : :
1656 : : /* Create a new Array node. */
1657 : 96993 : clone = (PyHamtNode_Array *)hamt_node_array_new(node->a_count);
1658 [ - + ]: 96993 : if (clone == NULL) {
1659 : 0 : return NULL;
1660 : : }
1661 : :
1662 : : /* Copy all elements from the current Array node to the new one. */
1663 [ + + ]: 3200769 : for (i = 0; i < HAMT_ARRAY_NODE_SIZE; i++) {
1664 : 3103776 : Py_XINCREF(node->a_array[i]);
1665 : 3103776 : clone->a_array[i] = node->a_array[i];
1666 : : }
1667 : :
1668 : : VALIDATE_ARRAY_NODE(clone)
1669 : 96993 : return clone;
1670 : : }
1671 : :
1672 : : static PyHamtNode *
1673 : 39679 : 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 : 39679 : uint32_t idx = hamt_mask(hash, shift);
1685 : 39679 : PyHamtNode *node = self->a_array[idx];
1686 : : PyHamtNode *child_node;
1687 : : PyHamtNode_Array *new_node;
1688 : : Py_ssize_t i;
1689 : :
1690 [ + + ]: 39679 : if (node == NULL) {
1691 : : /* There's no child node for the given hash. Create a new
1692 : : Bitmap node for this key. */
1693 : :
1694 : 1482 : PyHamtNode_Bitmap *empty = NULL;
1695 : :
1696 : : /* Get an empty Bitmap node to work with. */
1697 : 1482 : empty = (PyHamtNode_Bitmap *)hamt_node_bitmap_new(0);
1698 [ - + ]: 1482 : 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 : 1482 : child_node = hamt_node_bitmap_assoc(
1705 : : empty,
1706 : : shift + 5, hash, key, val, added_leaf);
1707 : 1482 : Py_DECREF(empty);
1708 [ - + ]: 1482 : if (child_node == NULL) {
1709 : 0 : return NULL;
1710 : : }
1711 : :
1712 : : /* Create a new Array node. */
1713 : 1482 : new_node = (PyHamtNode_Array *)hamt_node_array_new(self->a_count + 1);
1714 [ - + ]: 1482 : 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 [ + + ]: 48906 : for (i = 0; i < HAMT_ARRAY_NODE_SIZE; i++) {
1722 : 47424 : Py_XINCREF(self->a_array[i]);
1723 : 47424 : new_node->a_array[i] = self->a_array[i];
1724 : : }
1725 : :
1726 : : assert(new_node->a_array[idx] == NULL);
1727 : 1482 : new_node->a_array[idx] = child_node; /* borrow */
1728 : : 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 : 38197 : child_node = hamt_node_assoc(
1734 : : node, shift + 5, hash, key, val, added_leaf);
1735 [ - + ]: 38197 : if (child_node == NULL) {
1736 : 0 : return NULL;
1737 : : }
1738 [ - + ]: 38197 : else if (child_node == (PyHamtNode *)self) {
1739 : 0 : Py_DECREF(child_node);
1740 : 0 : return (PyHamtNode *)self;
1741 : : }
1742 : :
1743 : 38197 : new_node = hamt_node_array_clone(self);
1744 [ - + ]: 38197 : if (new_node == NULL) {
1745 : 0 : Py_DECREF(child_node);
1746 : 0 : return NULL;
1747 : : }
1748 : :
1749 : 38197 : Py_SETREF(new_node->a_array[idx], child_node); /* borrow */
1750 : : VALIDATE_ARRAY_NODE(new_node)
1751 : : }
1752 : :
1753 : 39679 : return (PyHamtNode *)new_node;
1754 : : }
1755 : :
1756 : : static hamt_without_t
1757 : 83643 : hamt_node_array_without(PyHamtNode_Array *self,
1758 : : uint32_t shift, int32_t hash,
1759 : : PyObject *key,
1760 : : PyHamtNode **new_node)
1761 : : {
1762 : 83643 : uint32_t idx = hamt_mask(hash, shift);
1763 : 83643 : PyHamtNode *node = self->a_array[idx];
1764 : :
1765 [ + + ]: 83643 : if (node == NULL) {
1766 : 367 : return W_NOT_FOUND;
1767 : : }
1768 : :
1769 : 83276 : PyHamtNode *sub_node = NULL;
1770 : 83276 : hamt_without_t res = hamt_node_without(
1771 : : (PyHamtNode *)node,
1772 : : shift + 5, hash, key, &sub_node);
1773 : :
1774 [ + + + - ]: 83276 : switch (res) {
1775 : 24281 : case W_NOT_FOUND:
1776 : : case W_ERROR:
1777 : : assert(sub_node == NULL);
1778 : 24281 : return res;
1779 : :
1780 : 55736 : case W_NEWNODE: {
1781 : : /* We need to replace a node at the `idx` index.
1782 : : Clone this node and replace.
1783 : : */
1784 : : assert(sub_node != NULL);
1785 : :
1786 : 55736 : PyHamtNode_Array *clone = hamt_node_array_clone(self);
1787 [ - + ]: 55736 : if (clone == NULL) {
1788 : 0 : Py_DECREF(sub_node);
1789 : 0 : return W_ERROR;
1790 : : }
1791 : :
1792 : 55736 : Py_SETREF(clone->a_array[idx], sub_node); /* borrow */
1793 : 55736 : *new_node = (PyHamtNode*)clone; /* borrow */
1794 : 55736 : return W_NEWNODE;
1795 : : }
1796 : :
1797 : 3259 : case W_EMPTY: {
1798 : : 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 : 3259 : Py_ssize_t new_count = self->a_count - 1;
1803 : :
1804 [ - + ]: 3259 : if (new_count == 0) {
1805 : 0 : return W_EMPTY;
1806 : : }
1807 : :
1808 [ + + ]: 3259 : 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 : 3060 : PyHamtNode_Array *new = hamt_node_array_clone(self);
1817 [ - + ]: 3060 : if (new == NULL) {
1818 : 0 : return W_ERROR;
1819 : : }
1820 : 3060 : new->a_count = new_count;
1821 [ + - ]: 3060 : Py_CLEAR(new->a_array[idx]);
1822 : :
1823 : 3060 : *new_node = (PyHamtNode*)new; /* borrow */
1824 : 3060 : 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 [ + + ]: 2175 : child->b_array[0] != NULL)
1859 : 2153 : {
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 : 2153 : PyObject *key = child->b_array[0];
1869 : 2153 : PyObject *val = child->b_array[1];
1870 : :
1871 : 2153 : Py_INCREF(key);
1872 : 2153 : new->b_array[new_i] = key;
1873 : 2153 : Py_INCREF(val);
1874 : 2153 : new->b_array[new_i + 1] = val;
1875 : : }
1876 : : else {
1877 : 832 : new->b_array[new_i] = NULL;
1878 : 832 : Py_INCREF(node);
1879 : 832 : new->b_array[new_i + 1] = (PyObject*)node;
1880 : : }
1881 : : }
1882 : : else {
1883 : :
1884 : : #ifdef Py_DEBUG
1885 : : if (IS_COLLISION_NODE(node)) {
1886 : : Py_ssize_t child_count = hamt_node_collision_count(
1887 : : (PyHamtNode_Collision*)node);
1888 : : assert(child_count > 1);
1889 : : }
1890 : : else if (IS_ARRAY_NODE(node)) {
1891 : : 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 : 146406 : 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 : 146406 : uint32_t idx = hamt_mask(hash, shift);
1923 : : PyHamtNode *node;
1924 : :
1925 : 146406 : node = self->a_array[idx];
1926 [ + + ]: 146406 : if (node == NULL) {
1927 : 3426 : return F_NOT_FOUND;
1928 : : }
1929 : :
1930 : : /* Dispatch to the generic hamt_node_find */
1931 : 142980 : return hamt_node_find(node, shift + 5, hash, key, val);
1932 : : }
1933 : :
1934 : : static int
1935 : 4918 : 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 [ + + ]: 162294 : for (i = 0; i < HAMT_ARRAY_NODE_SIZE; i++) {
1943 [ + + - + ]: 157376 : Py_VISIT(self->a_array[i]);
1944 : : }
1945 : :
1946 : 4918 : return 0;
1947 : : }
1948 : :
1949 : : static void
1950 : 98575 : hamt_node_array_dealloc(PyHamtNode_Array *self)
1951 : : {
1952 : : /* Array's tp_dealloc */
1953 : :
1954 : : Py_ssize_t i;
1955 : :
1956 : 98575 : PyObject_GC_UnTrack(self);
1957 [ + - - + ]: 98575 : Py_TRASHCAN_BEGIN(self, hamt_node_array_dealloc)
1958 : :
1959 [ + + ]: 3252975 : for (i = 0; i < HAMT_ARRAY_NODE_SIZE; i++) {
1960 : 3154400 : Py_XDECREF(self->a_array[i]);
1961 : : }
1962 : :
1963 : 98575 : Py_TYPE(self)->tp_free((PyObject *)self);
1964 [ + - ]: 98575 : Py_TRASHCAN_END
1965 : 98575 : }
1966 : :
1967 : : #ifdef Py_DEBUG
1968 : : static int
1969 : : 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 : : if (_hamt_dump_ident(writer, level + 1)) {
1977 : : goto error;
1978 : : }
1979 : :
1980 : : if (_hamt_dump_format(writer, "ArrayNode(id=%p):\n", node)) {
1981 : : goto error;
1982 : : }
1983 : :
1984 : : for (i = 0; i < HAMT_ARRAY_NODE_SIZE; i++) {
1985 : : if (node->a_array[i] == NULL) {
1986 : : continue;
1987 : : }
1988 : :
1989 : : if (_hamt_dump_ident(writer, level + 2)) {
1990 : : goto error;
1991 : : }
1992 : :
1993 : : if (_hamt_dump_format(writer, "%zd::\n", i)) {
1994 : : goto error;
1995 : : }
1996 : :
1997 : : if (hamt_node_dump(node->a_array[i], writer, level + 1)) {
1998 : : goto error;
1999 : : }
2000 : :
2001 : : if (_hamt_dump_format(writer, "\n")) {
2002 : : goto error;
2003 : : }
2004 : : }
2005 : :
2006 : : return 0;
2007 : : error:
2008 : : return -1;
2009 : : }
2010 : : #endif /* Py_DEBUG */
2011 : :
2012 : :
2013 : : /////////////////////////////////// Node Dispatch
2014 : :
2015 : :
2016 : : static PyHamtNode *
2017 : 74551 : 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 [ + + ]: 74551 : if (IS_BITMAP_NODE(node)) {
2033 : 34860 : return hamt_node_bitmap_assoc(
2034 : : (PyHamtNode_Bitmap *)node,
2035 : : shift, hash, key, val, added_leaf);
2036 : : }
2037 [ + + ]: 39691 : else if (IS_ARRAY_NODE(node)) {
2038 : 39679 : return hamt_node_array_assoc(
2039 : : (PyHamtNode_Array *)node,
2040 : : shift, hash, key, val, added_leaf);
2041 : : }
2042 : : else {
2043 : : 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 : 131292 : hamt_node_without(PyHamtNode *node,
2052 : : uint32_t shift, int32_t hash,
2053 : : PyObject *key,
2054 : : PyHamtNode **new_node)
2055 : : {
2056 [ + + ]: 131292 : if (IS_BITMAP_NODE(node)) {
2057 : 47645 : return hamt_node_bitmap_without(
2058 : : (PyHamtNode_Bitmap *)node,
2059 : : shift, hash, key,
2060 : : new_node);
2061 : : }
2062 [ + + ]: 83647 : else if (IS_ARRAY_NODE(node)) {
2063 : 83643 : return hamt_node_array_without(
2064 : : (PyHamtNode_Array *)node,
2065 : : shift, hash, key,
2066 : : new_node);
2067 : : }
2068 : : else {
2069 : : 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 : 241424 : 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 [ + + ]: 241424 : if (IS_BITMAP_NODE(node)) {
2096 : 94998 : return hamt_node_bitmap_find(
2097 : : (PyHamtNode_Bitmap *)node,
2098 : : shift, hash, key, val);
2099 : :
2100 : : }
2101 [ + + ]: 146426 : else if (IS_ARRAY_NODE(node)) {
2102 : 146406 : return hamt_node_array_find(
2103 : : (PyHamtNode_Array *)node,
2104 : : shift, hash, key, val);
2105 : : }
2106 : : else {
2107 : : 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 : : 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 : : if (IS_BITMAP_NODE(node)) {
2126 : : return hamt_node_bitmap_dump(
2127 : : (PyHamtNode_Bitmap *)node, writer, level);
2128 : : }
2129 : : else if (IS_ARRAY_NODE(node)) {
2130 : : return hamt_node_array_dump(
2131 : : (PyHamtNode_Array *)node, writer, level);
2132 : : }
2133 : : else {
2134 : : assert(IS_COLLISION_NODE(node));
2135 : : 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 : 318488 : hamt_iterator_bitmap_next(PyHamtIteratorState *iter,
2165 : : PyObject **key, PyObject **val)
2166 : : {
2167 : 318488 : int8_t level = iter->i_level;
2168 : :
2169 : 318488 : PyHamtNode_Bitmap *node = (PyHamtNode_Bitmap *)(iter->i_nodes[level]);
2170 : 318488 : Py_ssize_t pos = iter->i_pos[level];
2171 : :
2172 [ + + ]: 318488 : if (pos + 1 >= Py_SIZE(node)) {
2173 : : #ifdef Py_DEBUG
2174 : : assert(iter->i_level >= 0);
2175 : : iter->i_nodes[iter->i_level] = NULL;
2176 : : #endif
2177 : 71487 : iter->i_level--;
2178 : 71487 : return hamt_iterator_next(iter, key, val);
2179 : : }
2180 : :
2181 [ + + ]: 247001 : if (node->b_array[pos] == NULL) {
2182 : 15967 : iter->i_pos[level] = pos + 2;
2183 : :
2184 : 15967 : int8_t next_level = level + 1;
2185 : : assert(next_level < _Py_HAMT_MAX_TREE_DEPTH);
2186 : 15967 : iter->i_level = next_level;
2187 : 15967 : iter->i_pos[next_level] = 0;
2188 : 15967 : iter->i_nodes[next_level] = (PyHamtNode *)
2189 : 15967 : node->b_array[pos + 1];
2190 : :
2191 : 15967 : 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 : : assert(iter->i_level >= 0);
2212 : : 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 : 59292 : hamt_iterator_array_next(PyHamtIteratorState *iter,
2226 : : PyObject **key, PyObject **val)
2227 : : {
2228 : 59292 : int8_t level = iter->i_level;
2229 : :
2230 : 59292 : PyHamtNode_Array *node = (PyHamtNode_Array *)(iter->i_nodes[level]);
2231 : 59292 : Py_ssize_t pos = iter->i_pos[level];
2232 : :
2233 [ + + ]: 59292 : if (pos >= HAMT_ARRAY_NODE_SIZE) {
2234 : : #ifdef Py_DEBUG
2235 : : assert(iter->i_level >= 0);
2236 : : iter->i_nodes[iter->i_level] = NULL;
2237 : : #endif
2238 : 1846 : iter->i_level--;
2239 : 1846 : return hamt_iterator_next(iter, key, val);
2240 : : }
2241 : :
2242 [ + + ]: 62042 : for (Py_ssize_t i = pos; i < HAMT_ARRAY_NODE_SIZE; i++) {
2243 [ + + ]: 61952 : if (node->a_array[i] != NULL) {
2244 : 57356 : iter->i_pos[level] = i + 1;
2245 : :
2246 : 57356 : int8_t next_level = level + 1;
2247 : : assert(next_level < _Py_HAMT_MAX_TREE_DEPTH);
2248 : 57356 : iter->i_pos[next_level] = 0;
2249 : 57356 : iter->i_nodes[next_level] = node->a_array[i];
2250 : 57356 : iter->i_level = next_level;
2251 : :
2252 : 57356 : return hamt_iterator_next(iter, key, val);
2253 : : }
2254 : : }
2255 : :
2256 : : #ifdef Py_DEBUG
2257 : : assert(iter->i_level >= 0);
2258 : : iter->i_nodes[iter->i_level] = NULL;
2259 : : #endif
2260 : :
2261 : 90 : iter->i_level--;
2262 : 90 : return hamt_iterator_next(iter, key, val);
2263 : : }
2264 : :
2265 : : static hamt_iter_t
2266 : 377931 : hamt_iterator_next(PyHamtIteratorState *iter, PyObject **key, PyObject **val)
2267 : : {
2268 [ + + ]: 377931 : if (iter->i_level < 0) {
2269 : 118 : return I_END;
2270 : : }
2271 : :
2272 : : assert(iter->i_level < _Py_HAMT_MAX_TREE_DEPTH);
2273 : :
2274 : 377813 : PyHamtNode *current = iter->i_nodes[iter->i_level];
2275 : :
2276 [ + + ]: 377813 : if (IS_BITMAP_NODE(current)) {
2277 : 318488 : return hamt_iterator_bitmap_next(iter, key, val);
2278 : : }
2279 [ + + ]: 59325 : else if (IS_ARRAY_NODE(current)) {
2280 : 59292 : return hamt_iterator_array_next(iter, key, val);
2281 : : }
2282 : : else {
2283 : : 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 : : 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 : : 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 : 92540 : hamt_find(PyHamtObject *o, PyObject *key, PyObject **val)
2374 : : {
2375 [ + + ]: 92540 : if (o->h_count == 0) {
2376 : 87 : return F_NOT_FOUND;
2377 : : }
2378 : :
2379 : 92453 : int32_t key_hash = hamt_hash(key);
2380 [ + + ]: 92453 : if (key_hash == -1) {
2381 : 2 : return F_ERROR;
2382 : : }
2383 : :
2384 : 92451 : return hamt_node_find(o->h_root, 0, key_hash, key, val);
2385 : : }
2386 : :
2387 : :
2388 : : int
2389 : 17066 : _PyHamt_Find(PyHamtObject *o, PyObject *key, PyObject **val)
2390 : : {
2391 : 17066 : hamt_find_t res = hamt_find(o, key, val);
2392 [ + + + - ]: 17066 : switch (res) {
2393 : 2 : case F_ERROR:
2394 : 2 : return -1;
2395 : 1000 : case F_NOT_FOUND:
2396 : 1000 : return 0;
2397 : 16064 : case F_FOUND:
2398 : 16064 : 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 : 324 : _PyHamt_New(void)
2475 : : {
2476 [ + + ]: 324 : if (_empty_hamt != NULL) {
2477 : : /* HAMT is an immutable object so we can easily cache an
2478 : : empty instance. */
2479 : 289 : Py_INCREF(_empty_hamt);
2480 : 289 : 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 : : hamt_dump(PyHamtObject *self)
2507 : : {
2508 : : _PyUnicodeWriter writer;
2509 : :
2510 : : _PyUnicodeWriter_Init(&writer);
2511 : :
2512 : : if (_hamt_dump_format(&writer, "HAMT(len=%zd):\n", self->h_count)) {
2513 : : goto error;
2514 : : }
2515 : :
2516 : : if (hamt_node_dump(self->h_root, &writer, 0)) {
2517 : : goto error;
2518 : : }
2519 : :
2520 : : return _PyUnicodeWriter_Finish(&writer);
2521 : :
2522 : : error:
2523 : : _PyUnicodeWriter_Dealloc(&writer);
2524 : : 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 : 11940 : hamt_tp_traverse(PyHamtObject *self, visitproc visit, void *arg)
2712 : : {
2713 [ + - - + ]: 11940 : Py_VISIT(self->h_root);
2714 : 11940 : 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 : : hamt_py_dump(PyHamtObject *self, PyObject *Py_UNUSED(args))
2860 : : {
2861 : : 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 : 3125 : _PyHamt_Fini(PyInterpreterState *interp)
2962 : : {
2963 [ + + ]: 3125 : Py_CLEAR(_empty_hamt);
2964 [ + + ]: 3125 : Py_CLEAR(_empty_bitmap_node);
2965 : 3125 : }
|