tiny.choir.egraph.pattern

1const std = @import("std");2const alloc_arena = @import("alloc_arena");3const ir = @import("../core/root.zig");4const graph_mod = @import("graph.zig");56const ClassId = graph_mod.ClassId;7const Node = graph_mod.Node;8const Graph = graph_mod.Graph;910pub const max_variables = 4;1112pub const Scalar = union(enum) {13    int: i64,14    float: f64,15    boolean: bool,1617    pub fn eql(self: Scalar, other: Scalar) bool {18        if (std.meta.activeTag(self) != std.meta.activeTag(other)) return false;19        return switch (self) {20            .int => |value| value == other.int,21            .float => |value| @as(u64, @bitCast(value)) == @as(u64, @bitCast(other.float)),22            .boolean => |value| value == other.boolean,23        };24    }25};2627pub const TypeClass = enum {28    integer,29    float,30    boolean,31    other,32};3334pub const ClassifyFn = *const fn (ty: ir.Type) TypeClass;35pub const ReadFn = *const fn (attr: ir.Attribute) ?Scalar;36pub const BuildFn = *const fn (ir_ctx: *ir.Context, value: Scalar) anyerror!ir.Attribute;3738pub const ConstantStorage = enum {39    attribute,40    properties,41};4243pub const ConstantModel = struct {44    op_name: []const u8,45    attr_name: []const u8 = "value",46    classify: ClassifyFn,47    read: ReadFn = readScalar,48    build: BuildFn = buildScalar,49    storage: ConstantStorage = .attribute,50};5152pub fn readScalar(attr: ir.Attribute) ?Scalar {53    if (attr.cast(ir.Attribute.IntegerAttr)) |int_attr| return .{ .int = int_attr.getValue() };54    if (attr.cast(ir.Attribute.FloatAttr)) |float_attr| return .{ .float = float_attr.getValue() };55    if (attr.cast(ir.Attribute.BoolAttr)) |bool_attr| return .{ .boolean = bool_attr.getValue() };56    return null;57}5859pub fn buildScalar(ir_ctx: *ir.Context, value: Scalar) anyerror!ir.Attribute {60    return switch (value) {61        .int => |int_value| try ir_ctx.getI64Attr(int_value),62        .float => |float_value| try ir_ctx.getF64Attr(float_value),63        .boolean => |bool_value| try ir_ctx.getBoolAttr(bool_value),64    };65}6667pub const OpShape = struct {68    name: []const u8,69    operands: []const Pattern,70};7172pub const Pattern = union(enum) {73    variable: u8,74    constant: ?Scalar,75    operation: OpShape,76};7778pub const TemplateShape = struct {79    name: []const u8,80    operands: []const Template,81};8283pub const Template = union(enum) {84    variable: u8,85    constant: Scalar,86    operation: TemplateShape,87};8889pub const Rule = struct {90    name: []const u8,91    benefit: u32 = 1,92    classes: ?[]const TypeClass = null,93    lhs: OpShape,94    rhs: Template,95};9697pub const Bindings = struct {98    slots: [max_variables]?ClassId = @as([max_variables]?ClassId, @splat(null)),99100    pub fn get(self: *const Bindings, variable: u8) ?ClassId {101        return self.slots[variable];102    }103104    pub fn set(self: *Bindings, variable: u8, class: ClassId) void {105        self.slots[variable] = class;106    }107};108109pub fn classConstant(graph: *Graph, model: *const ConstantModel, class: ClassId) ?Scalar {110    for (graph.nodes(class)) |*node| {111        if (node.kind != .operation) continue;112        if (!std.mem.eql(u8, node.op_name, model.op_name)) continue;113        const attr = node.getAttr(model.attr_name) orelse continue;114        if (model.read(attr)) |value| return value;115    }116    return null;117}118119pub fn matchNode(120    graph: *Graph,121    model: *const ConstantModel,122    rule: *const Rule,123    node: *const Node,124    bindings: *Bindings,125) bool {126    if (node.kind != .operation) return false;127    if (!std.mem.eql(u8, node.op_name, rule.lhs.name)) return false;128    if (rule.classes) |classes| {129        if (node.result_types.len == 0) return false;130        const type_class = model.classify(node.result_types[0]);131        if (std.mem.indexOfScalar(TypeClass, classes, type_class) == null) return false;132    }133    return matchOperands(graph, model, rule.lhs.operands, node, bindings);134}135136fn matchOperands(137    graph: *Graph,138    model: *const ConstantModel,139    patterns: []const Pattern,140    node: *const Node,141    bindings: *Bindings,142) bool {143    if (patterns.len != node.operands.len) return false;144145    const saved = bindings.*;146    var swapped: [2]ClassId = undefined;147    const attempt_count: usize = if (node.commutative and node.operands.len == 2) 2 else 1;148149    attempts: for (0..attempt_count) |attempt| {150        bindings.* = saved;151        const operands = if (attempt == 0) node.operands else operands: {152            swapped = .{ node.operands[1], node.operands[0] };153            break :operands swapped[0..];154        };155156        for (patterns, operands) |pattern, operand| {157            switch (pattern) {158                .variable => |variable| {159                    if (bindings.get(variable)) |bound| {160                        if (!graph.find(bound).eql(graph.find(operand))) continue :attempts;161                    } else {162                        bindings.set(variable, graph.find(operand));163                    }164                },165                .constant => |expected| {166                    const actual = classConstant(graph, model, operand) orelse continue :attempts;167                    if (expected) |value| {168                        if (!value.eql(actual)) continue :attempts;169                    }170                },171                .operation => |shape| {172                    var matched = false;173                    const class_nodes = graph.nodes(operand);174                    for (class_nodes) |*candidate| {175                        if (candidate.kind != .operation) continue;176                        if (!std.mem.eql(u8, candidate.op_name, shape.name)) continue;177                        if (matchOperands(graph, model, shape.operands, candidate, bindings)) {178                            matched = true;179                            break;180                        }181                    }182                    if (!matched) continue :attempts;183                },184            }185        }186        return true;187    }188189    bindings.* = saved;190    return false;191}192193pub fn instantiate(194    graph: *Graph,195    ir_ctx: *ir.Context,196    model: *const ConstantModel,197    template: Template,198    bindings: *const Bindings,199    result_types: []ir.Type,200) anyerror!ClassId {201    switch (template) {202        .variable => |variable| {203            const bound = bindings.get(variable) orelse return error.UnboundPatternVariable;204            return graph.find(bound);205        },206        .constant => |value| {207            const attr = try model.build(ir_ctx, value);208            var attributes = [_]ir.NamedAttribute{.{ .name = model.attr_name, .value = attr }};209            var node = Node{210                .kind = .operation,211                .op_name = model.op_name,212                .operands = &.{},213                .result_types = result_types,214                .attributes = attributes[0..],215                .raw_attributes = if (model.storage == .attribute) attributes[0..] else &.{},216                .properties = if (model.storage == .properties) attr else null,217                .commutative = false,218            };219            return graph.addNode(&node);220        },221        .operation => |shape| {222            var operand_classes: [max_variables]ClassId = undefined;223            if (shape.operands.len > operand_classes.len) return error.PatternTooWide;224            for (shape.operands, 0..) |operand_template, index| {225                operand_classes[index] = try instantiate(graph, ir_ctx, model, operand_template, bindings, result_types);226            }227228            const commutative = if (ir_ctx.lookupOperation(shape.name)) |info| info.traits.is_commutative else false;229            var node = Node{230                .kind = .operation,231                .op_name = shape.name,232                .operands = operand_classes[0..shape.operands.len],233                .result_types = result_types,234                .attributes = &.{},235                .raw_attributes = &.{},236                .properties = null,237                .commutative = commutative,238            };239            node.normalizeOperands();240            return graph.addNode(&node);241        },242    }243}244245pub fn applyRule(246    graph: *Graph,247    ir_ctx: *ir.Context,248    model: *const ConstantModel,249    rule: *const Rule,250    class: ClassId,251    node: *const Node,252) anyerror!bool {253    var bindings = Bindings{};254    if (!matchNode(graph, model, rule, node, &bindings)) return false;255    const rhs_class = try instantiate(graph, ir_ctx, model, rule.rhs, &bindings, node.result_types);256    return graph.merge(class, rhs_class);257}258259const test_dialect_mod = @import("../dialects/fixture/root.zig");260261fn testClassify(ty: ir.Type) TypeClass {262    const name = ty.getDialectTypeName() orelse return .other;263    if (std.mem.eql(u8, name, "test.i64")) return .integer;264    return .other;265}266267const test_model = ConstantModel{268    .op_name = "test.constant",269    .classify = testClassify,270};271272fn testGraphContext(allocator: std.mem.Allocator) !ir.Context {273    var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);274    errdefer ctx.deinit(allocator);275    try test_dialect_mod.registerTestDialect(&ctx);276    _ = try ctx.registerOperation("test.constant", .{});277    _ = try ctx.registerOperation("test.binary", .{ .is_commutative = true });278    return ctx;279}280281fn addTestConstant(graph: *Graph, ir_ctx: *ir.Context, i64_type: ir.Type, value: i64) !ClassId {282    const attr = try ir_ctx.getI64Attr(value);283    var attributes = [_]ir.NamedAttribute{.{ .name = "value", .value = attr }};284    var result_types = [_]ir.Type{i64_type};285    var node = Node{286        .kind = .operation,287        .op_name = "test.constant",288        .operands = &.{},289        .result_types = result_types[0..],290        .attributes = attributes[0..],291        .raw_attributes = attributes[0..],292        .properties = null,293        .commutative = false,294    };295    return graph.addNode(&node);296}297298fn addTestBinary(graph: *Graph, i64_type: ir.Type, lhs: ClassId, rhs: ClassId) !ClassId {299    var operands = [_]ClassId{ lhs, rhs };300    var result_types = [_]ir.Type{i64_type};301    var node = Node{302        .kind = .operation,303        .op_name = "test.binary",304        .operands = operands[0..],305        .result_types = result_types[0..],306        .attributes = &.{},307        .raw_attributes = &.{},308        .properties = null,309        .commutative = true,310    };311    node.normalizeOperands();312    return graph.addNode(&node);313}314315test "scalar float equality distinguishes signed zero" {316    const positive = Scalar{ .float = 0.0 };317    const negative = Scalar{ .float = -0.0 };318    try std.testing.expect(!positive.eql(negative));319    try std.testing.expect(positive.eql(.{ .float = 0.0 }));320    try std.testing.expect(!positive.eql(.{ .int = 0 }));321}322323const add_zero_rule = Rule{324    .name = "test-binary-zero",325    .classes = &.{.integer},326    .lhs = .{ .name = "test.binary", .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 0 } } } },327    .rhs = .{ .variable = 0 },328};329330const self_to_zero_rule = Rule{331    .name = "test-binary-self-zero",332    .classes = &.{.integer},333    .lhs = .{ .name = "test.binary", .operands = &.{ .{ .variable = 0 }, .{ .variable = 0 } } },334    .rhs = .{ .constant = .{ .int = 0 } },335};336337test "pattern rule merges class with bound variable" {338    const testing = std.testing;339    var arena = alloc_arena.Arena.init(testing.allocator);340    defer arena.deinit();341    const allocator = arena.allocator();342343    var ctx = try testGraphContext(allocator);344    defer ctx.deinit(allocator);345    const i64_type = try test_dialect_mod.TestDialect.getI64Type(&ctx);346347    var graph = Graph.init(allocator);348    defer graph.deinit();349350    const seven = try addTestConstant(&graph, &ctx, i64_type, 7);351    const zero = try addTestConstant(&graph, &ctx, i64_type, 0);352    const sum = try addTestBinary(&graph, i64_type, seven, zero);353354    try testing.expect(!graph.find(sum).eql(graph.find(seven)));355356    const sum_nodes = graph.nodes(sum);357    var fired = false;358    for (sum_nodes) |*node| {359        if (node.kind != .operation) continue;360        if (!std.mem.eql(u8, node.op_name, "test.binary")) continue;361        fired = try applyRule(&graph, &ctx, &test_model, &add_zero_rule, sum, node);362        break;363    }364365    try testing.expect(fired);366    try testing.expect(graph.find(sum).eql(graph.find(seven)));367}368369test "pattern rule with commutative swap matches constant on the left" {370    const testing = std.testing;371    var arena = alloc_arena.Arena.init(testing.allocator);372    defer arena.deinit();373    const allocator = arena.allocator();374375    var ctx = try testGraphContext(allocator);376    defer ctx.deinit(allocator);377    const i64_type = try test_dialect_mod.TestDialect.getI64Type(&ctx);378379    var graph = Graph.init(allocator);380    defer graph.deinit();381382    const zero = try addTestConstant(&graph, &ctx, i64_type, 0);383    const seven = try addTestConstant(&graph, &ctx, i64_type, 7);384    const sum = try addTestBinary(&graph, i64_type, zero, seven);385386    const sum_nodes = graph.nodes(sum);387    var fired = false;388    for (sum_nodes) |*node| {389        if (node.kind != .operation) continue;390        if (!std.mem.eql(u8, node.op_name, "test.binary")) continue;391        fired = try applyRule(&graph, &ctx, &test_model, &add_zero_rule, sum, node);392        break;393    }394395    try testing.expect(fired);396    try testing.expect(graph.find(sum).eql(graph.find(seven)));397}398399test "pattern rule instantiates a fresh constant node" {400    const testing = std.testing;401    var arena = alloc_arena.Arena.init(testing.allocator);402    defer arena.deinit();403    const allocator = arena.allocator();404405    var ctx = try testGraphContext(allocator);406    defer ctx.deinit(allocator);407    const i64_type = try test_dialect_mod.TestDialect.getI64Type(&ctx);408409    var graph = Graph.init(allocator);410    defer graph.deinit();411412    const seven = try addTestConstant(&graph, &ctx, i64_type, 7);413    const sum = try addTestBinary(&graph, i64_type, seven, seven);414415    const sum_nodes = graph.nodes(sum);416    var fired = false;417    for (sum_nodes) |*node| {418        if (node.kind != .operation) continue;419        if (!std.mem.eql(u8, node.op_name, "test.binary")) continue;420        fired = try applyRule(&graph, &ctx, &test_model, &self_to_zero_rule, sum, node);421        break;422    }423424    try testing.expect(fired);425    const merged = classConstant(&graph, &test_model, sum) orelse return error.TestExpectedResult;426    try testing.expect(merged.eql(.{ .int = 0 }));427}428429test "pattern rule respects type class guard" {430    const testing = std.testing;431    var arena = alloc_arena.Arena.init(testing.allocator);432    defer arena.deinit();433    const allocator = arena.allocator();434435    var ctx = try testGraphContext(allocator);436    defer ctx.deinit(allocator);437    const i64_type = try test_dialect_mod.TestDialect.getI64Type(&ctx);438439    var graph = Graph.init(allocator);440    defer graph.deinit();441442    const seven = try addTestConstant(&graph, &ctx, i64_type, 7);443    const zero = try addTestConstant(&graph, &ctx, i64_type, 0);444    const sum = try addTestBinary(&graph, i64_type, seven, zero);445446    const float_guarded = Rule{447        .name = "float-only",448        .classes = &.{.float},449        .lhs = add_zero_rule.lhs,450        .rhs = add_zero_rule.rhs,451    };452453    const sum_nodes = graph.nodes(sum);454    var fired = false;455    for (sum_nodes) |*node| {456        if (node.kind != .operation) continue;457        if (!std.mem.eql(u8, node.op_name, "test.binary")) continue;458        fired = try applyRule(&graph, &ctx, &test_model, &float_guarded, sum, node);459        break;460    }461462    try testing.expect(!fired);463    try testing.expect(!graph.find(sum).eql(graph.find(seven)));464}465466test "nested pattern matches through operand classes" {467    const testing = std.testing;468    var arena = alloc_arena.Arena.init(testing.allocator);469    defer arena.deinit();470    const allocator = arena.allocator();471472    var ctx = try testGraphContext(allocator);473    defer ctx.deinit(allocator);474    const i64_type = try test_dialect_mod.TestDialect.getI64Type(&ctx);475476    var graph = Graph.init(allocator);477    defer graph.deinit();478479    const seven = try addTestConstant(&graph, &ctx, i64_type, 7);480    const zero = try addTestConstant(&graph, &ctx, i64_type, 0);481    const inner = try addTestBinary(&graph, i64_type, seven, zero);482    const outer = try addTestBinary(&graph, i64_type, inner, zero);483484    const nested_rule = Rule{485        .name = "nested-binary-zero",486        .lhs = .{ .name = "test.binary", .operands = &.{487            .{ .operation = .{ .name = "test.binary", .operands = &.{ .{ .variable = 0 }, .{ .constant = .{ .int = 0 } } } } },488            .{ .constant = .{ .int = 0 } },489        } },490        .rhs = .{ .variable = 0 },491    };492493    const outer_nodes = graph.nodes(outer);494    var fired = false;495    for (outer_nodes) |*node| {496        if (node.kind != .operation) continue;497        if (!std.mem.eql(u8, node.op_name, "test.binary")) continue;498        fired = try applyRule(&graph, &ctx, &test_model, &nested_rule, outer, node);499        break;500    }501502    try testing.expect(fired);503    try testing.expect(graph.find(outer).eql(graph.find(seven)));504}

1pub const pattern = @import("pattern.zig");