Merge pull request #22397 from Techatrix/type-safe-ast

improve type safety of std.zig.Ast

1 files changed, 11 insertions, 17 deletions

diff --git a/lib/std/zon/parse.zig b/lib/std/zon/parse.zig
index daf83d0bbd..1edb37654c 100644
--- a/lib/std/zon/parse.zig
+++ b/lib/std/zon/parse.zig
@@ -196,16 +196,15 @@ pub const Error = union(enum) {
         return .{ .err = self, .status = status };
     }
 
-    fn zoirErrorLocation(ast: Ast, maybe_token: Ast.TokenIndex, node_or_offset: u32) Ast.Location {
-        if (maybe_token == Zoir.CompileError.invalid_token) {
-            const main_tokens = ast.nodes.items(.main_token);
-            const ast_node = node_or_offset;
-            const token = main_tokens[ast_node];
-            return ast.tokenLocation(0, token);
-        } else {
-            var location = ast.tokenLocation(0, maybe_token);
+    fn zoirErrorLocation(ast: Ast, maybe_token: Ast.OptionalTokenIndex, node_or_offset: u32) Ast.Location {
+        if (maybe_token.unwrap()) |token| {
+            var location = ast.tokenLocation(0, token);
             location.column += node_or_offset;
             return location;
+        } else {
+            const ast_node: Ast.Node.Index = @enumFromInt(node_or_offset);
+            const token = ast.nodeMainToken(ast_node);
+            return ast.tokenLocation(0, token);
         }
     }
 };
@@ -632,7 +631,7 @@ const Parser = struct {
         switch (try ZonGen.parseStrLit(self.ast, ast_node, buf.writer(self.gpa))) {
             .success => {},
             .failure => |err| {
-                const token = self.ast.nodes.items(.main_token)[ast_node];
+                const token = self.ast.nodeMainToken(ast_node);
                 const raw_string = self.ast.tokenSlice(token);
                 return self.failTokenFmt(token, @intCast(err.offset()), "{s}", .{err.fmt(raw_string)});
             },
@@ -1005,8 +1004,7 @@ const Parser = struct {
         args: anytype,
     ) error{ OutOfMemory, ParseZon } {
         @branchHint(.cold);
-        const main_tokens = self.ast.nodes.items(.main_token);
-        const token = main_tokens[node.getAstNode(self.zoir)];
+        const token = self.ast.nodeMainToken(node.getAstNode(self.zoir));
         return self.failTokenFmt(token, 0, fmt, args);
     }
 
@@ -1025,8 +1023,7 @@ const Parser = struct {
         message: []const u8,
     ) error{ParseZon} {
         @branchHint(.cold);
-        const main_tokens = self.ast.nodes.items(.main_token);
-        const token = main_tokens[node.getAstNode(self.zoir)];
+        const token = self.ast.nodeMainToken(node.getAstNode(self.zoir));
         return self.failToken(.{
             .token = token,
             .offset = 0,
@@ -1059,10 +1056,7 @@ const Parser = struct {
             const struct_init = self.ast.fullStructInit(&buf, node.getAstNode(self.zoir)).?;
             const field_node = struct_init.ast.fields[f];
             break :b self.ast.firstToken(field_node) - 2;
-        } else b: {
-            const main_tokens = self.ast.nodes.items(.main_token);
-            break :b main_tokens[node.getAstNode(self.zoir)];
-        };
+        } else self.ast.nodeMainToken(node.getAstNode(self.zoir));
         switch (@typeInfo(T)) {
             inline .@"struct", .@"union", .@"enum" => |info| {
                 const note: Error.TypeCheckFailure.Note = if (info.fields.len == 0) b: {