1 files changed, 80 insertions, 4 deletions

diff --git a/module/zfs/range_tree.c b/module/zfs/range_tree.c
index 391533b3f..5919236d9 100644
--- a/module/zfs/range_tree.c
+++ b/module/zfs/range_tree.c
@@ -23,7 +23,7 @@
  * Use is subject to license terms.
  */
 /*
- * Copyright (c) 2013, 2017 by Delphix. All rights reserved.
+ * Copyright (c) 2013, 2019 by Delphix. All rights reserved.
  */
 
 #include <sys/zfs_context.h>
@@ -578,10 +578,10 @@ range_tree_vacate(range_tree_t *rt, range_tree_func_t *func, void *arg)
 void
 range_tree_walk(range_tree_t *rt, range_tree_func_t *func, void *arg)
 {
-	range_seg_t *rs;
-
-	for (rs = avl_first(&rt->rt_root); rs; rs = AVL_NEXT(&rt->rt_root, rs))
+	for (range_seg_t *rs = avl_first(&rt->rt_root); rs;
+	    rs = AVL_NEXT(&rt->rt_root, rs)) {
 		func(arg, rs->rs_start, rs->rs_end - rs->rs_start);
+	}
 }
 
 range_seg_t *
@@ -596,6 +596,12 @@ range_tree_space(range_tree_t *rt)
 	return (rt->rt_space);
 }
 
+uint64_t
+range_tree_numsegs(range_tree_t *rt)
+{
+	return ((rt == NULL) ? 0 : avl_numnodes(&rt->rt_root));
+}
+
 boolean_t
 range_tree_is_empty(range_tree_t *rt)
 {
@@ -667,3 +673,73 @@ range_tree_span(range_tree_t *rt)
 {
 	return (range_tree_max(rt) - range_tree_min(rt));
 }
+
+/*
+ * Remove any overlapping ranges between the given segment [start, end)
+ * from removefrom. Add non-overlapping leftovers to addto.
+ */
+void
+range_tree_remove_xor_add_segment(uint64_t start, uint64_t end,
+    range_tree_t *removefrom, range_tree_t *addto)
+{
+	avl_index_t where;
+	range_seg_t starting_rs = {
+		.rs_start = start,
+		.rs_end = start + 1
+	};
+
+	range_seg_t *curr = avl_find(&removefrom->rt_root,
+	    &starting_rs, &where);
+
+	if (curr == NULL)
+		curr = avl_nearest(&removefrom->rt_root, where, AVL_AFTER);
+
+	range_seg_t *next;
+	for (; curr != NULL; curr = next) {
+		next = AVL_NEXT(&removefrom->rt_root, curr);
+
+		if (start == end)
+			return;
+		VERIFY3U(start, <, end);
+
+		/* there is no overlap */
+		if (end <= curr->rs_start) {
+			range_tree_add(addto, start, end - start);
+			return;
+		}
+
+		uint64_t overlap_start = MAX(curr->rs_start, start);
+		uint64_t overlap_end = MIN(curr->rs_end, end);
+		uint64_t overlap_size = overlap_end - overlap_start;
+		ASSERT3S(overlap_size, >, 0);
+		range_tree_remove(removefrom, overlap_start, overlap_size);
+
+		if (start < overlap_start)
+			range_tree_add(addto, start, overlap_start - start);
+
+		start = overlap_end;
+	}
+	VERIFY3P(curr, ==, NULL);
+
+	if (start != end) {
+		VERIFY3U(start, <, end);
+		range_tree_add(addto, start, end - start);
+	} else {
+		VERIFY3U(start, ==, end);
+	}
+}
+
+/*
+ * For each entry in rt, if it exists in removefrom, remove it
+ * from removefrom. Otherwise, add it to addto.
+ */
+void
+range_tree_remove_xor_add(range_tree_t *rt, range_tree_t *removefrom,
+    range_tree_t *addto)
+{
+	for (range_seg_t *rs = avl_first(&rt->rt_root); rs;
+	    rs = AVL_NEXT(&rt->rt_root, rs)) {
+		range_tree_remove_xor_add_segment(rs->rs_start, rs->rs_end,
+		    removefrom, addto);
+	}
+}