A registry of call-sites

I want to collect all call-sites where a particular function is called. Roughly, I want to sprinkle

mark(@src());

calls around my code, and get something like

pub const mark_locations: []Location = ???

This is the same shaper of the problem as @import("builtin").test_functions.

Are there some ready to use recipes for this in Zig?

I know that this problem can be solved at the linker level. Basically, when you declare a linker symbol, you can instruct the liner “please put this symbol to an array with all similar symbols”, and then get a handle to that array. linkme - Rust implements this for Rust.

Anything like that in Zig?

I think I have an implementation, but I’m only searching the AST by module_name.function_name.

https://codeberg.org/Maarrk/polpo/src/branch/master/src/xgettext.zig

This naive solution has been working for me since then, but I’m looking forward to replacing it with something less fragile.

That would be a perfect fit for function instrumentation. I wrote a profiler a while back, where I tracked the path through the code with --finstrument-functions, but that was C and Fortran code. Would be nice if it would be possible in zig as well.

Are you able to talk at all about why you want this? Feels like a potential XY problem.

I am writing a fuzz test, and I want the test to be able to reflect on the runtime behavior of the program under test. Essentially, coverage, but:

  • I only care about specific logical conditions in my code, not line coverage,
  • I want this to work with vanilla zig build test, without any extra runtime libraries, processes, or compilation flags
  • I want to feed the information into the running program as it runs (so, it needs to be an array of counters, rather than an .html file, and I need to get it from inside the process).
  • The result should be deterministic across targets and optimization modes

This is horrifying, I like it! My take on the idea:

const std = @import("std");
const SourceLocation = std.builtin.SourceLocation;

pub fn main(init: std.process.Init.Minimal) void {
    if (init.args.vector.len > 1) {
        hit(@src());
    } else {
        hit(@src());
    }
    for (SrcRegistry.global.entries) |entry| {
        std.debug.print("line {} hit {} times\n", .{ entry.line_number, entry.value });
    }
}

fn hit(comptime src: SourceLocation) void {
    SrcRegistry.global.resolve(src).value += 1;
}
λ ./zig/zig run main.zig -- hello
line 6 hit 1 times
line 8 hit 0 times
line 27 hit 0 times
const SrcRegistry = struct {
    entries: []Entry,

    var global: SrcRegistry = .{
        .entries = struct {
            var array: [array_count(this_file())]Entry = array_init(this_file());

            fn this_file() []const u8 {
                return @embedFile(@src().file);
            }

            fn array_count(source: []const u8) u32 {
                @setEvalBranchQuota(1_000_000);
                return std.mem.count(u8, source, "@" ++ "src");
            }

            fn array_init(source: []const u8) [array_count(source)]Entry {
                @setEvalBranchQuota(1_000_000);
                var result: [array_count(source)]Entry = undefined;
                var rest = source;
                var line_index: u32 = 0;
                for (&result) |*entry| {
                    const prefix, rest = std.mem.cut(u8, rest, "@" ++ "src()").?;
                    line_index += std.mem.countScalar(u8, prefix, '\n');
                    const line_number = line_index + 1;
                    entry.* = Entry.init(line_number);
                }
                return result;
            }
        }.array[0..],
    };

    const Value = u32;
    const Entry = struct {
        line_number: u32,
        value: Value,

        fn init(line_number: u32) Entry {
            return .{ .line_number = line_number, .value = 0 };
        }
    };

    fn resolve(registry: *SrcRegistry, comptime src: SourceLocation) *Entry {
        const Cache = struct {
            var global: ?*Entry = null;

            // Ensure that each `src` gets a unique Cache
            const capture: SourceLocation = src;
        };

        if (Cache.global) |cached| {
            return cached;
        } else {
            const computed = for (registry.entries) |*entry| {
                if (entry.line_number == src.line) {
                    break entry;
                }
            } else unreachable;
            Cache.global = computed;
            return computed;
        }
    }
};

Oh dear! Unlike @import, @embedFile doesn’t require the path to be a string literal, it can be computed:

pub fn my_path() []const u8 {
    return @src().file;
}

pub fn my_text(comptime path: []const u8) []const u8 {
    return @embedFile(path);
}

pub fn main() void {
    std.debug.print("{s}\n", .{my_text(comptime my_path())});
}

So actually you can put SrcRegistry in the root source file, and have it, at comptime, recursively @embedFile everything that’s imported, by looking for @import statements (which must use a string literal).

Have you checked if the entire file ends up in the binary? Or does it get removed at some stage?

also I guess that using @src doesn’t work so well with generic functions.

It’s cool you made it work though, but this make me wish for builtin support

It’s a shame @src()doesn’t work outside a function scope, if it did you’d be able to define a single const MyGenericT = MyGeneric(T, @src());once, and then use the two source locations together to uniquely identify a more specific combination.

I guess you could do such an alias as a parameterless type function though?

I have a similar mechanism to what you described in my project which I call a “linker set”, see here, the idea is that you define a linker set and at comptime you can insert objects into it, e.g:

const my_set = LinkerSet("name", *u8);
var my_variable: u8 = 0;

comptime {
    my_set.insert(&my_variable);
}

And you can then iterate over those easily via a slice.

Using a linker section and its encapsulation symbols is brilliant!

Make the callee noinline and then use @returnAddress. It’s remarkably portable and it gives you a usize corresponding to the call machine code instruction at the call site. You can later resolve these usize values into source locations (line, column, file) based on debug info.

The catch is, I want to know that a call-site exists even if it was never called at runtime.