1 /* SPDX-License-Identifier: GPL-2.0-or-later */
2 /*
3   Interval Trees
4   (C) 2012  Michel Lespinasse <[email protected]>
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 static inline ITTYPE ITPREFIX ## _compute_subtree_last(ITSTRUCT *node)	      \
34 {									      \
35 	ITTYPE max = ITLAST(node), subtree_last;			      \
36 	if (node->ITRB.rb_left) {					      \
37 		subtree_last = rb_entry(node->ITRB.rb_left,		      \
38 					ITSTRUCT, ITRB)->ITSUBTREE;	      \
39 		if (max < subtree_last)					      \
40 			max = subtree_last;				      \
41 	}								      \
42 	if (node->ITRB.rb_right) {					      \
43 		subtree_last = rb_entry(node->ITRB.rb_right,		      \
44 					ITSTRUCT, ITRB)->ITSUBTREE;	      \
45 		if (max < subtree_last)					      \
46 			max = subtree_last;				      \
47 	}								      \
48 	return max;							      \
49 }									      \
50 									      \
51 RB_DECLARE_CALLBACKS(static, ITPREFIX ## _augment, ITSTRUCT, ITRB,	      \
52 		     ITTYPE, ITSUBTREE, ITPREFIX ## _compute_subtree_last)    \
53 									      \
54 /* Insert / remove interval nodes from the tree */			      \
55 									      \
56 ITSTATIC void ITPREFIX ## _insert(ITSTRUCT *node,			      \
57 				  struct rb_root_cached *root)	 	      \
58 {									      \
59 	struct rb_node **link = &root->rb_root.rb_node, *rb_parent = NULL;    \
60 	ITTYPE start = ITSTART(node), last = ITLAST(node);		      \
61 	ITSTRUCT *parent;						      \
62 	bool leftmost = true;						      \
63 									      \
64 	while (*link) {							      \
65 		rb_parent = *link;					      \
66 		parent = rb_entry(rb_parent, ITSTRUCT, ITRB);		      \
67 		if (parent->ITSUBTREE < last)				      \
68 			parent->ITSUBTREE = last;			      \
69 		if (start < ITSTART(parent))				      \
70 			link = &parent->ITRB.rb_left;			      \
71 		else {							      \
72 			link = &parent->ITRB.rb_right;			      \
73 			leftmost = false;				      \
74 		}							      \
75 	}								      \
76 									      \
77 	node->ITSUBTREE = last;						      \
78 	rb_link_node(&node->ITRB, rb_parent, link);			      \
79 	rb_insert_augmented_cached(&node->ITRB, root,			      \
80 				   leftmost, &ITPREFIX ## _augment);	      \
81 }									      \
82 									      \
83 ITSTATIC void ITPREFIX ## _remove(ITSTRUCT *node,			      \
84 				  struct rb_root_cached *root)		      \
85 {									      \
86 	rb_erase_augmented_cached(&node->ITRB, root, &ITPREFIX ## _augment);  \
87 }									      \
88 									      \
89 /*									      \
90  * Iterate over intervals intersecting [start;last]			      \
91  *									      \
92  * Note that a node's interval intersects [start;last] iff:		      \
93  *   Cond1: ITSTART(node) <= last					      \
94  * and									      \
95  *   Cond2: start <= ITLAST(node)					      \
96  */									      \
97 									      \
98 static ITSTRUCT *							      \
99 ITPREFIX ## _subtree_search(ITSTRUCT *node, ITTYPE start, ITTYPE last)	      \
100 {									      \
101 	while (true) {							      \
102 		/*							      \
103 		 * Loop invariant: start <= node->ITSUBTREE		      \
104 		 * (Cond2 is satisfied by one of the subtree nodes)	      \
105 		 */							      \
106 		if (node->ITRB.rb_left) {				      \
107 			ITSTRUCT *left = rb_entry(node->ITRB.rb_left,	      \
108 						  ITSTRUCT, ITRB);	      \
109 			if (start <= left->ITSUBTREE) {			      \
110 				/*					      \
111 				 * Some nodes in left subtree satisfy Cond2.  \
112 				 * Iterate to find the leftmost such node N.  \
113 				 * If it also satisfies Cond1, that's the     \
114 				 * match we are looking for. Otherwise, there \
115 				 * is no matching interval as nodes to the    \
116 				 * right of N can't satisfy Cond1 either.     \
117 				 */					      \
118 				node = left;				      \
119 				continue;				      \
120 			}						      \
121 		}							      \
122 		if (ITSTART(node) <= last) {		/* Cond1 */	      \
123 			if (start <= ITLAST(node))	/* Cond2 */	      \
124 				return node;	/* node is leftmost match */  \
125 			if (node->ITRB.rb_right) {			      \
126 				node = rb_entry(node->ITRB.rb_right,	      \
127 						ITSTRUCT, ITRB);	      \
128 				if (start <= node->ITSUBTREE)		      \
129 					continue;			      \
130 			}						      \
131 		}							      \
132 		return NULL;	/* No match */				      \
133 	}								      \
134 }									      \
135 									      \
136 ITSTATIC ITSTRUCT *							      \
137 ITPREFIX ## _iter_first(struct rb_root_cached *root,			      \
138 			ITTYPE start, ITTYPE last)			      \
139 {									      \
140 	ITSTRUCT *node, *leftmost;					      \
141 									      \
142 	if (!root->rb_root.rb_node)					      \
143 		return NULL;						      \
144 									      \
145 	/*								      \
146 	 * Fastpath range intersection/overlap between A: [a0, a1] and	      \
147 	 * B: [b0, b1] is given by:					      \
148 	 *								      \
149 	 *         a0 <= b1 && b0 <= a1					      \
150 	 *								      \
151 	 *  ... where A holds the lock range and B holds the smallest	      \
152 	 * 'start' and largest 'last' in the tree. For the later, we	      \
153 	 * rely on the root node, which by augmented interval tree	      \
154 	 * property, holds the largest value in its last-in-subtree.	      \
155 	 * This allows mitigating some of the tree walk overhead for	      \
156 	 * for non-intersecting ranges, maintained and consulted in O(1).     \
157 	 */								      \
158 	node = rb_entry(root->rb_root.rb_node, ITSTRUCT, ITRB);		      \
159 	if (node->ITSUBTREE < start)					      \
160 		return NULL;						      \
161 									      \
162 	leftmost = rb_entry(root->rb_leftmost, ITSTRUCT, ITRB);		      \
163 	if (ITSTART(leftmost) > last)					      \
164 		return NULL;						      \
165 									      \
166 	return ITPREFIX ## _subtree_search(node, start, last);		      \
167 }									      \
168 									      \
169 ITSTATIC ITSTRUCT *							      \
170 ITPREFIX ## _iter_next(ITSTRUCT *node, ITTYPE start, ITTYPE last)	      \
171 {									      \
172 	struct rb_node *rb = node->ITRB.rb_right, *prev;		      \
173 									      \
174 	while (true) {							      \
175 		/*							      \
176 		 * Loop invariants:					      \
177 		 *   Cond1: ITSTART(node) <= last			      \
178 		 *   rb == node->ITRB.rb_right				      \
179 		 *							      \
180 		 * First, search right subtree if suitable		      \
181 		 */							      \
182 		if (rb) {						      \
183 			ITSTRUCT *right = rb_entry(rb, ITSTRUCT, ITRB);	      \
184 			if (start <= right->ITSUBTREE)			      \
185 				return ITPREFIX ## _subtree_search(right,     \
186 								start, last); \
187 		}							      \
188 									      \
189 		/* Move up the tree until we come from a node's left child */ \
190 		do {							      \
191 			rb = rb_parent(&node->ITRB);			      \
192 			if (!rb)					      \
193 				return NULL;				      \
194 			prev = &node->ITRB;				      \
195 			node = rb_entry(rb, ITSTRUCT, ITRB);		      \
196 			rb = node->ITRB.rb_right;			      \
197 		} while (prev == rb);					      \
198 									      \
199 		/* Check if the node intersects [start;last] */		      \
200 		if (last < ITSTART(node))		/* !Cond1 */	      \
201 			return NULL;					      \
202 		else if (start <= ITLAST(node))		/* Cond2 */	      \
203 			return node;					      \
204 	}								      \
205 }
206