lib/choir/src/core/inspection.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const alloc_arena = @import("alloc_arena");
3 const Operation = @import("operation/root.zig").Operation;
4 const Block = @import("block.zig").Block;
5 const interfaces = @import("interfaces/root.zig");
6
7 pub fn countOperationsNamed(root: *Operation, name: []const u8) usize {
8 var state = CountNamed{ .name = name };
9 _ = root.walk(.{ .order = .pre_order }, &state, CountNamed.visit) catch unreachable;
10 return state.count;
11 }
12
13 pub fn findOperationNamed(root: *Operation, name: []const u8) ?*Operation {
14 var state = FindNamed{ .name = name };
15 _ = root.walk(.{ .order = .pre_order }, &state, FindNamed.visit) catch unreachable;
16 return state.found;
17 }
18
19 pub fn moduleBodyBlock(module: *Operation) ?*Block {
20 const region = module.getRegion(0) orelse return null;
21 return region.getEntryBlock();
22 }
23
24 pub fn isFunctionDefinition(op: *Operation) bool {
25 const iface = functionInterface(op) orelse return false;
26 return iface.call(.hasBody, .{});
27 }
28
29 pub fn functionByName(module: *Operation, name: []const u8) ?*Operation {
30 const block = moduleBodyBlock(module) orelse return null;
31 return functionByNameInBlock(block, name);
32 }
33
34 pub fn functionByNameInBlock(block: *Block, name: []const u8) ?*Operation {
35 var ops = block.getOperations();
36 while (ops.next()) |op| {
37 if (functionInterface(op) == null) continue;
38 const symbol_name = symbolName(op) orelse continue;
39 if (std.mem.eql(u8, symbol_name, name)) return op;
40 }
41 return null;
42 }
43
44 pub fn functionDefinitionByName(module: *Operation, name: []const u8) ?*Operation {
45 const block = moduleBodyBlock(module) orelse return null;
46 return functionDefinitionByNameInBlock(block, name);
47 }
48
49 pub fn functionDefinitionByNameInBlock(block: *Block, name: []const u8) ?*Operation {
50 var ops = block.getOperations();
51 while (ops.next()) |op| {
52 if (!isFunctionDefinition(op)) continue;
53 const symbol_name = symbolName(op) orelse continue;
54 if (std.mem.eql(u8, symbol_name, name)) return op;
55 }
56 return null;
57 }
58
59 fn functionInterface(op: *Operation) ?Operation.interface_handle(interfaces.FunctionOpInterface) {
60 if (symbolName(op) == null) return null;
61 return op.interface(interfaces.FunctionOpInterface);
62 }
63
64 fn symbolName(op: *Operation) ?[]const u8 {
65 const iface = op.interface(interfaces.SymbolOpInterface) orelse return null;
66 return iface.call(.getSymbolName, .{});
67 }
68
69 const CountNamed = struct {
70 name: []const u8,
71 count: usize = 0,
72
73 fn visit(self: *CountNamed, op: *Operation) void {
74 if (std.mem.eql(u8, op.name.name, self.name)) {
75 self.count += 1;
76 }
77 }
78 };
79
80 const FindNamed = struct {
81 name: []const u8,
82 found: ?*Operation = null,
83
84 fn visit(self: *FindNamed, op: *Operation) void {
85 if (self.found != null) return;
86 if (std.mem.eql(u8, op.name.name, self.name)) {
87 self.found = op;
88 }
89 }
90 };
91
92 test "inspection counts operations by name" {
93 const testing = std.testing;
94 const dialects = @import("../dialects/root.zig");
95 const test_dialect = @import("../dialects/fixture/root.zig");
96 const ir = @import("root.zig");
97
98 var arena = alloc_arena.Arena.init(testing.allocator);
99 defer arena.deinit();
100
101 var ctx = try ir.Context.init(arena.allocator(), ir.Context.Limits.testing);
102 defer ctx.deinit(arena.allocator());
103 try ctx.allowUnregistered();
104 try test_dialect.registerTestDialect(&ctx);
105
106 const loc = ir.Location.getUnknown();
107 const module = try dialects.BuiltinDialect.ModuleOp.create(&ctx, loc);
108 const block = module.getBodyBlock();
109
110 const left = try ctx.createOperation(ir.Operation.State.init("test.leaf", loc));
111 const right = try ctx.createOperation(ir.Operation.State.init("test.leaf", loc));
112 const other = try ctx.createOperation(ir.Operation.State.init("test.other", loc));
113 try block.addOperation(left);
114 try block.addOperation(other);
115 try block.addOperation(right);
116
117 try testing.expectEqual(@as(usize, 2), countOperationsNamed(module.op, "test.leaf"));
118 try testing.expectEqual(@as(usize, 1), countOperationsNamed(module.op, "test.other"));
119 try testing.expectEqual(@as(usize, 0), countOperationsNamed(module.op, "test.missing"));
120 try testing.expect(findOperationNamed(module.op, "test.leaf") == left);
121 try testing.expect(findOperationNamed(module.op, "test.other") == other);
122 try testing.expect(findOperationNamed(module.op, "test.missing") == null);
123 }
124
125 test "inspection finds function definitions by symbol name" {
126 const testing = std.testing;
127 const dialects = @import("../dialects/root.zig");
128 const ir = @import("root.zig");
129
130 var arena = alloc_arena.Arena.init(testing.allocator);
131 defer arena.deinit();
132
133 var ctx = try ir.Context.init(arena.allocator(), ir.Context.Limits.testing);
134 defer ctx.deinit(arena.allocator());
135
136 const loc = ir.Location.getUnknown();
137 const module = try dialects.BuiltinDialect.ModuleOp.create(&ctx, loc);
138 const block = module.getBodyBlock();
139
140 const body = try dialects.FuncDialect.FuncOp.create(&ctx, loc, "body", &.{}, &.{});
141 const declaration = try dialects.FuncDialect.FuncOp.createDeclaration(&ctx, loc, "declaration", &.{}, &.{});
142 try block.addOperation(declaration.op);
143 try block.addOperation(body.op);
144
145 try testing.expect(functionByName(module.op, "body") == body.op);
146 try testing.expect(functionByNameInBlock(block, "body") == body.op);
147 try testing.expect(functionByName(module.op, "declaration") == declaration.op);
148 try testing.expect(functionDefinitionByName(module.op, "body") == body.op);
149 try testing.expect(functionDefinitionByNameInBlock(block, "body") == body.op);
150 try testing.expect(functionDefinitionByName(module.op, "declaration") == null);
151 try testing.expect(functionByName(module.op, "missing") == null);
152 }