lib/pluck/src/order.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const pexpr = @import("pexpr.zig");
3
4 const Allocator = std.mem.Allocator;
5 const PExpr = pexpr.PExpr;
6 const Definitions = pexpr.Definitions;
7 const Symbol = pexpr.Symbol;
8
9 pub const DefinitionOrderMode = enum {
10 none,
11 topological,
12 min_fill,
13 };
14
15 pub const DefinitionOrder = struct {
16 map: std.StringHashMapUnmanaged(i32) = .{},
17
18 pub fn getIndex(self: *const DefinitionOrder, name: Symbol) ?i32 {
19 return self.map.get(name);
20 }
21
22 pub fn deinit(self: *DefinitionOrder, allocator: Allocator) void {
23 self.map.deinit(allocator);
24 }
25 };
26
27 pub fn modeToString(mode: DefinitionOrderMode) []const u8 {
28 return switch (mode) {
29 .none => "none",
30 .topological => "topological",
31 .min_fill => "min-fill",
32 };
33 }
34
35 pub fn buildDefinitionOrder(
36 allocator: Allocator,
37 definitions: *const Definitions,
38 mode: DefinitionOrderMode,
39 ) !?*DefinitionOrder {
40 if (mode == .none) return null;
41
42 const def_names = try collectUserDefNamesSorted(allocator, definitions);
43 defer allocator.free(def_names);
44 if (def_names.len == 0) return null;
45
46 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);
52
53 const order = try allocator.create(DefinitionOrder);
54 order.* = DefinitionOrder{};
55
56 for (ordered, 0..) |name, idx| {
57 try order.map.put(allocator, name, @intCast(idx));
58 }
59
60 return order;
61 }
62
63 const UserDefinitionNameOrder = struct {
64 fn lessThan(_: void, a: Symbol, b: Symbol) bool {
65 return std.mem.lessThan(u8, a, b);
66 }
67 };
68
69 fn collectUserDefNamesSorted(allocator: Allocator, definitions: *const Definitions) ![]Symbol {
70 var names: std.ArrayList(Symbol) = .empty;
71 defer names.deinit(allocator);
72
73 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 }
78
79 const slice = try names.toOwnedSlice(allocator);
80 std.sort.heap(Symbol, slice, {}, UserDefinitionNameOrder.lessThan);
81
82 return slice;
83 }
84
85 fn 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{};
92
93 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;
99
100 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 }
105
106 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 }
115
116 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 }
124
125 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 }
130
131 var result = try allocator.alloc(Symbol, n);
132 var result_idx: usize = 0;
133
134 while (queue.items.len > 0) {
135 const idx = queue.orderedRemove(0);
136 result[result_idx] = def_names[idx];
137 result_idx += 1;
138
139 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 }
146
147 if (result_idx < n) {
148 @memcpy(result, def_names);
149 }
150
151 return result;
152 }
153
154 fn 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{};
161
162 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);
166
167 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 }
172
173 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);
178
179 for (deps.items) |dep_idx| {
180 adj[i * n + dep_idx] = 1;
181 adj[dep_idx * n + i] = 1;
182 }
183
184 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 }
191
192 var eliminated = try allocator.alloc(bool, n);
193 defer allocator.free(eliminated);
194 @memset(eliminated, false);
195
196 var neighbors_buf = try allocator.alloc(usize, n);
197 defer allocator.free(neighbors_buf);
198
199 var order = try allocator.alloc(Symbol, n);
200 var out_idx: usize = 0;
201
202 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);
206
207 for (0..n) |v| {
208 if (eliminated[v]) continue;
209
210 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 }
218
219 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 }
225
226 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 }
233
234 const chosen = best_var orelse break;
235
236 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 }
244
245 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 }
251
252 for (neighbors_buf[0..degree]) |u| {
253 adj[chosen * n + u] = 0;
254 adj[u * n + chosen] = 0;
255 }
256
257 eliminated[chosen] = true;
258 order[out_idx] = def_names[chosen];
259 }
260
261 return order;
262 }
263
264 fn 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 }
281
282 for (expr.args) |arg| {
283 collectDefinedRefs(arg, name_to_idx, deps, allocator);
284 }
285 }
286
287 test "definition order topological respects dependencies" {
288 const allocator = std.testing.allocator;
289 var defs = Definitions.init(allocator);
290 defer defs.deinit();
291
292 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);
296
297 const expr_a = try PExpr.initWithArgs(allocator, .{ .defined = .{ .name = "b" } }, &[_]*PExpr{});
298 try defs.define("a", expr_a);
299
300 const order = (try buildDefinitionOrder(allocator, &defs, .topological)).?;
301 defer {
302 order.deinit(allocator);
303 allocator.destroy(order);
304 }
305
306 const idx_b = order.getIndex("b").?;
307 const idx_a = order.getIndex("a").?;
308 try std.testing.expect(idx_b < idx_a);
309 }
310
311 test "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();
315
316 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);
320
321 const expr_b = try PExpr.initWithArgs(allocator, .{ .defined = .{ .name = "d" } }, &[_]*PExpr{});
322 try defs.define("b", expr_b);
323
324 const expr_c = try PExpr.initWithArgs(allocator, .{ .defined = .{ .name = "d" } }, &[_]*PExpr{});
325 try defs.define("c", expr_c);
326
327 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);
332
333 const order = (try buildDefinitionOrder(allocator, &defs, .min_fill)).?;
334 defer {
335 order.deinit(allocator);
336 allocator.destroy(order);
337 }
338
339 const idx_a = order.getIndex("a").?;
340 try std.testing.expectEqual(@as(i32, 0), idx_a);
341 }