lib/choir/src/product/hashing/walk.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const alloc_arena = @import("alloc_arena");
3 const ir = @import("../../core/root.zig");
4 const hashing = @import("root.zig");
5 const capacity = hashing.capacity;
6
7 pub const Event = union(enum) {
8 operation: *ir.Operation,
9 region: *ir.Region,
10 block: *ir.Block,
11 };
12
13 pub const Iterator = struct {
14 frames: []capacity.OperationFrame,
15 depth: usize,
16
17 pub fn init(frames: []capacity.OperationFrame, root: *ir.Operation) Iterator {
18 if (frames.len == 0) @panic("operation traversal has no frame capacity");
19 frames[0] = .{ .operation = root };
20 return .{ .frames = frames, .depth = 1 };
21 }
22
23 pub fn next(self: *Iterator) ?Event {
24 while (self.depth > 0) {
25 const frame = &self.frames[self.depth - 1];
26 switch (frame.state) {
27 .operation => {
28 frame.state = .region;
29 return .{ .operation = frame.operation };
30 },
31 .region => {
32 if (frame.region_index >= frame.operation.regions.items.len) {
33 self.depth -= 1;
34 continue;
35 }
36 const region = &frame.operation.regions.items[frame.region_index];
37 frame.region_index += 1;
38 frame.block = if (region.blocks.head) |block_ptr|
39 @ptrCast(@alignCast(block_ptr))
40 else
41 null;
42 frame.state = .block;
43 return .{ .region = region };
44 },
45 .block => {
46 const block = frame.block orelse {
47 frame.state = .region;
48 continue;
49 };
50 frame.block = block.next;
51 frame.child = if (block.operations.head) |operation_ptr|
52 @ptrCast(@alignCast(operation_ptr))
53 else
54 null;
55 frame.state = .child;
56 return .{ .block = block };
57 },
58 .child => {
59 const child = frame.child orelse {
60 frame.state = .block;
61 continue;
62 };
63 frame.child = child.next_op;
64 if (self.depth >= self.frames.len) {
65 @panic("operation traversal exceeded inspected depth");
66 }
67 self.frames[self.depth] = .{ .operation = child };
68 self.depth += 1;
69 },
70 }
71 }
72 return null;
73 }
74 };
75
76 test "iterator preserves multi-region and multi-block preorder" {
77 const allocator = std.testing.allocator;
78 var arena = alloc_arena.Arena.init(allocator);
79 defer arena.deinit();
80 var context = try ir.Context.init(arena.allocator(), ir.Context.Limits.testing);
81 defer context.deinit(arena.allocator());
82 try context.allowUnregistered();
83 const location = ir.Location.getUnknown();
84
85 var first_region = ir.context.initRegion(&context);
86 defer first_region.deinit();
87 const first_block = try first_region.addBlock();
88 const first = try context.createOperation(ir.Operation.State.init("first", location));
89 try first_block.addOperation(first);
90 const second_block = try first_region.addBlock();
91 const second = try context.createOperation(ir.Operation.State.init("second", location));
92 try second_block.addOperation(second);
93
94 var second_region = ir.context.initRegion(&context);
95 defer second_region.deinit();
96 const third_block = try second_region.addBlock();
97 const third = try context.createOperation(ir.Operation.State.init("third", location));
98 try third_block.addOperation(third);
99
100 var root_state = ir.Operation.State.init("root", location);
101 root_state.addRegionBodies(&.{ &first_region, &second_region });
102 const root = try context.createOperation(root_state);
103 const limits = try capacity.Limits.inspect(root);
104 try std.testing.expectEqual(@as(usize, 2), limits.facts.operation_depth);
105
106 var frames: [2]capacity.OperationFrame = undefined;
107 var iterator = Iterator.init(&frames, root);
108 try std.testing.expectEqual(Event{ .operation = root }, iterator.next().?);
109 try std.testing.expectEqual(Event{ .region = &root.regions.items[0] }, iterator.next().?);
110 try std.testing.expectEqual(Event{ .block = first_block }, iterator.next().?);
111 try std.testing.expectEqual(Event{ .operation = first }, iterator.next().?);
112 try std.testing.expectEqual(Event{ .block = second_block }, iterator.next().?);
113 try std.testing.expectEqual(Event{ .operation = second }, iterator.next().?);
114 try std.testing.expectEqual(Event{ .region = &root.regions.items[1] }, iterator.next().?);
115 try std.testing.expectEqual(Event{ .block = third_block }, iterator.next().?);
116 try std.testing.expectEqual(Event{ .operation = third }, iterator.next().?);
117 try std.testing.expect(iterator.next() == null);
118 }
119
120 test "iterator consumes the maximum inspected operation depth" {
121 const allocator = std.testing.allocator;
122 var arena = alloc_arena.Arena.init(allocator);
123 defer arena.deinit();
124 var context = try ir.Context.init(arena.allocator(), ir.Context.Limits.testing);
125 defer context.deinit(arena.allocator());
126 try context.allowUnregistered();
127 const location = ir.Location.getUnknown();
128
129 var root_state = ir.Operation.State.init("root", location);
130 root_state.addRegion();
131 const root = try context.createOperation(root_state);
132 var current = root;
133 for (1..capacity.maximum_operation_depth) |_| {
134 var child_state = ir.Operation.State.init("child", location);
135 child_state.addRegion();
136 const child = try context.createOperation(child_state);
137 const block = try current.getRegion(0).?.addBlock();
138 try block.addOperation(child);
139 current = child;
140 }
141
142 const limits = try capacity.Limits.inspect(root);
143 try std.testing.expectEqual(
144 capacity.maximum_operation_depth,
145 limits.facts.operation_depth,
146 );
147 var frames: [capacity.maximum_operation_depth]capacity.OperationFrame = undefined;
148 var iterator = Iterator.init(&frames, root);
149 var operation_count: usize = 0;
150 while (iterator.next()) |event| switch (event) {
151 .operation => operation_count += 1,
152 .region, .block => {},
153 };
154 try std.testing.expectEqual(capacity.maximum_operation_depth, operation_count);
155 }