1 | /* SPDX-License-Identifier: GPL-2.0-or-later */ |
2 | /* |
3 | Interval Trees |
4 | (C) 2012 Michel Lespinasse <walken@google.com> |
5 | |
6 | |
7 | include/linux/interval_tree_generic.h |
8 | */ |
9 | |
10 | #include <linux/rbtree_augmented.h> |
11 | |
12 | /* |
13 | * Template for implementing interval trees |
14 | * |
15 | * ITSTRUCT: struct type of the interval tree nodes |
16 | * ITRB: name of struct rb_node field within ITSTRUCT |
17 | * ITTYPE: type of the interval endpoints |
18 | * ITSUBTREE: name of ITTYPE field within ITSTRUCT holding last-in-subtree |
19 | * ITSTART(n): start endpoint of ITSTRUCT node n |
20 | * ITLAST(n): last endpoint of ITSTRUCT node n |
21 | * ITSTATIC: 'static' or empty |
22 | * ITPREFIX: prefix to use for the inline tree definitions |
23 | * |
24 | * Note - before using this, please consider if generic version |
25 | * (interval_tree.h) would work for you... |
26 | */ |
27 | |
28 | #define INTERVAL_TREE_DEFINE(ITSTRUCT, ITRB, ITTYPE, ITSUBTREE, \ |
29 | ITSTART, ITLAST, ITSTATIC, ITPREFIX) \ |
30 | \ |
31 | /* Callbacks for augmented rbtree insert and remove */ \ |
32 | \ |
33 | RB_DECLARE_CALLBACKS_MAX(static, ITPREFIX ## _augment, \ |
34 | ITSTRUCT, ITRB, ITTYPE, ITSUBTREE, ITLAST) \ |
35 | \ |
36 | /* Insert / remove interval nodes from the tree */ \ |
37 | \ |
38 | ITSTATIC void ITPREFIX ## _insert(ITSTRUCT *node, \ |
39 | struct rb_root_cached *root) \ |
40 | { \ |
41 | struct rb_node **link = &root->rb_root.rb_node, *rb_parent = NULL; \ |
42 | ITTYPE start = ITSTART(node), last = ITLAST(node); \ |
43 | ITSTRUCT *parent; \ |
44 | bool leftmost = true; \ |
45 | \ |
46 | while (*link) { \ |
47 | rb_parent = *link; \ |
48 | parent = rb_entry(rb_parent, ITSTRUCT, ITRB); \ |
49 | if (parent->ITSUBTREE < last) \ |
50 | parent->ITSUBTREE = last; \ |
51 | if (start < ITSTART(parent)) \ |
52 | link = &parent->ITRB.rb_left; \ |
53 | else { \ |
54 | link = &parent->ITRB.rb_right; \ |
55 | leftmost = false; \ |
56 | } \ |
57 | } \ |
58 | \ |
59 | node->ITSUBTREE = last; \ |
60 | rb_link_node(&node->ITRB, rb_parent, link); \ |
61 | rb_insert_augmented_cached(&node->ITRB, root, \ |
62 | leftmost, &ITPREFIX ## _augment); \ |
63 | } \ |
64 | \ |
65 | ITSTATIC void ITPREFIX ## _remove(ITSTRUCT *node, \ |
66 | struct rb_root_cached *root) \ |
67 | { \ |
68 | rb_erase_augmented_cached(&node->ITRB, root, &ITPREFIX ## _augment); \ |
69 | } \ |
70 | \ |
71 | /* \ |
72 | * Iterate over intervals intersecting [start;last] \ |
73 | * \ |
74 | * Note that a node's interval intersects [start;last] iff: \ |
75 | * Cond1: ITSTART(node) <= last \ |
76 | * and \ |
77 | * Cond2: start <= ITLAST(node) \ |
78 | */ \ |
79 | \ |
80 | static ITSTRUCT * \ |
81 | ITPREFIX ## _subtree_search(ITSTRUCT *node, ITTYPE start, ITTYPE last) \ |
82 | { \ |
83 | while (true) { \ |
84 | /* \ |
85 | * Loop invariant: start <= node->ITSUBTREE \ |
86 | * (Cond2 is satisfied by one of the subtree nodes) \ |
87 | */ \ |
88 | if (node->ITRB.rb_left) { \ |
89 | ITSTRUCT *left = rb_entry(node->ITRB.rb_left, \ |
90 | ITSTRUCT, ITRB); \ |
91 | if (start <= left->ITSUBTREE) { \ |
92 | /* \ |
93 | * Some nodes in left subtree satisfy Cond2. \ |
94 | * Iterate to find the leftmost such node N. \ |
95 | * If it also satisfies Cond1, that's the \ |
96 | * match we are looking for. Otherwise, there \ |
97 | * is no matching interval as nodes to the \ |
98 | * right of N can't satisfy Cond1 either. \ |
99 | */ \ |
100 | node = left; \ |
101 | continue; \ |
102 | } \ |
103 | } \ |
104 | if (ITSTART(node) <= last) { /* Cond1 */ \ |
105 | if (start <= ITLAST(node)) /* Cond2 */ \ |
106 | return node; /* node is leftmost match */ \ |
107 | if (node->ITRB.rb_right) { \ |
108 | node = rb_entry(node->ITRB.rb_right, \ |
109 | ITSTRUCT, ITRB); \ |
110 | if (start <= node->ITSUBTREE) \ |
111 | continue; \ |
112 | } \ |
113 | } \ |
114 | return NULL; /* No match */ \ |
115 | } \ |
116 | } \ |
117 | \ |
118 | ITSTATIC ITSTRUCT * \ |
119 | ITPREFIX ## _iter_first(struct rb_root_cached *root, \ |
120 | ITTYPE start, ITTYPE last) \ |
121 | { \ |
122 | ITSTRUCT *node, *leftmost; \ |
123 | \ |
124 | if (!root->rb_root.rb_node) \ |
125 | return NULL; \ |
126 | \ |
127 | /* \ |
128 | * Fastpath range intersection/overlap between A: [a0, a1] and \ |
129 | * B: [b0, b1] is given by: \ |
130 | * \ |
131 | * a0 <= b1 && b0 <= a1 \ |
132 | * \ |
133 | * ... where A holds the lock range and B holds the smallest \ |
134 | * 'start' and largest 'last' in the tree. For the later, we \ |
135 | * rely on the root node, which by augmented interval tree \ |
136 | * property, holds the largest value in its last-in-subtree. \ |
137 | * This allows mitigating some of the tree walk overhead for \ |
138 | * for non-intersecting ranges, maintained and consulted in O(1). \ |
139 | */ \ |
140 | node = rb_entry(root->rb_root.rb_node, ITSTRUCT, ITRB); \ |
141 | if (node->ITSUBTREE < start) \ |
142 | return NULL; \ |
143 | \ |
144 | leftmost = rb_entry(root->rb_leftmost, ITSTRUCT, ITRB); \ |
145 | if (ITSTART(leftmost) > last) \ |
146 | return NULL; \ |
147 | \ |
148 | return ITPREFIX ## _subtree_search(node, start, last); \ |
149 | } \ |
150 | \ |
151 | ITSTATIC ITSTRUCT * \ |
152 | ITPREFIX ## _iter_next(ITSTRUCT *node, ITTYPE start, ITTYPE last) \ |
153 | { \ |
154 | struct rb_node *rb = node->ITRB.rb_right, *prev; \ |
155 | \ |
156 | while (true) { \ |
157 | /* \ |
158 | * Loop invariants: \ |
159 | * Cond1: ITSTART(node) <= last \ |
160 | * rb == node->ITRB.rb_right \ |
161 | * \ |
162 | * First, search right subtree if suitable \ |
163 | */ \ |
164 | if (rb) { \ |
165 | ITSTRUCT *right = rb_entry(rb, ITSTRUCT, ITRB); \ |
166 | if (start <= right->ITSUBTREE) \ |
167 | return ITPREFIX ## _subtree_search(right, \ |
168 | start, last); \ |
169 | } \ |
170 | \ |
171 | /* Move up the tree until we come from a node's left child */ \ |
172 | do { \ |
173 | rb = rb_parent(&node->ITRB); \ |
174 | if (!rb) \ |
175 | return NULL; \ |
176 | prev = &node->ITRB; \ |
177 | node = rb_entry(rb, ITSTRUCT, ITRB); \ |
178 | rb = node->ITRB.rb_right; \ |
179 | } while (prev == rb); \ |
180 | \ |
181 | /* Check if the node intersects [start;last] */ \ |
182 | if (last < ITSTART(node)) /* !Cond1 */ \ |
183 | return NULL; \ |
184 | else if (start <= ITLAST(node)) /* Cond2 */ \ |
185 | return node; \ |
186 | } \ |
187 | } |
188 | |