Coverage Report

Created: 2022-07-08 09:39

/home/mdboom/Work/builds/cpython/Python/hamt.c
Line
Count
Source (jump to first uncovered line)
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
hamt_hash(PyObject *o)
401
{
402
    Py_hash_t hash = PyObject_Hash(o);
403
404
#if SIZEOF_PY_HASH_T <= 4
405
    return hash;
406
#else
407
    if (hash == -1) {
  Branch (407:9): [True: 440, False: 170k]
408
        /* exception */
409
        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
    int32_t xored = (int32_t)(hash & 0xffffffffl) ^ (int32_t)(hash >> 32);
430
    return xored == -1 ? 
-20
: xored;
  Branch (430:12): [True: 0, False: 170k]
431
#endif
432
}
433
434
static inline uint32_t
435
hamt_mask(int32_t hash, uint32_t shift)
436
{
437
    return (((uint32_t)hash >> shift) & 0x01f);
438
}
439
440
static inline uint32_t
441
hamt_bitpos(int32_t hash, uint32_t shift)
442
{
443
    return (uint32_t)1 << hamt_mask(hash, shift);
444
}
445
446
static inline uint32_t
447
hamt_bitindex(uint32_t bitmap, uint32_t bit)
448
{
449
    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
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
    if (size == 0 && 
_empty_bitmap_node != NULL4.35k
) {
  Branch (527:9): [True: 4.35k, False: 70.0k]
  Branch (527:22): [True: 4.35k, False: 4]
528
        Py_INCREF(_empty_bitmap_node);
529
        return (PyHamtNode *)_empty_bitmap_node;
530
    }
531
532
    /* No freelist; allocate a new bitmap node */
533
    node = PyObject_GC_NewVar(
534
        PyHamtNode_Bitmap, &_PyHamt_BitmapNode_Type, size);
535
    if (node == NULL) {
  Branch (535:9): [True: 0, False: 70.0k]
536
        return NULL;
537
    }
538
539
    Py_SET_SIZE(node, size);
540
541
    for (i = 0; i < size; 
i++515k
) {
  Branch (541:17): [True: 515k, False: 70.0k]
542
        node->b_array[i] = NULL;
543
    }
544
545
    node->b_bitmap = 0;
546
547
    _PyObject_GC_TRACK(node);
548
549
    if (size == 0 && 
_empty_bitmap_node == NULL4
) {
  Branch (549:9): [True: 4, False: 70.0k]
  Branch (549:22): [True: 4, False: 0]
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
        _empty_bitmap_node = node;
554
        Py_INCREF(_empty_bitmap_node);
555
    }
556
557
    return (PyHamtNode *)node;
558
}
559
560
static inline Py_ssize_t
561
hamt_node_bitmap_count(PyHamtNode_Bitmap *node)
562
{
563
    return Py_SIZE(node) / 2;
564
}
565
566
static PyHamtNode_Bitmap *
567
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
    clone = (PyHamtNode_Bitmap *)hamt_node_bitmap_new(Py_SIZE(node));
575
    if (clone == NULL) {
  Branch (575:9): [True: 0, False: 15.8k]
576
        return NULL;
577
    }
578
579
    
for (i = 0; 15.8k
i < Py_SIZE(node);
i++124k
) {
  Branch (579:17): [True: 124k, False: 15.8k]
580
        Py_XINCREF(node->b_array[i]);
581
        clone->b_array[i] = node->b_array[i];
582
    }
583
584
    clone->b_bitmap = node->b_bitmap;
585
    return clone;
586
}
587
588
static PyHamtNode_Bitmap *
589
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
    PyHamtNode_Bitmap *new = (PyHamtNode_Bitmap *)hamt_node_bitmap_new(
595
        Py_SIZE(o) - 2);
596
    if (new == NULL) {
  Branch (596:9): [True: 0, False: 28.2k]
597
        return NULL;
598
    }
599
600
    uint32_t idx = hamt_bitindex(o->b_bitmap, bit);
601
    uint32_t key_idx = 2 * idx;
602
    uint32_t val_idx = key_idx + 1;
603
    uint32_t i;
604
605
    for (i = 0; i < key_idx; 
i++95.8k
) {
  Branch (605:17): [True: 95.8k, False: 28.2k]
606
        Py_XINCREF(o->b_array[i]);
607
        new->b_array[i] = o->b_array[i];
608
    }
609
610
    assert(Py_SIZE(o) >= 0 && Py_SIZE(o) <= 32);
611
    for (i = val_idx + 1; i < (uint32_t)Py_SIZE(o); 
i++95.3k
) {
  Branch (611:27): [True: 95.3k, False: 28.2k]
612
        Py_XINCREF(o->b_array[i]);
613
        new->b_array[i - 2] = o->b_array[i];
614
    }
615
616
    new->b_bitmap = o->b_bitmap & ~bit;
617
    return new;
618
}
619
620
static PyHamtNode *
621
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
    int32_t key1_hash = hamt_hash(key1);
635
    if (key1_hash == -1) {
  Branch (635:9): [True: 0, False: 2.78k]
636
        return NULL;
637
    }
638
639
    if (key1_hash == key2_hash) {
  Branch (639:9): [True: 9, False: 2.77k]
640
        PyHamtNode_Collision *n;
641
        n = (PyHamtNode_Collision *)hamt_node_collision_new(key1_hash, 4);
642
        if (n == NULL) {
  Branch (642:13): [True: 0, False: 9]
643
            return NULL;
644
        }
645
646
        Py_INCREF(key1);
647
        n->c_array[0] = key1;
648
        Py_INCREF(val1);
649
        n->c_array[1] = val1;
650
651
        Py_INCREF(key2);
652
        n->c_array[2] = key2;
653
        Py_INCREF(val2);
654
        n->c_array[3] = val2;
655
656
        return (PyHamtNode *)n;
657
    }
658
    else {
659
        int added_leaf = 0;
660
        PyHamtNode *n = hamt_node_bitmap_new(0);
661
        if (n == NULL) {
  Branch (661:13): [True: 0, False: 2.77k]
662
            return NULL;
663
        }
664
665
        PyHamtNode *n2 = hamt_node_assoc(
666
            n, shift, key1_hash, key1, val1, &added_leaf);
667
        Py_DECREF(n);
668
        if (n2 == NULL) {
  Branch (668:13): [True: 0, False: 2.77k]
669
            return NULL;
670
        }
671
672
        n = hamt_node_assoc(n2, shift, key2_hash, key2, val2, &added_leaf);
673
        Py_DECREF(n2);
674
        if (n == NULL) {
  Branch (674:13): [True: 0, False: 2.77k]
675
            return NULL;
676
        }
677
678
        return n;
679
    }
680
}
681
682
static PyHamtNode *
683
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
    uint32_t bit = hamt_bitpos(hash, shift);
