Proposal: restricted function types

[jacobly0]

Preamble

Indirect calls through function pointers are a major blocker for both #1006 and #6025. While a solution can be designed for the generic case, there always seems to be some core function pointer usage, often related to interfaces, that the breaks the generic solution. This proposal is intended to provide a way to represent some common usages of function pointers, particularly in interfaces, in a way that is both more optimizable and decoupled from various difficult design problems..

Motivating Example

Let's consider some simple interface for demonstration purposes:

pub const Operation = struct {
    vtable: *const VTable,

    pub const VTable = struct {
        operate: *const fn (x: i32) i32,
    };

    pub fn operate(op: Operation, x: i32) i32 {
        return op.vtable.operate(x);
    }
};

pub const addOne: Operation = .{ .vtable = &.{
    .operate = &addOneOperate,
} };
fn addOneOperate(x: i32) i32 {
    return x + 1;
}

pub const double: Operation = .{ .vtable = &.{
    .operate = &doubleOperate,
} };
fn doubleOperate(x: i32) i32 {
    return x * 2;
}

fn negate(x: i32) i32 {
    return -x;
}

fn testOperation(name: []const u8, op: Operation) void {
    for ([_]i32{ 1, 7, 12 }) |x| std.debug.print("{s}({d}) = {d}\n", .{ name, x, op.operate(x) });
}

const std = @import("std");
pub fn main() void {
    testOperation("addOne", addOne);
    testOperation("double", double);
    var negate_op: Operation = undefined;
    negate_op.vtable = &.{
        .operate = &negate,
    };
    testOperation("negate", negate_op);
}

$ zig run example.zig
addOne(1) = 2
addOne(7) = 8
addOne(12) = 13
double(1) = 2
double(7) = 14
double(12) = 24
negate(1) = -1
negate(7) = -7
negate(12) = -12

When the optimizer sees the indirect call op.vtable.operate(x), it would be nice if it could know that all of the possible callees are pure, or other useful properties that hold for all callees. This would be easily determinable by iterating all of the callees, but that requires knowing all of the callees in the first place. In a more complicated example, it would be easy for an Operation to come from somewhere where the optimizer is not sure which vtable it holds, and then it is forced to assume that op.vtable.operate(x); could mutate any and all global state.

Proposal

In pursuit of these goals, I propose having a language syntax that allows the compiler to be able to collect a complete list of all possible callees and communicate these to the optimizer at each call site. Here's a tentative syntax proposal, I'm sure it could be improved, but this is just to demonstrate the idea:

 pub const Operation = struct {
     vtable: *const VTable,
 
     pub const VTable = struct {
-        operate: *const fn (x: i32) i32,
+        operate: *const Operate,
+
+        pub const Operate = @RestrictCallees(fn (x: i32) i32);
     };
 
     pub fn operate(op: Operation, x: i32) i32 {
-        return op.vtable.operate(x);
+        return op.vtable.operate(x); // This may only call `addOneOperate`, `doubleOperate`, or `negate`
     }
 };
 
 pub const addOne: Operation = .{ .vtable = &.{
-    .operate = &addOneOperate,
+    .operate = &@allowCallee(addOneOperate),
 } };
 fn addOneOperate(x: i32) i32 {
     return x + 1;
 }
 
 pub const double: Operation = .{ .vtable = &.{
-    .operate = &doubleOperate,
+    .operate = &@allowCallee(doubleOperate),
 } };
 fn doubleOperate(x: i32) i32 {
     return x * 2;
 }
 
 fn negate(x: i32) i32 {
     return -x;
 }
 
 fn testOperation(name: []const u8, op: Operation) void {
     for ([_]i32{ 1, 7, 12 }) |x| std.debug.print("{s}({d}) = {d}\n", .{ name, x, op.operate(x) });
 }
 
 const std = @import("std");
 pub fn main() void {
     testOperation("addOne", addOne);
     testOperation("double", double);
     var negate_op: Operation = undefined;
     negate_op.vtable = &.{
-        .operate = &negate,
+        .operate = &@allowCallee(negate),
     };
     testOperation("negate", negate_op);
 }

