rangemap: implement free_range

src/util/btree.zig

@@ -16,7 +16,7 @@ pub fn BTree(comptime N: type, comptime compare_fn: CompareFn(N), comptime deini
         gpa: Allocator,
         root: ?*Node = null,
 
-        pub const Error = error{ already_exists, does_not_exist } || Allocator.Error;
+        pub const Error = error{ AlreadyExists, DoesNotExist } || Allocator.Error;
 
         pub fn WalkFn(comptime C: type) type {
             return fn (*const Node, C) void;
@@ -25,7 +25,7 @@ pub fn BTree(comptime N: type, comptime compare_fn: CompareFn(N), comptime deini
         pub const Iterator = struct {
             current: ?*Node,
 
-            pub fn next(self: *Iterator) ?*const Node {
+            pub fn next(self: *Iterator) ?*Node {
                 while (self.current) |n| {
                     const v = n;
 
@@ -94,7 +94,7 @@ pub fn BTree(comptime N: type, comptime compare_fn: CompareFn(N), comptime deini
                             child.parent = n;
                             n.left = child;
                         },
-                        .eq => return error.already_exists,
+                        .eq => return error.AlreadyExists,
                     }
                     return .{ n, inserted };
                 } else {
@@ -161,7 +161,7 @@ pub fn BTree(comptime N: type, comptime compare_fn: CompareFn(N), comptime deini
                     }
                     return node;
                 } else {
-                    return error.does_not_exist;
+                    return error.DoesNotExist;
                 }
             }
 

src/util/range.zig

@@ -4,19 +4,44 @@ const std = @import("std");
 
 /// Non-inclusive range type over `T`.
 pub fn Range(comptime T: type) type {
+    switch (@typeInfo(T)) {
+        .int => {},
+        else => {
+            @compileError("Range(T) only support integer ranges");
+        },
+    }
     return struct {
         /// Range start.
         start: T,
         /// Range length.
         len: T,
 
+        const Self = @This();
+
         /// Returns `start + len` of the range.
-        pub fn end(self: *const @This()) T {
+        pub fn end(self: Self) T {
             return self.start + self.len;
         }
 
+        /// Returns a range representing a subtraction of `other` from `self`
+        pub fn subtract(self: Self, other: Self) ?Self {
+            // No overlap
+            if (self.end() <= other.start or other.end() <= self.start)
+                return self;
+
+            // Self is fully in other
+            if (other.start <= self.start and other.end() >= self.end())
+                return null;
+
+            // Partial overlap
+            return if (self.start < other.start)
+                .{ .start = self.start, .len = other.start - self.start }
+            else
+                .{ .start = other.end(), .len = self.end() - other.end() };
+        }
+
         /// Returns a range representing an intersections of `self` with an`other` range.
-        pub fn intersect(self: *const @This(), other: *const @This()) ?Range(T) {
+        pub fn intersect(self: Self, other: Self) ?Self {
             if (self.start < other.start) {
                 const p = other.start - self.start;
                 if (p < self.len) {
@@ -32,11 +57,11 @@ pub fn Range(comptime T: type) type {
             return null;
         }
 
-        pub fn contains(self: *const @This(), scalar: T) bool {
+        pub fn contains(self: Self, scalar: T) bool {
             return scalar >= self.start and scalar - self.start < self.len;
         }
 
-        pub fn compare_disjoint(a: *const @This(), b: *const @This()) std.math.Order {
+        pub fn compare_disjoint(a: Self, b: Self) std.math.Order {
             if (a.start >= b.end()) {
                 return .gt;
             } else if (b.start >= a.end()) {

src/util/rangemap.zig

@@ -11,11 +11,44 @@ pub fn RangeMap(
     comptime K: type,
     comptime V: type,
     comptime ops: struct {
-        deinit_fn: ?fn(*V) void = null,
+        deinit_fn: ?fn (*V) void = null,
         merge_fn: ?fn (*const V, *const V) bool = null,
     },
 ) type {
     return struct {
+        btree: Tree,
+
+        const Self = @This();
+
+        pub const Error = error{
+            ScalarOutOfRange,
+            RangeOutOfBounds,
+        } || Tree.Error;
+
+        pub const Tree = BTree(Node, compare_fn, deinit_node_fn);
+
+        pub const WalkFn = fn (*const Node) void;
+
+        pub const Iterator = struct {
+            inner: Tree.Iterator,
+
+            pub fn next(self: *Iterator) ?*Node {
+                if (self.inner.next()) |n| {
+                    return &n.key;
+                } else {
+                    return null;
+                }
+            }
+        };
+
+        pub fn init(gpa: Allocator) @This() {
+            return .{ .btree = Tree.init(gpa) };
+        }
+
+        pub fn deinit(self: *@This()) void {
+            self.btree.deinit();
+        }
+
         pub const Node = struct {
             key: Range(K),
             value: V,
@@ -25,29 +58,8 @@ pub fn RangeMap(
             }
         };
 
-        pub const WalkFn = fn (*const Node) void;
-
-        pub const Iterator = struct {
-            inner: Tree.Iterator,
-
-            pub fn next(self: *Iterator) ?*const Node {
-                if (self.inner.next()) |n| {
-                    return &n.key;
-                } else {
-                    return null;
-                }
-            }
-        };
-
-        pub const Tree = BTree(Node, compare_fn, deinit_node_fn);
-
-        pub const Error = error{
-            scalar_out_of_range,
-            range_out_of_bounds,
-        } || Tree.Error;
-
         fn compare_fn(a: *const Node, b: *const Node) Order {
-            return Range(K).compare_disjoint(&a.key, &b.key);
+            return Range(K).compare_disjoint(a.key, b.key);
         }
 
         fn deinit_node_fn(n: *Node) void {
@@ -56,16 +68,6 @@ pub fn RangeMap(
             }
         }
 
-        btree: Tree,
-
-        pub fn init(gpa: Allocator) @This() {
-            return .{ .btree = Tree.init(gpa) };
-        }
-
-        pub fn deinit(self: *@This()) void {
-            self.btree.deinit();
-        }
-
         /// Returns the value at a given scalar point, along with the full range it belongs to.
         pub fn get_scalar(self: *const @This(), scalar: K) ?*Node {
             return if (self.get_scalar_node(scalar)) |n| &n.key else null;
@@ -86,6 +88,27 @@ pub fn RangeMap(
             }.call);
         }
 
+        /// Frees a range of the range map.
+        /// If there are any partially freed ranges, they are resized.
+        pub fn free_range(self: *Self, range: Range(K)) !void {
+            var it = self.iterator();
+            while (it.next()) |node| {
+                // If the node is fully in range, we can remove it
+                // If the node partially overlaps, we remove it and reinsert with the new range
+                // TODO: make the iterator actually work with it
+                if (range.intersect(node.key)) |intersection| {
+                    if (node.key.subtract(intersection)) |range_resized| {
+                        const value = node.value;
+                        try self.btree.remove(node.*);
+                        _ = try self.btree.insert(.{ .key = range_resized, .value = value });
+                    } else {
+                        try self.btree.remove(node.*);
+                    }
+                    it = self.iterator();
+                }
+            }
+        }
+
         /// Splits a given node at a scalar point inside its interval.
         ///
         /// The part of the interval before `at` is considered a "left" half, the remaining
@@ -98,14 +121,14 @@ pub fn RangeMap(
         ///
         /// # Errors
         ///
-        /// * `scalar_out_of_range` if the given `at` value is not inside the node's interval.
+        /// * `ScalarOutOfRange` if the given `at` value is not inside the node's interval.
         pub fn split_node(
             self: *@This(),
             node: *Tree.Node,
             at: K,
         ) Error!?struct { *Tree.Node, *Tree.Node } {
             if (!node.key.key.contains(at)) {
-                return error.scalar_out_of_range;
+                return error.ScalarOutOfRange;
             }
 
             const start = node.key.key.start;
@@ -195,7 +218,7 @@ pub fn RangeMap(
         fn validate_range(start: K, end: K) Error!void {
             // Check for addition overflowing the K's bit size
             if (std.math.add(K, start, end) == error.Overflow) {
-                return error.range_out_of_bounds;
+                return error.RangeOutOfBounds;
             }
         }
     };
@@ -210,7 +233,7 @@ test "Range map insertion" {
     _ = try map.insert(0, 10, "Range 1");
     _ = try map.insert(20, 10, "Range 3");
 
-    try std.testing.expectError(error.already_exists, map.insert(5, 10, "Invalid range"));
+    try std.testing.expectError(error.AlreadyExists, map.insert(5, 10, "Invalid range"));
 
     _ = try map.insert(1000, 10, "Range 4");
 }
@@ -244,11 +267,11 @@ test "Range map merging insertion" {
     {
         var it = map.iterator();
         const n0 = it.next().?;
-        try std.testing.expectEqual(Range(u32) { .start = 0, .len = 10 }, n0.key);
+        try std.testing.expectEqual(Range(u32){ .start = 0, .len = 10 }, n0.key);
         try std.testing.expectEqual(true, n0.value);
 
         const n1 = it.next().?;
-        try std.testing.expectEqual(Range(u32) { .start = 10, .len = 30 }, n1.key);
+        try std.testing.expectEqual(Range(u32){ .start = 10, .len = 30 }, n1.key);
         try std.testing.expectEqual(false, n1.value);
 
         try std.testing.expectEqual(null, it.next());
@@ -261,19 +284,19 @@ test "Range map merging insertion" {
     {
         var it = map.iterator();
         const n0 = it.next().?;
-        try std.testing.expectEqual(Range(u32) { .start = 0, .len = 10 }, n0.key);
+        try std.testing.expectEqual(Range(u32){ .start = 0, .len = 10 }, n0.key);
         try std.testing.expectEqual(true, n0.value);
 
         const n1 = it.next().?;
-        try std.testing.expectEqual(Range(u32) { .start = 10, .len = 30 }, n1.key);
+        try std.testing.expectEqual(Range(u32){ .start = 10, .len = 30 }, n1.key);
         try std.testing.expectEqual(false, n1.value);
 
         const n2 = it.next().?;
-        try std.testing.expectEqual(Range(u32) { .start = 40, .len = 10 }, n2.key);
+        try std.testing.expectEqual(Range(u32){ .start = 40, .len = 10 }, n2.key);
         try std.testing.expectEqual(true, n2.value);
 
         const n3 = it.next().?;
-        try std.testing.expectEqual(Range(u32) { .start = 50, .len = 10 }, n3.key);
+        try std.testing.expectEqual(Range(u32){ .start = 50, .len = 10 }, n3.key);
         try std.testing.expectEqual(false, n3.value);
 
         try std.testing.expectEqual(null, it.next());
@@ -286,23 +309,23 @@ test "Range map merging insertion" {
     {
         var it = map.iterator();
         const n0 = it.next().?;
-        try std.testing.expectEqual(Range(u32) { .start = 0, .len = 10 }, n0.key);
+        try std.testing.expectEqual(Range(u32){ .start = 0, .len = 10 }, n0.key);
         try std.testing.expectEqual(true, n0.value);
 
         const n1 = it.next().?;
-        try std.testing.expectEqual(Range(u32) { .start = 10, .len = 30 }, n1.key);
+        try std.testing.expectEqual(Range(u32){ .start = 10, .len = 30 }, n1.key);
         try std.testing.expectEqual(false, n1.value);
 
         const n2 = it.next().?;
-        try std.testing.expectEqual(Range(u32) { .start = 40, .len = 10 }, n2.key);
+        try std.testing.expectEqual(Range(u32){ .start = 40, .len = 10 }, n2.key);
         try std.testing.expectEqual(true, n2.value);
 
         const n3 = it.next().?;
-        try std.testing.expectEqual(Range(u32) { .start = 50, .len = 20 }, n3.key);
+        try std.testing.expectEqual(Range(u32){ .start = 50, .len = 20 }, n3.key);
         try std.testing.expectEqual(false, n3.value);
 
         const n4 = it.next().?;
-        try std.testing.expectEqual(Range(u32) { .start = 71, .len = 9 }, n4.key);
+        try std.testing.expectEqual(Range(u32){ .start = 71, .len = 9 }, n4.key);
         try std.testing.expectEqual(false, n4.value);
 
         try std.testing.expectEqual(null, it.next());
@@ -314,19 +337,19 @@ test "Range map merging insertion" {
     {
         var it = map.iterator();
         const n0 = it.next().?;
-        try std.testing.expectEqual(Range(u32) { .start = 0, .len = 10 }, n0.key);
+        try std.testing.expectEqual(Range(u32){ .start = 0, .len = 10 }, n0.key);
         try std.testing.expectEqual(true, n0.value);
 
         const n1 = it.next().?;
-        try std.testing.expectEqual(Range(u32) { .start = 10, .len = 30 }, n1.key);
+        try std.testing.expectEqual(Range(u32){ .start = 10, .len = 30 }, n1.key);
         try std.testing.expectEqual(false, n1.value);
 
         const n2 = it.next().?;
-        try std.testing.expectEqual(Range(u32) { .start = 40, .len = 10 }, n2.key);
+        try std.testing.expectEqual(Range(u32){ .start = 40, .len = 10 }, n2.key);
         try std.testing.expectEqual(true, n2.value);
 
         const n3 = it.next().?;
-        try std.testing.expectEqual(Range(u32) { .start = 50, .len = 30 }, n3.key);
+        try std.testing.expectEqual(Range(u32){ .start = 50, .len = 30 }, n3.key);
         try std.testing.expectEqual(false, n3.value);
 
         try std.testing.expectEqual(null, it.next());
@@ -339,23 +362,23 @@ test "Range map merging insertion" {
     {
         var it = map.iterator();
         const n0 = it.next().?;
-        try std.testing.expectEqual(Range(u32) { .start = 0, .len = 10 }, n0.key);
+        try std.testing.expectEqual(Range(u32){ .start = 0, .len = 10 }, n0.key);
         try std.testing.expectEqual(true, n0.value);
 
         const n1 = it.next().?;
-        try std.testing.expectEqual(Range(u32) { .start = 10, .len = 30 }, n1.key);
+        try std.testing.expectEqual(Range(u32){ .start = 10, .len = 30 }, n1.key);
         try std.testing.expectEqual(false, n1.value);
 
         const n2 = it.next().?;
-        try std.testing.expectEqual(Range(u32) { .start = 40, .len = 10 }, n2.key);
+        try std.testing.expectEqual(Range(u32){ .start = 40, .len = 10 }, n2.key);
         try std.testing.expectEqual(true, n2.value);
 
         const n3 = it.next().?;
-        try std.testing.expectEqual(Range(u32) { .start = 50, .len = 30 }, n3.key);
+        try std.testing.expectEqual(Range(u32){ .start = 50, .len = 30 }, n3.key);
         try std.testing.expectEqual(false, n3.value);
 
         const n4 = it.next().?;
-        try std.testing.expectEqual(Range(u32) { .start = 100, .len = 20 }, n4.key);
+        try std.testing.expectEqual(Range(u32){ .start = 100, .len = 20 }, n4.key);
         try std.testing.expectEqual(false, n4.value);
 
         try std.testing.expectEqual(null, it.next());
@@ -446,5 +469,25 @@ test "Range map should disallow overflowing ranges" {
     var map = Map.init(std.testing.allocator);
     defer map.deinit();
 
-    try std.testing.expectError(error.range_out_of_bounds, map.insert(0xF0000000, 0x20000000, false));
+    try std.testing.expectError(error.RangeOutOfBounds, map.insert(0xF0000000, 0x20000000, false));
+}
+
+test "Range map free" {
+    const Map = RangeMap(u32, bool, .{});
+    var map = Map.init(std.testing.allocator);
+    defer map.deinit();
+
+    _ = try map.insert(0x1000, 0x1000, true);
+    _ = try map.insert(0x3000, 0x1000, true);
+    _ = try map.insert(0x5000, 0x1000, true);
+    _ = try map.insert(0x7000, 0x1000, true);
+
+    try map.free_range(.{ .start = 0, .len = 0x3800 });
+
+    try std.testing.expectEqual(null, map.get_scalar(0x1000));
+    try std.testing.expectEqual(null, map.get_scalar(0x1800));
+    try std.testing.expectEqual(null, map.get_scalar(0x2000));
+    try std.testing.expectEqual(null, map.get_scalar(0x3000));
+
+    try std.testing.expectEqual(Range(u32){ .start = 0x3800, .len = 0x800 }, map.get_scalar(0x3800).?.key);
 }