lib/ui/src/tree/walk.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const std = @import("std");
 2 
 3 const abi = @import("../abi/root.zig");
 4 
 5 const Node = abi.Node;
 6 
 7 /// `Children` iterates over the direct children of one node, in preorder.
 8 pub const Children = struct {
 9     nodes: []const Node,
10     /// `Children.last` is the index of the last descendant of the node whose children are walked.
11     last: u32,
12     /// `Children.cursor` is the index the next call examines.
13     cursor: u32,
14 
15     /// `Children.next` returns the index of the next direct child, or null after the last one.
16     /// Each call jumps over that child's whole subtree through its `subtree_count`,
17     /// so a walk takes one step per child and never visits a grandchild.
18     pub fn next(self: *Children) ?u32 {
19         if (self.cursor > self.last) return null;
20         const found = self.cursor;
21         std.debug.assert(found < self.nodes.len);
22         const span = self.nodes[found].subtree_count;
23         std.debug.assert(found + span <= self.last);
24         self.cursor = found + 1 + span;
25         return found;
26     }
27 };
28 
29 /// `children` returns the `Children` iterator for node `index`.
30 /// It asserts that `index` and that node's subtree lie inside `nodes`.
31 pub fn children(nodes: []const Node, index: u32) Children {
32     std.debug.assert(index < nodes.len);
33     const span = nodes[index].subtree_count;
34     std.debug.assert(index + span < nodes.len);
35     return .{ .nodes = nodes, .last = index + span, .cursor = index + 1 };
36 }
37 
38 /// `subtree` returns node `index` followed by all of its descendants,
39 /// as one slice of `subtree_count + 1` nodes.
40 pub fn subtree(nodes: []const Node, index: u32) []const Node {
41     std.debug.assert(index < nodes.len);
42     const span = nodes[index].subtree_count;
43     std.debug.assert(index + span < nodes.len);
44     return nodes[index .. index + span + 1];
45 }
46 
47 /// `childCount` returns the number of direct children of node `index` by walking them.
48 pub fn childCount(nodes: []const Node, index: u32) u32 {
49     var walk = children(nodes, index);
50     var count: u32 = 0;
51     while (walk.next()) |_| count += 1;
52     return count;
53 }
54 
55 test "a preorder child walk names children and a child span names descendants" {
56     const nodes = [_]Node{
57         .{ .id = 1, .parent = 0, .subtree_count = 3 },
58         .{ .id = 2, .parent = 0, .subtree_count = 1 },
59         .{ .id = 3, .parent = 1, .subtree_count = 0 },
60         .{ .id = 4, .parent = 0, .subtree_count = 0 },
61     };
62     var walk = children(&nodes, 0);
63     try std.testing.expectEqual(@as(?u32, 1), walk.next());
64     try std.testing.expectEqual(@as(?u32, 3), walk.next());
65     try std.testing.expectEqual(@as(?u32, null), walk.next());
66     try std.testing.expectEqual(@as(u32, 2), childCount(&nodes, 0));
67     try std.testing.expectEqual(@as(u32, 0), nodes[3].parent);
68     try std.testing.expectEqual(@as(usize, 2), subtree(&nodes, 1).len);
69     try std.testing.expectEqual(@as(u64, 3), subtree(&nodes, 1)[1].id);
70 }
71 
72 test "a leaf enumerates no children and spans only itself" {
73     const nodes = [_]Node{
74         .{ .id = 1, .parent = 0, .subtree_count = 1 },
75         .{ .id = 2, .parent = 0, .subtree_count = 0 },
76     };
77     var walk = children(&nodes, 1);
78     try std.testing.expectEqual(@as(?u32, null), walk.next());
79     try std.testing.expectEqual(@as(u32, 0), childCount(&nodes, 1));
80     try std.testing.expectEqual(@as(usize, 1), subtree(&nodes, 1).len);
81 }
82 
83 test "every child of a wide node is enumerated exactly once" {
84     var nodes: [9]Node = undefined;
85     nodes[0] = .{ .id = 100, .parent = 0, .subtree_count = 8 };
86     for (1..9) |index| {
87         nodes[index] = .{ .id = 100 + index, .parent = 0, .subtree_count = 0 };
88     }
89     var seen: u32 = 0;
90     var walk = children(&nodes, 0);
91     while (walk.next()) |child| seen += 1 + child;
92     try std.testing.expectEqual(@as(u32, 8 + 36), seen);
93     try std.testing.expectEqual(@as(u32, 8), childCount(&nodes, 0));
94 }