compiler: handle eval branch quota in memoized calls

In a `memoized_call`, store how many backwards braches the call
performs. Add this to `sema.branch_count` when using a memoized call. If
this exceeds the quota, perform a non-memoized call to get a correct
"exceeded X backwards branches" error.

Also, do not memoize calls which do `@setEvalBranchQuota` or similar, as
this affects global state which must apply to the caller.

Change some eval branch quotas so that the compiler itself still builds correctly.

This commit manually changes a file in Aro which is automatically
generated. The sources which generate the file are not in this repo.
Upstream Aro should make the suitable changes on their end before the
next sync of Aro sources into the Zig repo.

Author: mlugg

lib/compiler/aro/aro/Builtins/Builtin.zig

@@ -5165,7 +5165,7 @@ const dafsa = [_]Node{
     .{ .char = 'e', .end_of_word = false, .end_of_list = true, .number = 1, .child_index = 4913 },
 };
 pub const data = blk: {
-    @setEvalBranchQuota(3986);
+    @setEvalBranchQuota(30_000);
     break :blk [_]@This(){
         // _Block_object_assign
         .{ .tag = @enumFromInt(0), .properties = .{ .param_str = "vv*vC*iC", .header = .blocks, .attributes = .{ .lib_function_without_prefix = true } } },

lib/compiler/aro/aro/Parser.zig

@@ -4802,6 +4802,7 @@ const CallExpr = union(enum) {
     }
 
     fn shouldPromoteVarArg(self: CallExpr, arg_idx: u32) bool {
+        @setEvalBranchQuota(2000);
         return switch (self) {
             .standard => true,
             .builtin => |builtin| switch (builtin.tag) {
@@ -4902,6 +4903,7 @@ const CallExpr = union(enum) {
     }
 
     fn returnType(self: CallExpr, p: *Parser, callable_ty: Type) Type {
+        @setEvalBranchQuota(6000);
         return switch (self) {
             .standard => callable_ty.returnType(),
             .builtin => |builtin| switch (builtin.tag) {

lib/std/crypto/pcurves/p384.zig

@@ -393,7 +393,7 @@ pub const P384 = struct {
     }
 
     const basePointPc = pc: {
-        @setEvalBranchQuota(50000);
+        @setEvalBranchQuota(70000);
         break :pc precompute(P384.basePoint, 15);
     };
 

src/InternPool.zig

@@ -2391,6 +2391,7 @@ pub const Key = union(enum) {
         func: Index,
         arg_values: []const Index,
         result: Index,
+        branch_count: u32,
     };
 
     pub fn hash32(key: Key, ip: *const InternPool) u32 {
@@ -6157,6 +6158,7 @@ pub const MemoizedCall = struct {
     func: Index,
     args_len: u32,
     result: Index,
+    branch_count: u32,
 };
 
 pub fn init(ip: *InternPool, gpa: Allocator, available_threads: usize) !void {
@@ -6785,6 +6787,7 @@ pub fn indexToKey(ip: *const InternPool, index: Index) Key {
                 .func = extra.data.func,
                 .arg_values = @ptrCast(extra_list.view().items(.@"0")[extra.end..][0..extra.data.args_len]),
                 .result = extra.data.result,
+                .branch_count = extra.data.branch_count,
             } };
         },
     };
@@ -7955,6 +7958,7 @@ pub fn get(ip: *InternPool, gpa: Allocator, tid: Zcu.PerThread.Id, key: Key) All
                     .func = memoized_call.func,
                     .args_len = @intCast(memoized_call.arg_values.len),
                     .result = memoized_call.result,
+                    .branch_count = memoized_call.branch_count,
                 }),
             });
             extra.appendSliceAssumeCapacity(.{@ptrCast(memoized_call.arg_values)});

src/Sema.zig

@@ -113,6 +113,11 @@ type_references: std.AutoArrayHashMapUnmanaged(InternPool.Index, void) = .{},
 /// `AnalUnit` multiple times.
 dependencies: std.AutoArrayHashMapUnmanaged(InternPool.Dependee, void) = .{},
 
+/// Whether memoization of this call is permitted. Operations with side effects global
+/// to the `Sema`, such as `@setEvalBranchQuota`, set this to `false`. It is observed
+/// by `analyzeCall`.
+allow_memoize: bool = true,
+
 const MaybeComptimeAlloc = struct {
     /// The runtime index of the `alloc` instruction.
     runtime_index: Value.RuntimeIndex,
@@ -5524,6 +5529,7 @@ fn zirSetEvalBranchQuota(sema: *Sema, block: *Block, inst: Zir.Inst.Index) Compi
         .needed_comptime_reason = "eval branch quota must be comptime-known",
     }));
     sema.branch_quota = @max(sema.branch_quota, quota);
+    sema.allow_memoize = false;
 }
 
 fn zirStoreNode(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!void {
@@ -6416,6 +6422,7 @@ fn zirSetAlignStack(sema: *Sema, block: *Block, extended: Zir.Inst.Extended.Inst
     }
 
     zcu.intern_pool.funcMaxStackAlignment(sema.func_index, alignment);
+    sema.allow_memoize = false;
 }
 
 fn zirSetCold(sema: *Sema, block: *Block, extended: Zir.Inst.Extended.InstData) CompileError!void {
@@ -6434,6 +6441,7 @@ fn zirSetCold(sema: *Sema, block: *Block, extended: Zir.Inst.Extended.InstData)
         .cau => return, // does nothing outside a function
     };
     ip.funcSetCold(func, is_cold);
+    sema.allow_memoize = false;
 }
 
 fn zirDisableInstrumentation(sema: *Sema) CompileError!void {
@@ -6445,6 +6453,7 @@ fn zirDisableInstrumentation(sema: *Sema) CompileError!void {
         .cau => return, // does nothing outside a function
     };
     ip.funcSetDisableInstrumentation(func);
+    sema.allow_memoize = false;
 }
 
 fn zirSetFloatMode(sema: *Sema, block: *Block, extended: Zir.Inst.Extended.InstData) CompileError!void {
@@ -7727,15 +7736,25 @@ fn analyzeCall(
         // This `res2` is here instead of directly breaking from `res` due to a stage1
         // bug generating invalid LLVM IR.
         const res2: Air.Inst.Ref = res2: {
-            if (should_memoize and is_comptime_call) {
-                if (zcu.intern_pool.getIfExists(.{ .memoized_call = .{
-                    .func = module_fn_index,
-                    .arg_values = memoized_arg_values,
-                    .result = .none,
-                } })) |memoized_call_index| {
-                    const memoized_call = zcu.intern_pool.indexToKey(memoized_call_index).memoized_call;
-                    break :res2 Air.internedToRef(memoized_call.result);
+            memoize: {
+                if (!should_memoize) break :memoize;
+                if (!is_comptime_call) break :memoize;
+                const memoized_call_index = ip.getIfExists(.{
+                    .memoized_call = .{
+                        .func = module_fn_index,
+                        .arg_values = memoized_arg_values,
+                        .result = undefined, // ignored by hash+eql
+                        .branch_count = undefined, // ignored by hash+eql
+                    },
+                }) orelse break :memoize;
+                const memoized_call = ip.indexToKey(memoized_call_index).memoized_call;
+                if (sema.branch_count + memoized_call.branch_count > sema.branch_quota) {
+                    // Let the call play out se we get the correct source location for the
+                    // "evaluation exceeded X backwards branches" error.
+                    break :memoize;
                 }
+                sema.branch_count += memoized_call.branch_count;
+                break :res2 Air.internedToRef(memoized_call.result);
             }
 
             new_fn_info.return_type = sema.fn_ret_ty.toIntern();
@@ -7773,6 +7792,17 @@ fn analyzeCall(
                 child_block.error_return_trace_index = error_return_trace_index;
             }
 
+            // We temporarily set `allow_memoize` to `true` to track this comptime call.
+            // It is restored after this call finishes analysis, so that a caller may
+            // know whether an in-progress call (containing this call) may be memoized.
+            const old_allow_memoize = sema.allow_memoize;
+            defer sema.allow_memoize = old_allow_memoize and sema.allow_memoize;
+            sema.allow_memoize = true;
+
+            // Store the current eval branch count so we can find out how many eval branches
+            // the comptime call caused.
+            const old_branch_count = sema.branch_count;
+
             const result = result: {
                 sema.analyzeFnBody(&child_block, fn_info.body) catch |err| switch (err) {
                     error.ComptimeReturn => break :result inlining.comptime_result,
@@ -7792,11 +7822,12 @@ fn analyzeCall(
                 // a reference to `comptime_allocs` so is not stable across instances of `Sema`.
                 // TODO: check whether any external comptime memory was mutated by the
                 // comptime function call. If so, then do not memoize the call here.
-                if (should_memoize and !Value.fromInterned(result_interned).canMutateComptimeVarState(zcu)) {
+                if (should_memoize and sema.allow_memoize and !Value.fromInterned(result_interned).canMutateComptimeVarState(zcu)) {
                     _ = try pt.intern(.{ .memoized_call = .{
                         .func = module_fn_index,
                         .arg_values = memoized_arg_values,
                         .result = result_transformed,
+                        .branch_count = sema.branch_count - old_branch_count,
                     } });
                 }
 

src/register_manager.zig

@@ -93,6 +93,8 @@ pub fn RegisterManager(
             comptime set: []const Register,
             reg: Register,
         ) ?std.math.IntFittingRange(0, set.len - 1) {
+            @setEvalBranchQuota(3000);
+
             const Id = @TypeOf(reg.id());
             comptime var min_id: Id = std.math.maxInt(Id);
             comptime var max_id: Id = std.math.minInt(Id);