1#include <kassert.h>
2#include <mem/alloc.h>
3#include <string.h>
4#include <thread/tid.h>
5
6static size_t tid_space_get_data(struct rbt_node *node) {
7 return container_of(node, struct tid_range, node)->start;
8}
9
10static int32_t tid_space_cmp(const struct rbt_node *a,
11 const struct rbt_node *b) {
12 int32_t l = tid_space_get_data(node: (void *) a);
13 int32_t r = tid_space_get_data(node: (void *) b);
14 return l - r;
15}
16
17struct tid_space *tid_space_init(uint64_t max_id) {
18 struct tid_space *ts = kmalloc(sizeof(*ts), ALLOC_FLAGS_ZERO);
19 if (!ts)
20 return NULL;
21
22 rbt_init(t: &ts->tree, get_data: tid_space_get_data, compare: tid_space_cmp);
23 spinlock_init(lock: &ts->lock);
24
25 ts->reserve_free = NULL;
26 for (int i = 0; i < TID_RANGE_RESERVE_COUNT; i++) {
27 ts->reserve_pool[i].node.left = NULL;
28 ts->reserve_pool[i].node.right = NULL;
29 ts->reserve_pool[i].node.parent = NULL;
30 ts->reserve_pool[i].start = 0;
31 ts->reserve_pool[i].length = 0;
32 ts->reserve_pool[i].next = ts->reserve_free;
33 ts->reserve_free = &ts->reserve_pool[i];
34 }
35
36 struct tid_range *r = kmalloc(sizeof(*r), ALLOC_FLAGS_ZERO);
37 if (!r)
38 return ts;
39
40 r->start = 1;
41 r->length = max_id;
42 rbt_insert(tree: &ts->tree, new_node: &r->node);
43
44 return ts;
45}
46
47static struct tid_range *tid_range_alloc(struct tid_space *ts) {
48 SPINLOCK_ASSERT_HELD(&ts->lock);
49 struct tid_range *r = kmalloc(sizeof(*r), ALLOC_FLAGS_ZERO);
50 if (r)
51 return r;
52
53 if (ts->reserve_free) {
54 r = ts->reserve_free;
55 ts->reserve_free = r->next;
56 memset(r, 0, sizeof(*r));
57 return r;
58 }
59
60 return NULL;
61}
62
63static void tid_range_free(struct tid_space *ts, struct tid_range *r) {
64 SPINLOCK_ASSERT_HELD(&ts->lock);
65 if ((uintptr_t) r >= (uintptr_t) &ts->reserve_pool[0] &&
66 (uintptr_t) r <
67 (uintptr_t) &ts->reserve_pool[TID_RANGE_RESERVE_COUNT]) {
68 r->next = ts->reserve_free;
69 ts->reserve_free = r;
70 } else {
71 kfree(r);
72 }
73}
74
75uint64_t tid_alloc(struct tid_space *ts) {
76 enum irql irql = spin_lock(lock: &ts->lock);
77
78 struct rbt_node *node = rbt_min(tree: &ts->tree);
79 if (!node) {
80 spin_unlock(lock: &ts->lock, old: irql);
81 return 0;
82 }
83
84 struct tid_range *range = rbt_entry(node, struct tid_range, node);
85 uint64_t id = range->start;
86
87 if (range->length == 1) {
88 rbt_delete(tree: &ts->tree, z: node);
89 tid_range_free(ts, r: range);
90 } else {
91 rbt_delete(tree: &ts->tree, z: &range->node);
92 range->start++;
93 range->length--;
94 rbt_insert(tree: &ts->tree, new_node: &range->node);
95 }
96
97 spin_unlock(lock: &ts->lock, old: irql);
98 return id;
99}
100
101void tid_free(struct tid_space *ts, uint64_t id) {
102 enum irql irql = spin_lock(lock: &ts->lock);
103
104 struct rbt_node *node = ts->tree.root;
105 struct tid_range *prev = NULL;
106 struct tid_range *next = NULL;
107
108 while (node) {
109 struct tid_range *r = rbt_entry(node, struct tid_range, node);
110 if (id < r->start) {
111 next = r;
112 node = node->left;
113 } else if (id > r->start + r->length - 1) {
114 prev = r;
115 node = node->right;
116 } else {
117 goto out;
118 }
119 }
120
121 bool merged_prev = false, merged_next = false;
122
123 if (prev && prev->start + prev->length == id) {
124 prev->length++;
125 merged_prev = true;
126 }
127
128 if (next && next->start == id + 1) {
129 if (merged_prev) {
130 prev->length += next->length;
131 rbt_delete(tree: &ts->tree, z: &next->node);
132 tid_range_free(ts, r: next);
133 } else {
134 rbt_delete(tree: &ts->tree, z: &next->node);
135 next->start = id;
136 next->length++;
137 rbt_insert(tree: &ts->tree, new_node: &next->node);
138 }
139 merged_next = true;
140 }
141
142 if (!merged_prev && !merged_next) {
143 struct tid_range *new_range = tid_range_alloc(ts);
144
145 if (!new_range)
146 goto out;
147
148 new_range->start = id;
149 new_range->length = 1;
150 rbt_insert(tree: &ts->tree, new_node: &new_range->node);
151 }
152
153out:
154 spin_unlock(lock: &ts->lock, old: irql);
155}
156