The way this works, is that some syntax like @RestrictCallees creates a function type that does not compare equal to other function types with the same signature. Instead, this produces a unique type in the same way as container type definitions, where a combination of the source location and captured values determines equality. Syntax like @allowCallee casts a comptime-known function value into the result restricted function type, checking that the signatures are compatible and adding a callee to the set of possible callees for that restricted function type when analyzed.

It would be checked illegal behavior to make an indirect call through a pointer to a restricted function type when the value of that pointer is not in the set of possible callees that were analyzed during compilation. Additionally, restricted function types would be required to be callconv(.auto), because this feature is not designed to handle function pointers passed to and from external code, and to facilitate the aforementioned safety checks.

As an example implementation in the llvm backend, a metadata node forward reference could be reserved the first time an indirect call to a given restricted function type is encountered, and attached to every indirect call through that restricted function type. Then, during flush, these could be iterated and each metadata node filled in with all of the callees collected by the frontend.

Open Questions

• Does this prevent an interface such as std.mem.Allocator from being passed to another zig compilation unit? I believe this wouldn't work in the future anyway, because non-extern struct types would not be compatible between different ZCUs, due to the not-yet-implemented type tags not matching. This is indeed already illegal behavior.
• How does this interact with incremental? Incrementally adding callees seems simple enough, but can we incrementally remove a callee when all instances of @allowCallee for a given function become no longer referenced? Just leaving it in the callee set wouldn't work since that would cause a false negative a safety check. Using restricted function pointers that are distinct from the original function pointer makes this a non-issue.
• Is it useful to introduce a "restricted function pointer cast" that casts a runtime known normal function pointer into a compatible restricted function pointer and safety checks that the runtime function pointer is one of the analyzed allowed callees? Depending on how the indirect call safety check is implemented, this could be trivial or difficult. Not considered useful at this time.

[mlugg]

Does this prevent an interface such as std.mem.Allocator from being passed to another zig compilation unit?

This is considered IB anyway since the auto-layout struct is internal to a ZCU, as are auto-callconv function pointers. It'd technically work today, but yeah, the type tags in future will definitely be one thing that stops it from working.

How does this interact with incremental?

Removing functions from the set is indeed problematic -- so, here's an idea. What if the return of @allowCallee, rather than pointing to the same function at runtime, pointed to a trampoline to the original function? This trampoline could legally stay in the set across incremental updates, because the only way to get to it (without invoking unchecked IB) is through @allowCallee. Not only does this solve the incremental problem, but it also makes the safety check better when you consider this case:

const Operate = @restrictCallees(fn (u32) void);
fn foo(x: u32) void {
    _ = x; // (this function's body really doesn't matter)
}
test {
    const fooOp: Operate = @allowCallee(foo);
    fooOp(123);
}
test {
    // I'm evil, and really feel like writing some bad code. Let's do an unsafe cast to `Operate`!
    const fooOp: *const Operate = @ptrCast(&foo);
    fooOp(123);
    // We haven't gone through `@allowCallee`, so in isolation, this test would fail, as `foo` isn't in the set.
    // But under this proposal as written, this would actually pass, because the other test in this file adds `foo` to the set.
    // My idea of a trampoline would mean that `foo` and `@allowCallee(foo)` actually have different addresses, so then
    // this test triggers the safety check. That seems better to me.
}

Is it useful to introduce a "restricted function pointer cast" that casts a runtime known normal function pointer into a compatible restricted function pointer and safety checks that the runtime function pointer is one of the analyzed allowed callees?

