summaryrefslogtreecommitdiffstats
path: root/Modules/rotatingtree.c
blob: 217e495b3d2a9d95b34637177ba91f64b4ceca16 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
#ifndef Py_BUILD_CORE_BUILTIN
#  define Py_BUILD_CORE_MODULE 1
#endif

#include "Python.h"
#include "pycore_lock.h"
#include "rotatingtree.h"

#define KEY_LOWER_THAN(key1, key2)  ((char*)(key1) < (char*)(key2))

/* The randombits() function below is a fast-and-dirty generator that
 * is probably irregular enough for our purposes.  Note that it's biased:
 * I think that ones are slightly more probable than zeroes.  It's not
 * important here, though.
 */

static unsigned int random_value = 1;
static unsigned int random_stream = 0;
static PyMutex random_mutex = {0};

static int
randombits(int bits)
{
    int result;
    PyMutex_Lock(&random_mutex);
    if (random_stream < (1U << bits)) {
        random_value *= 1082527;
        random_stream = random_value;
    }
    result = random_stream & ((1<<bits)-1);
    random_stream >>= bits;
    PyMutex_Unlock(&random_mutex);
    return result;
}


/* Insert a new node into the tree.
   (*root) is modified to point to the new root. */
void
RotatingTree_Add(rotating_node_t **root, rotating_node_t *node)
{
    while (*root != NULL) {
        if (KEY_LOWER_THAN(node->key, (*root)->key))
            root = &((*root)->left);
        else
            root = &((*root)->right);
    }
    node->left = NULL;
    node->right = NULL;
    *root = node;
}

/* Locate the node with the given key.  This is the most complicated
   function because it occasionally rebalances the tree to move the
   resulting node closer to the root. */
rotating_node_t *
RotatingTree_Get(rotating_node_t **root, void *key)
{
    if (randombits(3) != 4) {
        /* Fast path, no rebalancing */
        rotating_node_t *node = *root;
        while (node != NULL) {
            if (node->key == key)
                return node;
            if (KEY_LOWER_THAN(key, node->key))
                node = node->left;
            else
                node = node->right;
        }
        return NULL;
    }
    else {
        rotating_node_t **pnode = root;
        rotating_node_t *node = *pnode;
        rotating_node_t *next;
        int rotate;
        if (node == NULL)
            return NULL;
        while (1) {
            if (node->key == key)
                return node;
            rotate = !randombits(1);
            if (KEY_LOWER_THAN(key, node->key)) {
                next = node->left;
                if (next == NULL)
                    return NULL;
                if (rotate) {
                    node->left = next->right;
                    next->right = node;
                    *pnode = next;
                }
                else
                    pnode = &(node->left);
            }
            else {
                next = node->right;
                if (next == NULL)
                    return NULL;
                if (rotate) {
                    node->right = next->left;
                    next->left = node;
                    *pnode = next;
                }
                else
                    pnode = &(node->right);
            }
            node = next;
        }
    }
}

/* Enumerate all nodes in the tree.  The callback enumfn() should return
   zero to continue the enumeration, or non-zero to interrupt it.
   A non-zero value is directly returned by RotatingTree_Enum(). */
int
RotatingTree_Enum(rotating_node_t *root, rotating_tree_enum_fn enumfn,
                  void *arg)
{
    int result;
    rotating_node_t *node;
    while (root != NULL) {
        result = RotatingTree_Enum(root->left, enumfn, arg);
        if (result != 0) return result;
        node = root->right;
        result = enumfn(root, arg);
        if (result != 0) return result;
        root = node;
    }
    return 0;
}