diff mbox series

[v3,03/15] interval-tree: Add a utility to iterate over spans in an interval tree

Message ID 3-v3-402a7d6459de+24b-iommufd_jgg@nvidia.com
State Superseded
Headers show
Series IOMMUFD Generic interface | expand

Commit Message

Jason Gunthorpe Oct. 25, 2022, 6:12 p.m. UTC
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/.

Signed-off-by: Jason Gunthorpe <jgg@nvidia.com>
---
 .clang-format                 |   1 +
 include/linux/interval_tree.h |  50 +++++++++++++
 lib/Kconfig                   |   4 ++
 lib/interval_tree.c           | 132 ++++++++++++++++++++++++++++++++++
 4 files changed, 187 insertions(+)

Comments

Tian, Kevin Nov. 3, 2022, 5:31 a.m. UTC | #1
> From: Jason Gunthorpe <jgg@nvidia.com>
> Sent: Wednesday, October 26, 2022 2:12 AM
> +/*
> + * 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.
> + *
> + * Only is_hole, start_hole/used and last_hole/used are part of the external
> + * interface.

slightly better readability if moving this sentence into the structure as
the break line

> +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;
> +	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) {

"<=" to "<"

> +		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

Apart from those,

Reviewed-by: Kevin Tian <kevin.tian@intel.com>
Jason Gunthorpe Nov. 4, 2022, 7:38 p.m. UTC | #2
On Thu, Nov 03, 2022 at 05:31:17AM +0000, Tian, Kevin wrote:
> > From: Jason Gunthorpe <jgg@nvidia.com>
> > Sent: Wednesday, October 26, 2022 2:12 AM
> > +/*
> > + * 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.
> > + *
> > + * Only is_hole, start_hole/used and last_hole/used are part of the external
> > + * interface.
> 
> slightly better readability if moving this sentence into the structure as
> the break line

Do you mean like this?

@@ -37,13 +37,16 @@ interval_tree_iter_next(struct interval_tree_node *node,
  * The iterator is greedy, it always returns the largest hole or used possible,
  * consolidating all consecutive nodes.
- *
- * Only is_hole, start_hole/used and last_hole/used are part of the external
- * interface.
  */
 struct interval_tree_span_iter {
 	struct interval_tree_node *nodes[2];
 	unsigned long first_index;
 	unsigned long last_index;
+
+	/*
+	 * Only is_hole, start_hole/used and last_hole/used are part of the
+	 * external interface.
+	 */
 	union {
 		unsigned long start_hole;
 		unsigned long start_used;

> > +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;
> > +	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) {
> 
> "<=" to "<"
> 
> > +		iter->start_hole = new_index;
> > +		return;
> > +	}

Hrm, I don't think so..  This is testing if new_index is valid for the
current span. new_index == start_XXX is a valid index of the span and
new_index == last_XX is also a valid index, so either of them should
just update start_XX and do nothing else.

Thanks,
Jason
Tian, Kevin Nov. 5, 2022, 1:32 a.m. UTC | #3
> From: Jason Gunthorpe <jgg@nvidia.com>
> Sent: Saturday, November 5, 2022 3:38 AM
> 
> On Thu, Nov 03, 2022 at 05:31:17AM +0000, Tian, Kevin wrote:
> > > From: Jason Gunthorpe <jgg@nvidia.com>
> > > Sent: Wednesday, October 26, 2022 2:12 AM
> > > +/*
> > > + * 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.
> > > + *
> > > + * Only is_hole, start_hole/used and last_hole/used are part of the
> external
> > > + * interface.
> >
> > slightly better readability if moving this sentence into the structure as
> > the break line
> 
> Do you mean like this?
> 
> @@ -37,13 +37,16 @@ interval_tree_iter_next(struct interval_tree_node
> *node,
>   * The iterator is greedy, it always returns the largest hole or used possible,
>   * consolidating all consecutive nodes.
> - *
> - * Only is_hole, start_hole/used and last_hole/used are part of the external
> - * interface.
>   */
>  struct interval_tree_span_iter {
>  	struct interval_tree_node *nodes[2];
>  	unsigned long first_index;
>  	unsigned long last_index;
> +
> +	/*
> +	 * Only is_hole, start_hole/used and last_hole/used are part of the
> +	 * external interface.
> +	 */

/* Members below here are part of the external interface */

>  	union {
>  		unsigned long start_hole;
>  		unsigned long start_used;
> 
> > > +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;
> > > +	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) {
> >
> > "<=" to "<"
> >
> > > +		iter->start_hole = new_index;
> > > +		return;
> > > +	}
> 
> Hrm, I don't think so..  This is testing if new_index is valid for the
> current span. new_index == start_XXX is a valid index of the span and
> new_index == last_XX is also a valid index, so either of them should
> just update start_XX and do nothing else.
> 

You are right.
Matthew Wilcox Nov. 5, 2022, 1:48 a.m. UTC | #4
On Fri, Nov 04, 2022 at 04:38:19PM -0300, Jason Gunthorpe wrote:
> On Thu, Nov 03, 2022 at 05:31:17AM +0000, Tian, Kevin wrote:
> > > From: Jason Gunthorpe <jgg@nvidia.com>
> > > Sent: Wednesday, October 26, 2022 2:12 AM
> > > +/*
> > > + * 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.
> > > + *
> > > + * Only is_hole, start_hole/used and last_hole/used are part of the external
> > > + * interface.
> > 
> > slightly better readability if moving this sentence into the structure as
> > the break line
> 
> Do you mean like this?
> 
> @@ -37,13 +37,16 @@ interval_tree_iter_next(struct interval_tree_node *node,
>   * The iterator is greedy, it always returns the largest hole or used possible,
>   * consolidating all consecutive nodes.
> - *
> - * Only is_hole, start_hole/used and last_hole/used are part of the external
> - * interface.
>   */
>  struct interval_tree_span_iter {
>  	struct interval_tree_node *nodes[2];
>  	unsigned long first_index;
>  	unsigned long last_index;
> +
> +	/*
> +	 * Only is_hole, start_hole/used and last_hole/used are part of the
> +	 * external interface.
> +	 */
>  	union {
>  		unsigned long start_hole;
>  		unsigned long start_used;

Or you could kernel-doc it ...

/**
 * struct interval_tree_span_iter - Find used and unused spans.
 * @start_hole: ...
 * @start_used: ...
 ...
 *
 * 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 as used, which is fully covered by nodes. Each iteration step
 * alternates 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 span possible,
 * consolidating all consecutive nodes.
 */
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;
...
};
diff mbox series

Patch

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..3c803379b83cd8 100644
--- a/include/linux/interval_tree.h
+++ b/include/linux/interval_tree.h
@@ -27,4 +27,54 @@  extern struct interval_tree_node *
 interval_tree_iter_next(struct interval_tree_node *node,
 			unsigned long start, unsigned long last);
 
+/*
+ * 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.
+ *
+ * Only is_hole, start_hole/used and last_hole/used are part of the external
+ * interface.
+ */
+struct interval_tree_span_iter {
+	struct interval_tree_node *nodes[2];
+	unsigned long first_index;
+	unsigned long last_index;
+	union {
+		unsigned long start_hole;
+		unsigned long start_used;
+	};
+	union {
+		unsigned long last_hole;
+		unsigned long last_used;
+	};
+	/* 0 == used, 1 == is_hole, -1 == done iteration */
+	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..6edfb83894af8e 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
+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
+	 * nodes[1]->start. nodes[1] is now the next node and a hole is between
+	 * 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;
+	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