tiny.choir.inspection

1const std = @import("std");2const alloc_arena = @import("alloc_arena");3const Operation = @import("operation/root.zig").Operation;4const Block = @import("block.zig").Block;5const interfaces = @import("interfaces/root.zig");67pub 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}1213pub 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}1819pub fn moduleBodyBlock(module: *Operation) ?*Block {20    const region = module.getRegion(0) orelse return null;21    return region.getEntryBlock();22}2324pub fn isFunctionDefinition(op: *Operation) bool {25    const iface = functionInterface(op) orelse return false;26    return iface.call(.hasBody, .{});27}2829pub fn functionByName(module: *Operation, name: []const u8) ?*Operation {30    const block = moduleBodyBlock(module) orelse return null;31    return functionByNameInBlock(block, name);32}3334pub 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}4344pub fn functionDefinitionByName(module: *Operation, name: []const u8) ?*Operation {45    const block = moduleBodyBlock(module) orelse return null;46    return functionDefinitionByNameInBlock(block, name);47}4849pub 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}5859fn functionInterface(op: *Operation) ?Operation.interface_handle(interfaces.FunctionOpInterface) {60    if (symbolName(op) == null) return null;61    return op.interface(interfaces.FunctionOpInterface);62}6364fn symbolName(op: *Operation) ?[]const u8 {65    const iface = op.interface(interfaces.SymbolOpInterface) orelse return null;66    return iface.call(.getSymbolName, .{});67}6869const CountNamed = struct {70    name: []const u8,71    count: usize = 0,7273    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};7980const FindNamed = struct {81    name: []const u8,82    found: ?*Operation = null,8384    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};9192test "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");9798    var arena = alloc_arena.Arena.init(testing.allocator);99    defer arena.deinit();100101    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);105106    const loc = ir.Location.getUnknown();107    const module = try dialects.BuiltinDialect.ModuleOp.create(&ctx, loc);108    const block = module.getBodyBlock();109110    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);116117    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}124125test "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");129130    var arena = alloc_arena.Arena.init(testing.allocator);131    defer arena.deinit();132133    var ctx = try ir.Context.init(arena.allocator(), ir.Context.Limits.testing);134    defer ctx.deinit(arena.allocator());135136    const loc = ir.Location.getUnknown();137    const module = try dialects.BuiltinDialect.ModuleOp.create(&ctx, loc);138    const block = module.getBodyBlock();139140    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);144145    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}

1pub const inspection = ir.inspection;