One downside of my idea above is that it would make this somewhere between "difficult" and "essentially impossible". However, I don't think this is a desirable feature. The whole point here is to determine characteristics of the program which can be statically known at compile-time (but are more complex than we know from the direct function call graph alone); personally I'd need to see a quite compelling use case to add this kind of runtime cast operation into the mix.

[jacobly0]

Is it useful to introduce a "restricted function pointer cast" that casts a runtime known normal function pointer into a compatible restricted function pointer and safety checks that the runtime function pointer is one of the analyzed allowed callees?

One downside of my idea above is that it would make this somewhere between "difficult" and "essentially impossible". However, I don't think this is a desirable feature. The whole point here is to determine characteristics of the program which can be statically known at compile-time (but are more complex than we know from the direct function call graph alone); personally I'd need to see a quite compelling use case to add this kind of runtime cast operation into the mix.

I think that if you want a @ptrCast that hasn't gone through @allowCallee to fail the safety check, which sounds like a good idea, then the concept of a runtime cast doesn't really make sense anymore. Therefore, this features should not be taken into consideration when designing the implementation.

[mlugg]

The proposal mentions that this could go a long way towards solving issues with #1006 and #6025; let me elaborate on that, and discuss how I personally see this working.

#1006 is essentially the same problem as #157; we want to be able to statically determine the amount of stack space required of a function. This issue is also a major blocker for #6025 or similar language-level coroutine ideas. The problem is, function pointers hide a part of the call graph from the compiler, so it can't be statically known. That means we can't statically know the stack size in these cases. Perhaps that wouldn't be a problem if we didn't use (runtime) function pointers much; however, we use them for vtables, a lot! Most notably, std.mem.Allocator is based on vtables. This is a huge problem; for other vtable interfaces we could at least think about having to allocate new stacks to do those calls, but to allocate you need an Allocator, and we can't require allocation before calling into an Allocator.

Enter this proposal. Let's redefine Allocator.VTable to look like this (comments omitted for brevity):

pub const VTable = struct {
    pub const Alloc = @restrictCallees(fn (*anyopaque, len: usize, alignment: Alignment, ret_addr: usize) ?[*]u8);
    pub const Resize = @restrictCallees(fn (*anyopaque, memory: []u8, alignment: Alignment, new_len: usize, ret_addr: usize) bool);
    pub const Remap = @restrictCallees(fn (*anyopaque, memory: []u8, alignment: Alignment, new_len: usize, ret_addr: usize) ?[*]u8);
    pub const Free = @restrictCallees(fn (*anyopaque, memory: []u8, alignment: Alignment, ret_addr: usize) void);
    alloc: *const Alloc,
    resize: *const Resize,
    remap: *const Remap,
    free: *const Free,
};

