Hi Jason, On 11/8/22 01:48, Jason Gunthorpe wrote: > The span iterator travels over the indexes of the interval_tree, not the > nodes, and classifies spans of indexes as either 'used' or 'hole'. > > 'used' spans are fully covered by nodes in the tree and 'hole' spans have > no node intersecting the span. > > This is done greedily such that spans are maximally sized and every > iteration step switches between used/hole. > > As an example a trivial allocator can be written as: > > for (interval_tree_span_iter_first(&span, itree, 0, ULONG_MAX); > !interval_tree_span_iter_done(&span); > interval_tree_span_iter_next(&span)) > if (span.is_hole && > span.last_hole - span.start_hole >= allocation_size - 1) > return span.start_hole; > > With all the tricky boundary conditions handled by the library code. > > The following iommufd patches have several algorithms for its overlapping > node interval trees that are significantly simplified with this kind of > iteration primitive. As it seems generally useful, put it into lib/. > > Tested-by: Nicolin Chen <nicolinc@xxxxxxxxxx> > Reviewed-by: Kevin Tian <kevin.tian@xxxxxxxxx> > Signed-off-by: Jason Gunthorpe <jgg@xxxxxxxxxx> > --- > .clang-format | 1 + > include/linux/interval_tree.h | 58 +++++++++++++++ > lib/Kconfig | 4 ++ > lib/interval_tree.c | 132 ++++++++++++++++++++++++++++++++++ > 4 files changed, 195 insertions(+) > > diff --git a/.clang-format b/.clang-format > index 1247d54f9e49fa..96d07786dcfb46 100644 > --- a/.clang-format > +++ b/.clang-format > @@ -440,6 +440,7 @@ ForEachMacros: > - 'inet_lhash2_for_each_icsk' > - 'inet_lhash2_for_each_icsk_continue' > - 'inet_lhash2_for_each_icsk_rcu' > + - 'interval_tree_for_each_span' > - 'intlist__for_each_entry' > - 'intlist__for_each_entry_safe' > - 'kcore_copy__for_each_phdr' > diff --git a/include/linux/interval_tree.h b/include/linux/interval_tree.h > index 288c26f50732d7..2b8026a3990645 100644 > --- a/include/linux/interval_tree.h > +++ b/include/linux/interval_tree.h > @@ -27,4 +27,62 @@ extern struct interval_tree_node * > interval_tree_iter_next(struct interval_tree_node *node, > unsigned long start, unsigned long last); > > +/** > + * struct interval_tree_span_iter - Find used and unused spans. > + * @start_hole: Start of an interval for a hole when is_hole == 1 > + * @last_hole: Inclusive end of an interval for a hole when is_hole == 1 > + * @start_used: Start of a used interval when is_hole == 0 > + * @last_used: Inclusive end of a used interval when is_hole == 0 > + * @is_hole: 0 == used, 1 == is_hole, -1 == done iteration > + * > + * This iterator travels over spans in an interval tree. It does not return > + * nodes but classifies each span as either a hole, where no nodes intersect, or > + * a used, which is fully covered by nodes. Each iteration step toggles between > + * hole and used until the entire range is covered. The returned spans always > + * fully cover the requested range. > + * > + * The iterator is greedy, it always returns the largest hole or used possible, > + * consolidating all consecutive nodes. > + * > + * Use interval_tree_span_iter_done() to detect end of iteration. > + */ > +struct interval_tree_span_iter { > + /* private: not for use by the caller */ > + struct interval_tree_node *nodes[2]; > + unsigned long first_index; > + unsigned long last_index; > + > + /* public: */ > + union { > + unsigned long start_hole; > + unsigned long start_used; > + }; > + union { > + unsigned long last_hole; > + unsigned long last_used; > + }; > + int is_hole; > +}; > + > +void interval_tree_span_iter_first(struct interval_tree_span_iter *state, > + struct rb_root_cached *itree, > + unsigned long first_index, > + unsigned long last_index); > +void interval_tree_span_iter_advance(struct interval_tree_span_iter *iter, > + struct rb_root_cached *itree, > + unsigned long new_index); > +void interval_tree_span_iter_next(struct interval_tree_span_iter *state); > + > +static inline bool > +interval_tree_span_iter_done(struct interval_tree_span_iter *state) > +{ > + return state->is_hole == -1; > +} > + > +#define interval_tree_for_each_span(span, itree, first_index, last_index) \ > + for (interval_tree_span_iter_first(span, itree, \ > + first_index, last_index); \ > + !interval_tree_span_iter_done(span); \ > + interval_tree_span_iter_next(span)) > + > #endif /* _LINUX_INTERVAL_TREE_H */ > diff --git a/lib/Kconfig b/lib/Kconfig > index 9bbf8a4b2108e6..c6c323fd251721 100644 > --- a/lib/Kconfig > +++ b/lib/Kconfig > @@ -479,6 +479,10 @@ config INTERVAL_TREE > > for more information. > > +config INTERVAL_TREE_SPAN_ITER > + bool > + depends on INTERVAL_TREE > + > config XARRAY_MULTI > bool > help > diff --git a/lib/interval_tree.c b/lib/interval_tree.c > index 593ce56ece5050..d2882db8fa2a07 100644 > --- a/lib/interval_tree.c > +++ b/lib/interval_tree.c > @@ -15,3 +15,135 @@ EXPORT_SYMBOL_GPL(interval_tree_insert); > EXPORT_SYMBOL_GPL(interval_tree_remove); > EXPORT_SYMBOL_GPL(interval_tree_iter_first); > EXPORT_SYMBOL_GPL(interval_tree_iter_next); > + > +#ifdef CONFIG_INTERVAL_TREE_SPAN_ITER Maybe add in a kernel doc that a prerequisite is state.nodes[1] must be populated > +static void > +interval_tree_span_iter_next_gap(struct interval_tree_span_iter *state) > +{ > + struct interval_tree_node *cur = state->nodes[1]; > + > + /* > + * Roll nodes[1] into nodes[0] by advancing nodes[1] to the end of a > + * contiguous span of nodes. This makes nodes[0]->last the end of that > + * contiguous span of valid indexes that started at the original I would suggest s/contiguous span/contiguous used span and remove "of valid indexes" > + * nodes[1]->start. nodes[1] is now the next node and a hole is between nodes[1] is now the first node starting the next used span. A hole span is between nodes[0]->last and nodes[1]->start > + * nodes[0] and [1]. > + */ > + state->nodes[0] = cur; > + do { > + if (cur->last > state->nodes[0]->last) > + state->nodes[0] = cur; > + cur = interval_tree_iter_next(cur, state->first_index, > + state->last_index); > + } while (cur && (state->nodes[0]->last >= cur->start || > + state->nodes[0]->last + 1 == cur->start)); > + state->nodes[1] = cur; > +} > + > +void interval_tree_span_iter_first(struct interval_tree_span_iter *iter, > + struct rb_root_cached *itree, > + unsigned long first_index, > + unsigned long last_index) > +{ > + iter->first_index = first_index; > + iter->last_index = last_index; > + iter->nodes[0] = NULL; > + iter->nodes[1] = > + interval_tree_iter_first(itree, first_index, last_index); > + if (!iter->nodes[1]) { > + /* No nodes intersect the span, whole span is hole */ > + iter->start_hole = first_index; > + iter->last_hole = last_index; > + iter->is_hole = 1; > + return; > + } > + if (iter->nodes[1]->start > first_index) { > + /* Leading hole on first iteration */ > + iter->start_hole = first_index; > + iter->last_hole = iter->nodes[1]->start - 1; > + iter->is_hole = 1; > + interval_tree_span_iter_next_gap(iter); > + return; > + } > + > + /* Starting inside a used */ > + iter->start_used = first_index; > + iter->is_hole = 0; > + interval_tree_span_iter_next_gap(iter); > + iter->last_used = iter->nodes[0]->last; > + if (iter->last_used >= last_index) { > + iter->last_used = last_index; > + iter->nodes[0] = NULL; > + iter->nodes[1] = NULL; > + } > +} > +EXPORT_SYMBOL_GPL(interval_tree_span_iter_first); > + > +void interval_tree_span_iter_next(struct interval_tree_span_iter *iter) > +{ > + if (!iter->nodes[0] && !iter->nodes[1]) { > + iter->is_hole = -1; > + return; > + } > + > + if (iter->is_hole) { > + iter->start_used = iter->last_hole + 1; > + iter->last_used = iter->nodes[0]->last; > + if (iter->last_used >= iter->last_index) { > + iter->last_used = iter->last_index; > + iter->nodes[0] = NULL; > + iter->nodes[1] = NULL; > + } > + iter->is_hole = 0; > + return; > + } > + > + if (!iter->nodes[1]) { > + /* Trailing hole */ > + iter->start_hole = iter->nodes[0]->last + 1; > + iter->last_hole = iter->last_index; > + iter->nodes[0] = NULL; > + iter->is_hole = 1; > + return; > + } > + > + /* must have both nodes[0] and [1], interior hole */ > + iter->start_hole = iter->nodes[0]->last + 1; > + iter->last_hole = iter->nodes[1]->start - 1; > + iter->is_hole = 1; > + interval_tree_span_iter_next_gap(iter); > +} > +EXPORT_SYMBOL_GPL(interval_tree_span_iter_next); > + > +/* > + * Advance the iterator index to a specific position. The returned used/hole is > + * updated to start at new_index. This is faster than calling > + * interval_tree_span_iter_first() as it can avoid full searches in several > + * cases where the iterator is already set. > + */ > +void interval_tree_span_iter_advance(struct interval_tree_span_iter *iter, > + struct rb_root_cached *itree, > + unsigned long new_index) > +{ > + if (iter->is_hole == -1) > + return; > + > + iter->first_index = new_index; check new_index > iter->first_index? > + if (new_index > iter->last_index) { > + iter->is_hole = -1; > + return; > + } > + > + /* Rely on the union aliasing hole/used */ > + if (iter->start_hole <= new_index && new_index <= iter->last_hole) { > + iter->start_hole = new_index; > + return; > + } > + if (new_index == iter->last_hole + 1) > + interval_tree_span_iter_next(iter); > + else > + interval_tree_span_iter_first(iter, itree, new_index, > + iter->last_index); > +} > +EXPORT_SYMBOL_GPL(interval_tree_span_iter_advance); > +#endif Besides, looks good to me Reviewed-by: Eric Auger <eric.auger@xxxxxxxxxx> Eric