From 5433e0438ce464adea53e4849a80d3854b2a036d Mon Sep 17 00:00:00 2001
From: Jacob Young <jacobly0@users.noreply.github.com>
Date: Sun, 9 Feb 2025 08:10:20 -0500
Subject: [PATCH] cbe: fix ub triggered by mulw overflowing the promoted type

Closes #21914
---
 lib/zig.h              | 44 ++++++++++++++++++++++++++++-----
 test/behavior/math.zig | 56 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 94 insertions(+), 6 deletions(-)

diff --git a/lib/zig.h b/lib/zig.h
index f8aa3fd24e..863230d705 100644
--- a/lib/zig.h
+++ b/lib/zig.h
@@ -650,7 +650,7 @@ typedef ptrdiff_t intptr_t;
     zig_operator(Type,    Type, operation, operator)
 #define zig_shift_operator(Type, operation, operator) \
     zig_operator(Type, uint8_t, operation, operator)
-#define zig_int_helpers(w) \
+#define zig_int_helpers(w, PromotedUnsigned) \
     zig_basic_operator(uint##w##_t, and_u##w,  &) \
     zig_basic_operator( int##w##_t, and_i##w,  &) \
     zig_basic_operator(uint##w##_t,  or_u##w,  |) \
@@ -726,16 +726,48 @@ typedef ptrdiff_t intptr_t;
     } \
 \
     static inline uint##w##_t zig_mulw_u##w(uint##w##_t lhs, uint##w##_t rhs, uint8_t bits) { \
-        return zig_wrap_u##w(lhs * rhs, bits); \
+        return zig_wrap_u##w((PromotedUnsigned)lhs * rhs, bits); \
     } \
 \
     static inline int##w##_t zig_mulw_i##w(int##w##_t lhs, int##w##_t rhs, uint8_t bits) { \
         return zig_wrap_i##w((int##w##_t)((uint##w##_t)lhs * (uint##w##_t)rhs), bits); \
     }
-zig_int_helpers(8)
-zig_int_helpers(16)
-zig_int_helpers(32)
-zig_int_helpers(64)
+#if UINT8_MAX <= UINT_MAX
+zig_int_helpers(8, unsigned int)
+#elif UINT8_MAX <= ULONG_MAX
+zig_int_helpers(8, unsigned long)
+#elif UINT8_MAX <= ULLONG_MAX
+zig_int_helpers(8, unsigned long long)
+#else
+zig_int_helpers(8, uint8_t)
+#endif
+#if UINT16_MAX <= UINT_MAX
+zig_int_helpers(16, unsigned int)
+#elif UINT16_MAX <= ULONG_MAX
+zig_int_helpers(16, unsigned long)
+#elif UINT16_MAX <= ULLONG_MAX
+zig_int_helpers(16, unsigned long long)
+#else
+zig_int_helpers(16, uint16_t)
+#endif
+#if UINT32_MAX <= UINT_MAX
+zig_int_helpers(32, unsigned int)
+#elif UINT32_MAX <= ULONG_MAX
+zig_int_helpers(32, unsigned long)
+#elif UINT32_MAX <= ULLONG_MAX
+zig_int_helpers(32, unsigned long long)
+#else
+zig_int_helpers(32, uint32_t)
+#endif
+#if UINT64_MAX <= UINT_MAX
+zig_int_helpers(64, unsigned int)
+#elif UINT64_MAX <= ULONG_MAX
+zig_int_helpers(64, unsigned long)
+#elif UINT64_MAX <= ULLONG_MAX
+zig_int_helpers(64, unsigned long long)
+#else
+zig_int_helpers(64, uint64_t)
+#endif
 
 static inline bool zig_addo_u32(uint32_t *res, uint32_t lhs, uint32_t rhs, uint8_t bits) {
 #if zig_has_builtin(add_overflow) || defined(zig_gnuc)
diff --git a/test/behavior/math.zig b/test/behavior/math.zig
index fb325a95b6..a658a0559c 100644
--- a/test/behavior/math.zig
+++ b/test/behavior/math.zig
@@ -778,6 +778,62 @@ fn should_not_be_zero(x: f128) !void {
     try expect(x != 0.0);
 }
 
+test "umax wrapped squaring" {
+    if (builtin.zig_backend == .stage2_riscv64) return error.SkipZigTest;
+    if (builtin.zig_backend == .stage2_x86_64) return error.SkipZigTest;
+
+    {
+        var x: u4 = maxInt(u4);
+        x *%= x;
+        try expect(x == 1);
+    }
+    {
+        var x: u8 = maxInt(u8);
+        x *%= x;
+        try expect(x == 1);
+    }
+    {
+        var x: u12 = maxInt(u12);
+        x *%= x;
+        try expect(x == 1);
+    }
+    {
+        var x: u16 = maxInt(u16);
+        x *%= x;
+        try expect(x == 1);
+    }
+    {
+        var x: u24 = maxInt(u24);
+        x *%= x;
+        try expect(x == 1);
+    }
+    {
+        var x: u32 = maxInt(u32);
+        x *%= x;
+        try expect(x == 1);
+    }
+    {
+        var x: u48 = maxInt(u48);
+        x *%= x;
+        try expect(x == 1);
+    }
+    {
+        var x: u64 = maxInt(u64);
+        x *%= x;
+        try expect(x == 1);
+    }
+    {
+        var x: u96 = maxInt(u96);
+        x *%= x;
+        try expect(x == 1);
+    }
+    {
+        var x: u128 = maxInt(u128);
+        x *%= x;
+        try expect(x == 1);
+    }
+}
+
 test "128-bit multiplication" {
     if (builtin.zig_backend == .stage2_aarch64) return error.SkipZigTest; // TODO
     if (builtin.zig_backend == .stage2_arm) return error.SkipZigTest; // TODO