interval_tree_generic.h

  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									      \
 33RB_DECLARE_CALLBACKS_MAX(static, ITPREFIX ## _augment,			      \
 34			 ITSTRUCT, ITRB, ITTYPE, ITSUBTREE, ITLAST)	      \
 35									      \
 36/* Insert / remove interval nodes from the tree */			      \
 37									      \
 38ITSTATIC 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									      \
 65ITSTATIC 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									      \
 80ITSTATIC ITSTRUCT *							      \
 81ITPREFIX ## _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			node = rb_entry(node->ITRB.rb_right, ITSTRUCT, ITRB); \
108			continue;					      \
109		}							      \
110		return NULL;	/* No match */				      \
111	}								      \
112}									      \
113									      \
114ITSTATIC ITSTRUCT *							      \
115ITPREFIX ## _iter_first(struct rb_root_cached *root,			      \
116			ITTYPE start, ITTYPE last)			      \
117{									      \
118	ITSTRUCT *node, *leftmost;					      \
119									      \
120	if (!root->rb_root.rb_node)					      \
121		return NULL;						      \
122									      \
123	/*								      \
124	 * Fastpath range intersection/overlap between A: [a0, a1] and	      \
125	 * B: [b0, b1] is given by:					      \
126	 *								      \
127	 *         a0 <= b1 && b0 <= a1					      \
128	 *								      \
129	 *  ... where A holds the lock range and B holds the smallest	      \
130	 * 'start' and largest 'last' in the tree. For the later, we	      \
131	 * rely on the root node, which by augmented interval tree	      \
132	 * property, holds the largest value in its last-in-subtree.	      \
133	 * This allows mitigating some of the tree walk overhead for	      \
134	 * for non-intersecting ranges, maintained and consulted in O(1).     \
135	 */								      \
136	node = rb_entry(root->rb_root.rb_node, ITSTRUCT, ITRB);		      \
137	if (node->ITSUBTREE < start)					      \
138		return NULL;						      \
139									      \
140	leftmost = rb_entry(root->rb_leftmost, ITSTRUCT, ITRB);		      \
141	if (ITSTART(leftmost) > last)					      \
142		return NULL;						      \
143									      \
144	return ITPREFIX ## _subtree_search(node, start, last);		      \
145}									      \
146									      \
147ITSTATIC ITSTRUCT *							      \
148ITPREFIX ## _iter_next(ITSTRUCT *node, ITTYPE start, ITTYPE last)	      \
149{									      \
150	struct rb_node *rb = node->ITRB.rb_right, *prev;		      \
151									      \
152	while (true) {							      \
153		/*							      \
154		 * Loop invariants:					      \
155		 *   Cond1: ITSTART(node) <= last			      \
156		 *   rb == node->ITRB.rb_right				      \
157		 *							      \
158		 * First, search right subtree if suitable		      \
159		 */							      \
160		if (rb) {						      \
161			ITSTRUCT *right = rb_entry(rb, ITSTRUCT, ITRB);	      \
162			if (start <= right->ITSUBTREE)			      \
163				return ITPREFIX ## _subtree_search(right,     \
164								start, last); \
165		}							      \
166									      \
167		/* Move up the tree until we come from a node's left child */ \
168		do {							      \
169			rb = rb_parent(&node->ITRB);			      \
170			if (!rb)					      \
171				return NULL;				      \
172			prev = &node->ITRB;				      \
173			node = rb_entry(rb, ITSTRUCT, ITRB);		      \
174			rb = node->ITRB.rb_right;			      \
175		} while (prev == rb);					      \
176									      \
177		/* Check if the node intersects [start;last] */		      \
178		if (last < ITSTART(node))		/* !Cond1 */	      \
179			return NULL;					      \
180		else if (start <= ITLAST(node))		/* Cond2 */	      \
181			return node;					      \
182	}								      \
183}