162306a36Sopenharmony_ci/* SPDX-License-Identifier: GPL-2.0 */
262306a36Sopenharmony_ci/* Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.
362306a36Sopenharmony_ci */
462306a36Sopenharmony_ci#ifndef __IOMMUFD_DOUBLE_SPAN_H
562306a36Sopenharmony_ci#define __IOMMUFD_DOUBLE_SPAN_H
662306a36Sopenharmony_ci
762306a36Sopenharmony_ci#include <linux/interval_tree.h>
862306a36Sopenharmony_ci
962306a36Sopenharmony_ci/*
1062306a36Sopenharmony_ci * This is a variation of the general interval_tree_span_iter that computes the
1162306a36Sopenharmony_ci * spans over the union of two different interval trees. Used ranges are broken
1262306a36Sopenharmony_ci * up and reported based on the tree that provides the interval. The first span
1362306a36Sopenharmony_ci * always takes priority. Like interval_tree_span_iter it is greedy and the same
1462306a36Sopenharmony_ci * value of is_used will not repeat on two iteration cycles.
1562306a36Sopenharmony_ci */
1662306a36Sopenharmony_cistruct interval_tree_double_span_iter {
1762306a36Sopenharmony_ci	struct rb_root_cached *itrees[2];
1862306a36Sopenharmony_ci	struct interval_tree_span_iter spans[2];
1962306a36Sopenharmony_ci	union {
2062306a36Sopenharmony_ci		unsigned long start_hole;
2162306a36Sopenharmony_ci		unsigned long start_used;
2262306a36Sopenharmony_ci	};
2362306a36Sopenharmony_ci	union {
2462306a36Sopenharmony_ci		unsigned long last_hole;
2562306a36Sopenharmony_ci		unsigned long last_used;
2662306a36Sopenharmony_ci	};
2762306a36Sopenharmony_ci	/* 0 = hole, 1 = used span[0], 2 = used span[1], -1 done iteration */
2862306a36Sopenharmony_ci	int is_used;
2962306a36Sopenharmony_ci};
3062306a36Sopenharmony_ci
3162306a36Sopenharmony_civoid interval_tree_double_span_iter_update(
3262306a36Sopenharmony_ci	struct interval_tree_double_span_iter *iter);
3362306a36Sopenharmony_civoid interval_tree_double_span_iter_first(
3462306a36Sopenharmony_ci	struct interval_tree_double_span_iter *iter,
3562306a36Sopenharmony_ci	struct rb_root_cached *itree1, struct rb_root_cached *itree2,
3662306a36Sopenharmony_ci	unsigned long first_index, unsigned long last_index);
3762306a36Sopenharmony_civoid interval_tree_double_span_iter_next(
3862306a36Sopenharmony_ci	struct interval_tree_double_span_iter *iter);
3962306a36Sopenharmony_ci
4062306a36Sopenharmony_cistatic inline bool
4162306a36Sopenharmony_ciinterval_tree_double_span_iter_done(struct interval_tree_double_span_iter *state)
4262306a36Sopenharmony_ci{
4362306a36Sopenharmony_ci	return state->is_used == -1;
4462306a36Sopenharmony_ci}
4562306a36Sopenharmony_ci
4662306a36Sopenharmony_ci#define interval_tree_for_each_double_span(span, itree1, itree2, first_index, \
4762306a36Sopenharmony_ci					   last_index)                        \
4862306a36Sopenharmony_ci	for (interval_tree_double_span_iter_first(span, itree1, itree2,       \
4962306a36Sopenharmony_ci						  first_index, last_index);   \
5062306a36Sopenharmony_ci	     !interval_tree_double_span_iter_done(span);                      \
5162306a36Sopenharmony_ci	     interval_tree_double_span_iter_next(span))
5262306a36Sopenharmony_ci
5362306a36Sopenharmony_ci#endif
54