lib/choir/src/egraph/node.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const ir = @import("../core/root.zig");
3
4 pub const ClassId = struct {
5 index: u32,
6
7 pub fn eql(self: ClassId, other: ClassId) bool {
8 return self.index == other.index;
9 }
10 };
11
12 fn class_id_less_than(_: void, lhs: ClassId, rhs: ClassId) bool {
13 return lhs.index < rhs.index;
14 }
15
16 pub const ValueEntry = struct {
17 value: *ir.Value,
18 cost: u32,
19 order: usize,
20 };
21
22 pub const NodeKind = enum {
23 value,
24 operation,
25 };
26
27 pub const Node = struct {
28 kind: NodeKind,
29 value: ?*ir.Value = null,
30 op_name: []const u8 = "",
31 operands: []ClassId = &.{},
32 result_types: []ir.Type = &.{},
33 attributes: []ir.NamedAttribute = &.{},
34 raw_attributes: []ir.NamedAttribute = &.{},
35 properties: ?ir.Attribute = null,
36 commutative: bool = false,
37
38 pub fn valueNode(value: *ir.Value) Node {
39 return .{ .kind = .value, .value = value };
40 }
41
42 pub fn operationNode(
43 allocator: std.mem.Allocator,
44 op: *ir.Operation,
45 operands: []const ClassId,
46 commutative: bool,
47 ) !Node {
48 const owned_operands = try allocator.alloc(ClassId, operands.len);
49 errdefer allocator.free(owned_operands);
50 @memcpy(owned_operands, operands);
51
52 const result_types = op.getResultTypes();
53 const owned_result_types = try allocator.alloc(ir.Type, result_types.len);
54 errdefer allocator.free(owned_result_types);
55 @memcpy(owned_result_types, result_types);
56
57 const owned_attributes = try allocator.alloc(ir.NamedAttribute, op.getNumAttrs());
58 errdefer allocator.free(owned_attributes);
59 var attributes = op.getAttrs();
60 var attribute_index: usize = 0;
61 while (attributes.next()) |attr| : (attribute_index += 1) {
62 if (attribute_index >= owned_attributes.len) {
63 return error.OperationAttributesChanged;
64 }
65 owned_attributes[attribute_index] = attr;
66 }
67 if (attribute_index != owned_attributes.len) return error.OperationAttributesChanged;
68
69 const properties = try op.getPropertiesAsAttr();
70 const raw_attributes = op.getRawDictionaryAttrs();
71 const owned_raw_attributes = if (namedAttributesEqual(owned_attributes, raw_attributes))
72 owned_attributes
73 else blk: {
74 const owned = try allocator.alloc(ir.NamedAttribute, raw_attributes.len);
75 @memcpy(owned, raw_attributes);
76 break :blk owned;
77 };
78 errdefer if (owned_raw_attributes.ptr != owned_attributes.ptr) {
79 allocator.free(owned_raw_attributes);
80 };
81
82 var node = Node{
83 .kind = .operation,
84 .op_name = op.name.name,
85 .operands = owned_operands,
86 .result_types = owned_result_types,
87 .attributes = owned_attributes,
88 .raw_attributes = owned_raw_attributes,
89 .properties = properties,
90 .commutative = commutative,
91 };
92 node.normalizeOperands();
93 return node;
94 }
95
96 pub fn clone(self: Node, allocator: std.mem.Allocator) !Node {
97 switch (self.kind) {
98 .value => return valueNode(self.value.?),
99 .operation => {
100 const owned_operands = try allocator.alloc(ClassId, self.operands.len);
101 errdefer allocator.free(owned_operands);
102 @memcpy(owned_operands, self.operands);
103
104 const owned_result_types = try allocator.alloc(ir.Type, self.result_types.len);
105 errdefer allocator.free(owned_result_types);
106 @memcpy(owned_result_types, self.result_types);
107
108 const owned_attributes = try allocator.alloc(ir.NamedAttribute, self.attributes.len);
109 errdefer allocator.free(owned_attributes);
110 @memcpy(owned_attributes, self.attributes);
111
112 const owned_raw_attributes = if (namedAttributesEqual(
113 self.attributes,
114 self.raw_attributes,
115 ))
116 owned_attributes
117 else blk: {
118 const owned = try allocator.alloc(
119 ir.NamedAttribute,
120 self.raw_attributes.len,
121 );
122 @memcpy(owned, self.raw_attributes);
123 break :blk owned;
124 };
125 errdefer if (owned_raw_attributes.ptr != owned_attributes.ptr) {
126 allocator.free(owned_raw_attributes);
127 };
128
129 return .{
130 .kind = .operation,
131 .op_name = self.op_name,
132 .operands = owned_operands,
133 .result_types = owned_result_types,
134 .attributes = owned_attributes,
135 .raw_attributes = owned_raw_attributes,
136 .properties = self.properties,
137 .commutative = self.commutative,
138 };
139 },
140 }
141 }
142
143 pub fn deinit(self: *Node, allocator: std.mem.Allocator) void {
144 switch (self.kind) {
145 .value => {},
146 .operation => {
147 allocator.free(self.operands);
148 allocator.free(self.result_types);
149 if (self.raw_attributes.ptr != self.attributes.ptr) {
150 allocator.free(self.raw_attributes);
151 }
152 allocator.free(self.attributes);
153 },
154 }
155 self.* = .{ .kind = .value, .value = undefined };
156 }
157
158 pub fn getAttr(self: *const Node, name: []const u8) ?ir.Attribute {
159 if (self.kind != .operation) return null;
160 for (self.attributes) |attribute| {
161 if (std.mem.eql(u8, attribute.name, name)) return attribute.value;
162 }
163 return null;
164 }
165
166 pub fn normalizeOperands(self: *Node) void {
167 if (self.kind != .operation) return;
168 if (!self.commutative) return;
169 std.mem.sort(ClassId, self.operands, {}, class_id_less_than);
170 }
171
172 pub fn hash(self: *const Node) u64 {
173 var hasher = std.hash.Wyhash.init(0);
174 const tag: u8 = @backingInt(self.kind);
175 hasher.update(std.mem.asBytes(&tag));
176
177 switch (self.kind) {
178 .value => {
179 const ptr = @intFromPtr(self.value.?);
180 hasher.update(std.mem.asBytes(&ptr));
181 const type_id = self.value.?.type.uniqueId();
182 hasher.update(std.mem.asBytes(&type_id));
183 },
184 .operation => {
185 hasher.update(self.op_name);
186 hasher.update(std.mem.asBytes(&self.commutative));
187 for (self.operands) |operand| {
188 hasher.update(std.mem.asBytes(&operand.index));
189 }
190 for (self.result_types) |result_type| {
191 const type_id = result_type.uniqueId();
192 hasher.update(std.mem.asBytes(&type_id));
193 }
194 for (self.attributes) |attribute| {
195 hasher.update(attribute.name);
196 hashAttribute(&hasher, attribute.value);
197 }
198 },
199 }
200
201 return hasher.final();
202 }
203
204 pub fn eql(self: *const Node, other: *const Node) bool {
205 if (self.kind != other.kind) return false;
206
207 switch (self.kind) {
208 .value => return self.value.? == other.value.?,
209 .operation => {
210 if (!std.mem.eql(u8, self.op_name, other.op_name)) return false;
211 if (self.commutative != other.commutative) return false;
212 if (self.operands.len != other.operands.len) return false;
213 if (self.result_types.len != other.result_types.len) return false;
214 if (self.attributes.len != other.attributes.len) return false;
215
216 for (self.operands, other.operands) |lhs, rhs| {
217 if (!lhs.eql(rhs)) return false;
218 }
219 for (self.result_types, other.result_types) |lhs, rhs| {
220 if (!lhs.eql(rhs)) return false;
221 }
222 for (self.attributes, other.attributes) |lhs, rhs| {
223 if (!std.mem.eql(u8, lhs.name, rhs.name)) return false;
224 if (!lhs.value.eql(rhs.value)) return false;
225 }
226 return true;
227 },
228 }
229 }
230 };
231
232 fn namedAttributesEqual(
233 lhs: []const ir.NamedAttribute,
234 rhs: []const ir.NamedAttribute,
235 ) bool {
236 if (lhs.len != rhs.len) return false;
237 for (lhs, rhs) |left, right| {
238 if (!std.mem.eql(u8, left.name, right.name)) return false;
239 if (!left.value.eql(right.value)) return false;
240 }
241 return true;
242 }
243
244 fn hashAttribute(hasher: *std.hash.Wyhash, attr: ir.Attribute) void {
245 hasher.update(attr.abstract.name);
246
247 if (std.mem.eql(u8, attr.abstract.name, ir.builtin_attr_names.integer)) {
248 const value = attr.cast(ir.Attribute.IntegerAttr).?;
249 hasher.update(std.mem.asBytes(&value.value));
250 hasher.update(std.mem.asBytes(&value.width));
251 hasher.update(std.mem.asBytes(&value.is_signed));
252 return;
253 }
254
255 if (std.mem.eql(u8, attr.abstract.name, ir.builtin_attr_names.float_)) {
256 const value = attr.cast(ir.Attribute.FloatAttr).?;
257 const bits: u64 = @bitCast(value.value);
258 hasher.update(std.mem.asBytes(&bits));
259 hasher.update(std.mem.asBytes(&value.width));
260 return;
261 }
262
263 if (std.mem.eql(u8, attr.abstract.name, ir.builtin_attr_names.bool_)) {
264 const value = attr.cast(ir.Attribute.BoolAttr).?;
265 hasher.update(std.mem.asBytes(&value.value));
266 return;
267 }
268
269 if (std.mem.eql(u8, attr.abstract.name, ir.builtin_attr_names.string)) {
270 const value = attr.cast(ir.Attribute.StringAttr).?;
271 hasher.update(value.value);
272 return;
273 }
274
275 if (std.mem.eql(u8, attr.abstract.name, ir.builtin_attr_names.symbol_ref)) {
276 const value = attr.cast(ir.Attribute.SymbolRefAttr).?;
277 const root_len = value.root_reference.len;
278 hasher.update(std.mem.asBytes(&root_len));
279 hasher.update(value.root_reference);
280 const nested_count = value.nested_references.len;
281 hasher.update(std.mem.asBytes(&nested_count));
282 for (value.nested_references) |item| {
283 const item_len = item.len;
284 hasher.update(std.mem.asBytes(&item_len));
285 hasher.update(item);
286 }
287 return;
288 }
289
290 if (std.mem.eql(u8, attr.abstract.name, ir.builtin_attr_names.string_list)) {
291 const value = attr.cast(ir.Attribute.StringListAttr).?;
292 for (value.values) |item| {
293 hasher.update(item);
294 }
295 return;
296 }
297
298 if (std.mem.eql(u8, attr.abstract.name, ir.builtin_attr_names.type_list)) {
299 const value = attr.cast(ir.Attribute.TypeListAttr).?;
300 for (value.values) |item| {
301 const type_id = item.uniqueId();
302 hasher.update(std.mem.asBytes(&type_id));
303 }
304 return;
305 }
306
307 if (std.mem.eql(u8, attr.abstract.name, ir.builtin_attr_names.array)) {
308 const value = attr.cast(ir.Attribute.ArrayAttr).?;
309 const count = value.values.len;
310 hasher.update(std.mem.asBytes(&count));
311 for (value.values) |item| {
312 hashAttribute(hasher, item);
313 }
314 return;
315 }
316
317 if (attr.cast(ir.Attribute.DialectAttr)) |value| {
318 hasher.update(value.payload);
319 return;
320 }
321
322 const ptr = @intFromPtr(attr.impl);
323 hasher.update(std.mem.asBytes(&ptr));
324 }