tigerbeetle-node 0.10.0 → 0.11.0

package/src/tigerbeetle/src/lsm/segmented_array.zig

@@ -3,31 +3,72 @@ const std = @import("std");
 const assert = std.debug.assert;
 const math = std.math;
 const mem = std.mem;
-const meta = std.meta;
 
 const div_ceil = @import("../util.zig").div_ceil;
-
+const binary_search_values_raw = @import("binary_search.zig").binary_search_values_raw;
+const binary_search_keys = @import("binary_search.zig").binary_search_keys;
 const Direction = @import("direction.zig").Direction;
 
-pub const Cursor = struct {
-    node: u32,
-    relative_index: u32,
+/// A "segmented array" is an array with efficient (amortized) random-insert/remove operations.
+/// Also known as an "unrolled linked list": https://en.wikipedia.org/wiki/Unrolled_linked_list
+///
+/// The structure consists of an array list of "nodes". Each node is a non-empty array of T.
+/// When a node fills, it is split into two adjacent, partially-full nodes.
+/// When a node empties, it is joined with a nearby node.
+///
+/// An absolute index is offset from the start of the segmented array.
+/// A relative index is offset from the start of a node.
+pub fn SegmentedArray(
+    comptime T: type,
+    comptime NodePool: type,
+    comptime element_count_max: u32,
+    comptime options: Options,
+) type {
+    return SegmentedArrayType(T, NodePool, element_count_max, null, {}, {}, options);
+}
+
+pub fn SortedSegmentedArray(
+    comptime T: type,
+    comptime NodePool: type,
+    comptime element_count_max: u32,
+    comptime Key: type,
+    comptime key_from_value: fn (*const T) callconv(.Inline) Key,
+    comptime compare_keys: fn (Key, Key) callconv(.Inline) math.Order,
+    comptime options: Options,
+) type {
+    return SegmentedArrayType(T, NodePool, element_count_max, Key, key_from_value, compare_keys, options);
+}
+
+pub const Options = struct {
+    /// Assert all invariants before/after every public function.
+    /// Very expensive - only enable for debugging/fuzzing.
+    verify: bool = false,
 };
 
-pub fn SegmentedArray(
+fn SegmentedArrayType(
     comptime T: type,
     comptime NodePool: type,
     comptime element_count_max: u32,
+    // Set when the SegmentedArray is ordered:
+    comptime Key: ?type,
+    comptime key_from_value: if (Key) |K| (fn (*const T) callconv(.Inline) K) else void,
+    comptime compare_keys: if (Key) |K| (fn (K, K) callconv(.Inline) math.Order) else void,
+    comptime options: Options,
 ) type {
     return struct {
         const Self = @This();
 
+        pub const Cursor = struct {
+            node: u32,
+            relative_index: u32,
+        };
+
         // We can't use @divExact() here as we store TableInfo structs of various sizes in this
         // data structure. This means that there may be padding at the end of the node.
         pub const node_capacity = blk: {
-            const max = NodePool.node_size / @sizeOf(T);
+            const max = @divFloor(NodePool.node_size, @sizeOf(T));
 
-            // We require that the node is evenly divisible by 2 to simplify our code
+            // We require that the node capacity is evenly divisible by 2 to simplify our code
             // that splits/joins nodes at the midpoint.
             const capacity = if (max % 2 == 0) max else max - 1;
 
@@ -39,9 +80,15 @@ pub fn SegmentedArray(
         comptime {
             // If this assert fails, we should be using a non-segmented array instead!
             assert(element_count_max > node_capacity);
+
+            // We use u32 for indexes and counts.
+            assert(element_count_max <= std.math.maxInt(u32));
+
+            // The buffers returned from the node_pool must be able to store T with correct alignment.
+            assert(NodePool.node_alignment >= @alignOf(T));
         }
 
-        pub const node_count_max = blk: {
+        pub const node_count_max_naive = blk: {
             // If a node fills up it is divided into two new nodes. Therefore,
             // the worst possible space overhead is when all nodes are half full.
             // This uses flooring division, we want to examine the worst case here.
@@ -49,8 +96,23 @@ pub fn SegmentedArray(
             break :blk div_ceil(element_count_max, elements_per_node_min);
         };
 
+        // We can't always actually reach node_count_max_naive in all configurations.
+        // If we're at node_count_max_naive-1 nodes, in order to split one more node we need:
+        pub const node_count_max = if (element_count_max >=
+            // * The node that we split must be full.
+            node_capacity +
+            // * The last node must have at least one element.
+            1 +
+            // * All other nodes must be at least half-full.
+            ((node_count_max_naive -| 3) * @divExact(node_capacity, 2)) +
+            // * And then we insert one more element into the full node.
+            1)
+            node_count_max_naive
+        else
+            node_count_max_naive - 1;
+
         node_count: u32 = 0,
-        /// This is the segmented array, the first key_node_count pointers are non-null.
+        /// This is the segmented array. The first node_count pointers are non-null.
         /// The rest are null. We only use optional pointers here to get safety checks.
         nodes: *[node_count_max]?*[node_capacity]T,
         /// Since nodes in a segmented array are usually not full, computing the absolute index
@@ -72,13 +134,19 @@ pub fn SegmentedArray(
             mem.set(?*[node_capacity]T, nodes, null);
             indexes[0] = 0;
 
-            return Self{
+            const array = Self{
                 .nodes = nodes,
                 .indexes = indexes,
             };
+
+            if (options.verify) array.verify();
+
+            return array;
         }
 
         pub fn deinit(array: Self, allocator: mem.Allocator, node_pool: ?*NodePool) void {
+            if (options.verify) array.verify();
+
             for (array.nodes[0..array.node_count]) |node| {
                 node_pool.?.release(
                     @ptrCast(NodePool.Node, @alignCast(NodePool.node_alignment, node.?)),
@@ -88,7 +156,78 @@ pub fn SegmentedArray(
             allocator.free(array.indexes);
         }
 
-        pub fn insert_elements(
+        pub fn verify(array: Self) void {
+            assert(array.node_count <= node_count_max);
+            for (array.nodes) |node, node_index| {
+                if (node_index < array.node_count) {
+                    // The first node_count pointers are non-null.
+                    assert(node != null);
+                } else {
+                    // The rest are non-null.
+                    assert(node == null);
+                }
+            }
+            for (array.nodes[0..array.node_count]) |_, node_index| {
+                const c = array.count(@intCast(u32, node_index));
+                // Every node is at most full.
+                assert(c <= node_capacity);
+                // Every node is at least half-full, except the last.
+                if (node_index < array.node_count - 1) {
+                    assert(c >= @divTrunc(node_capacity, 2));
+                }
+            }
+            if (Key) |K| {
+                // If Key is not null then the elements must be sorted by key_from_value (but not necessarily unique).
+                var key_prior_or_null: ?K = null;
+                for (array.nodes[0..array.node_count]) |_, node_index| {
+                    for (array.node_elements(@intCast(u32, node_index))) |*value| {
+                        const key = key_from_value(value);
+                        if (key_prior_or_null) |key_prior| {
+                            assert(compare_keys(key_prior, key) != .gt);
+                        }
+                        key_prior_or_null = key;
+                    }
+                }
+            }
+        }
+
+        pub usingnamespace if (Key) |_| struct {
+            /// Returns the absolute index of the element being inserted.
+            pub fn insert_element(
+                array: *Self,
+                node_pool: *NodePool,
+                element: T,
+            ) u32 {
+                if (options.verify) array.verify();
+
+                const cursor = array.search(key_from_value(&element));
+                const absolute_index = array.absolute_index_for_cursor(cursor);
+                array.insert_elements_at_absolute_index(node_pool, absolute_index, &[_]T{element});
+
+                if (options.verify) array.verify();
+
+                return absolute_index;
+            }
+        } else struct {
+            pub fn insert_elements(
+                array: *Self,
+                node_pool: *NodePool,
+                absolute_index: u32,
+                elements: []const T,
+            ) void {
+                if (options.verify) array.verify();
+
+                array.insert_elements_at_absolute_index(
+                    node_pool,
+                    absolute_index,
+                    elements,
+                );
+
+                if (options.verify) array.verify();
+            }
+        };
+
+        fn insert_elements_at_absolute_index(
             array: *Self,
             node_pool: *NodePool,
             absolute_index: u32,
@@ -136,6 +275,7 @@ pub fn SegmentedArray(
 
             const a = cursor.node;
             const a_pointer = array.nodes[a].?;
+            assert(cursor.relative_index <= array.count(a));
 
             const total = array.count(a) + @intCast(u32, elements.len);
             if (total <= node_capacity) {
@@ -164,86 +304,99 @@ pub fn SegmentedArray(
             // The 5th case can be seen as a special case of the 4th.
             //
             // elements: [yyyyyy], relative_index: 0
-            // [xxxxx_]
-            // [yyyyyy][xxxxx_]
+            // [xxxxx_][______]
+            // [______][xxxxx_] // after first copy_backwards
+            // [______][xxxxx_] // skip mem.copyBackwards (a_half >= relative_index)
+            // [yyyyyy][xxxxx_] // after second copy_backwards
             //
             // elements: [yy], relative_index: 1
-            // [xxxxx_]
-            // [xyyx__][xxx___]
+            // [xxxxx_][______]
+            // [x__x__][xxx___] // after first copy_backwards
+            // [x__x__][xxx___] // skip mem.copyBackwards (a_half >= relative_index)
+            // [xyyx__][xxx___] // after second copy_backwards
             //
             // elements: [yy], relative_index: 2
-            // [xxx_]
-            // [xxy_][yx__]
+            // [xxx_][____]
+            // [xx__][_x__] // after first copy_backwards
+            // [xx__][_x__] // skip mem.copyBackwards (a_half >= relative_index)
+            // [xxy_][yx__] // after second copy_backwards
             //
             // elements: [yy], relative_index: 5
-            // [xxxxxx]
-            // [xxxx__][xyyx__]
+            // [xxxxxx][______]
+            // [xxxxx_][___x__] // after first copy_backwards
+            // [xxxx__][x__x__] // after mem.copyBackwards (a_half < relative_index)
+            // [xxxx__][xyyx__] // after second copy_backwards
             //
             // elements: [yyyyy_], relative_index: 5
-            // [xxxxx_]
-            // [xxxxx_][yyyyy_]
-
-            assert(cursor.relative_index <= array.count(a));
+            // [xxxxx_][______]
+            // [xxxxx_][______] // after first copy_backwards
+            // [xxxxx_][______] // skip mem.copyBackwards (a_half >= relative_index)
+            // [xxxxx_][yyyyy_] // after second copy_backwards
+
+            const a_half_pointer = a_pointer[0..a_half];
+            const b_half_pointer = b_pointer[0..b_half];
+
+            // Move part of `a` forwards to make space for elements.
+            copy_backwards(
+                a_half_pointer,
+                b_half_pointer,
+                cursor.relative_index + elements.len,
+                a_pointer[cursor.relative_index..array.count(a)],
+            );
 
-            const existing_a_head = a_pointer[0..@minimum(a_half, cursor.relative_index)];
-            const existing_b_head = a_pointer[existing_a_head.len..cursor.relative_index];
-
-            const existing_a_tail = a_pointer[existing_a_head.len..][0..a_half -|
-                (cursor.relative_index + elements.len)];
-            const existing_b_tail = a_pointer[existing_a_head.len + existing_a_tail.len +
-                existing_b_head.len .. array.count(a)];
-
-            const elements_a = elements[0 .. a_half - existing_a_head.len - existing_a_tail.len];
-            const elements_b = elements[elements_a.len..];
-
-            assert(array.count(a) == existing_a_head.len + existing_b_head.len +
-                existing_a_tail.len + existing_b_tail.len);
-            assert(elements.len == elements_a.len + elements_b.len);
-            assert(a_half == existing_a_head.len + elements_a.len + existing_a_tail.len);
-            assert(b_half == existing_b_head.len + elements_b.len + existing_b_tail.len);
-            assert(total == a_half + b_half);
-
-            if (existing_a_tail.len > 0) assert(existing_b_head.len == 0 and elements_b.len == 0);
-            if (existing_b_head.len > 0) assert(elements_a.len == 0 and existing_a_tail.len == 0);
-            if (elements_a.len > 0) assert(existing_b_head.len == 0);
-            if (elements_b.len > 0) assert(existing_a_tail.len == 0);
-
-            assert(existing_a_head.ptr == a_pointer);
-            assert(existing_a_head.ptr + existing_a_head.len == existing_a_tail.ptr);
-            assert(existing_a_head.ptr + existing_a_head.len == existing_b_head.ptr);
-            if (existing_b_head.len > 0) {
-                assert(existing_b_head.ptr + existing_b_head.len == existing_b_tail.ptr);
-            } else {
-                assert(existing_a_tail.ptr + existing_a_tail.len == existing_b_tail.ptr);
+            if (a_half < cursor.relative_index) {
+                // Move the part of `a` that is past the half-way point into `b`.
+                mem.copyBackwards(
+                    T,
+                    b_half_pointer,
+                    a_pointer[a_half..cursor.relative_index],
+                );
             }
 
-            mem.copy(T, b_pointer[existing_b_head.len + elements_b.len ..], existing_b_tail);
-            mem.copy(T, b_pointer[existing_b_head.len..], elements_b);
-            mem.copy(T, b_pointer, existing_b_head);
-
-            mem.copyBackwards(
-                T,
-                a_pointer[existing_a_head.len + elements_a.len ..],
-                existing_a_tail,
+            // Move `elements` into `a` and/or `b`.
+            copy_backwards(
+                a_half_pointer,
+                b_half_pointer,
+                cursor.relative_index,
+                elements,
             );
-            mem.copy(T, a_pointer[existing_a_head.len..], elements_a);
 
             array.indexes[b] = array.indexes[a] + a_half;
             array.increment_indexes_after(b, @intCast(u32, elements.len));
         }
 
+        /// Behaves like mem.copyBackwards, but as if `a` and `b` were a single contiguous slice.
+        /// `target` is the destination index within the concatenation of `a` and `b`.
+        fn copy_backwards(
+            a: []T,
+            b: []T,
+            target: usize,
+            source: []const T,
+        ) void {
+            assert(target + source.len <= a.len + b.len);
+            const target_a = a[@minimum(target, a.len)..@minimum(target + source.len, a.len)];
+            const target_b = b[target -| a.len..(target + source.len) -| a.len];
+            assert(target_a.len + target_b.len == source.len);
+            const source_a = source[0..target_a.len];
+            const source_b = source[target_a.len..];
+            if (target_b.ptr != source_b.ptr) {
+                mem.copyBackwards(T, target_b, source_b);
+            }
+            if (target_a.ptr != source_a.ptr) {
+                mem.copyBackwards(T, target_a, source_a);
+            }
+        }
+
         /// Insert an empty node at index `node`.
         fn insert_empty_node_at(array: *Self, node_pool: *NodePool, node: u32) void {
-            assert(array.node_count + 1 < node_count_max);
-
             assert(node <= array.node_count);
-            if (node < array.node_count) {
-                mem.copyBackwards(
-                    ?*[node_capacity]T,
-                    array.nodes[node + 1 .. array.node_count + 1],
-                    array.nodes[node..array.node_count],
-                );
-            }
+            assert(array.node_count + 1 <= node_count_max);
+
+            mem.copyBackwards(
+                ?*[node_capacity]T,
+                array.nodes[node + 1 .. array.node_count + 1],
+                array.nodes[node..array.node_count],
+            );
             mem.copyBackwards(
                 u32,
                 array.indexes[node + 1 .. array.node_count + 2],
@@ -251,7 +404,13 @@ pub fn SegmentedArray(
             );
 
             array.node_count += 1;
-            array.nodes[node] = @ptrCast(*[node_capacity]T, node_pool.acquire());
+            const node_pointer = node_pool.acquire();
+            comptime {
+                // @ptrCast does not check that the size or alignment agree
+                assert(std.meta.alignment(@TypeOf(node_pointer)) >= @alignOf(T));
+                assert(@sizeOf(@TypeOf(node_pointer.*)) >= @sizeOf([node_capacity]T));
+            }
+            array.nodes[node] = @ptrCast(*[node_capacity]T, node_pointer);
             assert(array.indexes[node] == array.indexes[node + 1]);
         }
 
@@ -261,6 +420,8 @@ pub fn SegmentedArray(
             absolute_index: u32,
             remove_count: u32,
         ) void {
+            if (options.verify) array.verify();
+
             assert(array.node_count > 0);
             assert(remove_count > 0);
             assert(absolute_index + remove_count <= element_count_max);
@@ -274,6 +435,8 @@ pub fn SegmentedArray(
                 array.remove_elements_batch(node_pool, absolute_index, batch);
                 i -= batch;
             }
+
+            if (options.verify) array.verify();
         }
 
         fn remove_elements_batch(
@@ -351,7 +514,11 @@ pub fn SegmentedArray(
                 array.decrement_indexes_after(b, remove_count);
 
                 array.remove_empty_node_at(node_pool, b);
-                array.maybe_remove_or_merge_node_with_next(node_pool, a);
+
+                // Either:
+                // * `b` was the last node so now `a` is the last node
+                // * both `a` and `b` were at least half-full so now `a` is at least half-full
+                assert(b == array.node_count or array.count(a) >= half);
             }
         }
 
@@ -442,14 +609,11 @@ pub fn SegmentedArray(
                 @ptrCast(NodePool.Node, @alignCast(NodePool.node_alignment, array.nodes[node].?)),
             );
 
-            assert(node <= array.node_count - 1);
-            if (node < array.node_count - 1) {
-                mem.copy(
-                    ?*[node_capacity]T,
-                    array.nodes[node .. array.node_count - 1],
-                    array.nodes[node + 1 .. array.node_count],
-                );
-            }
+            mem.copy(
+                ?*[node_capacity]T,
+                array.nodes[node .. array.node_count - 1],
+                array.nodes[node + 1 .. array.node_count],
+            );
             mem.copy(
                 u32,
                 array.indexes[node..array.node_count],
@@ -535,27 +699,31 @@ pub fn SegmentedArray(
             assert(absolute_index < element_count_max);
             assert(absolute_index <= array.len());
 
-            var node: u32 = 0;
-            while (node + 1 < array.node_count and
-                absolute_index >= array.indexes[node + 1])
-            {
-                node += 1;
-            }
-            assert(node < array.node_count);
-
-            const relative_index = absolute_index - array.indexes[node];
+            const result = binary_search_keys(u32, struct {
+                inline fn compare(a: u32, b: u32) math.Order {
+                    return math.order(a, b);
+                }
+            }.compare, array.indexes[0..array.node_count], absolute_index, .{});
 
-            if (node == array.node_count - 1) {
-                // Insertion may target the index one past the end of the array.
-                assert(relative_index <= array.count(node));
+            if (result.exact) {
+                return .{
+                    .node = result.index,
+                    .relative_index = 0,
+                };
             } else {
-                assert(relative_index < array.count(node));
+                const node = result.index - 1;
+                const relative_index = absolute_index - array.indexes[node];
+                if (node == array.node_count - 1) {
+                    // Insertion may target the index one past the end of the array.
+                    assert(relative_index <= array.count(node));
+                } else {
+                    assert(relative_index < array.count(node));
+                }
+                return .{
+                    .node = node,
+                    .relative_index = relative_index,
+                };
             }
-
-            return .{
-                .node = node,
-                .relative_index = relative_index,
-            };
         }
 
         pub const Iterator = struct {
@@ -607,19 +775,15 @@ pub fn SegmentedArray(
             }
         };
 
-        pub fn iterator(
+        pub fn iterator_from_cursor(
             array: *const Self,
-            /// Absolute index to start iteration at.
-            absolute_index: u32,
-            /// The start node allows us to skip over nodes as an optimization.
-            /// If ascending start from the first element of the start node and ascend.
-            /// If descending start from the last element of the start node and descend.
-            start_node: u32,
+            /// First element of iteration.
+            cursor: Cursor,
             direction: Direction,
         ) Iterator {
             if (array.node_count == 0) {
-                assert(absolute_index == 0);
-                assert(start_node == 0);
+                assert(cursor.node == 0);
+                assert(cursor.relative_index == 0);
 
                 return Iterator{
                     .array = array,
@@ -627,58 +791,111 @@ pub fn SegmentedArray(
                     .cursor = .{ .node = 0, .relative_index = 0 },
                     .done = true,
                 };
+            } else {
+                assert(cursor.node < array.node_count);
+                assert(cursor.relative_index < array.count(cursor.node));
+
+                return .{
+                    .array = array,
+                    .direction = direction,
+                    .cursor = cursor,
+                };
             }
+        }
 
-            assert(start_node < array.node_count);
+        pub fn iterator_from_index(
+            array: *const Self,
+            /// First element of iteration.
+            absolute_index: u32,
+            direction: Direction,
+        ) Iterator {
             assert(absolute_index < element_count_max);
-            assert(absolute_index < array.indexes[array.node_count]);
 
-            switch (direction) {
-                .ascending => {
-                    assert(absolute_index >= array.indexes[start_node]);
+            if (array.node_count == 0) {
+                assert(absolute_index == 0);
 
-                    var node = start_node;
-                    while (node + 1 < array.node_count and
-                        absolute_index >= array.indexes[node + 1])
-                    {
-                        node += 1;
-                    }
-                    assert(node < array.node_count);
+                return Iterator{
+                    .array = array,
+                    .direction = direction,
+                    .cursor = .{ .node = 0, .relative_index = 0 },
+                    .done = true,
+                };
+            } else {
+                assert(absolute_index < array.len());
 
-                    const relative_index = absolute_index - array.indexes[node];
-                    assert(relative_index < array.count(node));
+                return Iterator{
+                    .array = array,
+                    .direction = direction,
+                    .cursor = array.cursor_for_absolute_index(absolute_index),
+                };
+            }
+        }
 
+        pub usingnamespace if (Key) |K| struct {
+            /// Returns a cursor to the index of the key either exactly equal to the target key or,
+            /// if there is no exact match, the next greatest key.
+            pub fn search(
+                array: *const Self,
+                key: K,
+            ) Cursor {
+                if (array.node_count == 0) {
                     return .{
-                        .array = array,
-                        .direction = direction,
-                        .cursor = .{
-                            .node = node,
-                            .relative_index = relative_index,
-                        },
+                        .node = 0,
+                        .relative_index = 0,
                     };
-                },
-                .descending => {
-                    assert(absolute_index < array.indexes[start_node + 1]);
+                }
 
-                    var node = start_node;
-                    while (node > 0 and absolute_index < array.indexes[node]) {
-                        node -= 1;
-                    }
+                var offset: usize = 0;
+                var length: usize = array.node_count;
+                while (length > 1) {
+                    const half = length / 2;
+                    const mid = offset + half;
 
-                    const relative_index = absolute_index - array.indexes[node];
-                    assert(relative_index < array.count(node));
+                    const node = &array.nodes[mid].?[0];
+                    // This trick seems to be what's needed to get llvm to emit branchless code for this,
+                    // a ternary-style if expression was generated as a jump here for whatever reason.
+                    const next_offsets = [_]usize{ offset, mid };
+                    offset = next_offsets[@boolToInt(compare_keys(key_from_value(node), key) == .lt)];
 
+                    length -= half;
+                }
+
+                // Unlike a normal binary search, don't increment the offset when "key" is higher
+                // than the element — "round down" to the previous node.
+                // This guarantees that the node result is never "== node_count".
+                //
+                // (If there are two adjacent nodes starting with keys A and C, and we search B,
+                // we want to pick the A node.)
+                const node = @intCast(u32, offset);
+                assert(node < array.node_count);
+
+                const relative_index = binary_search_values_raw(
+                    K,
+                    T,
+                    key_from_value,
+                    compare_keys,
+                    array.node_elements(node),
+                    key,
+                    .{},
+                );
+
+                // Follow the same rule as absolute_index_for_cursor:
+                // only return relative_index==array.count() at the last node.
+                if (node + 1 < array.node_count and
+                    relative_index == array.count(node))
+                {
                     return .{
-                        .array = array,
-                        .direction = direction,
-                        .cursor = .{
-                            .node = node,
-                            .relative_index = relative_index,
-                        },
+                        .node = node + 1,
+                        .relative_index = 0,
                     };
-                },
+                } else {
+                    return .{
+                        .node = node,
+                        .relative_index = relative_index,
+                    };
+                }
             }
-        }
+        } else struct {};
     };
 }
 
@@ -686,19 +903,26 @@ fn TestContext(
     comptime T: type,
     comptime node_size: u32,
     comptime element_count_max: u32,
+    comptime Key: type,
+    comptime key_from_value: fn (*const T) callconv(.Inline) Key,
+    comptime compare_keys: fn (Key, Key) callconv(.Inline) math.Order,
+    comptime element_order: enum { sorted, unsorted },
+    comptime options: Options,
 ) type {
     return struct {
         const Self = @This();
 
         const testing = std.testing;
-
         const log = false;
 
         const NodePool = @import("node_pool.zig").NodePool;
 
         // Test overaligned nodes to catch compile errors for missing @alignCast()
         const TestPool = NodePool(node_size, 2 * @alignOf(T));
-        const TestArray = SegmentedArray(T, TestPool, element_count_max);
+        const TestArray = switch (element_order) {
+            .sorted => SortedSegmentedArray(T, TestPool, element_count_max, Key, key_from_value, compare_keys, options),
+            .unsorted => SegmentedArray(T, TestPool, element_count_max, options),
+        };
 
         random: std.rand.Random,
 
@@ -761,6 +985,27 @@ fn TestContext(
             }
 
             try context.remove_all();
+
+            if (element_order == .unsorted) {
+                // Insert at the beginning of the array until the array is full.
+                while (context.array.len() < element_count_max) {
+                    try context.insert_before_first();
+                }
+                assert(context.array.node_count >= TestArray.node_count_max - 1);
+
+                // Remove all-but-one elements from the last node and insert them into the first node.
+                const element_count_last = context.array.count(context.array.node_count - 1);
+                var element_index: usize = 0;
+                while (element_index < element_count_last - 1) : (element_index += 1) {
+                    try context.remove_last();
+                    try context.insert_before_first();
+                }
+
+                // We should now have maxed out our node count.
+                assert(context.array.node_count == TestArray.node_count_max);
+
+                try context.remove_all();
+            }
         }
 
         fn insert(context: *Self) !void {
@@ -770,19 +1015,29 @@ fn TestContext(
             if (count_free == 0) return;
 
             var buffer: [TestArray.node_capacity * 3]T = undefined;
-
             const count_max = @minimum(count_free, TestArray.node_capacity * 3);
             const count = context.random.uintAtMostBiased(u32, count_max - 1) + 1;
             context.random.bytes(mem.sliceAsBytes(buffer[0..count]));
 
             assert(context.reference.items.len <= element_count_max);
-            const index = context.random.uintAtMostBiased(u32, reference_len);
-
-            context.array.insert_elements(&context.pool, index, buffer[0..count]);
 
-            // TODO the standard library could use an AssumeCapacity variant of this.
-            context.reference.insertSlice(index, buffer[0..count]) catch unreachable;
+            switch (element_order) {
+                .unsorted => {
+                    const index = context.random.uintAtMostBiased(u32, reference_len);
 
+                    context.array.insert_elements(&context.pool, index, buffer[0..count]);
+                    // TODO the standard library could use an AssumeCapacity variant of this.
+                    context.reference.insertSlice(index, buffer[0..count]) catch unreachable;
+                },
+                .sorted => {
+                    for (buffer[0..count]) |value| {
+                        const index_actual = context.array.insert_element(&context.pool, value);
+                        const index_expect = context.reference_index(key_from_value(&value));
+                        context.reference.insert(index_expect, value) catch unreachable;
+                        try std.testing.expectEqual(index_expect, index_actual);
+                    }
+                },
+            }
             context.inserts += count;
 
             try context.verify();
@@ -807,6 +1062,35 @@ fn TestContext(
             try context.verify();
         }
 
+        fn insert_before_first(context: *Self) !void {
+            assert(element_order == .unsorted);
+
+            const insert_index = context.array.absolute_index_for_cursor(context.array.first());
+
+            var element: T = undefined;
+            context.random.bytes(mem.asBytes(&element));
+
+            context.array.insert_elements(&context.pool, insert_index, &.{element});
+            context.reference.insert(insert_index, element) catch unreachable;
+
+            context.inserts += 1;
+
+            try context.verify();
+        }
+
+        fn remove_last(context: *Self) !void {
+            assert(element_order == .unsorted);
+
+            const remove_index = context.array.absolute_index_for_cursor(context.array.last());
+
+            context.array.remove_elements(&context.pool, remove_index, 1);
+            context.reference.replaceRange(remove_index, 1, &[0]T{}) catch unreachable;
+
+            context.removes += 1;
+
+            try context.verify();
+        }
+
         fn remove_all(context: *Self) !void {
             while (context.reference.items.len > 0) try context.remove();
 
@@ -830,7 +1114,7 @@ fn TestContext(
                 for (context.reference.items) |i| std.debug.print("{}, ", .{i});
 
                 std.debug.print("\nactual: ", .{});
-                var it = context.array.iterator(0, 0, .ascending);
+                var it = context.array.iterator_from_index(0, .ascending);
                 while (it.next()) |i| std.debug.print("{}, ", .{i.*});
                 std.debug.print("\n", .{});
             }
@@ -838,7 +1122,7 @@ fn TestContext(
             try testing.expectEqual(context.reference.items.len, context.array.len());
 
             {
-                var it = context.array.iterator(0, 0, .ascending);
+                var it = context.array.iterator_from_index(0, .ascending);
 
                 for (context.reference.items) |expect| {
                     const actual = it.next() orelse return error.TestUnexpectedResult;
@@ -848,9 +1132,8 @@ fn TestContext(
             }
 
             {
-                var it = context.array.iterator(
+                var it = context.array.iterator_from_index(
                     @intCast(u32, context.reference.items.len) -| 1,
-                    context.array.last().node,
                     .descending,
                 );
 
@@ -865,7 +1148,26 @@ fn TestContext(
                 try testing.expectEqual(@as(?*const T, null), it.next());
             }
 
-            try testing.expectEqual(context.reference.items.len, context.array.len());
+            {
+                for (context.reference.items) |_, i| {
+                    try testing.expect(std.meta.eql(
+                        i,
+                        context.array.absolute_index_for_cursor(
+                            context.array.cursor_for_absolute_index(@intCast(u32, i)),
+                        ),
+                    ));
+                }
+            }
+
+            if (element_order == .sorted) {
+                for (context.reference.items) |*expect, i| {
+                    if (i == 0) continue;
+                    try testing.expect(compare_keys(
+                        key_from_value(&context.reference.items[i - 1]),
+                        key_from_value(expect),
+                    ) != .gt);
+                }
+            }
 
             if (context.array.len() == 0) {
                 try testing.expectEqual(@as(u32, 0), context.array.node_count);
@@ -875,45 +1177,6 @@ fn TestContext(
                 try testing.expectEqual(@as(?*[TestArray.node_capacity]T, null), node);
             }
 
-            if (context.reference.items.len > 0) {
-                const reference_len = @intCast(u32, context.reference.items.len);
-                const index = context.random.uintLessThanBiased(u32, reference_len);
-                const cursor = context.array.cursor_for_absolute_index(index);
-
-                {
-                    const start_node = context.random.uintAtMostBiased(u32, cursor.node);
-
-                    var it = context.array.iterator(index, start_node, .ascending);
-
-                    for (context.reference.items[index..]) |expect| {
-                        const actual = it.next() orelse return error.TestUnexpectedResult;
-                        try testing.expectEqual(expect, actual.*);
-                    }
-                    try testing.expectEqual(@as(?*const T, null), it.next());
-                }
-
-                {
-                    const start_node = cursor.node + context.random.uintAtMostBiased(
-                        u32,
-                        context.array.node_count - 1 - cursor.node,
-                    );
-                    assert(start_node >= cursor.node);
-                    assert(start_node < context.array.node_count);
-
-                    var it = context.array.iterator(index, start_node, .descending);
-
-                    var i = index + 1;
-                    while (i > 0) {
-                        i -= 1;
-
-                        const expect = context.reference.items[i];
-                        const actual = it.next() orelse return error.TestUnexpectedResult;
-                        try testing.expectEqual(expect, actual.*);
-                    }
-                    try testing.expectEqual(@as(?*const T, null), it.next());
-                }
-            }
-
             {
                 var i: u32 = 0;
                 while (i < context.array.node_count -| 1) : (i += 1) {
@@ -921,38 +1184,83 @@ fn TestContext(
                         @divExact(TestArray.node_capacity, 2));
                 }
             }
+            if (element_order == .sorted) try context.verify_search();
+        }
+
+        fn verify_search(context: *Self) !void {
+            var queries: [20]Key = undefined;
+            context.random.bytes(mem.sliceAsBytes(&queries));
+
+            // Test min/max exceptional values on different SegmentedArray shapes.
+            queries[0] = 0;
+            queries[1] = math.maxInt(Key);
+
+            for (queries) |query| {
+                try testing.expectEqual(
+                    context.reference_index(query),
+                    context.array.absolute_index_for_cursor(context.array.search(query)),
+                );
+            }
+        }
+
+        fn reference_index(context: *const Self, key: Key) u32 {
+            return binary_search_values_raw(
+                Key,
+                T,
+                key_from_value,
+                compare_keys,
+                context.reference.items,
+                key,
+                .{},
+            );
         }
     };
 }
 
-test "SegmentedArray" {
-    const seed = 42;
-
+pub fn run_tests(seed: u64, comptime options: Options) !void {
     var prng = std.rand.DefaultPrng.init(seed);
     const random = prng.random();
 
-    // TODO Import this type from lsm/tree.zig.
-    const TableInfo = extern struct {
-        checksum: u128,
-        key_min: u128,
-        key_max: u128,
+    const Key = @import("composite_key.zig").CompositeKey(u64);
+    const TableType = @import("table.zig").TableType;
+    const TableInfoType = @import("manifest.zig").TableInfoType;
+    const TableInfo = TableInfoType(TableType(
+        Key,
+        Key.Value,
+        Key.compare_keys,
+        Key.key_from_value,
+        Key.sentinel_key,
+        Key.tombstone,
+        Key.tombstone_from_key,
+    ));
+
+    const CompareInt = struct {
+        inline fn compare_keys(a: u32, b: u32) std.math.Order {
+            return std.math.order(a, b);
+        }
+
+        inline fn key_from_value(value: *const u32) u32 {
+            return value.*;
+        }
+    };
 
-        address: u64,
+    const CompareTable = struct {
+        inline fn compare_keys(a: u64, b: u64) std.math.Order {
+            return std.math.order(a, b);
+        }
 
-        /// Set to the current snapshot tick on creation.
-        snapshot_min: u64,
-        /// Initially math.maxInt(u64) on creation, set to the current
-        /// snapshot tick on deletion.
-        snapshot_max: u64,
-        flags: u64 = 0,
+        inline fn key_from_value(value: *const TableInfo) u64 {
+            return value.address;
+        }
     };
 
     comptime {
         assert(@sizeOf(TableInfo) == 48 + @sizeOf(u128) * 2);
         assert(@alignOf(TableInfo) == 16);
+        assert(@bitSizeOf(TableInfo) == @sizeOf(TableInfo) * 8);
     }
 
-    const Options = struct {
+    const TestOptions = struct {
         element_type: type,
         node_size: u32,
         element_count_max: u32,
@@ -964,35 +1272,47 @@ test "SegmentedArray" {
     // We want to explore not just the bottom boundary but also the surrounding area
     // as it may also have interesting edge cases.
     inline for (.{
-        Options{ .element_type = u32, .node_size = 8, .element_count_max = 3 },
-        Options{ .element_type = u32, .node_size = 8, .element_count_max = 4 },
-        Options{ .element_type = u32, .node_size = 8, .element_count_max = 5 },
-        Options{ .element_type = u32, .node_size = 8, .element_count_max = 6 },
-        Options{ .element_type = u32, .node_size = 8, .element_count_max = 1024 },
-        Options{ .element_type = u32, .node_size = 16, .element_count_max = 1024 },
-        Options{ .element_type = u32, .node_size = 32, .element_count_max = 1024 },
-        Options{ .element_type = u32, .node_size = 64, .element_count_max = 1024 },
-        Options{ .element_type = TableInfo, .node_size = 256, .element_count_max = 3 },
-        Options{ .element_type = TableInfo, .node_size = 256, .element_count_max = 4 },
-        Options{ .element_type = TableInfo, .node_size = 256, .element_count_max = 1024 },
-        Options{ .element_type = TableInfo, .node_size = 512, .element_count_max = 1024 },
-        Options{ .element_type = TableInfo, .node_size = 1024, .element_count_max = 1024 },
-    }) |options| {
-        const Context = TestContext(
-            options.element_type,
-            options.node_size,
-            options.element_count_max,
-        );
-
-        var context = try Context.init(random);
-        defer context.deinit();
-
-        try context.run();
-
-        if (options.node_size % @sizeOf(options.element_type) != 0) tested_padding = true;
-        if (Context.TestArray.node_capacity == 2) tested_node_capacity_min = true;
+        TestOptions{ .element_type = u32, .node_size = 8, .element_count_max = 3 },
+        TestOptions{ .element_type = u32, .node_size = 8, .element_count_max = 4 },
+        TestOptions{ .element_type = u32, .node_size = 8, .element_count_max = 5 },
+        TestOptions{ .element_type = u32, .node_size = 8, .element_count_max = 6 },
+        TestOptions{ .element_type = u32, .node_size = 8, .element_count_max = 1024 },
+        TestOptions{ .element_type = u32, .node_size = 16, .element_count_max = 1024 },
+        TestOptions{ .element_type = u32, .node_size = 32, .element_count_max = 1024 },
+        TestOptions{ .element_type = u32, .node_size = 64, .element_count_max = 1024 },
+        TestOptions{ .element_type = TableInfo, .node_size = 256, .element_count_max = 3 },
+        TestOptions{ .element_type = TableInfo, .node_size = 256, .element_count_max = 4 },
+        TestOptions{ .element_type = TableInfo, .node_size = 256, .element_count_max = 1024 },
+        TestOptions{ .element_type = TableInfo, .node_size = 512, .element_count_max = 1024 },
+        TestOptions{ .element_type = TableInfo, .node_size = 1024, .element_count_max = 1024 },
+    }) |test_options| {
+        inline for (.{ .sorted, .unsorted }) |order| {
+            const Context = TestContext(
+                test_options.element_type,
+                test_options.node_size,
+                test_options.element_count_max,
+                if (test_options.element_type == u32) u32 else u64,
+                if (test_options.element_type == u32) CompareInt.key_from_value else CompareTable.key_from_value,
+                if (test_options.element_type == u32) CompareInt.compare_keys else CompareTable.compare_keys,
+                order,
+                options,
+            );
+
+            var context = try Context.init(random);
+            defer context.deinit();
+
+            try context.run();
+
+            if (test_options.node_size % @sizeOf(test_options.element_type) != 0) tested_padding = true;
+            if (Context.TestArray.node_capacity == 2) tested_node_capacity_min = true;
+        }
     }
 
     assert(tested_padding);
     assert(tested_node_capacity_min);
 }
+
+test "SegmentedArray" {
+    const seed = 42;
+    try run_tests(seed, .{ .verify = true });
+}