Re: [PATCH v4 03/17] interval-tree: Add a utility to iterate over spans in an interval tree

[Eric Auger <eric.auger@redhat.com>]

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@nvidia.com>
> Reviewed-by: Kevin Tian <kevin.tian@intel.com>
> Signed-off-by: Jason Gunthorpe <jgg@nvidia.com>
> ---
>  .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@redhat.com>

Eric