lib/choir/src/egraph/test.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const alloc_arena = @import("alloc_arena");
3 const egraph = @import("root.zig");
4 const ir = @import("../core/root.zig");
5 const rewrite = ir.rewrite;
6 const dialects = @import("../dialects/root.zig");
7 const test_support = @import("../dialects/fixture/root.zig");
8 const ClassId = egraph.ClassId;
9 const Extraction = egraph.Extraction;
10 const Graph = egraph.Graph;
11 const Node = egraph.Node;
12 const RewriteContext = egraph.RewriteContext;
13 const RewriteSet = egraph.RewriteSet;
14 const Scalar = egraph.Scalar;
15 const TypeClass = egraph.TypeClass;
16 test {
17 std.testing.refAllDecls(egraph.node);
18 std.testing.refAllDecls(egraph.graph);
19 std.testing.refAllDecls(egraph.rules);
20 std.testing.refAllDecls(egraph.pattern);
21 std.testing.refAllDecls(egraph.extract);
22 }
23
24 fn buildTestContext(allocator: std.mem.Allocator) !ir.Context {
25 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
26 errdefer ctx.deinit(allocator);
27 try test_support.registerTestDialect(&ctx);
28 _ = try ctx.registerOperation("test.constant", .{});
29 _ = try ctx.registerOperation("test.binary", .{ .is_commutative = true });
30 return ctx;
31 }
32
33 const EGraphProperties = struct {
34 payload: ?ir.Attribute = null,
35
36 fn from(storage: *anyopaque) *@This() {
37 return @ptrCast(@alignCast(storage));
38 }
39
40 fn fromConst(storage: *const anyopaque) *const @This() {
41 return @ptrCast(@alignCast(storage));
42 }
43
44 fn init(storage: *anyopaque, _: std.mem.Allocator) anyerror!void {
45 from(storage).* = .{};
46 }
47
48 fn deinit(_: *anyopaque, _: std.mem.Allocator) void {}
49
50 fn get(_: *const ir.Operation, storage: *const anyopaque, name: []const u8) ?ir.Attribute {
51 const payload = fromConst(storage).payload orelse return null;
52 const values = payload.cast(ir.Attribute.ArrayAttr).?.values;
53 if (std.mem.eql(u8, name, "left")) return values[0];
54 if (std.mem.eql(u8, name, "right")) return values[1];
55 return null;
56 }
57
58 fn getProperties(_: *const ir.Operation, storage: *const anyopaque) ?ir.Attribute {
59 return fromConst(storage).payload;
60 }
61
62 fn setProperties(_: *ir.Operation, storage: *anyopaque, attr: ir.Attribute) anyerror!void {
63 const array = attr.cast(ir.Attribute.ArrayAttr) orelse return error.InvalidTestProperties;
64 if (array.values.len != 2) return error.InvalidTestProperties;
65 from(storage).payload = attr;
66 }
67
68 fn copyProperties(dest: *anyopaque, source: *const anyopaque) anyerror!void {
69 from(dest).* = fromConst(source).*;
70 }
71
72 const model = ir.OperationPropertiesModel{
73 .name = "test.egraph.properties",
74 .size = @sizeOf(@This()),
75 .alignment = std.mem.Alignment.fromByteUnits(@alignOf(@This())),
76 .init = init,
77 .deinit = deinit,
78 .getInherentAttr = get,
79 .getPropertiesAsAttr = getProperties,
80 .setPropertiesFromAttr = setProperties,
81 .copyProperties = copyProperties,
82 };
83 };
84
85 fn registerEGraphProperties(context: *ir.Context) !void {
86 _ = try context.registerOperation("test.egraph_properties", .{});
87 try context.registerOperationInherentAttributeNames(
88 "test.egraph_properties",
89 &.{ "left", "right" },
90 );
91 try context.registerOperationPropertiesModel(
92 "test.egraph_properties",
93 EGraphProperties.model,
94 );
95 }
96
97 test "egraph canonical dedup retains the first storage witness" {
98 const testing = std.testing;
99 var arena = alloc_arena.Arena.init(testing.allocator);
100 defer arena.deinit();
101 var context = try ir.Context.init(arena.allocator(), ir.Context.Limits.testing);
102 defer context.deinit(arena.allocator());
103
104 const visible = try context.getI64Attr(11);
105 const first_raw = try context.getI64Attr(99);
106 const second_raw = try context.getI64Attr(100);
107 const first_properties = try context.getI64Attr(7);
108 const second_properties = try context.getI64Attr(8);
109 var canonical = [_]ir.NamedAttribute{.{ .name = "value", .value = visible }};
110 var first_storage = [_]ir.NamedAttribute{.{ .name = "value", .value = first_raw }};
111 var second_storage = [_]ir.NamedAttribute{.{ .name = "value", .value = second_raw }};
112 const first = Node{
113 .kind = .operation,
114 .op_name = "test.canonical",
115 .attributes = &canonical,
116 .raw_attributes = &first_storage,
117 .properties = first_properties,
118 };
119 const second = Node{
120 .kind = .operation,
121 .op_name = "test.canonical",
122 .attributes = &canonical,
123 .raw_attributes = &second_storage,
124 .properties = second_properties,
125 };
126
127 var graph = Graph.init(testing.allocator);
128 defer graph.deinit();
129 const first_class = try graph.addNode(&first);
130 const second_class = try graph.addNode(&second);
131 try testing.expect(first_class.eql(second_class));
132 const retained = graph.nodes(first_class);
133 try testing.expectEqual(@as(usize, 1), retained.len);
134 try testing.expect(retained[0].properties.?.eql(first_properties));
135 try testing.expect(retained[0].raw_attributes[0].value.eql(first_raw));
136 }
137
138 test "egraph extraction refuses unqualified properties and preserves source shadows" {
139 const testing = std.testing;
140 var arena = alloc_arena.Arena.init(testing.allocator);
141 defer arena.deinit();
142 const allocator = arena.allocator();
143
144 var context = try buildTestContext(allocator);
145 defer context.deinit(allocator);
146 try registerEGraphProperties(&context);
147
148 const location = ir.Location.getUnknown();
149 const result_type = try test_support.TestDialect.getI64Type(&context);
150 const left = try context.getI64Attr(11);
151 const right = try context.getI64Attr(22);
152 const payload = try context.getArrayAttr(&.{ left, right });
153 const raw_left = try context.getI64Attr(99);
154 const note = try context.getStringAttr("kept");
155 const raw_attributes = [_]ir.NamedAttribute{
156 .{ .name = "left", .value = raw_left },
157 .{ .name = "debug.note", .value = note },
158 };
159 var source_state = ir.Operation.State.init("test.egraph_properties", location);
160 source_state.addTypes(&.{result_type});
161 source_state.addRawAttributes(&raw_attributes);
162 try source_state.setPropertiesAttr(payload);
163 const source = try context.createOperation(source_state);
164
165 const module = try test_support.TestDialect.ModuleOp.create(&context, location);
166 const block = module.getBodyBlock();
167 try block.addOperation(source);
168 const terminator = try test_support.TestDialect.ReturnOp.create(
169 &context,
170 location,
171 &.{source.getResult(0).?},
172 );
173 try block.addOperation(terminator.op);
174
175 var graph = Graph.init(allocator);
176 defer graph.deinit();
177 const class = try graph.addOperation(source, &.{}, 1, 0);
178 var extraction = try Extraction.init(allocator, &graph, .{});
179 defer extraction.deinit();
180 extraction.analyze();
181 var rewriter = rewrite.PatternRewriter.init(allocator, &context);
182 defer rewriter.deinit();
183 const result = try extraction.materialize(
184 &rewriter,
185 class,
186 terminator.op,
187 0,
188 source.getResult(0).?,
189 );
190 try testing.expect(result == null);
191 const materialized = source;
192
193 try testing.expectEqual(
194 @as(i64, 11),
195 materialized.getAttr("left").?.cast(ir.Attribute.IntegerAttr).?.value,
196 );
197 try testing.expectEqual(
198 @as(i64, 22),
199 materialized.getAttr("right").?.cast(ir.Attribute.IntegerAttr).?.value,
200 );
201 try testing.expectEqual(@as(usize, 2), materialized.getRawDictionaryAttrs().len);
202 try testing.expectEqual(
203 @as(i64, 99),
204 materialized.raw_dictionary_attrs.get("left").?.cast(ir.Attribute.IntegerAttr).?.value,
205 );
206 try testing.expectEqualStrings(
207 "kept",
208 materialized.getAttrAs(ir.Attribute.StringAttr, "debug.note").?.value,
209 );
210 }