lib/choir/src/dialects/arith/scalar.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 
  3 pub fn maskForBits(bits: u8) u64 {
  4     if (bits >= 64) return std.math.maxInt(u64);
  5     const shift: u6 = @intCast(bits);
  6     return (@as(u64, 1) << shift) - 1;
  7 }
  8 
  9 pub fn maskToBits(value: i64, bits: u8) u64 {
 10     return @as(u64, @bitCast(value)) & maskForBits(bits);
 11 }
 12 
 13 pub fn signExtend(value: u64, bits: u8) i64 {
 14     if (bits >= 64) return @as(i64, @bitCast(value));
 15     const shift: u6 = @intCast(bits - 1);
 16     const sign_bit = @as(u64, 1) << shift;
 17     const mask = maskForBits(bits);
 18     if ((value & sign_bit) != 0) {
 19         return @as(i64, @bitCast(value | ~mask));
 20     }
 21     return @as(i64, @bitCast(value));
 22 }
 23 
 24 pub const Limits = struct { min: i128, max: i128 };
 25 
 26 pub fn intLimits(bits: u8) Limits {
 27     if (bits >= 64) {
 28         return .{ .min = @as(i128, std.math.minInt(i64)), .max = @as(i128, std.math.maxInt(i64)) };
 29     }
 30     const shift: u7 = @intCast(bits - 1);
 31     const max = (@as(i128, 1) << shift) - 1;
 32     const min = -(@as(i128, 1) << shift);
 33     return .{ .min = min, .max = max };
 34 }
 35 
 36 pub fn fitsSigned(value: i128, bits: u8) bool {
 37     const limits = intLimits(bits);
 38     return value >= limits.min and value <= limits.max;
 39 }
 40 
 41 pub fn truncate(value: i64, bits: u8) i64 {
 42     return signExtend(maskToBits(value, bits), bits);
 43 }
 44 
 45 pub fn addWrap(lhs: i64, rhs: i64, bits: u8) i64 {
 46     return truncate(lhs +% rhs, bits);
 47 }
 48 
 49 pub fn subWrap(lhs: i64, rhs: i64, bits: u8) i64 {
 50     return truncate(lhs -% rhs, bits);
 51 }
 52 
 53 pub fn mulWrap(lhs: i64, rhs: i64, bits: u8) i64 {
 54     return truncate(lhs *% rhs, bits);
 55 }
 56 
 57 pub fn negWrap(value: i64, bits: u8) i64 {
 58     return truncate(0 -% value, bits);
 59 }
 60 
 61 pub fn absWrap(value: i64, bits: u8) i64 {
 62     if (value >= 0) return truncate(value, bits);
 63     return negWrap(value, bits);
 64 }
 65 
 66 pub fn addChecked(lhs: i64, rhs: i64, bits: u8) ?i64 {
 67     const res = @as(i128, lhs) + rhs;
 68     if (!fitsSigned(res, bits)) return null;
 69     return @as(i64, @intCast(res));
 70 }
 71 
 72 pub fn subChecked(lhs: i64, rhs: i64, bits: u8) ?i64 {
 73     const res = @as(i128, lhs) - rhs;
 74     if (!fitsSigned(res, bits)) return null;
 75     return @as(i64, @intCast(res));
 76 }
 77 
 78 pub fn mulChecked(lhs: i64, rhs: i64, bits: u8) ?i64 {
 79     const res = @as(i128, lhs) * rhs;
 80     if (!fitsSigned(res, bits)) return null;
 81     return @as(i64, @intCast(res));
 82 }
 83 
 84 pub fn negChecked(value: i64, bits: u8) ?i64 {
 85     const limits = intLimits(bits);
 86     if (@as(i128, value) == limits.min) return null;
 87     return -value;
 88 }
 89 
 90 pub fn absChecked(value: i64, bits: u8) ?i64 {
 91     if (value >= 0) return value;
 92     return negChecked(value, bits);
 93 }
 94 
 95 pub fn divTruncChecked(lhs: i64, rhs: i64, bits: u8) ?i64 {
 96     const left = truncate(lhs, bits);
 97     const right = truncate(rhs, bits);
 98     if (right == 0) return null;
 99     const limits = intLimits(bits);
100     if (@as(i128, left) == limits.min and right == -1) return null;
101     return truncate(@divTrunc(left, right), bits);
102 }
103 
104 pub fn remTruncChecked(lhs: i64, rhs: i64, bits: u8) ?i64 {
105     const left = truncate(lhs, bits);
106     const right = truncate(rhs, bits);
107     if (right == 0) return null;
108     const limits = intLimits(bits);
109     if (@as(i128, left) == limits.min and right == -1) return null;
110     return truncate(@rem(left, right), bits);
111 }
112 
113 pub fn unsignedResult(value: u64, bits: u8) i64 {
114     if (bits >= 64) return @as(i64, @bitCast(value));
115     return @intCast(value & maskForBits(bits));
116 }
117 
118 pub fn divTruncUnsignedChecked(lhs: i64, rhs: i64, bits: u8) ?i64 {
119     const left = maskToBits(lhs, bits);
120     const right = maskToBits(rhs, bits);
121     if (right == 0) return null;
122     return unsignedResult(@divTrunc(left, right), bits);
123 }
124 
125 pub fn remTruncUnsignedChecked(lhs: i64, rhs: i64, bits: u8) ?i64 {
126     const left = maskToBits(lhs, bits);
127     const right = maskToBits(rhs, bits);
128     if (right == 0) return null;
129     return unsignedResult(@rem(left, right), bits);
130 }
131 
132 pub fn bitAnd(lhs: i64, rhs: i64, bits: u8) i64 {
133     return signExtend(maskToBits(lhs, bits) & maskToBits(rhs, bits), bits);
134 }
135 
136 pub fn bitOr(lhs: i64, rhs: i64, bits: u8) i64 {
137     return signExtend(maskToBits(lhs, bits) | maskToBits(rhs, bits), bits);
138 }
139 
140 pub fn bitXor(lhs: i64, rhs: i64, bits: u8) i64 {
141     return signExtend(maskToBits(lhs, bits) ^ maskToBits(rhs, bits), bits);
142 }
143 
144 pub fn bitNot(value: i64, bits: u8) i64 {
145     return signExtend(~maskToBits(value, bits) & maskForBits(bits), bits);
146 }
147 
148 pub fn shiftCount(value: i64, bits: u8) ?u6 {
149     if (value < 0) return null;
150     const count: u64 = @intCast(value);
151     if (count >= bits) return null;
152     return @intCast(count);
153 }
154 
155 pub fn shiftLeftWrap(value: i64, count: u6, bits: u8) i64 {
156     const shifted = maskToBits(value, bits) *% (@as(u64, 1) << count);
157     return signExtend(shifted & maskForBits(bits), bits);
158 }
159 
160 pub fn shiftRightArithmetic(value: i64, count: u6, bits: u8) i64 {
161     return truncate(truncate(value, bits) >> count, bits);
162 }
163 
164 pub fn shiftRightLogical(value: i64, count: u6, bits: u8) i64 {
165     return signExtend(maskToBits(value, bits) >> count, bits);
166 }
167 
168 test "wrapping arithmetic at 64-bit boundaries" {
169     const testing = std.testing;
170     try testing.expectEqual(std.math.minInt(i64), addWrap(std.math.maxInt(i64), 1, 64));
171     try testing.expectEqual(std.math.maxInt(i64), subWrap(std.math.minInt(i64), 1, 64));
172     try testing.expectEqual(@as(i64, 0), mulWrap(1 << 32, 1 << 32, 64));
173     try testing.expectEqual(std.math.minInt(i64), negWrap(std.math.minInt(i64), 64));
174     try testing.expectEqual(@as(i64, 5), addWrap(2, 3, 64));
175 }
176 
177 test "wrapping arithmetic truncates to declared width" {
178     const testing = std.testing;
179     try testing.expectEqual(@as(i64, -128), addWrap(127, 1, 8));
180     try testing.expectEqual(@as(i64, 127), addWrap(-128, -1, 8));
181     try testing.expectEqual(@as(i64, -1), signExtend(0xff, 8));
182     try testing.expectEqual(@as(i64, 127), signExtend(0x7f, 8));
183     try testing.expectEqual(@as(i64, -2), mulWrap(0x7fff_ffff, 2, 32));
184 }
185 
186 test "checked arithmetic refuses out-of-range results" {
187     const testing = std.testing;
188     try testing.expectEqual(@as(?i64, null), addChecked(std.math.maxInt(i64), 1, 64));
189     try testing.expectEqual(@as(?i64, 5), addChecked(2, 3, 64));
190     try testing.expectEqual(@as(?i64, 120), addChecked(60, 60, 8));
191     try testing.expectEqual(@as(?i64, null), addChecked(64, 64, 8));
192     try testing.expectEqual(@as(?i64, null), negChecked(std.math.minInt(i64), 64));
193     try testing.expectEqual(@as(?i64, null), mulChecked(std.math.maxInt(i64), 2, 64));
194 }
195 
196 test "division and remainder guard against zero and overflow" {
197     const testing = std.testing;
198     try testing.expectEqual(@as(?i64, 3), divTruncChecked(7, 2, 64));
199     try testing.expectEqual(@as(?i64, -3), divTruncChecked(-7, 2, 64));
200     try testing.expectEqual(@as(?i64, null), divTruncChecked(1, 0, 64));
201     try testing.expectEqual(@as(?i64, null), divTruncChecked(std.math.minInt(i64), -1, 64));
202     try testing.expectEqual(@as(?i64, 1), remTruncChecked(7, 3, 64));
203     try testing.expectEqual(@as(?i64, -1), remTruncChecked(-7, 3, 64));
204     try testing.expectEqual(@as(?i64, null), remTruncChecked(1, 0, 64));
205     try testing.expectEqual(@as(?i64, null), remTruncChecked(std.math.minInt(i64), -1, 64));
206 }
207 
208 test "unsigned division and remainder use masked ordering" {
209     const testing = std.testing;
210     try testing.expectEqual(@as(?i64, @as(i64, @bitCast(@as(u64, 0x7fff_ffff_ffff_ffff)))), divTruncUnsignedChecked(-1, 2, 64));
211     try testing.expectEqual(@as(?i64, 1), remTruncUnsignedChecked(-1, 2, 64));
212     try testing.expectEqual(@as(?i64, 127), divTruncUnsignedChecked(-1, 2, 8));
213     try testing.expectEqual(@as(?i64, 1), remTruncUnsignedChecked(-1, 2, 8));
214     try testing.expectEqual(@as(?i64, null), divTruncUnsignedChecked(1, 0, 64));
215     try testing.expectEqual(@as(?i64, null), remTruncUnsignedChecked(1, 0, 64));
216 }
217 
218 test "bitwise operations respect width" {
219     const testing = std.testing;
220     try testing.expectEqual(@as(i64, -1), bitNot(0, 64));
221     try testing.expectEqual(@as(i64, -1), bitNot(0, 8));
222     try testing.expectEqual(@as(i64, 0), bitAnd(-1, 0, 64));
223     try testing.expectEqual(@as(i64, -1), bitOr(0, -1, 64));
224     try testing.expectEqual(@as(i64, 0), bitXor(0x55, 0x55, 8));
225     try testing.expectEqual(@as(i64, -86), bitXor(0x55, -1, 8));
226 }
227 
228 test "shift helpers require bounded counts and preserve bit semantics" {
229     const testing = std.testing;
230     try testing.expectEqual(@as(?u6, 0), shiftCount(0, 64));
231     try testing.expectEqual(@as(?u6, 63), shiftCount(63, 64));
232     try testing.expectEqual(@as(?u6, null), shiftCount(64, 64));
233     try testing.expectEqual(@as(?u6, null), shiftCount(-1, 64));
234     try testing.expectEqual(@as(i64, 16), shiftLeftWrap(1, 4, 64));
235     try testing.expectEqual(@as(i64, std.math.minInt(i64)), shiftLeftWrap(1, 63, 64));
236     try testing.expectEqual(@as(i64, -4), shiftRightArithmetic(-16, 2, 64));
237     try testing.expectEqual(@as(i64, std.math.maxInt(i64)), shiftRightLogical(-1, 1, 64));
238     try testing.expectEqual(@as(i64, 127), shiftRightLogical(-1, 1, 8));
239 }
240 
241 /// Ruling 48: test the finite truncation against exact exclusive upper bounds.
242 pub fn floatToInt(value: f64, bits: u8, signed: bool) ?i64 {
243     std.debug.assert(bits > 0);
244     std.debug.assert(bits <= 64);
245     if (!std.math.isFinite(value)) return null;
246     const truncated = @trunc(value);
247     const exponent: u7 = @intCast(bits - @as(u8, if (signed) 1 else 0));
248     const upper: f64 = @floatFromInt(@as(u128, 1) << exponent);
249     const lower: f64 = if (signed) -upper else 0;
250     if (truncated < lower or truncated >= upper) return null;
251     return if (signed) @intFromFloat(truncated) else @bitCast(@as(u64, @intFromFloat(truncated)));
252 }
253 
254 /// Round f64 directly to bfloat16, correcting ties introduced by the f32 intermediate.
255 pub fn roundBfloat(value: f64) f64 {
256     if (std.math.isNan(value)) return std.math.nan(f64);
257     const narrow: f32 = @floatCast(value);
258     const bits: u32 = @bitCast(narrow);
259     if (!std.math.isFinite(narrow)) return narrow;
260     const remainder = bits & 0xffff;
261     var high = bits >> 16;
262     const up = if (remainder == 0x8000 and @abs(value) != @abs(@as(f64, narrow)))
263         @abs(value) > @abs(@as(f64, narrow))
264     else
265         remainder > 0x8000 or (remainder == 0x8000 and high & 1 != 0);
266     if (up) high += 1;
267     return @as(f32, @bitCast(high << 16));
268 }
269 
270 test "Precision1 bfloat rounding preserves ties and avoids an intermediate f32 tie" {
271     const testing = std.testing;
272     try testing.expectEqual(@as(f64, 256), roundBfloat(257));
273     try testing.expectEqual(@as(f64, 260), roundBfloat(259));
274     try testing.expectEqual(@as(f64, 258), roundBfloat(257 + 0x1p-20));
275     try testing.expectEqual(@as(f64, 256), roundBfloat(257 - 0x1p-20));
276     try testing.expectEqual(@as(f64, -258), roundBfloat(-257 - 0x1p-20));
277     try testing.expectEqual(@as(f64, 0), roundBfloat(0x1p-134));
278     try testing.expectEqual(@as(f64, 0x1p-133), roundBfloat(0x1p-134 + 0x1p-150));
279     try testing.expect(std.math.isNan(roundBfloat(std.math.nan(f64))));
280     try testing.expectEqual(std.math.inf(f64), roundBfloat(std.math.inf(f64)));
281 }