tiny.choir.egraph.graph

1const std = @import("std");2const ir = @import("../core/root.zig");3const node_mod = @import("node.zig");45pub const ClassId = node_mod.ClassId;6pub const ValueEntry = node_mod.ValueEntry;7pub const Node = node_mod.Node;89const Class = struct {10    parent: ClassId,11    rank: u8 = 0,12    nodes: std.ArrayListUnmanaged(Node) = .empty,13    values: std.ArrayListUnmanaged(ValueEntry) = .empty,1415    fn init(id: ClassId) Class {16        return .{ .parent = id };17    }1819    fn deinit(self: *Class, allocator: std.mem.Allocator) void {20        for (self.nodes.items) |*node| {21            node.deinit(allocator);22        }23        self.nodes.deinit(allocator);24        self.values.deinit(allocator);25    }26};2728pub const GraphStats = struct {29    classes_created: usize = 0,30    nodes_added: usize = 0,31    unions: usize = 0,32    rebuilds: usize = 0,33};3435pub const Graph = struct {36    allocator: std.mem.Allocator,37    classes: std.ArrayListUnmanaged(Class) = .empty,38    memo: std.AutoHashMapUnmanaged(u64, std.ArrayListUnmanaged(ClassId)) = .empty,39    stats: GraphStats = .{},4041    /// Cumulative allocation traffic for a graph whose rules only merge existing42    /// classes. Counts cover all blocks together; each class and collision bucket43    /// may grow to the entire population. Rebuild passes include terminal scans.44    pub fn mergeStorageBound(45        node_count: u64,46        atoms: u64,47        graphs: u64,48        rebuild_passes: u64,49    ) !u64 {50        if (node_count > std.math.maxInt(u32)) return error.CapacityOverflow;51        const classes = try mul(graphs, try listStorage(Class, node_count));52        const contents = try mul(node_count, try sum(53            try listStorage(Node, node_count),54            try listStorage(ValueEntry, node_count),55        ));56        const payload_item = @sizeOf(ClassId) + @sizeOf(ir.Type) +57            2 * @sizeOf(ir.NamedAttribute) + 64;58        const payload = try mul(try mul(2, atoms), payload_item);59        const table = try mul(graphs, try memoStorage(node_count));60        const buckets = try mul(node_count, try listStorage(ClassId, 1));61        const pairs = try listStorage(struct { lhs: ClassId, rhs: ClassId }, node_count);62        const rebuilds = try mul(rebuild_passes, try sum(buckets, pairs));63        return sum(try sum(classes, contents), try sum(payload, try sum(64            table,65            try sum(buckets, rebuilds),66        )));67    }6869    pub fn init(allocator: std.mem.Allocator) Graph {70        return .{ .allocator = allocator };71    }7273    pub fn deinit(self: *Graph) void {74        self.clearMemo();75        self.memo.deinit(self.allocator);76        for (self.classes.items) |*class| {77            class.deinit(self.allocator);78        }79        self.classes.deinit(self.allocator);80    }8182    pub fn find(self: *Graph, id: ClassId) ClassId {83        const idx: usize = @intCast(id.index);84        const parent = self.classes.items[idx].parent;85        if (parent.eql(id)) return id;86        const root = self.find(parent);87        self.classes.items[idx].parent = root;88        return root;89    }9091    pub fn classCount(self: *const Graph) usize {92        return self.classes.items.len;93    }9495    pub fn addValue(self: *Graph, value: *ir.Value, cost: u32, order: usize) !ClassId {96        var node = Node.valueNode(value);97        const id = try self.addNode(&node);98        try self.attachValue(id, value, cost, order);99        return id;100    }101102    pub fn addOperation(103        self: *Graph,104        op: *ir.Operation,105        operands: []const ClassId,106        cost: u32,107        order: usize,108    ) !ClassId {109        var node = try Node.operationNode(self.allocator, op, operands, op.hasTrait("is_commutative"));110        defer node.deinit(self.allocator);111        self.canonicalizeNode(&node);112        const id = try self.addNode(&node);113        if (op.getResult(0)) |value| {114            try self.attachValue(id, value, cost, order);115        }116        return id;117    }118119    pub fn addNode(self: *Graph, node: *const Node) !ClassId {120        const hash = node.hash();121        if (self.memo.get(hash)) |ids| {122            for (ids.items) |candidate| {123                const root = self.find(candidate);124                const class = &self.classes.items[@intCast(root.index)];125                for (class.nodes.items) |*existing| {126                    if (existing.eql(node)) return root;127                }128            }129        }130131        const id = ClassId{ .index = @intCast(self.classes.items.len) };132        var class = Class.init(id);133        var class_owned = true;134        errdefer if (class_owned) class.deinit(self.allocator);135136        var owned = try node.clone(self.allocator);137        var owned_in_class = false;138        errdefer if (!owned_in_class) owned.deinit(self.allocator);139        try class.nodes.append(self.allocator, owned);140        owned_in_class = true;141142        try self.classes.append(self.allocator, class);143        class_owned = false;144        try self.insertMemo(hash, id);145        self.stats.classes_created += 1;146        self.stats.nodes_added += 1;147        return id;148    }149150    pub fn attachValue(self: *Graph, id: ClassId, value: *ir.Value, cost: u32, order: usize) !void {151        const root = self.find(id);152        try self.attachValueToRoot(root, .{ .value = value, .cost = cost, .order = order });153    }154155    pub fn representativeValue(self: *Graph, id: ClassId) ?ValueEntry {156        const root = self.find(id);157        const class = &self.classes.items[@intCast(root.index)];158        if (class.values.items.len == 0) return null;159160        var best = class.values.items[0];161        for (class.values.items[1..]) |candidate| {162            if (candidate.cost < best.cost or163                (candidate.cost == best.cost and candidate.order < best.order))164            {165                best = candidate;166            }167        }168        return best;169    }170171    pub fn nodes(self: *Graph, id: ClassId) []const Node {172        const root = self.find(id);173        return self.classes.items[@intCast(root.index)].nodes.items;174    }175176    pub fn classValues(self: *Graph, id: ClassId) []const ValueEntry {177        const root = self.find(id);178        return self.classes.items[@intCast(root.index)].values.items;179    }180181    pub fn merge(self: *Graph, lhs: ClassId, rhs: ClassId) !bool {182        var lhs_root = self.find(lhs);183        var rhs_root = self.find(rhs);184        if (lhs_root.eql(rhs_root)) return false;185186        var lhs_class = &self.classes.items[@intCast(lhs_root.index)];187        var rhs_class = &self.classes.items[@intCast(rhs_root.index)];188        if (lhs_class.rank < rhs_class.rank) {189            const tmp_root = lhs_root;190            lhs_root = rhs_root;191            rhs_root = tmp_root;192193            const tmp_class = lhs_class;194            lhs_class = rhs_class;195            rhs_class = tmp_class;196        }197198        try lhs_class.nodes.appendSlice(self.allocator, rhs_class.nodes.items);199        rhs_class.nodes.clearRetainingCapacity();200201        for (rhs_class.values.items) |entry| {202            try self.attachValueToRoot(lhs_root, entry);203        }204        rhs_class.values.clearRetainingCapacity();205206        rhs_class.parent = lhs_root;207        if (lhs_class.rank == rhs_class.rank) {208            lhs_class.rank += 1;209        }210211        self.stats.unions += 1;212        return true;213    }214215    pub fn rebuild(self: *Graph) !bool {216        var any_changed = false;217        while (try self.rebuildOnce()) {218            any_changed = true;219            self.stats.rebuilds += 1;220        }221        return any_changed;222    }223224    fn attachValueToRoot(self: *Graph, root: ClassId, entry: ValueEntry) !void {225        const class = &self.classes.items[@intCast(root.index)];226        for (class.values.items) |*existing| {227            if (existing.value == entry.value) {228                if (entry.cost < existing.cost or229                    (entry.cost == existing.cost and entry.order < existing.order))230                {231                    existing.cost = entry.cost;232                    existing.order = entry.order;233                }234                return;235            }236        }237        try class.values.append(self.allocator, entry);238    }239240    fn canonicalizeNode(self: *Graph, node: *Node) void {241        if (node.kind != .operation) return;242        for (node.operands) |*operand| {243            operand.* = self.find(operand.*);244        }245        node.normalizeOperands();246    }247248    fn rebuildOnce(self: *Graph) !bool {249        self.clearMemo();250251        var pairs: std.ArrayListUnmanaged(struct { lhs: ClassId, rhs: ClassId }) = .empty;252        defer pairs.deinit(self.allocator);253254        for (self.classes.items, 0..) |*class, index| {255            const id = ClassId{ .index = @intCast(index) };256            const root = self.find(id);257            if (!root.eql(id)) continue;258259            for (class.nodes.items) |*node| {260                self.canonicalizeNode(node);261                const hash = node.hash();262                var found: ?ClassId = null;263                if (self.memo.get(hash)) |ids| {264                    for (ids.items) |candidate| {265                        const candidate_root = self.find(candidate);266                        const candidate_class = &self.classes.items[@intCast(candidate_root.index)];267                        for (candidate_class.nodes.items) |*existing| {268                            if (existing.eql(node)) {269                                found = candidate_root;270                                break;271                            }272                        }273                        if (found != null) break;274                    }275                }276277                if (found) |existing| {278                    if (!existing.eql(root)) {279                        try pairs.append(self.allocator, .{ .lhs = root, .rhs = existing });280                    }281                } else {282                    try self.insertMemo(hash, root);283                }284            }285        }286287        var changed = false;288        for (pairs.items) |pair| {289            if (try self.merge(pair.lhs, pair.rhs)) {290                changed = true;291            }292        }293        return changed;294    }295296    fn insertMemo(self: *Graph, hash: u64, id: ClassId) !void {297        var gop = try self.memo.getOrPut(self.allocator, hash);298        if (!gop.found_existing) {299            gop.value_ptr.* = .empty;300        }301        const root = self.find(id);302        for (gop.value_ptr.items) |existing| {303            if (self.find(existing).eql(root)) return;304        }305        try gop.value_ptr.append(self.allocator, root);306    }307308    fn clearMemo(self: *Graph) void {309        var iter = self.memo.valueIterator();310        while (iter.next()) |ids| {311            ids.deinit(self.allocator);312        }313        self.memo.clearRetainingCapacity();314    }315};316317fn sum(a: u64, b: u64) !u64 {318    return std.math.add(u64, a, b) catch error.CapacityOverflow;319}320321fn mul(a: u64, b: u64) !u64 {322    return std.math.mul(u64, a, b) catch error.CapacityOverflow;323}324325fn listStorage(comptime Item: type, count: u64) !u64 {326    if (count == 0) return 0;327    return mul(try mul(4, try sum(try mul(4, count), 64)), @sizeOf(Item) + @alignOf(Item));328}329330fn memoStorage(count: u64) !u64 {331    if (count == 0) return 0;332    const required = try sum(try mul(count, 100) / 80, 1);333    if (required > std.math.maxInt(u32)) return error.CapacityOverflow;334    const capacity = std.math.ceilPowerOfTwo(u32, @intCast(required)) catch335        return error.CapacityOverflow;336    if (capacity > std.math.maxInt(u32) / 80) return error.CapacityOverflow;337    const bytes = 1 + @sizeOf(u64) + @sizeOf(std.ArrayListUnmanaged(ClassId)) +338        4 * @sizeOf(usize) + 3 * @alignOf(usize);339    return mul(try mul(2, @max(8, capacity)), bytes);340}341342test "saturation graph storage covers collisions merges and repeated rebuilds" {343    for ([_]u32{ 1, 8, 64 }) |count| try checkMergeStorage(count);344    try std.testing.expectError(345        error.CapacityOverflow,346        Graph.mergeStorageBound(std.math.maxInt(u64), 1, 1, 1),347    );348}349350fn checkMergeStorage(count: u32) !void {351    const allocator = std.testing.allocator;352    const fixed = @import("alloc_fixed");353    var context = try ir.Context.init(allocator, ir.Context.Limits.testing);354    defer context.deinit(allocator);355    const bound = try Graph.mergeStorageBound(count, count, 1, count + 2);356    const storage = try allocator.alignedAlloc(u8, .@"64", @intCast(bound));357    defer allocator.free(storage);358    var backing = fixed.Tracked.init(storage);359    var retained = fixed.Monotonic.init(backing.allocator(), @intCast(bound));360    var graph = Graph.init(retained.allocator());361    defer graph.deinit();362    for (0..count) |index| {363        var attrs = [_]ir.NamedAttribute{.{364            .name = "value",365            .value = try context.getI64Attr(@intCast(index)),366        }};367        const node = Node{ .kind = .operation, .op_name = "test.node", .attributes = &attrs };368        _ = try graph.addNode(&node);369    }370    graph.clearMemo();371    for (0..count) |index| try graph.insertMemo(0, .{ .index = @intCast(index) });372    try std.testing.expectEqual(count, graph.memo.get(0).?.items.len);373    _ = try graph.rebuild();374    for (1..count) |index| {375        try std.testing.expect(try graph.merge(.{ .index = 0 }, .{ .index = @intCast(index) }));376        _ = try graph.rebuild();377    }378    try std.testing.expectEqual(count, graph.nodes(.{ .index = 0 }).len);379    try std.testing.expect(!backing.exhausted);380    const used = if (retained.current) |*current| fixed.used(current) else 0;381    try std.testing.expect(used <= bound);382}

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