tiny.pluck.definition_order

1const std = @import("std");2const pexpr = @import("pexpr.zig");34const Allocator = std.mem.Allocator;5const PExpr = pexpr.PExpr;6const Definitions = pexpr.Definitions;7const Symbol = pexpr.Symbol;89pub const DefinitionOrderMode = enum {10    none,11    topological,12    min_fill,13};1415pub const DefinitionOrder = struct {16    map: std.StringHashMapUnmanaged(i32) = .{},1718    pub fn getIndex(self: *const DefinitionOrder, name: Symbol) ?i32 {19        return self.map.get(name);20    }2122    pub fn deinit(self: *DefinitionOrder, allocator: Allocator) void {23        self.map.deinit(allocator);24    }25};2627pub fn modeToString(mode: DefinitionOrderMode) []const u8 {28    return switch (mode) {29        .none => "none",30        .topological => "topological",31        .min_fill => "min-fill",32    };33}3435pub fn buildDefinitionOrder(36    allocator: Allocator,37    definitions: *const Definitions,38    mode: DefinitionOrderMode,39) !?*DefinitionOrder {40    if (mode == .none) return null;4142    const def_names = try collectUserDefNamesSorted(allocator, definitions);43    defer allocator.free(def_names);44    if (def_names.len == 0) return null;4546    const ordered = switch (mode) {47        .topological => try topologicalSortDefs(allocator, definitions, def_names),48        .min_fill => try minFillOrderDefs(allocator, definitions, def_names),49        .none => unreachable,50    };51    defer allocator.free(ordered);5253    const order = try allocator.create(DefinitionOrder);54    order.* = DefinitionOrder{};5556    for (ordered, 0..) |name, idx| {57        try order.map.put(allocator, name, @intCast(idx));58    }5960    return order;61}6263const UserDefinitionNameOrder = struct {64    fn lessThan(_: void, a: Symbol, b: Symbol) bool {65        return std.mem.lessThan(u8, a, b);66    }67};6869fn collectUserDefNamesSorted(allocator: Allocator, definitions: *const Definitions) ![]Symbol {70    var names: std.ArrayList(Symbol) = .empty;71    defer names.deinit(allocator);7273    var iter = definitions.defs.iterator();74    while (iter.next()) |entry| {75        if (entry.value_ptr.is_stdlib) continue;76        try names.append(allocator, entry.value_ptr.name);77    }7879    const slice = try names.toOwnedSlice(allocator);80    std.sort.heap(Symbol, slice, {}, UserDefinitionNameOrder.lessThan);8182    return slice;83}8485fn topologicalSortDefs(86    allocator: Allocator,87    definitions: *const Definitions,88    def_names: []const Symbol,89) ![]Symbol {90    const n = def_names.len;91    if (n == 0) return &[_]Symbol{};9293    var edges = try allocator.alloc(std.ArrayList(usize), n);94    defer {95        for (edges) |*e| e.deinit(allocator);96        allocator.free(edges);97    }98    for (edges) |*e| e.* = .empty;99100    var name_to_idx = std.StringHashMap(usize).init(allocator);101    defer name_to_idx.deinit();102    for (def_names, 0..) |name, i| {103        try name_to_idx.put(name, i);104    }105106    for (def_names, 0..) |name, i| {107        const def = definitions.defs.get(name) orelse continue;108        var deps: std.ArrayList(usize) = .empty;109        defer deps.deinit(allocator);110        collectDefinedRefs(def.expr, &name_to_idx, &deps, allocator);111        for (deps.items) |dep_idx| {112            try edges[dep_idx].append(allocator, i);113        }114    }115116    var in_degree = try allocator.alloc(usize, n);117    defer allocator.free(in_degree);118    @memset(in_degree, 0);119    for (edges) |e| {120        for (e.items) |dep_idx| {121            in_degree[dep_idx] += 1;122        }123    }124125    var queue: std.ArrayList(usize) = .empty;126    defer queue.deinit(allocator);127    for (in_degree, 0..) |deg, i| {128        if (deg == 0) try queue.append(allocator, i);129    }130131    var result = try allocator.alloc(Symbol, n);132    var result_idx: usize = 0;133134    while (queue.items.len > 0) {135        const idx = queue.orderedRemove(0);136        result[result_idx] = def_names[idx];137        result_idx += 1;138139        for (edges[idx].items) |dep_idx| {140            in_degree[dep_idx] -= 1;141            if (in_degree[dep_idx] == 0) {142                try queue.append(allocator, dep_idx);143            }144        }145    }146147    if (result_idx < n) {148        @memcpy(result, def_names);149    }150151    return result;152}153154fn minFillOrderDefs(155    allocator: Allocator,156    definitions: *const Definitions,157    def_names: []const Symbol,158) ![]Symbol {159    const n = def_names.len;160    if (n == 0) return &[_]Symbol{};161162    const adj_size = std.math.mul(usize, n, n) catch return error.OutOfMemory;163    var adj = try allocator.alloc(u8, adj_size);164    defer allocator.free(adj);165    @memset(adj, 0);166167    var name_to_idx = std.StringHashMap(usize).init(allocator);168    defer name_to_idx.deinit();169    for (def_names, 0..) |name, i| {170        try name_to_idx.put(name, i);171    }172173    for (def_names, 0..) |name, i| {174        const def = definitions.defs.get(name) orelse continue;175        var deps: std.ArrayList(usize) = .empty;176        defer deps.deinit(allocator);177        collectDefinedRefs(def.expr, &name_to_idx, &deps, allocator);178179        for (deps.items) |dep_idx| {180            adj[i * n + dep_idx] = 1;181            adj[dep_idx * n + i] = 1;182        }183184        for (deps.items, 0..) |dep_i, idx_i| {185            for (deps.items[idx_i + 1 ..]) |dep_j| {186                adj[dep_i * n + dep_j] = 1;187                adj[dep_j * n + dep_i] = 1;188            }189        }190    }191192    var eliminated = try allocator.alloc(bool, n);193    defer allocator.free(eliminated);194    @memset(eliminated, false);195196    var neighbors_buf = try allocator.alloc(usize, n);197    defer allocator.free(neighbors_buf);198199    var order = try allocator.alloc(Symbol, n);200    var out_idx: usize = 0;201202    while (out_idx < n) : (out_idx += 1) {203        var best_var: ?usize = null;204        var best_fill: usize = std.math.maxInt(usize);205        var best_degree: usize = std.math.maxInt(usize);206207        for (0..n) |v| {208            if (eliminated[v]) continue;209210            var degree: usize = 0;211            for (0..n) |u| {212                if (u == v or eliminated[u]) continue;213                if (adj[v * n + u] != 0) {214                    neighbors_buf[degree] = u;215                    degree += 1;216                }217            }218219            var fill: usize = 0;220            for (neighbors_buf[0..degree], 0..) |u, idx_i| {221                for (neighbors_buf[idx_i + 1 .. degree]) |w| {222                    if (adj[u * n + w] == 0) fill += 1;223                }224            }225226            const better = fill < best_fill or (fill == best_fill and (degree < best_degree or (degree == best_degree and (best_var == null or v < best_var.?))));227            if (better) {228                best_var = v;229                best_fill = fill;230                best_degree = degree;231            }232        }233234        const chosen = best_var orelse break;235236        var degree: usize = 0;237        for (0..n) |u| {238            if (u == chosen or eliminated[u]) continue;239            if (adj[chosen * n + u] != 0) {240                neighbors_buf[degree] = u;241                degree += 1;242            }243        }244245        for (neighbors_buf[0..degree], 0..) |u, idx_i| {246            for (neighbors_buf[idx_i + 1 .. degree]) |w| {247                adj[u * n + w] = 1;248                adj[w * n + u] = 1;249            }250        }251252        for (neighbors_buf[0..degree]) |u| {253            adj[chosen * n + u] = 0;254            adj[u * n + chosen] = 0;255        }256257        eliminated[chosen] = true;258        order[out_idx] = def_names[chosen];259    }260261    return order;262}263264fn collectDefinedRefs(265    expr: *PExpr,266    name_to_idx: *const std.StringHashMap(usize),267    deps: *std.ArrayList(usize),268    allocator: Allocator,269) void {270    switch (expr.head) {271        .defined => |d| {272            if (name_to_idx.get(d.name)) |idx| {273                for (deps.items) |existing| {274                    if (existing == idx) return;275                }276                deps.append(allocator, idx) catch {};277            }278        },279        else => {},280    }281282    for (expr.args) |arg| {283        collectDefinedRefs(arg, name_to_idx, deps, allocator);284    }285}286287test "definition order topological respects dependencies" {288    const allocator = std.testing.allocator;289    var defs = Definitions.init(allocator);290    defer defs.deinit();291292    const expr_b = try PExpr.initWithArgs(allocator, .{ .flip = {} }, &[_]*PExpr{293        try PExpr.initWithArgs(allocator, .{ .const_native = .{ .float = 0.5 } }, &[_]*PExpr{}),294    });295    try defs.define("b", expr_b);296297    const expr_a = try PExpr.initWithArgs(allocator, .{ .defined = .{ .name = "b" } }, &[_]*PExpr{});298    try defs.define("a", expr_a);299300    const order = (try buildDefinitionOrder(allocator, &defs, .topological)).?;301    defer {302        order.deinit(allocator);303        allocator.destroy(order);304    }305306    const idx_b = order.getIndex("b").?;307    const idx_a = order.getIndex("a").?;308    try std.testing.expect(idx_b < idx_a);309}310311test "definition order min-fill chooses low-fill variable first" {312    const allocator = std.testing.allocator;313    var defs = Definitions.init(allocator);314    defer defs.deinit();315316    const expr_d = try PExpr.initWithArgs(allocator, .{ .flip = {} }, &[_]*PExpr{317        try PExpr.initWithArgs(allocator, .{ .const_native = .{ .float = 0.4 } }, &[_]*PExpr{}),318    });319    try defs.define("d", expr_d);320321    const expr_b = try PExpr.initWithArgs(allocator, .{ .defined = .{ .name = "d" } }, &[_]*PExpr{});322    try defs.define("b", expr_b);323324    const expr_c = try PExpr.initWithArgs(allocator, .{ .defined = .{ .name = "d" } }, &[_]*PExpr{});325    try defs.define("c", expr_c);326327    const expr_a = try PExpr.initWithArgs(allocator, .{ .app = {} }, &[_]*PExpr{328        try PExpr.initWithArgs(allocator, .{ .defined = .{ .name = "b" } }, &[_]*PExpr{}),329        try PExpr.initWithArgs(allocator, .{ .defined = .{ .name = "c" } }, &[_]*PExpr{}),330    });331    try defs.define("a", expr_a);332333    const order = (try buildDefinitionOrder(allocator, &defs, .min_fill)).?;334    defer {335        order.deinit(allocator);336        allocator.destroy(order);337    }338339    const idx_a = order.getIndex("a").?;340    try std.testing.expectEqual(@as(i32, 0), idx_a);341}

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