authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2023-02-17 11:51:22-07:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2023-02-18 19:17:20-07:00
logfaa44e2e5875036b105d8b7d38ccb2e93757a3c5
treefc384ad77143ae6a5c7df5a19da4c2ccfef364d9
parent6733e43d87d4fe7b9d89948ebb95a72515c44fee

AstGen: rework multi-object for loop

* Allow unbounded looping. * Lower by incrementing raw pointers for each iterable rather than incrementing a single index variable. This elides safety checks without any analysis required thanks to the length assertion and lowers to decent machine code even in debug builds. - An "end" value is selected, prioritizing a counter if possible, falling back to a runtime calculation of ptr+len on a slice input. * Specialize on the pattern `0..`, avoiding an unnecessary subtraction instruction being emitted. * Add the `for_check_lens` ZIR instruction.

4 files changed, 156 insertions(+), 55 deletions(-)

src/AstGen.zig+114-55
......@@ -2666,6 +2666,7 @@ fn addEnsureResult(gz: *GenZir, maybe_unused_result: Zir.Inst.Ref, statement: As
26662666 .validate_deref,
26672667 .save_err_ret_index,
26682668 .restore_err_ret_index,
2669 .for_check_lens,
26692670 => break :b true,
26702671
26712672 .@"defer" => unreachable,
......@@ -6294,37 +6295,35 @@ fn forExpr(
62946295 try astgen.checkLabelRedefinition(scope, label_token);
62956296 }
62966297
6297 // Set up variables and constants.
62986298 const is_inline = parent_gz.force_comptime or for_full.inline_token != null;
62996299 const tree = astgen.tree;
63006300 const token_tags = tree.tokens.items(.tag);
63016301 const node_tags = tree.nodes.items(.tag);
63026302 const node_data = tree.nodes.items(.data);
6303 const gpa = astgen.gpa;
63036304
6304 // Check for unterminated ranges.
6305 {
6306 var unterminated: ?Ast.Node.Index = null;
6307 for (for_full.ast.inputs) |input| {
6308 if (node_tags[input] != .for_range) break;
6309 if (node_data[input].rhs != 0) break;
6310 unterminated = unterminated orelse input;
6311 } else {
6312 return astgen.failNode(unterminated.?, "unterminated for range", .{});
6313 }
6314 }
6315
6316 var lens = astgen.gpa.alloc(Zir.Inst.Ref, for_full.ast.inputs.len);
6317 defer astgen.gpa.free(lens);
6318 var indexables = astgen.gpa.alloc(Zir.Inst.Ref, for_full.ast.inputs.len);
6319 defer astgen.gpa.free(indexables);
6320 var counters = std.ArrayList(Zir.Inst.Ref).init(astgen.gpa);
6321 defer counters.deinit();
6305 const allocs = try gpa.alloc(Zir.Inst.Ref, for_full.ast.inputs.len);
6306 defer gpa.free(allocs);
6307 // elements of this array can be `none`, indicating no length check.
6308 const lens = try gpa.alloc(Zir.Inst.Ref, for_full.ast.inputs.len);
6309 defer gpa.free(lens);
63226310
63236311 const counter_alloc_tag: Zir.Inst.Tag = if (is_inline) .alloc_comptime_mut else .alloc;
63246312
6313 // Tracks the index of allocs/lens that has a length to be checked and is
6314 // used for the end value.
6315 // If this is null, there are no len checks.
6316 var end_input_index: ?u32 = null;
6317 // This is a value to use to find out if the for loop has reached the end
6318 // yet. It prefers to use a counter since the end value is provided directly,
6319 // and otherwise falls back to adding ptr+len of a slice to compute end.
6320 // Corresponds to end_input_index and will be .none in case that value is null.
6321 var cond_end_val: Zir.Inst.Ref = .none;
6322
63256323 {
63266324 var payload = for_full.payload_token;
6327 for (for_full.ast.inputs) |input, i| {
6325 for (for_full.ast.inputs) |input, i_usize| {
6326 const i = @intCast(u32, i_usize);
63286327 const payload_is_ref = token_tags[payload] == .asterisk;
63296328 const ident_tok = payload + @boolToInt(payload_is_ref);
63306329
......@@ -6339,59 +6338,101 @@ fn forExpr(
63396338 return astgen.failTok(ident_tok, "cannot capture reference to range", .{});
63406339 }
63416340 const counter_ptr = try parent_gz.addUnNode(counter_alloc_tag, .usize_type, node);
6342 const start_val = try expr(parent_gz, scope, node_data[input].lhs, input);
6341 const start_node = node_data[input].lhs;
6342 const start_val = try expr(parent_gz, scope, .{ .rl = .none }, start_node);
63436343 _ = try parent_gz.addBin(.store, counter_ptr, start_val);
6344 indexables[i] = counter_ptr;
6345 try counters.append(counter_ptr);
63466344
63476345 const end_node = node_data[input].rhs;
6348 const end_val = if (end_node != 0) try expr(parent_gz, scope, node_data[input].rhs, input) else .none;
6349 const range_len = try parent_gz.addPlNode(.for_range_len, input, Zir.Inst.Bin{
6350 .lhs = start_val,
6351 .rhs = end_val,
6352 });
6346 const end_val = if (end_node != 0)
6347 try expr(parent_gz, scope, .{ .rl = .none }, node_data[input].rhs)
6348 else
6349 .none;
6350
6351 const range_len = if (end_val == .none or nodeIsTriviallyZero(tree, start_node))
6352 end_val
6353 else
6354 try parent_gz.addPlNode(.sub, input, Zir.Inst.Bin{
6355 .lhs = end_val,
6356 .rhs = start_val,
6357 });
6358
6359 if (range_len != .none and cond_end_val == .none) {
6360 end_input_index = i;
6361 cond_end_val = end_val;
6362 }
6363
6364 allocs[i] = counter_ptr;
63536365 lens[i] = range_len;
63546366 } else {
63556367 const cond_ri: ResultInfo = .{ .rl = if (payload_is_ref) .ref else .none };
63566368 const indexable = try expr(parent_gz, scope, cond_ri, input);
6357 indexables[i] = indexable;
6369 const base_ptr = try parent_gz.addPlNode(.elem_ptr_imm, input, Zir.Inst.ElemPtrImm{
6370 .ptr = indexable,
6371 .index = 0,
6372 });
63586373
6359 const indexable_len = try parent_gz.addUnNode(.indexable_ptr_len, indexable, input);
6360 lens[i] = indexable_len;
6374 if (end_input_index == null) {
6375 end_input_index = i;
6376 assert(cond_end_val == .none);
6377 }
6378
6379 allocs[i] = base_ptr;
6380 lens[i] = try parent_gz.addUnNode(.indexable_ptr_len, indexable, input);
63616381 }
63626382 }
63636383 }
63646384
6365 const len = "check_for_lens";
6385 // In case there are no counters which already have an end computed, we
6386 // compute an end from base pointer plus length.
6387 if (end_input_index) |i| {
6388 if (cond_end_val == .none) {
6389 cond_end_val = try parent_gz.addPlNode(.add, for_full.ast.inputs[i], Zir.Inst.Bin{
6390 .lhs = allocs[i],
6391 .rhs = lens[i],
6392 });
6393 }
6394 }
63666395
6367 const index_ptr = blk: {
6368 // Future optimization:
6369 // for loops with only ranges don't need a separate index variable.
6370 const index_ptr = try parent_gz.addUnNode(counter_alloc_tag, .usize_type, node);
6371 // initialize to zero
6372 _ = try parent_gz.addBin(.store, index_ptr, .zero_usize);
6373 try counters.append(index_ptr);
6374 break :blk index_ptr;
6375 };
6396 // We use a dedicated ZIR instruction to assert the lengths to assist with
6397 // nicer error reporting as well as fewer ZIR bytes emitted.
6398 if (end_input_index != null) {
6399 const lens_len = @intCast(u32, lens.len);
6400 try astgen.extra.ensureUnusedCapacity(gpa, @typeInfo(Zir.Inst.MultiOp).Struct.fields.len + lens_len);
6401 _ = try parent_gz.addPlNode(.for_check_lens, node, Zir.Inst.MultiOp{
6402 .operands_len = lens_len,
6403 });
6404 appendRefsAssumeCapacity(astgen, lens);
6405 }
63766406
63776407 const loop_tag: Zir.Inst.Tag = if (is_inline) .block_inline else .loop;
63786408 const loop_block = try parent_gz.makeBlockInst(loop_tag, node);
6379 try parent_gz.instructions.append(astgen.gpa, loop_block);
6409 try parent_gz.instructions.append(gpa, loop_block);
63806410
63816411 var loop_scope = parent_gz.makeSubBlock(scope);
63826412 loop_scope.is_inline = is_inline;
63836413 loop_scope.setBreakResultInfo(ri);
63846414 defer loop_scope.unstack();
6385 defer loop_scope.labeled_breaks.deinit(astgen.gpa);
6415 defer loop_scope.labeled_breaks.deinit(gpa);
63866416
63876417 var cond_scope = parent_gz.makeSubBlock(&loop_scope.base);
63886418 defer cond_scope.unstack();
63896419
6390 // check condition i < array_expr.len
6391 const index = try cond_scope.addUnNode(.load, index_ptr, for_full.ast.cond_expr);
6392 const cond = try cond_scope.addPlNode(.cmp_lt, for_full.ast.cond_expr, Zir.Inst.Bin{
6393 .lhs = index,
6394 .rhs = len,
6420 // Load all the iterables.
6421 const loaded_ptrs = try gpa.alloc(Zir.Inst.Ref, allocs.len);
6422 defer gpa.free(loaded_ptrs);
6423 for (allocs) |alloc, i| {
6424 loaded_ptrs[i] = try cond_scope.addUnNode(.load, alloc, for_full.ast.inputs[i]);
6425 }
6426
6427 // Check the condition.
6428 const input_index = end_input_index orelse {
6429 return astgen.failNode(node, "TODO: handle infinite for loop", .{});
6430 };
6431 assert(cond_end_val != .none);
6432
6433 const cond = try cond_scope.addPlNode(.cmp_neq, for_full.ast.inputs[input_index], Zir.Inst.Bin{
6434 .lhs = loaded_ptrs[input_index],
6435 .rhs = cond_end_val,
63956436 });
63966437
63976438 const condbr_tag: Zir.Inst.Tag = if (is_inline) .condbr_inline else .condbr;
......@@ -6400,16 +6441,15 @@ fn forExpr(
64006441 const cond_block = try loop_scope.makeBlockInst(block_tag, node);
64016442 try cond_scope.setBlockBody(cond_block);
64026443 // cond_block unstacked now, can add new instructions to loop_scope
6403 try loop_scope.instructions.append(astgen.gpa, cond_block);
6444 try loop_scope.instructions.append(gpa, cond_block);
64046445
6405 // Increment the index variable and ranges.
6406 for (counters) |counter_ptr| {
6407 const counter = try loop_scope.addUnNode(.load, counter_ptr, for_full.ast.cond_expr);
6408 const counter_plus_one = try loop_scope.addPlNode(.add, node, Zir.Inst.Bin{
6409 .lhs = counter,
6446 // Increment the loop variables.
6447 for (allocs) |alloc, i| {
6448 const incremented = try loop_scope.addPlNode(.add, node, Zir.Inst.Bin{
6449 .lhs = loaded_ptrs[i],
64106450 .rhs = .one_usize,
64116451 });
6412 _ = try loop_scope.addBin(.store, counter_ptr, counter_plus_one);
6452 _ = try loop_scope.addBin(.store, alloc, incremented);
64136453 }
64146454 const repeat_tag: Zir.Inst.Tag = if (is_inline) .repeat_inline else .repeat;
64156455 _ = try loop_scope.addNode(repeat_tag, node);
......@@ -8960,6 +9000,25 @@ comptime {
89609000 }
89619001}
89629002
9003fn nodeIsTriviallyZero(tree: *const Ast, node: Ast.Node.Index) bool {
9004 const node_tags = tree.nodes.items(.tag);
9005 const main_tokens = tree.nodes.items(.main_token);
9006
9007 switch (node_tags[node]) {
9008 .number_literal => {
9009 const ident = main_tokens[node];
9010 return switch (std.zig.parseNumberLiteral(tree.tokenSlice(ident))) {
9011 .int => |number| switch (number) {
9012 0 => true,
9013 else => false,
9014 },
9015 else => false,
9016 };
9017 },
9018 else => return false,
9019 }
9020}
9021
89639022fn nodeMayNeedMemoryLocation(tree: *const Ast, start_node: Ast.Node.Index, have_res_ty: bool) bool {
89649023 const node_tags = tree.nodes.items(.tag);
89659024 const node_datas = tree.nodes.items(.data);
src/Sema.zig+15
......@@ -1386,6 +1386,11 @@ fn analyzeBodyInner(
13861386 i += 1;
13871387 continue;
13881388 },
1389 .for_check_lens => {
1390 try sema.zirForCheckLens(block, inst);
1391 i += 1;
1392 continue;
1393 },
13891394
13901395 // Special case instructions to handle comptime control flow.
13911396 .@"break" => {
......@@ -17096,6 +17101,16 @@ fn zirRestoreErrRetIndex(sema: *Sema, start_block: *Block, inst: Zir.Inst.Index)
1709617101 return sema.popErrorReturnTrace(start_block, src, operand, saved_index);
1709717102}
1709817103
17104fn zirForCheckLens(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!void {
17105 const inst_data = sema.code.instructions.items(.data)[inst].pl_node;
17106 const extra = sema.code.extraData(Zir.Inst.MultiOp, inst_data.payload_index);
17107 const args = sema.code.refSlice(extra.end, extra.data.operands_len);
17108 const src = inst_data.src();
17109
17110 _ = args;
17111 return sema.fail(block, src, "TODO implement zirForCheckLens", .{});
17112}
17113
1709917114fn addToInferredErrorSet(sema: *Sema, uncasted_operand: Air.Inst.Ref) !void {
1710017115 assert(sema.fn_ret_ty.zigTypeTag() == .ErrorUnion);
1710117116
src/Zir.zig+12
......@@ -497,6 +497,15 @@ pub const Inst = struct {
497497 /// Sends comptime control flow back to the beginning of the current block.
498498 /// Uses the `node` field.
499499 repeat_inline,
500 /// Asserts that all the lengths provided match. Used to build a for loop.
501 /// Return value is always void.
502 /// Uses the `pl_node` field with payload `MultiOp`.
503 /// There is exactly one item corresponding to each AST node inside the for
504 /// loop condition. Each item may be `none`, indicating an unbounded range.
505 /// Illegal behaviors:
506 /// * If all lengths are unbounded ranges (always a compile error).
507 /// * If any two lengths do not match each other.
508 for_check_lens,
500509 /// Merge two error sets into one, `E1 || E2`.
501510 /// Uses the `pl_node` field with payload `Bin`.
502511 merge_error_sets,
......@@ -1242,6 +1251,7 @@ pub const Inst = struct {
12421251 .defer_err_code,
12431252 .save_err_ret_index,
12441253 .restore_err_ret_index,
1254 .for_check_lens,
12451255 => false,
12461256
12471257 .@"break",
......@@ -1309,6 +1319,7 @@ pub const Inst = struct {
13091319 .memcpy,
13101320 .memset,
13111321 .check_comptime_control_flow,
1322 .for_check_lens,
13121323 .@"defer",
13131324 .defer_err_code,
13141325 .restore_err_ret_index,
......@@ -1588,6 +1599,7 @@ pub const Inst = struct {
15881599 .@"break" = .@"break",
15891600 .break_inline = .@"break",
15901601 .check_comptime_control_flow = .un_node,
1602 .for_check_lens = .pl_node,
15911603 .call = .pl_node,
15921604 .cmp_lt = .pl_node,
15931605 .cmp_lte = .pl_node,
src/print_zir.zig+15
......@@ -355,6 +355,8 @@ const Writer = struct {
355355 .array_type,
356356 => try self.writePlNodeBin(stream, inst),
357357
358 .for_check_lens => try self.writePlNodeMultiOp(stream, inst),
359
358360 .elem_ptr_imm => try self.writeElemPtrImm(stream, inst),
359361
360362 .@"export" => try self.writePlNodeExport(stream, inst),
......@@ -868,6 +870,19 @@ const Writer = struct {
868870 try self.writeSrc(stream, inst_data.src());
869871 }
870872
873 fn writePlNodeMultiOp(self: *Writer, stream: anytype, inst: Zir.Inst.Index) !void {
874 const inst_data = self.code.instructions.items(.data)[inst].pl_node;
875 const extra = self.code.extraData(Zir.Inst.MultiOp, inst_data.payload_index);
876 const args = self.code.refSlice(extra.end, extra.data.operands_len);
877 try stream.writeAll("{");
878 for (args) |arg, i| {
879 if (i != 0) try stream.writeAll(", ");
880 try self.writeInstRef(stream, arg);
881 }
882 try stream.writeAll("}) ");
883 try self.writeSrc(stream, inst_data.src());
884 }
885
871886 fn writeElemPtrImm(self: *Writer, stream: anytype, inst: Zir.Inst.Index) !void {
872887 const inst_data = self.code.instructions.items(.data)[inst].pl_node;
873888 const extra = self.code.extraData(Zir.Inst.ElemPtrImm, inst_data.payload_index).data;