Merge pull request #9679 from travisstaloch/sat-arith-operators

sat-arithmetic: add operator support

5 files changed, 189 insertions, 24 deletions

diff --git a/lib/std/zig/Ast.zig b/lib/std/zig/Ast.zig
index 5838dcd37a..4ee3a45221 100644
--- a/lib/std/zig/Ast.zig
+++ b/lib/std/zig/Ast.zig
@@ -395,14 +395,18 @@ pub fn firstToken(tree: Tree, node: Node.Index) TokenIndex {
         .assign_mod,
         .assign_add,
         .assign_sub,
-        .assign_bit_shift_left,
-        .assign_bit_shift_right,
+        .assign_shl,
+        .assign_shl_sat,
+        .assign_shr,
         .assign_bit_and,
         .assign_bit_xor,
         .assign_bit_or,
         .assign_mul_wrap,
         .assign_add_wrap,
         .assign_sub_wrap,
+        .assign_mul_sat,
+        .assign_add_sat,
+        .assign_sub_sat,
         .assign,
         .merge_error_sets,
         .mul,
@@ -410,13 +414,17 @@ pub fn firstToken(tree: Tree, node: Node.Index) TokenIndex {
         .mod,
         .array_mult,
         .mul_wrap,
+        .mul_sat,
         .add,
         .sub,
         .array_cat,
         .add_wrap,
         .sub_wrap,
-        .bit_shift_left,
-        .bit_shift_right,
+        .add_sat,
+        .sub_sat,
+        .shl,
+        .shl_sat,
+        .shr,
         .bit_and,
         .bit_xor,
         .bit_or,
@@ -651,14 +659,18 @@ pub fn lastToken(tree: Tree, node: Node.Index) TokenIndex {
         .assign_mod,
         .assign_add,
         .assign_sub,
-        .assign_bit_shift_left,
-        .assign_bit_shift_right,
+        .assign_shl,
+        .assign_shl_sat,
+        .assign_shr,
         .assign_bit_and,
         .assign_bit_xor,
         .assign_bit_or,
         .assign_mul_wrap,
         .assign_add_wrap,
         .assign_sub_wrap,
+        .assign_mul_sat,
+        .assign_add_sat,
+        .assign_sub_sat,
         .assign,
         .merge_error_sets,
         .mul,
@@ -666,13 +678,17 @@ pub fn lastToken(tree: Tree, node: Node.Index) TokenIndex {
         .mod,
         .array_mult,
         .mul_wrap,
+        .mul_sat,
         .add,
         .sub,
         .array_cat,
         .add_wrap,
         .sub_wrap,
-        .bit_shift_left,
-        .bit_shift_right,
+        .add_sat,
+        .sub_sat,
+        .shl,
+        .shl_sat,
+        .shr,
         .bit_and,
         .bit_xor,
         .bit_or,
@@ -2524,9 +2540,11 @@ pub const Node = struct {
         /// `lhs -= rhs`. main_token is op.
         assign_sub,
         /// `lhs <<= rhs`. main_token is op.
-        assign_bit_shift_left,
+        assign_shl,
+        /// `lhs <<|= rhs`. main_token is op.
+        assign_shl_sat,
         /// `lhs >>= rhs`. main_token is op.
-        assign_bit_shift_right,
+        assign_shr,
         /// `lhs &= rhs`. main_token is op.
         assign_bit_and,
         /// `lhs ^= rhs`. main_token is op.
@@ -2539,6 +2557,12 @@ pub const Node = struct {
         assign_add_wrap,
         /// `lhs -%= rhs`. main_token is op.
         assign_sub_wrap,
+        /// `lhs *|= rhs`. main_token is op.
+        assign_mul_sat,
+        /// `lhs +|= rhs`. main_token is op.
+        assign_add_sat,
+        /// `lhs -|= rhs`. main_token is op.
+        assign_sub_sat,
         /// `lhs = rhs`. main_token is op.
         assign,
         /// `lhs || rhs`. main_token is the `||`.
@@ -2553,6 +2577,8 @@ pub const Node = struct {
         array_mult,
         /// `lhs *% rhs`. main_token is the `*%`.
         mul_wrap,
+        /// `lhs *| rhs`. main_token is the `*|`.
+        mul_sat,
         /// `lhs + rhs`. main_token is the `+`.
         add,
         /// `lhs - rhs`. main_token is the `-`.
@@ -2563,10 +2589,16 @@ pub const Node = struct {
         add_wrap,
         /// `lhs -% rhs`. main_token is the `-%`.
         sub_wrap,
+        /// `lhs +| rhs`. main_token is the `+|`.
+        add_sat,
+        /// `lhs -| rhs`. main_token is the `-|`.
+        sub_sat,
         /// `lhs << rhs`. main_token is the `<<`.
-        bit_shift_left,
+        shl,
+        /// `lhs <<| rhs`. main_token is the `<<|`.
+        shl_sat,
         /// `lhs >> rhs`. main_token is the `>>`.
-        bit_shift_right,
+        shr,
         /// `lhs & rhs`. main_token is the `&`.
         bit_and,
         /// `lhs ^ rhs`. main_token is the `^`.
diff --git a/lib/std/zig/parse.zig b/lib/std/zig/parse.zig
index f7697027a3..021b028455 100644
--- a/lib/std/zig/parse.zig
+++ b/lib/std/zig/parse.zig
@@ -1268,14 +1268,18 @@ const Parser = struct {
             .percent_equal => .assign_mod,
             .plus_equal => .assign_add,
             .minus_equal => .assign_sub,
-            .angle_bracket_angle_bracket_left_equal => .assign_bit_shift_left,
-            .angle_bracket_angle_bracket_right_equal => .assign_bit_shift_right,
+            .angle_bracket_angle_bracket_left_equal => .assign_shl,
+            .angle_bracket_angle_bracket_left_pipe_equal => .assign_shl_sat,
+            .angle_bracket_angle_bracket_right_equal => .assign_shr,
             .ampersand_equal => .assign_bit_and,
             .caret_equal => .assign_bit_xor,
             .pipe_equal => .assign_bit_or,
             .asterisk_percent_equal => .assign_mul_wrap,
             .plus_percent_equal => .assign_add_wrap,
             .minus_percent_equal => .assign_sub_wrap,
+            .asterisk_pipe_equal => .assign_mul_sat,
+            .plus_pipe_equal => .assign_add_sat,
+            .minus_pipe_equal => .assign_sub_sat,
             .equal => .assign,
             else => return expr,
         };
@@ -1342,14 +1346,17 @@ const Parser = struct {
         .keyword_orelse = .{ .prec = 40, .tag = .@"orelse" },
         .keyword_catch = .{ .prec = 40, .tag = .@"catch" },
 
-        .angle_bracket_angle_bracket_left = .{ .prec = 50, .tag = .bit_shift_left },
-        .angle_bracket_angle_bracket_right = .{ .prec = 50, .tag = .bit_shift_right },
+        .angle_bracket_angle_bracket_left = .{ .prec = 50, .tag = .shl },
+        .angle_bracket_angle_bracket_left_pipe = .{ .prec = 50, .tag = .shl_sat },
+        .angle_bracket_angle_bracket_right = .{ .prec = 50, .tag = .shr },
 
         .plus = .{ .prec = 60, .tag = .add },
         .minus = .{ .prec = 60, .tag = .sub },
         .plus_plus = .{ .prec = 60, .tag = .array_cat },
         .plus_percent = .{ .prec = 60, .tag = .add_wrap },
         .minus_percent = .{ .prec = 60, .tag = .sub_wrap },
+        .plus_pipe = .{ .prec = 60, .tag = .add_sat },
+        .minus_pipe = .{ .prec = 60, .tag = .sub_sat },
 
         .pipe_pipe = .{ .prec = 70, .tag = .merge_error_sets },
         .asterisk = .{ .prec = 70, .tag = .mul },
@@ -1357,6 +1364,7 @@ const Parser = struct {
         .percent = .{ .prec = 70, .tag = .mod },
         .asterisk_asterisk = .{ .prec = 70, .tag = .array_mult },
         .asterisk_percent = .{ .prec = 70, .tag = .mul_wrap },
+        .asterisk_pipe = .{ .prec = 70, .tag = .mul_sat },
     });
 
     fn parseExprPrecedence(p: *Parser, min_prec: i32) Error!Node.Index {
diff --git a/lib/std/zig/parser_test.zig b/lib/std/zig/parser_test.zig
index 2f79cc175c..57f081decb 100644
--- a/lib/std/zig/parser_test.zig
+++ b/lib/std/zig/parser_test.zig
@@ -4739,6 +4739,26 @@ test "zig fmt: assignment with inline for and inline while" {
     );
 }
 
+test "zig fmt: saturating arithmetic" {
+    try testCanonical(
+        \\test {
+        \\    const actual = switch (op) {
+        \\        .add => a +| b,
+        \\        .sub => a -| b,
+        \\        .mul => a *| b,
+        \\        .shl => a <<| b,
+        \\    };
+        \\    switch (op) {
+        \\        .add => actual +|= b,
+        \\        .sub => actual -|= b,
+        \\        .mul => actual *|= b,
+        \\        .shl => actual <<|= b,
+        \\    }
+        \\}
+        \\
+    );
+}
+
 test "zig fmt: insert trailing comma if there are comments between switch values" {
     try testTransform(
         \\const a = switch (b) {
diff --git a/lib/std/zig/render.zig b/lib/std/zig/render.zig
index 3029d38cb9..4357960251 100644
--- a/lib/std/zig/render.zig
+++ b/lib/std/zig/render.zig
@@ -333,27 +333,33 @@ fn renderExpression(gpa: *Allocator, ais: *Ais, tree: Ast, node: Ast.Node.Index,
 
         .add,
         .add_wrap,
+        .add_sat,
         .array_cat,
         .array_mult,
         .assign,
         .assign_bit_and,
         .assign_bit_or,
-        .assign_bit_shift_left,
-        .assign_bit_shift_right,
+        .assign_shl,
+        .assign_shl_sat,
+        .assign_shr,
         .assign_bit_xor,
         .assign_div,
         .assign_sub,
         .assign_sub_wrap,
+        .assign_sub_sat,
         .assign_mod,
         .assign_add,
         .assign_add_wrap,
+        .assign_add_sat,
         .assign_mul,
         .assign_mul_wrap,
+        .assign_mul_sat,
         .bang_equal,
         .bit_and,
         .bit_or,
-        .bit_shift_left,
-        .bit_shift_right,
+        .shl,
+        .shl_sat,
+        .shr,
         .bit_xor,
         .bool_and,
         .bool_or,
@@ -367,8 +373,10 @@ fn renderExpression(gpa: *Allocator, ais: *Ais, tree: Ast, node: Ast.Node.Index,
         .mod,
         .mul,
         .mul_wrap,
+        .mul_sat,
         .sub,
         .sub_wrap,
+        .sub_sat,
         .@"orelse",
         => {
             const infix = datas[node];
@@ -2520,8 +2528,8 @@ fn nodeCausesSliceOpSpace(tag: Ast.Node.Tag) bool {
         .assign,
         .assign_bit_and,
         .assign_bit_or,
-        .assign_bit_shift_left,
-        .assign_bit_shift_right,
+        .assign_shl,
+        .assign_shr,
         .assign_bit_xor,
         .assign_div,
         .assign_sub,
@@ -2534,8 +2542,8 @@ fn nodeCausesSliceOpSpace(tag: Ast.Node.Tag) bool {
         .bang_equal,
         .bit_and,
         .bit_or,
-        .bit_shift_left,
-        .bit_shift_right,
+        .shl,
+        .shr,
         .bit_xor,
         .bool_and,
         .bool_or,
diff --git a/lib/std/zig/tokenizer.zig b/lib/std/zig/tokenizer.zig
index 3ef6c9a6ba..02fa3dd381 100644
--- a/lib/std/zig/tokenizer.zig
+++ b/lib/std/zig/tokenizer.zig
@@ -103,15 +103,21 @@ pub const Token = struct {
         plus_equal,
         plus_percent,
         plus_percent_equal,
+        plus_pipe,
+        plus_pipe_equal,
         minus,
         minus_equal,
         minus_percent,
         minus_percent_equal,
+        minus_pipe,
+        minus_pipe_equal,
         asterisk,
         asterisk_equal,
         asterisk_asterisk,
         asterisk_percent,
         asterisk_percent_equal,
+        asterisk_pipe,
+        asterisk_pipe_equal,
         arrow,
         colon,
         slash,
@@ -124,6 +130,8 @@ pub const Token = struct {
         angle_bracket_left_equal,
         angle_bracket_angle_bracket_left,
         angle_bracket_angle_bracket_left_equal,
+        angle_bracket_angle_bracket_left_pipe,
+        angle_bracket_angle_bracket_left_pipe_equal,
         angle_bracket_right,
         angle_bracket_right_equal,
         angle_bracket_angle_bracket_right,
@@ -227,15 +235,21 @@ pub const Token = struct {
                 .plus_equal => "+=",
                 .plus_percent => "+%",
                 .plus_percent_equal => "+%=",
+                .plus_pipe => "+|",
+                .plus_pipe_equal => "+|=",
                 .minus => "-",
                 .minus_equal => "-=",
                 .minus_percent => "-%",
                 .minus_percent_equal => "-%=",
+                .minus_pipe => "-|",
+                .minus_pipe_equal => "-|=",
                 .asterisk => "*",
                 .asterisk_equal => "*=",
                 .asterisk_asterisk => "**",
                 .asterisk_percent => "*%",
                 .asterisk_percent_equal => "*%=",
+                .asterisk_pipe => "*|",
+                .asterisk_pipe_equal => "*|=",
                 .arrow => "->",
                 .colon => ":",
                 .slash => "/",
@@ -248,6 +262,8 @@ pub const Token = struct {
                 .angle_bracket_left_equal => "<=",
                 .angle_bracket_angle_bracket_left => "<<",
                 .angle_bracket_angle_bracket_left_equal => "<<=",
+                .angle_bracket_angle_bracket_left_pipe => "<<|",
+                .angle_bracket_angle_bracket_left_pipe_equal => "<<|=",
                 .angle_bracket_right => ">",
                 .angle_bracket_right_equal => ">=",
                 .angle_bracket_angle_bracket_right => ">>",
@@ -352,8 +368,10 @@ pub const Tokenizer = struct {
         pipe,
         minus,
         minus_percent,
+        minus_pipe,
         asterisk,
         asterisk_percent,
+        asterisk_pipe,
         slash,
         line_comment_start,
         line_comment,
@@ -382,8 +400,10 @@ pub const Tokenizer = struct {
         percent,
         plus,
         plus_percent,
+        plus_pipe,
         angle_bracket_left,
         angle_bracket_angle_bracket_left,
+        angle_bracket_angle_bracket_left_pipe,
         angle_bracket_right,
         angle_bracket_angle_bracket_right,
         period,
@@ -584,6 +604,9 @@ pub const Tokenizer = struct {
                     '%' => {
                         state = .asterisk_percent;
                     },
+                    '|' => {
+                        state = .asterisk_pipe;
+                    },
                     else => {
                         result.tag = .asterisk;
                         break;
@@ -602,6 +625,18 @@ pub const Tokenizer = struct {
                     },
                 },
 
+                .asterisk_pipe => switch (c) {
+                    '=' => {
+                        result.tag = .asterisk_pipe_equal;
+                        self.index += 1;
+                        break;
+                    },
+                    else => {
+                        result.tag = .asterisk_pipe;
+                        break;
+                    },
+                },
+
                 .percent => switch (c) {
                     '=' => {
                         result.tag = .percent_equal;
@@ -628,6 +663,9 @@ pub const Tokenizer = struct {
                     '%' => {
                         state = .plus_percent;
                     },
+                    '|' => {
+                        state = .plus_pipe;
+                    },
                     else => {
                         result.tag = .plus;
                         break;
@@ -646,6 +684,18 @@ pub const Tokenizer = struct {
                     },
                 },
 
+                .plus_pipe => switch (c) {
+                    '=' => {
+                        result.tag = .plus_pipe_equal;
+                        self.index += 1;
+                        break;
+                    },
+                    else => {
+                        result.tag = .plus_pipe;
+                        break;
+                    },
+                },
+
                 .caret => switch (c) {
                     '=' => {
                         result.tag = .caret_equal;
@@ -903,6 +953,9 @@ pub const Tokenizer = struct {
                     '%' => {
                         state = .minus_percent;
                     },
+                    '|' => {
+                        state = .minus_pipe;
+                    },
                     else => {
                         result.tag = .minus;
                         break;
@@ -920,6 +973,17 @@ pub const Tokenizer = struct {
                         break;
                     },
                 },
+                .minus_pipe => switch (c) {
+                    '=' => {
+                        result.tag = .minus_pipe_equal;
+                        self.index += 1;
+                        break;
+                    },
+                    else => {
+                        result.tag = .minus_pipe;
+                        break;
+                    },
+                },
 
                 .angle_bracket_left => switch (c) {
                     '<' => {
@@ -942,12 +1006,27 @@ pub const Tokenizer = struct {
                         self.index += 1;
                         break;
                     },
+                    '|' => {
+                        state = .angle_bracket_angle_bracket_left_pipe;
+                    },
                     else => {
                         result.tag = .angle_bracket_angle_bracket_left;
                         break;
                     },
                 },
 
+                .angle_bracket_angle_bracket_left_pipe => switch (c) {
+                    '=' => {
+                        result.tag = .angle_bracket_angle_bracket_left_pipe_equal;
+                        self.index += 1;
+                        break;
+                    },
+                    else => {
+                        result.tag = .angle_bracket_angle_bracket_left_pipe;
+                        break;
+                    },
+                },
+
                 .angle_bracket_right => switch (c) {
                     '>' => {
                         state = .angle_bracket_angle_bracket_right;
@@ -1936,6 +2015,24 @@ test "tokenizer - invalid token with unfinished escape right before eof" {
     try testTokenize("'\\u", &.{.invalid});
 }
 
+test "tokenizer - saturating" {
+    try testTokenize("<<", &.{.angle_bracket_angle_bracket_left});
+    try testTokenize("<<|", &.{.angle_bracket_angle_bracket_left_pipe});
+    try testTokenize("<<|=", &.{.angle_bracket_angle_bracket_left_pipe_equal});
+
+    try testTokenize("*", &.{.asterisk});
+    try testTokenize("*|", &.{.asterisk_pipe});
+    try testTokenize("*|=", &.{.asterisk_pipe_equal});
+
+    try testTokenize("+", &.{.plus});
+    try testTokenize("+|", &.{.plus_pipe});
+    try testTokenize("+|=", &.{.plus_pipe_equal});
+
+    try testTokenize("-", &.{.minus});
+    try testTokenize("-|", &.{.minus_pipe});
+    try testTokenize("-|=", &.{.minus_pipe_equal});
+}
+
 fn testTokenize(source: [:0]const u8, expected_tokens: []const Token.Tag) !void {
     var tokenizer = Tokenizer.init(source);
     for (expected_tokens) |expected_token_id| {