#include <linux/init.h>
#include <linux/interval_tree.h>
/* Callbacks for augmented rbtree insert and remove */
static inline unsigned long
compute_subtree_last(struct interval_tree_node *node)
{
unsigned long max = node->last, subtree_last;
if (node->rb.rb_left) {
subtree_last = rb_entry(node->rb.rb_left,
struct interval_tree_node, rb)->__subtree_last;
if (max < subtree_last)
max = subtree_last;
}
if (node->rb.rb_right) {
subtree_last = rb_entry(node->rb.rb_right,
struct interval_tree_node, rb)->__subtree_last;
if (max < subtree_last)
max = subtree_last;
}
return max;
}
RB_DECLARE_CALLBACKS(static, augment_callbacks, struct interval_tree_node, rb,
unsigned long, __subtree_last, compute_subtree_last)
/* Insert / remove interval nodes from the tree */
void interval_tree_insert(struct interval_tree_node *node,
struct rb_root *root)
{
struct rb_node **link = &root->rb_node, *rb_parent = NULL;
unsigned long start = node->start, last = node->last;
struct interval_tree_node *parent;
while (*link) {
rb_parent = *link;
parent = rb_entry(rb_parent, struct interval_tree_node, rb);
if (parent->__subtree_last < last)
parent->__subtree_last = last;
if (start < parent->start)
link = &parent->rb.rb_left;
else
link = &parent->rb.rb_right;
}
node->__subtree_last = last;
rb_link_node(&node->rb, rb_parent, link);
rb_insert_augmented(&node->rb, root, &augment_callbacks);
}
void interval_tree_remove(struct interval_tree_node *node,
struct rb_root *root)
{
rb_erase_augmented(&node->rb, root, &augment_callbacks);
}
/*
* Iterate over intervals intersecting [start;last]
*
* Note that a node's interval intersects [start;last] iff:
* Cond1: node->start <= last
* and
* Cond2: start <= node->last
*/
static struct interval_tree_node *
subtree_search(struct interval_tree_node *node,
unsigned long start, unsigned long last)
{
while (true) {
/*
* Loop invariant: start <= node->__subtree_last
* (Cond2 is satisfied by one of the subtree nodes)
*/
if (node->rb.rb_left) {
struct interval_tree_node *left =
rb_entry(node->rb.rb_left,
struct interval_tree_node, rb);
if (start <= left->__subtree_last) {
/*
* Some nodes in left subtree satisfy Cond2.
* Iterate to find the leftmost such node N.
* If it also satisfies Cond1, that's the match
* we are looking for. Otherwise, there is no
* matching interval as nodes to the right of N
* can't satisfy Cond1 either.
*/
node = left;
continue;
}
}
if (node->start <= last) { /* Cond1 */
if (start <= node->last) /* Cond2 */
return node; /* node is leftmost match */
if (node->rb.rb_right) {
node = rb_entry(node->rb.rb_right,
struct interval_tree_node, rb);
if (start <= node->__subtree_last)
continue;
}
}
return NULL; /* No match */
}
}
struct interval_tree_node *
interval_tree_iter_first(struct rb_root *root,
unsigned long start, unsigned long last)
{
struct interval_tree_node *node;
if (!root->rb_node)
return NULL;
node = rb_entry(root->rb_node, struct interval_tree_node, rb);
if (node->__subtree_last < start)
return NULL;
return subtree_search(node, start, last);
}
struct interval_tree_node *
interval_tree_iter_next(struct interval_tree_node *node,
unsigned long start, unsigned long last)
{
struct rb_node *rb = node->rb.rb_right, *prev;
while (true) {
/*
* Loop invariants:
* Cond1: node->start <= last
* rb == node->rb.rb_right
*
* First, search right subtree if suitable
*/
if (rb) {
struct interval_tree_node *right =
rb_entry(rb, struct interval_tree_node, rb);
if (start <= right->__subtree_last)
return subtree_search(right, start, last);
}
/* Move up the tree until we come from a node's left child */
do {
rb = rb_parent(&node->rb);
if (!rb)
return NULL;
prev = &node->rb;
node = rb_entry(rb, struct interval_tree_node, rb);
rb = node->rb.rb_right;
} while (prev == rb);
/* Check if the node intersects [start;last] */
if (last < node->start) /* !Cond1 */
return NULL;
else if (start <= node->last) /* Cond2 */
return node;
}
}
