Commit c14a630041a196c6af4959c789d90178e0fe7ae1 | zig

author	lg <abbix@riseup.net> 2026-08-31 11:54:22+02:00
committer	Matthew Lugg <mlugg@mlugg.co.uk> 2026-08-31 11:54:22+02:00
log	c14a630041a196c6af4959c789d90178e0fe7ae1
tree	00932ac3c10efa411d3d76bc8556ddb13ceccbc0
parent	fd3048c5a91220102420ef51ba5193d6febca521

Sema: allow comptime call memoization when callee uses `@setEvalBranchQuota` (#36494)

Previously, memoization was disabled as soon as setEvalBranchQuota was
called. This commit keeps memoization enabled when quota is being set by keeping
track of the requested quota in the call itself and externally.

This has shown significant improvements when calling functions where the callee calls `@setEvalBranchQuota`

For example:
```zig
fn fib(n: u64) u64 {
    @setEvalBranchQuota(1_000_000_000); // before, this disabled memoization
    if (n < 2) return n;
    return fib(n - 1) + fib(n - 2);
}
pub export fn entry() u64 {
    var acc: u64 = 0;
    inline for (0..40) |_| {
        acc +%= comptime fib(25);
    }
    return acc;
}
```
This went from ~30s to near instant on my machine, admittedly this is the perfect case for memoization though.

A more "real-world" case like this:
```zig
  pub export fn entry() usize {
      @setEvalBranchQuota(10_000_000);
      var n: usize = 0;
      inline for (0..200) |_| {
         // comptimePrint calls @setEvalBranchQuota!
          const s = std.fmt.comptimePrint("value {d} of {s} at {x}", .{ 42, "abcdefgh", 7 });
          n += s.len;
      }
      return n;
  }
```
Showed around ~2.2x improvement.

Reviewed-on: https://codeberg.org/ziglang/zig/pulls/36494
Reviewed-by: mlugg <mlugg@mlugg.co.uk>

2 files changed, 17 insertions(+), 1 deletions(-)

src
- InternPool.zig
- Sema.zig

src/InternPool.zig+4

...	...	@@ -2596,6 +2596,7 @@ pub const Key = union(enum) {
2596	2596	arg_values: []const Index,
2597	2597	result: Index,
2598	2598	branch_count: u32,
	2599	branch_quota: u32,
2599	2600	};
2600	2601
2601	2602	pub fn hash32(key: Key, ip: *const InternPool) u32 {
...	...	@@ -6272,6 +6273,7 @@ pub const MemoizedCall = struct {
6272	6273	args_len: u32,
6273	6274	result: Index,
6274	6275	branch_count: u32,
	6276	branch_quota: u32,
6275	6277	};
6276	6278
6277	6279	pub fn init(ip: *InternPool, gpa: Allocator, io: Io, available_threads: usize) !void {
...	...	@@ -6913,6 +6915,7 @@ pub fn indexToKey(ip: *const InternPool, index: Index) Key {
6913	6915	.arg_values = @ptrCast(extra_list.view().items(.@"0")[extra.end..][0..extra.data.args_len]),
6914	6916	.result = extra.data.result,
6915	6917	.branch_count = extra.data.branch_count,
	6918	.branch_quota = extra.data.branch_quota,
6916	6919	} };
6917	6920	},
6918	6921	};
...	...	@@ -7992,6 +7995,7 @@ pub fn get(ip: *InternPool, gpa: Allocator, io: Io, tid: Zcu.PerThread.Id, key:
7992	7995	.args_len = @intCast(memoized_call.arg_values.len),
7993	7996	.result = memoized_call.result,
7994	7997	.branch_count = memoized_call.branch_count,
	7998	.branch_quota = memoized_call.branch_quota,
7995	7999	}),
7996	8000	});
7997	8001	extra.appendSliceAssumeCapacity(.{@ptrCast(memoized_call.arg_values)});

src/Sema.zig+13-1

...	...	@@ -132,6 +132,10 @@ dependencies: std.array_hash_map.Auto(InternPool.Dependee, void) = .empty,
132	132	/// by `analyzeCall`.
133	133	allow_memoize: bool = true,
134	134
	135	/// The largest quota requested by `@setEvalBranchQuota` within the comptime call
	136	/// currently being analyzed.
	137	quota_request: u32 = 0,
	138
135	139	/// The `BranchHint` for the current branch of runtime control flow.
136	140	/// This state is on `Sema` so that `cold` hints can be propagated up through blocks with less special handling.
137	141	branch_hint: ?std.lang.BranchHint = null,
...	...	@@ -4914,7 +4918,7 @@ fn zirSetEvalBranchQuota(sema: Sema, block: Block, inst: Zir.Inst.Index) Compi
4914	4918	const src = block.nodeOffset(inst_data.src_node);
4915	4919	const quota: u32 = @intCast(try sema.resolveInt(block, src, inst_data.operand, .u32, .{ .simple = .operand_setEvalBranchQuota }));
4916	4920	sema.branch_quota = @max(sema.branch_quota, quota);
4917		sema.allow_memoize = false;
	4921	sema.quota_request = @max(sema.quota_request, quota);
4918	4922	}
4919	4923
4920	4924	fn zirStoreNode(sema: Sema, block: Block, inst: Zir.Inst.Index) CompileError!void {
...	...	@@ -7218,6 +7222,7 @@ fn analyzeCall(
7218	7222	.arg_values = memoized_arg_values,
7219	7223	.result = undefined, // ignored by hash+eql
7220	7224	.branch_count = undefined, // ignored by hash+eql
	7225	.branch_quota = undefined, // ignored by hash+eql
7221	7226	},
7222	7227	}) orelse break :memoize;
7223	7228	const memoized_call = ip.indexToKey(memoized_call_index).memoized_call;
...	...	@@ -7227,6 +7232,8 @@ fn analyzeCall(
7227	7232	break :memoize;
7228	7233	}
7229	7234	sema.branch_count += memoized_call.branch_count;
	7235	sema.branch_quota = @max(sema.branch_quota, memoized_call.branch_quota);
	7236	sema.quota_request = @max(sema.quota_request, memoized_call.branch_quota);
7230	7237	const result = Air.internedToRef(memoized_call.result);
7231	7238	if (ensure_result_used) {
7232	7239	try sema.ensureResultUsed(block, sema.typeOf(result), call_src);
...	...	@@ -7353,6 +7360,10 @@ fn analyzeCall(
7353	7360	defer sema.allow_memoize = old_allow_memoize and sema.allow_memoize;
7354	7361	sema.allow_memoize = true;
7355	7362
	7363	const old_quota_request = sema.quota_request;
	7364	defer sema.quota_request = @max(old_quota_request, sema.quota_request);
	7365	sema.quota_request = 0;
	7366
7356	7367	// Store the current eval branch count so we can find out how many eval branches
7357	7368	// the comptime call caused.
7358	7369	const old_branch_count = sema.branch_count;
...	...	@@ -7388,6 +7399,7 @@ fn analyzeCall(
7388	7399	.arg_values = memoized_arg_values,
7389	7400	.result = result_val.toIntern(),
7390	7401	.branch_count = sema.branch_count - old_branch_count,
	7402	.branch_quota = sema.quota_request,
7391	7403	} });
7392	7404	}
7393	7405	}