697
    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
    if (self->b_bitmap & bit) {
  Branch (713:9): [True: 11.8k, False: 25.8k]
714
        /* The key is set in this node */
715
716
        uint32_t key_idx = 2 * idx;
717
        uint32_t val_idx = key_idx + 1;
718
719
        assert(val_idx < (size_t)Py_SIZE(self));
720
721
        PyObject *key_or_null = self->b_array[key_idx];
722
        PyObject *val_or_node = self->b_array[val_idx];
723
724
        if (key_or_null == NULL) {
  Branch (724:13): [True: 1.27k, False: 10.5k]
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
            PyHamtNode *sub_node = hamt_node_assoc(
731
                (PyHamtNode *)val_or_node,
732
                shift + 5, hash, key, val, added_leaf);
733
            if (sub_node == NULL) {
  Branch (733:17): [True: 0, False: 1.27k]
734
                return NULL;
735
            }
736
737
            if (val_or_node == (PyObject *)sub_node) {
  Branch (737:17): [True: 0, False: 1.27k]
738
                Py_DECREF(sub_node);
739
                Py_INCREF(self);
740
                return (PyHamtNode *)self;
741
            }
742
743
            PyHamtNode_Bitmap *ret = hamt_node_bitmap_clone(self);
744
            if (ret == NULL) {
  Branch (744:17): [True: 0, False: 1.27k]
745
                return NULL;
746
            }
747
            Py_SETREF(ret->b_array[val_idx], (PyObject*)sub_node);
748
            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
        int comp_err = PyObject_RichCompareBool(key, key_or_null, Py_EQ);
756
        if (comp_err < 0) {  /* exception in __eq__ */
  Branch (756:13): [True: 0, False: 10.5k]
757
            return NULL;
758
        }
759
        if (comp_err == 1) {  /* key == key_or_null */
  Branch (759:13): [True: 7.81k, False: 2.78k]
760
            if (val == val_or_node) {
  Branch (760:17): [True: 2, False: 7.81k]
761
                /* we already have the same key/val pair; return self. */
762
                Py_INCREF(self);
763
                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
            PyHamtNode_Bitmap *ret = hamt_node_bitmap_clone(self);
769
            if (ret == NULL) {
  Branch (769:17): [True: 0, False: 7.81k]
770
                return NULL;
771
            }
772
            Py_INCREF(val);
773
            Py_SETREF(ret->b_array[val_idx], val);
774
            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
        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
        if (sub_node == NULL) {
  Branch (791:13): [True: 0, False: 2.78k]
792
            return NULL;
793
        }
794
795
        PyHamtNode_Bitmap *ret = hamt_node_bitmap_clone(self);
796
        if (ret == NULL) {
  Branch (796:13): [True: 0, False: 2.78k]
797
            Py_DECREF(sub_node);
798
            return NULL;
799
        }
800
        Py_SETREF(ret->b_array[key_idx], NULL);
801
        Py_SETREF(ret->b_array[val_idx], (PyObject *)sub_node);
802
803
        *added_leaf = 1;
804
        return (PyHamtNode *)ret;
805
    }
806
    else {
807
        /* There was no key before with the same (shift,hash). */
808
809
        uint32_t n = (uint32_t)_Py_popcount32(self->b_bitmap);
810
811
        if (n >= 16) {
  Branch (811:13): [True: 100, False: 25.7k]
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
            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
            PyHamtNode *empty = NULL;
832
            PyHamtNode_Array *new_node = NULL;
833
            PyHamtNode *res = NULL;
834
835
            /* Create a new Array node. */
836
            new_node = (PyHamtNode_Array *)hamt_node_array_new(n + 1);
837
            if (new_node == NULL) {
  Branch (837:17): [True: 0, False: 100]
838
                goto fin;
839
            }
840
841
            /* Create an empty bitmap node for the next
842
               hamt_node_assoc call. */
843
            empty = hamt_node_bitmap_new(0);
844
            if (empty == NULL) {
  Branch (844:17): [True: 0, False: 100]
845
                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
            new_node->a_array[jdx] = hamt_node_assoc(
851
                empty, shift + 5, hash, key, val, added_leaf);
852
            if (new_node->a_array[jdx] == NULL) {
  Branch (852:17): [True: 0, False: 100]
853
                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
            for (i = 0, j = 0; i < HAMT_ARRAY_NODE_SIZE; 
i++3.20k
) {
  Branch (859:32): [True: 3.20k, False: 100]
860
                if (((self->b_bitmap >> i) & 1) != 0) {
  Branch (860:21): [True: 1.60k, False: 1.60k]
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
                    if (self->b_array[j] == NULL) {
  Branch (866:25): [True: 543, False: 1.05k]
867
                        new_node->a_array[i] =
868
                            (PyHamtNode *)self->b_array[j + 1];
869
                        Py_INCREF(new_node->a_array[i]);
870
                    }
871
                    else {
872
                        int32_t rehash = hamt_hash(self->b_array[j]);
873
                        if (rehash == -1) {
  Branch (873:29): [True: 0, False: 1.05k]
874
                            goto fin;
875
                        }
876
877
                        new_node->a_array[i] = hamt_node_assoc(
878
                            empty, shift + 5,
879
                            rehash,
880
                            self->b_array[j],
881
                            self->b_array[j + 1],
882
                            added_leaf);
883
884
                        if (new_node->a_array[i] == NULL) {
  Branch (884:29): [True: 0, False: 1.05k]
885
                            goto fin;
886
                        }
887
                    }
888
                    j += 2;
889
                }
890
            }
891
892
            VALIDATE_ARRAY_NODE(new_node)
893
894
            /* That's it! */
895
            res = (PyHamtNode *)new_node;
896
897
        fin:
898
            Py_XDECREF(empty);
899
            if (res == NULL) {
  Branch (899:17): [True: 0, False: 100]
900
                Py_XDECREF(new_node);
901
            }
902
            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
            uint32_t key_idx = 2 * idx;
910
            uint32_t val_idx = key_idx + 1;
911
            uint32_t i;
912
913
            *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
                (PyHamtNode_Bitmap *)hamt_node_bitmap_new(2 * (n + 1));
919
            if (new_node == NULL) {
  Branch (919:17): [True: 0, False: 25.7k]
920
                return NULL;
921
            }
922
923
            /* Copy all keys/values that will be before the new key/value
924
               we are adding. */
925
            
for (i = 0; 25.7k
i < key_idx;
i++71.2k
) {
  Branch (925:25): [True: 71.2k, False: 25.7k]
926
                Py_XINCREF(self->b_array[i]);
927
                new_node->b_array[i] = self->b_array[i];
928
            }
929
930
            /* Set the new key/value to the new Bitmap node. */
931
            Py_INCREF(key);
932
            new_node->b_array[key_idx] = key;
933
            Py_INCREF(val);
934
            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
            for (i = key_idx; i < (uint32_t)Py_SIZE(self); 
i++71.5k
) {
  Branch (939:31): [True: 71.5k, False: 25.7k]
940
                Py_XINCREF(self->b_array[i]);
941
                new_node->b_array[i + 2] = self->b_array[i];
942
            }
943
944
            new_node->b_bitmap = self->b_bitmap | bit;
945
            return (PyHamtNode *)new_node;
946
        }
947
    }
948
}
949
950
static hamt_without_t
951
hamt_node_bitmap_without(PyHamtNode_Bitmap *self,
952
                         uint32_t shift, int32_t hash,
953
                         PyObject *key,
954
                         PyHamtNode **new_node)
955
{
956
    uint32_t bit = hamt_bitpos(hash, shift);
957
    if ((self->b_bitmap & bit) == 0) {
  Branch (957:9): [True: 9.14k, False: 38.8k]
958
        return W_NOT_FOUND;
959
    }
960
961
    uint32_t idx = hamt_bitindex(self->b_bitmap, bit);
962
963
    uint32_t key_idx = 2 * idx;
964
    uint32_t val_idx = key_idx + 1;
965
966
    PyObject *key_or_null = self->b_array[key_idx];
967
    PyObject *val_or_node = self->b_array[val_idx];
968
969
    if (key_or_null == NULL) {
  Branch (969:9): [True: 4.35k, False: 34.5k]
970
        /* key == NULL means that 'value' is another tree node. */
971
972
        PyHamtNode *sub_node = NULL;
973
974
        hamt_without_t res = hamt_node_without(
975
            (PyHamtNode *)val_or_node,
976
            shift + 5, hash, key, &sub_node);
977
978
        switch (res) {
979
            case W_EMPTY:
  Branch (979:13): [True: 0, False: 4.35k]
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
            case W_NEWNODE: {
  Branch (996:13): [True: 3.99k, False: 366]
997
                assert(sub_node != NULL);
998
999
                if (IS_BITMAP_NODE(sub_node)) {
1000
                    PyHamtNode_Bitmap *sub_tree = (PyHamtNode_Bitmap *)sub_node;
1001
                    if (hamt_node_bitmap_count(sub_tree) == 1 &&
  Branch (1001:25): [True: 3.58k, False: 411]
1002
                            
sub_tree->b_array[0] != NULL3.58k
)
  Branch (1002:29): [True: 3.56k, False: 12]
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
                        PyHamtNode_Bitmap *clone = hamt_node_bitmap_clone(self);
1014
                        if (clone == NULL) {
  Branch (1014:29): [True: 0, False: 3.56k]
1015
                            Py_DECREF(sub_node);
1016
                            return W_ERROR;
1017
                        }
1018
1019
                        PyObject *key = sub_tree->b_array[0];
1020
                        PyObject *val = sub_tree->b_array[1];
1021
1022
                        Py_INCREF(key);
1023
                        Py_XSETREF(clone->b_array[key_idx], key);
1024
                        Py_INCREF(val);
1025
                        Py_SETREF(clone->b_array[val_idx], val);
1026
1027
                        Py_DECREF(sub_tree);
1028
1029
                        *new_node = (PyHamtNode *)clone;
1030
                        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
                PyHamtNode_Bitmap *clone = hamt_node_bitmap_clone(self);
1045
                if (clone == NULL) {
  Branch (1045:21): [True: 0, False: 424]
1046
                    return W_ERROR;
1047
                }
1048
1049
                Py_SETREF(clone->b_array[val_idx],
1050
                          (PyObject *)sub_node);  /* borrow */
1051
1052
                *new_node = (PyHamtNode *)clone;
1053
                return W_NEWNODE;
1054
            }
1055
1056
            case W_ERROR:
  Branch (1056:13): [True: 277, False: 4.08k]
1057
            case W_NOT_FOUND:
  Branch (1057:13): [True: 89, False: 4.26k]
1058
                assert(sub_node == NULL);
1059
                return res;
1060
1061
            default:
  Branch (1061:13): [True: 0, False: 4.35k]
1062
                Py_UNREACHABLE();
1063
        }
1064
    }
1065
    else {
1066
        /* We have a regular key/value pair */
1067
1068
        int cmp = PyObject_RichCompareBool(key_or_null, key, Py_EQ);
1069
        if (cmp < 0) {
  Branch (1069:13): [True: 1.91k, False: 32.6k]
1070
            return W_ERROR;
1071
        }
1072
        if (cmp == 0) {
  Branch (1072:13): [True: 1.08k, False: 31.5k]
1073
            return W_NOT_FOUND;
1074
        }
1075
1076
        if (hamt_node_bitmap_count(self) == 1) {
  Branch (1076:13): [True: 3.28k, False: 28.2k]
1077
            return W_EMPTY;
1078
        }
1079
1080
        *new_node = (PyHamtNode *)
1081
            hamt_node_bitmap_clone_without(self, bit);
1082
        if (*new_node == NULL) {
  Branch (1082:13): [True: 0, False: 28.2k]
1083
            return W_ERROR;
1084
        }
1085
1086
        return W_NEWNODE;
1087
    }
1088
}
1089
1090
static hamt_find_t
1091
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
    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
    if ((self->b_bitmap & bit) == 0) {
  Branch (1105:9): [True: 34.9k, False: 62.4k]
1106
        return F_NOT_FOUND;
1107
    }
1108
1109
    idx = hamt_bitindex(self->b_bitmap, bit);
1110
    key_idx = idx * 2;
1111
    val_idx = key_idx + 1;
1112
1113
    assert(val_idx < (size_t)Py_SIZE(self));
1114
1115
    key_or_null = self->b_array[key_idx];
1116
    val_or_node = self->b_array[val_idx];
1117
1118
    if (key_or_null == NULL) {
  Branch (1118:9): [True: 8.25k, False: 54.1k]
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
        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
    comp_err = PyObject_RichCompareBool(key, key_or_null, Py_EQ);
1130
    if (comp_err < 0) {  /* exception in __eq__ */
  Branch (1130:9): [True: 1.91k, False: 52.2k]
1131
        return F_ERROR;
1132
    }
1133
    if (comp_err == 1) {  /* key == key_or_null */
  Branch (1133:9): [True: 47.5k, False: 4.64k]
1134
        *val = val_or_node;
1135
        return F_FOUND;
1136
    }
1137
1138
    return F_NOT_FOUND;
1139
}
1140
1141
static int
1142
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
    for (i = 
Py_SIZE112k
(self); --i >= 0; ) {
  Branch (1148:29): [True: 627k, False: 112k]
1149
        Py_VISIT(self->b_array[i]);
1150
    }
1151
1152
    return 0;
1153
}
1154
1155
static void
1156
hamt_node_bitmap_dealloc(PyHamtNode_Bitmap *self)
1157
{
1158
    /* Bitmap's tp_dealloc */
1159
1160
    Py_ssize_t len = Py_SIZE(self);
1161
    Py_ssize_t i;
1162
1163
    PyObject_GC_UnTrack(self);
1164
    Py_TRASHCAN_BEGIN(self, hamt_node_bitmap_dealloc)
1165
1166
    if (len > 0) {
  Branch (1166:9): [True: 70.0k, False: 4]
1167
        i = len;
1168
        while (--i >= 0) {
  Branch (1168:16): [True: 515k, False: 70.0k]
1169
            Py_XDECREF(self->b_array[i]);
1170
        }
1171
    }
1172
1173
    Py_TYPE(self)->tp_free((PyObject *)self);
1174
    Py_TRASHCAN_END
1175
}
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
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
    node = PyObject_GC_NewVar(
1267
        PyHamtNode_Collision, &_PyHamt_CollisionNode_Type, size);
1268
    if (node == NULL) {
  Branch (1268:9): [True: 0, False: 16]
1269
        return NULL;
1270
    }
1271
1272
    
for (i = 0; 16
i < size;
i++72
) {
  Branch (1272:17): [True: 72, False: 16]
1273
        node->c_array[i] = NULL;
1274
    }
1275
1276
    Py_SET_SIZE(node, size);
1277
    node->c_hash = hash;
1278
1279
    _PyObject_GC_TRACK(node);
1280
1281
    return (PyHamtNode *)node;
1282
}
1283
1284
static hamt_find_t
1285
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
    for (i = 0; i < Py_SIZE(self); 
i += 227
) {
  Branch (1294:17): [True: 52, False: 5]
1295
        el = self->c_array[i];
1296
1297
        assert(el != NULL);
1298
        int cmp = PyObject_RichCompareBool(key, el, Py_EQ);
1299
        if (cmp < 0) {
  Branch (1299:13): [True: 0, False: 52]
1300
            return F_ERROR;
1301
        }
1302
        if (cmp == 1) {
  Branch (1302:13): [True: 25, False: 27]
1303
            *idx = i;
1304
            return F_FOUND;
1305
        }
1306
    }
1307
1308
    return F_NOT_FOUND;
1309
}
1310
1311
static PyHamtNode *
1312
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
    if (hash == self->c_hash) {
  Branch (1319:9): [True: 6, False: 6]
1320
        /* The hash of the 'key' we are adding matches the hash of
1321
           other keys in this Collision node. */
1322
1323
        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
        found = hamt_node_collision_find_index(self, key, &key_idx);
1330
        switch (found) {
1331
            case F_ERROR:
  Branch (1331:13): [True: 0, False: 6]
1332
                /* Exception. */
1333
                return NULL;
1334
1335
            case F_NOT_FOUND:
  Branch (1335:13): [True: 4, False: 2]
1336
                /* This is a totally new key.  Clone the current node,
1337
                   add a new key/value to the cloned node. */
1338
1339
                new_node = (PyHamtNode_Collision *)hamt_node_collision_new(
1340
                    self->c_hash, Py_SIZE(self) + 2);
1341
                if (new_node == NULL) {
  Branch (1341:21): [True: 0, False: 4]
1342
                    return NULL;
1343
                }
1344
1345
                
for (i = 0; 4
i < Py_SIZE(self);
i++16
) {
  Branch (1345:29): [True: 16, False: 4]
1346
                    Py_INCREF(self->c_array[i]);
1347
                    new_node->c_array[i] = self->c_array[i];
1348
                }
1349
1350
                Py_INCREF(key);
1351
                new_node->c_array[i] = key;
1352
                Py_INCREF(val);
1353
                new_node->c_array[i + 1] = val;
1354
1355
                *added_leaf = 1;
1356
                return (PyHamtNode *)new_node;
1357
1358
            case F_FOUND:
  Branch (1358:13): [True: 2, False: 4]
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
                Py_ssize_t val_idx = key_idx + 1;
1364
1365
                if (self->c_array[val_idx] == val) {
  Branch (1365:21): [True: 0, False: 2]
1366
                    /* We're setting a key/value pair that's already set. */
1367
                    Py_INCREF(self);
1368
                    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
                new_node = (PyHamtNode_Collision *)hamt_node_collision_new(
1374
                    self->c_hash, Py_SIZE(self));
1375
                if (new_node == NULL) {
  Branch (1375:21): [True: 0, False: 2]
1376
                    return NULL;
1377
                }
1378
1379
                /* Copy all elements of the old node to the new one. */
1380
                
for (i = 0; 2
i < Py_SIZE(self);
i++8
) {
  Branch (1380:29): [True: 8, False: 2]
1381
                    Py_INCREF(self->c_array[i]);
1382
                    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
                Py_DECREF(new_node->c_array[val_idx]);
1387
                Py_INCREF(val);
1388
                new_node->c_array[val_idx] = val;
1389
1390
                return (PyHamtNode *)new_node;
1391
1392
            default:
  Branch (1392:13): [True: 0, False: 6]
1393
                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
        new_node = (PyHamtNode_Bitmap *)hamt_node_bitmap_new(2);
1409
        if (new_node == NULL) {
  Branch (1409:13): [True: 0, False: 6]
1410
            return NULL;
1411
        }
1412
        new_node->b_bitmap = hamt_bitpos(self->c_hash, shift);
1413
        Py_INCREF(self);
1414
        new_node->b_array[1] = (PyObject*) self;
1415
1416
        assoc_res = hamt_node_bitmap_assoc(
1417
            new_node, shift, hash, key, val, added_leaf);
1418
        Py_DECREF(new_node);
1419
        return assoc_res;
1420
    }
1421
}
1422
1423
static inline Py_ssize_t
1424
hamt_node_collision_count(PyHamtNode_Collision *node)
1425
{
1426
    return Py_SIZE(node) / 2;
1427
}
1428
1429
static hamt_without_t
1430
hamt_node_collision_without(PyHamtNode_Collision *self,
1431
                            uint32_t shift, int32_t hash,
1432
                            PyObject *key,
1433
                            PyHamtNode **new_node)
1434
{
1435
    if (hash != self->c_hash) {
  Branch (1435:9): [True: 0, False: 4]
1436
        return W_NOT_FOUND;
1437
    }
1438
1439
    Py_ssize_t key_idx = -1;
1440
    hamt_find_t found = hamt_node_collision_find_index(self, key, &key_idx);
1441
1442
    switch (found) {
1443
        case F_ERROR:
  Branch (1443:9): [True: 0, False: 4]
1444
            return W_ERROR;
1445
1446
        case F_NOT_FOUND:
  Branch (1446:9): [True: 0, False: 4]
1447
            return W_NOT_FOUND;
1448
1449
        case F_FOUND:
  Branch (1449:9): [True: 4, False: 0]
1450
            assert(key_idx >= 0);
1451
            assert(key_idx < Py_SIZE(self));
1452
1453
            Py_ssize_t new_count = hamt_node_collision_count(self) - 1;
1454
1455
            if (new_count == 0) {
  Branch (1455:17): [True: 0, False: 4]
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
                return W_EMPTY;
1461
            }
1462
1463
            if (new_count == 1) {
  Branch (1463:17): [True: 3, False: 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
                    hamt_node_bitmap_new(2);
1471
                if (node == NULL) {
  Branch (1471:21): [True: 0, False: 3]
1472
                    return W_ERROR;
1473
                }
1474
1475
                if (key_idx == 0) {
  Branch (1475:21): [True: 0, False: 3]
1476
                    Py_INCREF(self->c_array[2]);
1477
                    node->b_array[0] = self->c_array[2];
1478
                    Py_INCREF(self->c_array[3]);
1479
                    node->b_array[1] = self->c_array[3];
1480
                }
1481
                else {
1482
                    assert(key_idx == 2);
1483
                    Py_INCREF(self->c_array[0]);
1484
                    node->b_array[0] = self->c_array[0];
1485
                    Py_INCREF(self->c_array[1]);
1486
                    node->b_array[1] = self->c_array[1];
1487
                }
1488
1489
                node->b_bitmap = hamt_bitpos(hash, shift);
1490
1491
                *new_node = (PyHamtNode *)node;
1492
                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
                hamt_node_collision_new(
1499
                    self->c_hash, Py_SIZE(self) - 2);
1500
            if (new == NULL) {
  Branch (1500:17): [True: 0, False: 1]
1501
                return W_ERROR;
1502
            }
1503
1504
            /* Copy all other keys from `self` to `new` */
1505
            Py_ssize_t i;
1506
            for (i = 0; i < key_idx; 
i++2
) {
  Branch (1506:25): [True: 2, False: 1]
1507
                Py_INCREF(self->c_array[i]);
1508
                new->c_array[i] = self->c_array[i];
1509
            }
1510
            for (i = key_idx + 2; i < Py_SIZE(self); 
i++2
) {
  Branch (1510:35): [True: 2, False: 1]
1511
                Py_INCREF(self->c_array[i]);
1512
                new->c_array[i - 2] = self->c_array[i];
1513
            }
1514
1515
            *new_node = (PyHamtNode*)new;
1516
            return W_NEWNODE;
1517
1518
        default:
  Branch (1518:9): [True: 0, False: 4]
1519
            Py_UNREACHABLE();
1520
    }
1521
}
1522
1523
static hamt_find_t
1524
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
    Py_ssize_t idx = -1;
1532
    hamt_find_t res;
1533
1534
    res = hamt_node_collision_find_index(self, key, &idx);
1535
    if (res == F_ERROR || res == F_NOT_FOUND) {
  Branch (1535:9): [True: 0, False: 20]
  Branch (1535:27): [True: 1, False: 19]
1536
        return res;
1537
    }
1538
1539
    assert(idx >= 0);
1540
    assert(idx + 1 < Py_SIZE(self));
1541
1542
    *val = self->c_array[idx + 1];
1543
    assert(*val != NULL);
1544
1545
    return F_FOUND;
1546
}
1547
1548
1549
static int
1550
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
    for (i = Py_SIZE(self); --i >= 0; ) {
  Branch (1557:29): [True: 0, False: 0]
1558
        Py_VISIT(self->c_array[i]);
1559
    }
1560
1561
    return 0;
1562
}
1563
1564
static void
1565
hamt_node_collision_dealloc(PyHamtNode_Collision *self)
1566
{
1567
    /* Collision's tp_dealloc */
1568
1569
    Py_ssize_t len = Py_SIZE(self);
1570
1571
    PyObject_GC_UnTrack(self);
1572
    Py_TRASHCAN_BEGIN(self, hamt_node_collision_dealloc)
1573
1574
    if (len > 0) {
  Branch (1574:9): [True: 16, False: 0]
1575
1576
        while (--len >= 0) {
  Branch (1576:16): [True: 72, False: 16]
1577
            Py_XDECREF(self->c_array[len]);
1578
        }
1579
    }
1580
1581
    Py_TYPE(self)->tp_free((PyObject *)self);
1582
    Py_TRASHCAN_END
1583
}
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
hamt_node_array_new(Py_ssize_t count)
1629
{
1630
    Py_ssize_t i;
1631
1632
    PyHamtNode_Array *node = PyObject_GC_New(
1633
        PyHamtNode_Array, &_PyHamt_ArrayNode_Type);
1634
    if (node == NULL) {
  Branch (1634:9): [True: 0, False: 98.4k]
1635
        return NULL;
1636
    }
1637
1638
    
for (i = 0; 98.4k
i < HAMT_ARRAY_NODE_SIZE;
i++3.14M
) {
  Branch (1638:17): [True: 3.14M, False: 98.4k]
1639
        node->a_array[i] = NULL;
1640
    }
1641
1642
    node->a_count = count;
1643
1644
    _PyObject_GC_TRACK(node);
1645
    return (PyHamtNode *)node;
1646
}
1647
1648
static PyHamtNode_Array *
1649
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
    clone = (PyHamtNode_Array *)hamt_node_array_new(node->a_count);
1658
    if (clone == NULL) {
  Branch (1658:9): [True: 0, False: 96.8k]
1659
        return NULL;
1660
    }
1661
1662
    /* Copy all elements from the current Array node to the new one. */
1663
    
for (i = 0; 96.8k
i < HAMT_ARRAY_NODE_SIZE;
i++3.09M
) {
  Branch (1663:17): [True: 3.09M, False: 96.8k]
1664
        Py_XINCREF(node->a_array[i]);
1665
        clone->a_array[i] = node->a_array[i];
1666
    }
1667
1668
    VALIDATE_ARRAY_NODE(clone)
1669
    return clone;
1670
}
1671
1672
static PyHamtNode *
1673
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
    uint32_t idx = hamt_mask(hash, shift);
1685
    PyHamtNode *node = self->a_array[idx];
1686
    PyHamtNode *child_node;
1687
    PyHamtNode_Array *new_node;
1688
    Py_ssize_t i;
1689
1690
    if (node == NULL) {
  Branch (1690:9): [True: 1.47k, False: 38.1k]
1691
        /* There's no child node for the given hash.  Create a new
1692
           Bitmap node for this key. */
1693
1694
        PyHamtNode_Bitmap *empty = NULL;
1695
1696
        /* Get an empty Bitmap node to work with. */
1697
        empty = (PyHamtNode_Bitmap *)hamt_node_bitmap_new(0);
1698
        if (empty == NULL) {
  Branch (1698:13): [True: 0, False: 1.47k]
1699
            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
        child_node = hamt_node_bitmap_assoc(
1705
            empty,
1706
            shift + 5, hash, key, val, added_leaf);
1707
        Py_DECREF(empty);
1708
        if (child_node == NULL) {
  Branch (1708:13): [True: 0, False: 1.47k]
1709
            return NULL;
1710
        }
1711
1712
        /* Create a new Array node. */
1713
        new_node = (PyHamtNode_Array *)hamt_node_array_new(self->a_count + 1);
1714
        if (new_node == NULL) {
  Branch (1714:13): [True: 0, False: 1.47k]
1715
            Py_DECREF(child_node);
1716
            return NULL;
1717
        }
1718
1719
        /* Copy all elements from the current Array node to the
1720
           new one. */
1721
        
for (i = 0; 1.47k
i < HAMT_ARRAY_NODE_SIZE;
i++47.3k
) {
  Branch (1721:21): [True: 47.3k, False: 1.47k]
1722
            Py_XINCREF(self->a_array[i]);
1723
            new_node->a_array[i] = self->a_array[i];
1724
        }
1725
1726
        assert(new_node->a_array[idx] == NULL);
1727
        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
        child_node = hamt_node_assoc(
1734
            node, shift + 5, hash, key, val, added_leaf);
1735
        if (child_node == NULL) {
  Branch (1735:13): [True: 0, False: 38.1k]
1736
            return NULL;
1737
        }
1738
        else if (child_node == (PyHamtNode *)self) {
  Branch (1738:18): [True: 0, False: 38.1k]
1739
            Py_DECREF(child_node);
1740
            return (PyHamtNode *)self;
1741
        }
1742
1743
        new_node = hamt_node_array_clone(self);
1744
        if (new_node == NULL) {
  Branch (1744:13): [True: 0, False: 38.1k]
1745
            Py_DECREF(child_node);
1746
            return NULL;
1747
        }
1748
1749
        Py_SETREF(new_node->a_array[idx], child_node);  /* borrow */
1750
        VALIDATE_ARRAY_NODE(new_node)
1751
    }
1752
1753
    return (PyHamtNode *)new_node;
1754
}
1755
1756
static hamt_without_t
1757
hamt_node_array_without(PyHamtNode_Array *self,
1758
                        uint32_t shift, int32_t hash,
1759
                        PyObject *key,
1760
                        PyHamtNode **new_node)
1761
{
1762
    uint32_t idx = hamt_mask(hash, shift);
1763
    PyHamtNode *node = self->a_array[idx];
1764
1765
    if (node == NULL) {
  Branch (1765:9): [True: 280, False: 83.2k]
1766
        return W_NOT_FOUND;
1767
    }
1768
1769
    PyHamtNode *sub_node = NULL;
1770
    hamt_without_t res = hamt_node_without(
1771
        (PyHamtNode *)node,
1772
        shift + 5, hash, key, &sub_node);
1773
1774
    switch (res) {
1775
        case W_NOT_FOUND:
  Branch (1775:9): [True: 20.7k, False: 62.5k]
1776
        case W_ERROR:
  Branch (1776:9): [True: 3.63k, False: 79.6k]
1777
            assert(sub_node == NULL);
1778
            return res;
1779
1780
        case W_NEWNODE: {
  Branch (1780:9): [True: 55.6k, False: 27.6k]
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
            PyHamtNode_Array *clone = hamt_node_array_clone(self);
1787
            if (clone == NULL) {
  Branch (1787:17): [True: 0, False: 55.6k]
1788
                Py_DECREF(sub_node);
1789
                return W_ERROR;
1790
            }
1791
1792
            Py_SETREF(clone->a_array[idx], sub_node);  /* borrow */
1793
            *new_node = (PyHamtNode*)clone;  /* borrow */
1794
            return W_NEWNODE;
1795
        }
1796
1797
        case W_EMPTY: {
  Branch (1797:9): [True: 3.27k, False: 79.9k]
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
            Py_ssize_t new_count = self->a_count - 1;
1803
1804
            if (new_count == 0) {
  Branch (1804:17): [True: 0, False: 3.27k]
1805
                return W_EMPTY;
1806
            }
1807
1808
            if (new_count >= 16) {
  Branch (1808:17): [True: 3.07k, False: 199]
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
                PyHamtNode_Array *new = hamt_node_array_clone(self);
1817
                if (new == NULL) {
  Branch (1817:21): [True: 0, False: 3.07k]
1818
                    return W_ERROR;
1819
                }
1820
                new->a_count = new_count;
1821
                Py_CLEAR(new->a_array[idx]);
1822
1823
                *new_node = (PyHamtNode*)new;  /* borrow */
1824
                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
            Py_ssize_t bitmap_size = new_count * 2;
1831
            uint32_t bitmap = 0;
1832
1833
            PyHamtNode_Bitmap *new = (PyHamtNode_Bitmap *)
1834
                hamt_node_bitmap_new(bitmap_size);
1835
            if (new == NULL) {
  Branch (1835:17): [True: 0, False: 199]
1836
                return W_ERROR;
1837
            }
1838
1839
            Py_ssize_t new_i = 0;
1840
            for (uint32_t i = 0; i < HAMT_ARRAY_NODE_SIZE; 
i++6.36k
) {
  Branch (1840:34): [True: 6.36k, False: 199]
1841
                if (i == idx) {
  Branch (1841:21): [True: 199, False: 6.16k]
1842
                    /* Skip the node we are deleting. */
1843
                    continue;
1844
                }
1845
1846
                PyHamtNode *node = self->a_array[i];
1847
                if (node == NULL) {
  Branch (1847:21): [True: 3.18k, False: 2.98k]
1848
                    /* Skip any missing nodes. */
1849
                    continue;
1850
                }
1851
1852
                bitmap |= 1U << i;
1853
1854
                if (IS_BITMAP_NODE(node)) {
1855
                    PyHamtNode_Bitmap *child = (PyHamtNode_Bitmap *)node;
1856
1857
                    if (hamt_node_bitmap_count(child) == 1 &&
  Branch (1857:25): [True: 2.09k, False: 886]
1858
                            
child->b_array[0] != NULL2.09k
)
  Branch (1858:29): [True: 2.07k, False: 20]
1859
                    {
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
                        PyObject *key = child->b_array[0];
1869
                        PyObject *val = child->b_array[1];
1870
1871
                        Py_INCREF(key);
1872
                        new->b_array[new_i] = key;
1873
                        Py_INCREF(val);
1874
                        new->b_array[new_i + 1] = val;
1875
                    }
1876
                    else {
1877
                        new->b_array[new_i] = NULL;
1878
                        Py_INCREF(node);
1879
                        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
                    new->b_array[new_i] = NULL;
1897
                    Py_INCREF(node);
1898
                    new->b_array[new_i + 1] = (PyObject*)node;
1899
                }
1900
1901
                new_i += 2;
1902
            }
1903
1904
            new->b_bitmap = bitmap;
1905
            *new_node = (PyHamtNode*)new;  /* borrow */
1906
            return W_NEWNODE;
1907
        }
1908
1909
        default:
  Branch (1909:9): [True: 0, False: 83.2k]
1910
            Py_UNREACHABLE();
1911
    }
1912
}
1913
1914
static hamt_find_t
1915
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
    uint32_t idx = hamt_mask(hash, shift);
1923
    PyHamtNode *node;
1924
1925
    node = self->a_array[idx];
1926
    if (node == NULL) {
  Branch (1926:9): [True: 3.35k, False: 142k]
1927
        return F_NOT_FOUND;
1928
    }
1929
1930
    /* Dispatch to the generic hamt_node_find */
1931
    return hamt_node_find(node, shift + 5, hash, key, val);
1932
}
1933
1934
static int
1935
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
    for (i = 0; i < HAMT_ARRAY_NODE_SIZE; 
i++158k
) {
  Branch (1942:17): [True: 158k, False: 4.95k]
1943
        Py_VISIT(self->a_array[i]);
1944
    }
1945
1946
    return 0;
1947
}
1948
1949
static void
1950
hamt_node_array_dealloc(PyHamtNode_Array *self)
1951
{
1952
    /* Array's tp_dealloc */
1953
1954
    Py_ssize_t i;
1955
1956
    PyObject_GC_UnTrack(self);
1957
    Py_TRASHCAN_BEGIN(self, hamt_node_array_dealloc)
1958
1959
    
for (i = 0; 98.4k
i < HAMT_ARRAY_NODE_SIZE;
i++3.14M
) {
  Branch (1959:17): [True: 3.14M, False: 98.4k]
1960
        Py_XDECREF(self->a_array[i]);
1961
    }
1962
1963
    Py_TYPE(self)->tp_free((PyObject *)self);
1964
    Py_TRASHCAN_END
1965
}
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
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
    if (IS_BITMAP_NODE(node)) {
2033
        return hamt_node_bitmap_assoc(
2034
            (PyHamtNode_Bitmap *)node,
2035
            shift, hash, key, val, added_leaf);
2036
    }
2037
    else if (IS_ARRAY_NODE(node)) {
2038
        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
        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
hamt_node_without(PyHamtNode *node,
2052
                  uint32_t shift, int32_t hash,
2053
                  PyObject *key,
2054
                  PyHamtNode **new_node)
2055
{
2056
    if (IS_BITMAP_NODE(node)) {
2057
        return hamt_node_bitmap_without(
2058
            (PyHamtNode_Bitmap *)node,
2059
            shift, hash, key,
2060
            new_node);
2061
    }
2062
    else if (IS_ARRAY_NODE(node)) {
2063
        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
        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
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
    if (IS_BITMAP_NODE(node)) {
2096
        return hamt_node_bitmap_find(
2097
            (PyHamtNode_Bitmap *)node,
2098
            shift, hash, key, val);
2099
2100
    }
2101
    else if (IS_ARRAY_NODE(node)) {
2102
        return hamt_node_array_find(
2103
            (PyHamtNode_Array *)node,
2104
            shift, hash, key, val);
2105
    }
2106
    else {
2107
        assert(IS_COLLISION_NODE(node));
2108
        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
hamt_iterator_init(PyHamtIteratorState *iter, PyHamtNode *root)
2151
{
2152
    for (uint32_t i = 0; i < _Py_HAMT_MAX_TREE_DEPTH; 
i++1.81k
) {
  Branch (2152:26): [True: 1.81k, False: 227]
2153
        iter->i_nodes[i] = NULL;
2154
        iter->i_pos[i] = 0;
2155
    }
2156
2157
    iter->i_level = 0;
2158
2159
    /* Note: we don't incref/decref nodes in i_nodes. */
2160
    iter->i_nodes[0] = root;
2161
}
2162
2163
static hamt_iter_t
2164
hamt_iterator_bitmap_next(PyHamtIteratorState *iter,
2165
                          PyObject **key, PyObject **val)
2166
{
2167
    int8_t level = iter->i_level;
2168
2169
    PyHamtNode_Bitmap *node = (PyHamtNode_Bitmap *)(iter->i_nodes[level]);
2170
    Py_ssize_t pos = iter->i_pos[level];
2171
2172
    if (pos + 1 >= Py_SIZE(node)) {
  Branch (2172:9): [True: 71.7k, False: 247k]
2173
#ifdef Py_DEBUG
2174
        assert(iter->i_level >= 0);
2175
        iter->i_nodes[iter->i_level] = NULL;
2176
#endif
2177
        iter->i_level--;
2178
        return hamt_iterator_next(iter, key, val);
2179
    }
2180
2181
    if (node->b_array[pos] == NULL) {
  Branch (2181:9): [True: 16.5k, False: 231k]
2182
        iter->i_pos[level] = pos + 2;
2183
2184
        int8_t next_level = level + 1;
2185
        assert(next_level < _Py_HAMT_MAX_TREE_DEPTH);
2186
        iter->i_level = next_level;
2187
        iter->i_pos[next_level] = 0;
2188
        iter->i_nodes[next_level] = (PyHamtNode *)
2189
            node->b_array[pos + 1];
2190
2191
        return hamt_iterator_next(iter, key, val);
2192
    }
2193
2194
    *key = node->b_array[pos];
2195
    *val = node->b_array[pos + 1];
2196
    iter->i_pos[level] = pos + 2;
2197
    return I_ITEM;
2198
}
2199
2200
static hamt_iter_t
2201
hamt_iterator_collision_next(PyHamtIteratorState *iter,
2202
                             PyObject **key, PyObject **val)
2203
{
2204
    int8_t level = iter->i_level;
2205
2206
    PyHamtNode_Collision *node = (PyHamtNode_Collision *)(iter->i_nodes[level]);
2207
    Py_ssize_t pos = iter->i_pos[level];
2208
2209
    if (pos + 1 >= Py_SIZE(node)) {
  Branch (2209:9): [True: 6, False: 27]
2210
#ifdef Py_DEBUG
2211
        assert(iter->i_level >= 0);
2212
        iter->i_nodes[iter->i_level] = NULL;
2213
#endif
2214
        iter->i_level--;
2215
        return hamt_iterator_next(iter, key, val);
2216
    }
2217
2218
    *key = node->c_array[pos];
2219
    *val = node->c_array[pos + 1];
2220
    iter->i_pos[level] = pos + 2;
2221
    return I_ITEM;
2222
}
2223
2224
static hamt_iter_t
2225
hamt_iterator_array_next(PyHamtIteratorState *iter,
2226
                         PyObject **key, PyObject **val)
2227
{
2228
    int8_t level = iter->i_level;
2229
2230
    PyHamtNode_Array *node = (PyHamtNode_Array *)(iter->i_nodes[level]);
2231
    Py_ssize_t pos = iter->i_pos[level];
2232
2233
    if (pos >= HAMT_ARRAY_NODE_SIZE) {
  Branch (2233:9): [True: 1.80k, False: 57.1k]
2234
#ifdef Py_DEBUG
2235
        assert(iter->i_level >= 0);
2236
        iter->i_nodes[iter->i_level] = NULL;
2237
#endif
2238
        iter->i_level--;
2239
        return hamt_iterator_next(iter, key, val);
2240
    }
2241
2242
    
for (Py_ssize_t i = pos; 57.1k
i < HAMT_ARRAY_NODE_SIZE;
i++4.71k
) {
  Branch (2242:30): [True: 61.6k, False: 123]
2243
        if (node->a_array[i] != NULL) {
  Branch (2243:13): [True: 56.9k, False: 4.71k]
2244
            iter->i_pos[level] = i + 1;
2245
2246
            int8_t next_level = level + 1;
2247
            assert(next_level < _Py_HAMT_MAX_TREE_DEPTH);
2248
            iter->i_pos[next_level] = 0;
2249
            iter->i_nodes[next_level] = node->a_array[i];
2250
            iter->i_level = next_level;
2251
2252
            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
    iter->i_level--;
2262
    return hamt_iterator_next(iter, key, val);
2263
}
2264
2265
static hamt_iter_t
2266
hamt_iterator_next(PyHamtIteratorState *iter, PyObject **key, PyObject **val)
2267
{
2268
    if (iter->i_level < 0) {
  Branch (2268:9): [True: 118, False: 378k]
2269
        return I_END;
2270
    }
2271
2272
    assert(iter->i_level < _Py_HAMT_MAX_TREE_DEPTH);
2273
2274
    PyHamtNode *current = iter->i_nodes[iter->i_level];
2275
2276
    if (IS_BITMAP_NODE(current)) {
2277
        return hamt_iterator_bitmap_next(iter, key, val);
2278
    }
2279
    else if (IS_ARRAY_NODE(current)) {
2280
        return hamt_iterator_array_next(iter, key, val);
2281
    }
2282
    else {
2283
        assert(IS_COLLISION_NODE(current));
2284
        return hamt_iterator_collision_next(iter, key, val);
2285
    }
2286
}
2287
2288
2289
/////////////////////////////////// HAMT high-level functions
2290
2291
2292
PyHamtObject *
2293
_PyHamt_Assoc(PyHamtObject *o, PyObject *key, PyObject *val)
2294
{
2295
    int32_t key_hash;
2296
    int added_leaf = 0;
2297
    PyHamtNode *new_root;
2298
    PyHamtObject *new_o;
2299
2300
    key_hash = hamt_hash(key);
2301
    if (key_hash == -1) {
  Branch (2301:9): [True: 219, False: 29.7k]
2302
        return NULL;
2303
    }
2304
2305
    new_root = hamt_node_assoc(
2306
        (PyHamtNode *)(o->h_root),
2307
        0, key_hash, key, val, &added_leaf);
2308
    if (new_root == NULL) {
  Branch (2308:9): [True: 0, False: 29.7k]
2309
        return NULL;
2310
    }
2311
2312
    if (new_root == o->h_root) {
  Branch (2312:9): [True: 2, False: 29.7k]
2313
        Py_DECREF(new_root);
2314
        Py_INCREF(o);
2315
        return o;
2316
    }
2317
2318
    new_o = hamt_alloc();
2319
    if (new_o == NULL) {
  Branch (2319:9): [True: 0, False: 29.7k]
2320
        Py_DECREF(new_root);
2321
        return NULL;
2322
    }
2323
2324
    new_o->h_root = new_root;  /* borrow */
2325
    new_o->h_count = added_leaf ? 
o->h_count + 121.9k
:
o->h_count7.81k
;
  Branch (2325:22): [True: 21.9k, False: 7.81k]
2326
2327
    return new_o;
2328
}
2329
2330
PyHamtObject *
2331
_PyHamt_Without(PyHamtObject *o, PyObject *key)
2332
{
2333
    int32_t key_hash = hamt_hash(key);
2334
    if (key_hash == -1) {
  Branch (2334:9): [True: 219, False: 43.9k]
2335
        return NULL;
2336
    }
2337
2338
    PyHamtNode *new_root = NULL;
2339
2340
    hamt_without_t res = hamt_node_without(
2341
        (PyHamtNode *)(o->h_root),
2342
        0, key_hash, key,
2343
        &new_root);
2344
2345
    switch (res) {
2346
        case W_ERROR:
  Branch (2346:9): [True: 1.91k, False: 42.0k]
2347
            return NULL;
2348
        case W_EMPTY:
  Branch (2348:9): [True: 12, False: 43.9k]
2349
            return _PyHamt_New();
2350
        case W_NOT_FOUND:
  Branch (2350:9): [True: 10.5k, False: 33.4k]
2351
            Py_INCREF(o);
2352
            return o;
2353
        case W_NEWNODE: {
  Branch (2353:9): [True: 31.5k, False: 12.4k]
2354
            assert(new_root != NULL);
2355
2356
            PyHamtObject *new_o = hamt_alloc();
2357
            if (new_o == NULL) {
  Branch (2357:17): [True: 0, False: 31.5k]
2358
                Py_DECREF(new_root);
2359
                return NULL;
2360
            }
2361
2362
            new_o->h_root = new_root;  /* borrow */
2363
            new_o->h_count = o->h_count - 1;
2364
            assert(new_o->h_count >= 0);
2365
            return new_o;
2366
        }
2367
        default:
  Branch (2367:9): [True: 0, False: 43.9k]
2368
            Py_UNREACHABLE();
2369
    }
2370
}
2371
2372
static hamt_find_t
2373
hamt_find(PyHamtObject *o, PyObject *key, PyObject **val)
2374
{
2375
    if (o->h_count == 0) {
  Branch (2375:9): [True: 54, False: 92.4k]
2376
        return F_NOT_FOUND;
2377
    }
2378
2379
    int32_t key_hash = hamt_hash(key);
2380
    if (key_hash == -1) {
  Branch (2380:9): [True: 2, False: 92.4k]
2381
        return F_ERROR;
2382
    }
2383
2384
    return hamt_node_find(o->h_root, 0, key_hash, key, val);
2385
}
2386
2387
2388
int
2389
_PyHamt_Find(PyHamtObject *o, PyObject *key, PyObject **val)
2390
{
2391
    hamt_find_t res = hamt_find(o, key, val);
2392
    switch (res) {
2393
        case F_ERROR:
  Branch (2393:9): [True: 2, False: 17.0k]
2394
            return -1;
2395
        case F_NOT_FOUND:
  Branch (2395:9): [True: 975, False: 16.0k]
2396
            return 0;
2397
        case F_FOUND:
  Branch (2397:9): [True: 16.0k, False: 977]
2398
            return 1;
2399
        default:
  Branch (2399:9): [True: 0, False: 17.0k]
2400
            Py_UNREACHABLE();
2401
    }
2402
}
2403
2404
2405
int
2406
_PyHamt_Eq(PyHamtObject *v, PyHamtObject *w)
2407
{
2408
    if (v == w) {
  Branch (2408:9): [True: 1, False: 18]
2409
        return 1;
2410
    }
2411
2412
    if (v->h_count != w->h_count) {
  Branch (2412:9): [True: 8, False: 10]
2413
        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
    hamt_iterator_init(&iter, v->h_root);
2424
2425
    do {
2426
        iter_res = hamt_iterator_next(&iter, &v_key, &v_val);
2427
        if (iter_res == I_ITEM) {
  Branch (2427:13): [True: 28, False: 2]
2428
            find_res = hamt_find(w, v_key, &w_val);
2429
            switch (find_res) {
  Branch (2429:21): [True: 0, False: 28]
2430
                case F_ERROR:
  Branch (2430:17): [True: 2, False: 26]
2431
                    return -1;
2432
2433
                case F_NOT_FOUND:
  Branch (2433:17): [True: 4, False: 24]
2434
                    return 0;
2435
2436
                case F_FOUND: {
  Branch (2436:17): [True: 22, False: 6]
2437
                    int cmp = PyObject_RichCompareBool(v_val, w_val, Py_EQ);
2438
                    if (cmp < 0) {
  Branch (2438:25): [True: 0, False: 22]
2439
                        return -1;
2440
                    }
2441
                    if (cmp == 0) {
  Branch (2441:25): [True: 2, False: 20]
2442
                        return 0;
2443
                    }
2444
                }
2445
            }
2446
        }
2447
    } while (
iter_res != I_END22
);
  Branch (2447:14): [True: 20, False: 2]
2448
2449
    return 1;
2450
}
2451
2452
Py_ssize_t
2453
_PyHamt_Len(PyHamtObject *o)
2454
{
2455
    return o->h_count;
2456
}
2457
2458
static PyHamtObject *
2459
hamt_alloc(void)
2460
{
2461
    PyHamtObject *o;
2462
    o = PyObject_GC_New(PyHamtObject, &_PyHamt_Type);
2463
    if (o == NULL) {
  Branch (2463:9): [True: 0, False: 61.3k]
2464
        return NULL;
2465
    }
2466
    o->h_count = 0;
2467
    o->h_root = NULL;
2468
    o->h_weakreflist = NULL;
2469
    PyObject_GC_Track(o);
2470
    return o;
2471
}
2472
2473
PyHamtObject *
2474
_PyHamt_New(void)
2475
{
2476
    if (_empty_hamt != NULL) {
  Branch (2476:9): [True: 280, False: 4]
2477
        /* HAMT is an immutable object so we can easily cache an
2478
           empty instance. */
2479
        Py_INCREF(_empty_hamt);
2480
        return _empty_hamt;
2481
    }
2482
2483
    PyHamtObject *o = hamt_alloc();
2484
    if (o == NULL) {
  Branch (2484:9): [True: 0, False: 4]
2485
        return NULL;
2486
    }
2487
2488
    o->h_root = hamt_node_bitmap_new(0);
2489
    if (o->h_root == NULL) {
  Branch (2489:9): [True: 0, False: 4]
2490
        Py_DECREF(o);
2491
        return NULL;
2492
    }
2493
2494
    o->h_count = 0;
2495
2496
    if (_empty_hamt == NULL) {
  Branch (2496:9): [True: 4, False: 0]
2497
        Py_INCREF(o);
2498
        _empty_hamt = o;
2499
    }
2500
2501
    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
hamt_baseiter_tp_clear(PyHamtIterator *it)
2534
{
2535
    Py_CLEAR(it->hi_obj);
2536
    return 0;
2537
}
2538
2539
static void
2540
hamt_baseiter_tp_dealloc(PyHamtIterator *it)
2541
{
2542
    PyObject_GC_UnTrack(it);
2543
    (void)hamt_baseiter_tp_clear(it);
2544
    PyObject_GC_Del(it);
2545
}
2546
2547
static int
2548
hamt_baseiter_tp_traverse(PyHamtIterator *it, visitproc visit, void *arg)
2549
{
2550
    Py_VISIT(it->hi_obj);
2551
    return 0;
2552
}
2553
2554
static PyObject *
2555
hamt_baseiter_tp_iternext(PyHamtIterator *it)
2556
{
2557
    PyObject *key;
2558
    PyObject *val;
2559
    hamt_iter_t res = hamt_iterator_next(&it->hi_iter, &key, &val);
2560
2561
    switch (res) {
2562
        case I_END:
  Branch (2562:9): [True: 116, False: 231k]
2563
            PyErr_SetNone(PyExc_StopIteration);
2564
            return NULL;
2565
2566
        case I_ITEM: {
  Branch (2566:9): [True: 231k, False: 116]
2567
            return (*(it->hi_yield))(key, val);
2568
        }
2569
2570
        default: {
  Branch (2570:9): [True: 0, False: 231k]
2571
            Py_UNREACHABLE();
2572
        }
2573
    }
2574
}
2575
2576
static Py_ssize_t
2577
hamt_baseiter_tp_len(PyHamtIterator *it)
2578
{
2579
    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
hamt_baseiter_new(PyTypeObject *type, binaryfunc yield, PyHamtObject *o)
2588
{
2589
    PyHamtIterator *it = PyObject_GC_New(PyHamtIterator, type);
2590
    if (it == NULL) {
  Branch (2590:9): [True: 0, False: 217]
2591
        return NULL;
2592
    }
2593
2594
    Py_INCREF(o);
2595
    it->hi_obj = o;
2596
    it->hi_yield = yield;
2597
2598
    hamt_iterator_init(&it->hi_iter, o->h_root);
2599
2600
    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
hamt_iter_yield_items(PyObject *key, PyObject *val)
2627
{
2628
    return PyTuple_Pack(2, key, val);
2629
}
2630
2631
PyObject *
2632
_PyHamt_NewIterItems(PyHamtObject *o)
2633
{
2634
    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
hamt_iter_yield_keys(PyObject *key, PyObject *val)
2650
{
2651
    Py_INCREF(key);
2652
    return key;
2653
}
2654
2655
PyObject *
2656
_PyHamt_NewIterKeys(PyHamtObject *o)
2657
{
2658
    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
hamt_iter_yield_values(PyObject *key, PyObject *val)
2674
{
2675
    Py_INCREF(val);
2676
    return val;
2677
}
2678
2679
PyObject *
2680
_PyHamt_NewIterValues(PyHamtObject *o)
2681
{
2682
    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
hamt_tp_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2698
{
2699
    return (PyObject*)_PyHamt_New();
2700
}
2701
2702
static int
2703
hamt_tp_clear(PyHamtObject *self)
2704
{
2705
    Py_CLEAR(self->h_root);
2706
    return 0;
2707
}
2708
2709
2710
static int
2711
hamt_tp_traverse(PyHamtObject *self, visitproc visit, void *arg)
2712
{
2713
    Py_VISIT(self->h_root);
2714
    return 0;
2715
}
2716
2717
static void
2718
hamt_tp_dealloc(PyHamtObject *self)
2719
{
2720
    PyObject_GC_UnTrack(self);
2721
    if (self->h_weakreflist != NULL) {
  Branch (2721:9): [True: 1, False: 61.3k]
2722
        PyObject_ClearWeakRefs((PyObject*)self);
2723
    }
2724
    (void)hamt_tp_clear(self);
2725
    Py_TYPE(self)->tp_free(self);
2726
}
2727
2728
2729
static PyObject *
2730
hamt_tp_richcompare(PyObject *v, PyObject *w, int op)
2731
{
2732
    if (!PyHamt_Check(v) || !PyHamt_Check(w) || (op != Py_EQ && 
op != 9
Py_NE9
)) {
  Branch (2732:9): [True: 0, False: 18]
  Branch (2732:29): [True: 0, False: 18]
  Branch (2732:50): [True: 9, False: 9]
  Branch (2732:65): [True: 0, False: 9]
2733
        Py_RETURN_NOTIMPLEMENTED;
2734
    }
2735
2736
    int res = _PyHamt_Eq((PyHamtObject *)v, (PyHamtObject *)w);
2737
    if (res < 0) {
  Branch (2737:9): [True: 2, False: 16]
2738
        return NULL;
2739
    }
2740
2741
    if (op == Py_NE) {
  Branch (2741:9): [True: 8, False: 8]
2742
        res = !res;
2743
    }
2744
2745
    if (res) {
  Branch (2745:9): [True: 8, False: 8]
2746
        Py_RETURN_TRUE;
2747
    }
2748
    else {
2749
        Py_RETURN_FALSE;
2750
    }
2751
}
2752
2753
static int
2754
hamt_tp_contains(PyHamtObject *self, PyObject *key)
2755
{
2756
    PyObject *val;
2757
    return _PyHamt_Find(self, key, &val);
2758
}
2759
2760
static PyObject *
2761
hamt_tp_subscript(PyHamtObject *self, PyObject *key)
2762
{
2763
    PyObject *val;
2764
    hamt_find_t res = hamt_find(self, key, &val);
2765
    switch (res) {
2766
        case F_ERROR:
  Branch (2766:9): [True: 2, False: 3]
2767
            return NULL;
2768
        case F_FOUND:
  Branch (2768:9): [True: 2, False: 3]
2769
            Py_INCREF(val);
2770
            return val;
2771
        case F_NOT_FOUND:
  Branch (2771:9): [True: 1, False: 4]
2772
            PyErr_SetObject(PyExc_KeyError, key);
2773
            return NULL;
2774
        default:
  Branch (2774:9): [True: 0, False: 5]
2775
            Py_UNREACHABLE();
2776
    }
2777
}
2778
2779
static Py_ssize_t
2780
hamt_tp_len(PyHamtObject *self)
2781
{
2782
    return _PyHamt_Len(self);
2783
}
2784
2785
static PyObject *
2786
hamt_tp_iter(PyHamtObject *self)
2787
{
2788
    return _PyHamt_NewIterKeys(self);
2789
}
2790
2791
static PyObject *
2792
hamt_py_set(PyHamtObject *self, PyObject *args)
2793
{
2794
    PyObject *key;
2795
    PyObject *val;
2796
2797
    if (!PyArg_UnpackTuple(args, "set", 2, 2, &key, &val)) {
  Branch (2797:9): [True: 0, False: 21.3k]
2798
        return NULL;
2799
    }
2800
2801
    return (PyObject *)_PyHamt_Assoc(self, key, val);
2802
}
2803
2804
static PyObject *
2805
hamt_py_get(PyHamtObject *self, PyObject *args)
2806
{
2807
    PyObject *key;
2808
    PyObject *def = NULL;
2809
2810
    if (!PyArg_UnpackTuple(args, "get", 1, 2, &key, &def)) {
  Branch (2810:9): [True: 0, False: 75.4k]
2811
        return NULL;
2812
    }
2813
2814
    PyObject *val = NULL;
2815
    hamt_find_t res = hamt_find(self, key, &val);
2816
    switch (res) {
2817
        case F_ERROR:
  Branch (2817:9): [True: 1.91k, False: 73.5k]
2818
            return NULL;
2819
        case F_FOUND:
  Branch (2819:9): [True: 31.5k, False: 43.9k]
2820
            Py_INCREF(val);
2821
            return val;
2822
        case F_NOT_FOUND:
  Branch (2822:9): [True: 42.0k, False: 33.4k]
2823
            if (def == NULL) {
  Branch (2823:17): [True: 8, False: 42.0k]
2824
                Py_RETURN_NONE;
2825
            }
2826
            Py_INCREF(def);
2827
            return def;
2828
        default:
  Branch (2828:9): [True: 0, False: 75.4k]
2829
            Py_UNREACHABLE();
2830
    }
2831
}
2832
2833
static PyObject *
2834
hamt_py_delete(PyHamtObject *self, PyObject *key)
2835
{
2836
    return (PyObject *)_PyHamt_Without(self, key);
2837
}
2838
2839
static PyObject *
2840
hamt_py_items(PyHamtObject *self, PyObject *args)
2841
{
2842
    return _PyHamt_NewIterItems(self);
2843
}
2844
2845
static PyObject *
2846
hamt_py_values(PyHamtObject *self, PyObject *args)
2847
{
2848
    return _PyHamt_NewIterValues(self);
2849
}
2850
2851
static PyObject *
2852
hamt_py_keys(PyHamtObject *self, PyObject *Py_UNUSED(args))
2853
{
2854
    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
_PyHamt_Fini(PyInterpreterState *interp)
2962
{
2963
    Py_CLEAR(_empty_hamt);
2964
    Py_CLEAR(_empty_bitmap_node);
2965
}