Anything creating an Allocator.VTable uses @allowCallee on these 4 functions (perhaps we provide a wrapper in Allocator for doing this, I don't know). This way, when the compiler sees a virtual call to a VTable.Alloc, it knows precisely the set of functions that call might be referring to.

Okay, so, let's consider #157. @stackSize on a function which (somewhere in its static call graph) calls into a Allocator.VTable.Free (i.e. by calling gpa.free or similar) would have to consider the full set of possible functions, and from that, determine the worst case of those functions' stack size. Essentially, we extend the call graph with edges going out to every possible callee. This means that if any of those functions can themselves call into Allocator.VTable.Free, we get an error about a recursive call (there is a major unresolved question here which I will cover below), just like you would for direct recursion. Because we need to know the full set of possible callees for this, @stackSize cannot be eagerly known at comptime; we need to wait for all analysis to conclude in order to gather the full set of potential callees to determine the maximum possible stack usage.

Given that @stackSize cannot be eagerly comptime-known, there are two options:

• It is runtime-known.
• It is a comptime-known lazy value, which it is a compile error to ever resolve.

Option 2 doesn't seem very useful in practice; what good really is a comptime-known value if you can't resolve it? Well, I suppose you can store it in a global, but that seems of dubious usefulness. For now, I'll assume the first option (we treat it as runtime-known), but the second might be a better idea, I've not put extensive thought into this.

Let's think about avoiding stack overflow then. First, consider the case where we don't want to use recursion. My knowledge in this area was a little shaky, so I've confirmed with @alexrp that the main thread stack size is basically set in one of two ways:

• We embed it in the binary, and the OS respects this and provides us a stack of this size; for instance, Windows does this.
• The OS provides us a default stack size, and we need to manually resize it at runtime (with the desired stack size which we pull from somewhere); for instance, Linux does this (with setrlimit).

In the first case, we can embed @stackSize(whatever_the_entry_point_is) into the binary; this is fine, because we can resolve this value at the end of semantic analysis and patch it into the binary. In the latter case, std.start code could even just literally use @stackSize(whatever_the_entry_point_is), and reserve that much space. Either way, everything is fine.

Now, let's think about recursion. Direct recursion under #157 / #1006 becomes a compile error, because the required stack size is theoretically infinite. Instead, you need to allocate a new stack for the call somewhere else (perhaps using async/await #6025, or some other mechanism). In this case, what you'd probably do is something like gpa.alloc(u8, @stackSize(recursiveFn)), and use that buffer as the new stack. Of course, alloc itself is a vtable call, but that's fine, because allocators are not recursive (again, there's an unsolved problem here I'll discuss below; if you've guessed it by now, have a cookie!).

So, this proposal allows implementing @stackSize sanely in common vtable cases, and consequently solves #1006. If you have a direct function pointer call for any reason, @stackSize could probably be made to work at runtime for arbitrary runtime function pointers so you can allocate the necessary stack at runtime, but that's definitely to do with #157 / #1006, not this proposal.

---

One major concern I have is how this interacts with arenas. std.heap.ArenaAllocator is an implementation of std.mem.Allocator which can itself call into other Allocator implementations. This is problematic, because it makes all calls through Allocator.VTable potentially recursive. That's not a deficiency of the system, either; you absolutely could nest an arena in an arena in an arena in [...]. (I don't think there's any reasonable use case for this; it would very likely waste a lot of memory due to the arenas nesting their nodes weirdly.)

I'm not quite sure how we could solve this, but I'm almost certain it would involve changing ArenaAllocator rather than amending this proposal in some way. For instance, perhaps ArenaAllocator could become generic over a concrete implementation (potentially concerning for code bloat, but maybe not in practice?); or perhaps it calls into a new ArenaBacking vtable, which is implemented only by "direct" allocators. It'll add a little complexity to the standard library, but I don't think that's necessarily a huge deal given the benefits we get from this proposal.

[mlugg]

After some discussion with @jacobly0, we've figured out a solution to the ArenaAllocator issue which doesn't require any significant API changes. It serves to demonstrate that this is truly a general-purpose mechanism which should work for all kinds of vtable.

Here's how I think about this:

• The general approach to recursion -- which is what arenas are actually effectively doing -- is heap-allocating a new frame for the recursive call
• So, we want to heap-allocate a new frame in e.g. ArenaAllocator.alloc for the recursive call
• However, we can't allocate that frame without calling into Allocator.alloc, so we seem to be stuck
• Key insight: we can allocate that frame ahead-of-time instead!

So, here's what the new ArenaAllocator.init looks like:

fn init(backing: Allocator) Allocator.Error!ArenaAllocator {
    const stack_buf_len = @max(
        @stackSize(Allocator.free), // whatever set of functions arena uses
        @stackSize(Allocator.alloc),
        @stackSize(Allocator.resize),
    );
    const stack_buf = try backing.alloc(u8, stack_buf_len);
    return .{
        .state = .{},
        .backing_allocator = backing,
        .stack_buf = stack_buf,
    };
}

The key insight here is that because ArenaAllocator.init isn't anywhere in any Allocator.VTable, we can safely do this allocation directly. It doesn't matter if backing is another ArenaAllocator, because that allocator already has its own stack_buf for calling its own backing allocator! So, if ArenaAllocator wants to alloc from its backing allocator, that looks something like this (depending on details of #1006):

fn allocFromBacking(arena: ArenaAllocator, n: usize) Allocator.Error![]u8 {
    try @newStackCall(arena.stack_buf, Allocator.alloc, .{ arena.backing_allocator, u8, n });
}

This eliminates the potential recursion loop in ArenaAllocator's vtable in just the same way you would for a simple call graph: heap-allocate the frame. Implementations of Allocator are by far the trickiest example here, because of the catch-22 issue of needing an Allocator to allocate a new stack in the first place; if we've solved this, I'm confident that we've solved essentially all vtables.

[andrewrk]

1. Why the existence of @allowCallee rather than having the same semantics apply upon coercion?
2. What about generalizing this to pointers rather than function bodies? Although I don't have a use case in mind, it may be a syntactically nicer and conceptually simpler place to draw the line.

Regarding the discussion around safe recursion: one unsolved problem that I observe here is that the stack allocation should happen amortized O(1) like how it does in ArrayList, or at least it should be in large batches. From what I see here, what's conceptually solved so far would require to make an allocation for every recursive cycle. For instance, in the classic fibonacci example, it would make an allocation for every function call. That's going to be significantly worse performance that programming languages that have yolo recursion.

[jacobly0]

1. Why the existence of @allowCallee rather than having the same semantics apply upon coercion?

That's certainly a possible implementation, but it doesn't feel much like "restricting" callees anymore, so it would probably go with a wholly different syntax altogether.

2. What about generalizing this to pointers rather than function bodies? Although I don't have a use case in mind, it may be a syntactically nicer and conceptually simpler place to draw the line.

Since it requires comptime-known functions to work, accepting a function body seems like a natural way to implement it. Since it only works with functions and not arbitrary pointers, I don't personally find that conceptually simpler.

Regarding the discussion around safe recursion: one unsolved problem that I observe here is that the stack allocation should happen amortized O(1) like how it does in ArrayList, or at least it should be in large batches. From what I see here, what's conceptually solved so far would require to make an allocation for every recursive cycle. For instance, in the classic fibonacci example, it would make an allocation for every function call. That's going to be significantly worse performance that programming languages that have yolo recursion.

That's just one possible naive allocation mechanism, arena allocators already provide continuous amortized O(1) allocation by allocating a larger chunk each time, and you could pass the same arena all the way through your recursive call graph.

[mlugg]

1. Why the existence of @allowCallee rather than having the same semantics apply upon coercion?

Personally I feel that this operation should be explicit, because it effectively has global side effects; it's not simply converting data from one form to another, which is usually all a coercion really is.

2. What about generalizing this to pointers rather than function bodies?

I mean, we could do, but I don't see any reason to. You can convert between the two incredibly trivially (.* and &), and it feels more conceptually correct for it to work on function bodies, although I don't have a strong justification there (EDIT: Jacob just posted a comment which I think justifies this). I wouldn't really mind if it worked on function pointers instead, but to be honest, this seems to me like a distinction without a difference.

Regarding the discussion around safe recursion: one unsolved problem that I observe here is that the stack allocation should happen amortized O(1) like how it does in ArrayList, or at least it should be in large batches. From what I see here, what's conceptually solved so far would require to make an allocation for every recursive cycle. For instance, in the classic fibonacci example, it would make an allocation for every function call. That's going to be significantly worse performance that programming languages that have yolo recursion.

The obvious thing here is that you can allocate your frames into an arena or similar to amortize the allocation cost. There is still obviously a minor overhead to the frame allocation (a few comparisons and a pointer bump I suppose?), and you could argue that this overhead is a problem. However, that argument doesn't seem very convincing to me, because fibonacci, while a typical example of recursion, isn't actually a realistic example. Any example of recursion which is that simple is usually better to implement in an iterative approach; it'll probably require less memory and be more performant. In my experience, most cases where recursion is a useful tool (for instance, traversing an AST) have relatively few points where recursion happens compared to the amount of logic in between, to the point that the overhead of an amortized arena allocation wouldn't be noticeable. (Bear in mind, you probably don't need much code in the middle to be more significant than a simple arena alloc!)

[andrewrk]

To illustrate the motivation behind my 2 points above, see if you prefer the before or after in these diffs:

        operate: *const Operate,

        pub const Operate = @restrictCallees(fn (x: i32) i32);

⬇

        operate: Operate,

        pub const Operate = @closed(*const fn (x: i32) i32);

and

    .operate = &@allowCallee(addOneOperate),

⬇

    .operate = &addOneOperate,

To me these rules seem easier to reason about:

@restrictCallees can be called on function bodies and creates a new function body type that is different but corresponds to the same function body. That new type can be also used as a function pointer. Later, @allowCallee uses result location syntax that propagates through addr-of operator to add elements to the set.

⬇

@closed makes a distinct pointer type from a pointer type that is the same, except the compiler tracks the set of possible values it can be. Coercion to the closed type requires the operand to be comptime known, adding the value to the set.

It took me 5m to type the first description, thinking hard, and 10s to type the second one, which felt relaxing.

[mlugg]

Ignoring the builtin naming (that can be bikeshed later, I certainly don't love the names in this proposal), I hugely prefer the former:

• This operation, while I suppose you could define it for other pointers, is not helpful in that context; AFAICT this really only aids the optimizer and the language for functions. Your definition just seems like a pointless over-generalization.
• The coercion is doing what amount to an implicit mutation of global state; that makes me a little uneasy. I like this being flagged up by the builtin.

I don't think you're being disingenuous with your definition of the builtins, but I do think you've unintentionally made the former far more complicated than it needs to be. You can follow the exact same scheme as the latter:

@restrictCallees takes a function body type and creates a distinct function body type that is the same, except the compiler tracks the set of all possible values it can be. Function bodies are given this type with the @allowCallee builtin, which requires the operand to be comptime-known and adds the value to the set.

Your definition, for some reason, just clarifies unrelated language semantics:

That new type can be also used as a function pointer.

i.e. "given a type T, there is a type *T" -- we already know this!

[...] uses result location syntax that propagates through addr-of operator [...]

i.e. "result types propagate through expressions" -- we already know this!

[mlugg]

Oh, also, your variant raises issues for our plan for runtime safety, because you might more reasonably expect @intFromPtr on the pointer to give the same address given that a) it looks like a trivial coercion and b) it applies to all pointers; whereas the safety plan, discussed in my first comment, is to make the result of @allowCallee actually point to a different address in safe builds.

[andrewrk]

Your definition, for some reason, just clarifies unrelated language semantics:

I think that's precisely my point. It interacts with other language features, which you have to keep in mind how they work in order to understand this feature. In other words the programmer faces a more mesh-like graph of language features, whereas my counter-proposal has only 1 edge between language features (function pointers and the proposed restricted type feature).

It's easier to reason about a smaller set of language features at the same time.

The coercion is doing what amount to an implicit mutation of global state; that makes me a little uneasy. I like this being flagged up by the builtin.

I don't really buy this reasoning. It's not actually doing that because there is no observable global state (at compile time) introduced by this feature. By the same reasoning, simply introducing another function into the compilation unit is mutating global state.

The proposed allow builtin adds intentional friction against adding a function to a set, but it's completely harmless and safe to add a function to the closed set. I see no reason for this friction. Can you come up with a more concrete reason for this friction to exist?

your variant raises issues for our plan for runtime safety

Have we considered this plan instead?

• Don't allow @ptrCast or @ptrFromInt on this kind of pointer
• No safety checks required