diff --git a/lib/std/math/big/int.zig b/lib/std/math/big/int.zig index 189093aa84..1040a6dc8a 100644 --- a/lib/std/math/big/int.zig +++ b/lib/std/math/big/int.zig @@ -1959,6 +1959,52 @@ pub const Const = struct { return bits; } + /// @popCount with two's complement semantics. + /// + /// This returns the number of 1 bits set when the value would be represented in + /// two's complement with the given integer width (bit_count). + /// This includes the leading sign bit, which will be set for negative values. + /// + /// Asserts that bit_count is enough to represent value in two's compliment + /// and that the final result fits in a usize. + /// Asserts that there are no trailing empty limbs on the most significant end, + /// i.e. that limb count matches `calcLimbLen()` and zero is not negative. + pub fn popCount(self: Const, bit_count: usize) usize { + var sum: usize = 0; + if (self.positive) { + for (self.limbs) |limb| { + sum += @popCount(limb); + } + } else { + assert(self.fitsInTwosComp(.signed, bit_count)); + assert(self.limbs[self.limbs.len - 1] != 0); + + var remaining_bits = bit_count; + var carry: u1 = 1; + var add_res: Limb = undefined; + + // All but the most significant limb. + for (self.limbs[0 .. self.limbs.len - 1]) |limb| { + carry = @boolToInt(@addWithOverflow(Limb, ~limb, carry, &add_res)); + sum += @popCount(add_res); + remaining_bits -= limb_bits; // Asserted not to undeflow by fitsInTwosComp + } + + // The most significant limb may have fewer than @bitSizeOf(Limb) meaningful bits, + // which we can detect with @clz(). + // There may also be fewer limbs than needed to fill bit_count. + const limb = self.limbs[self.limbs.len - 1]; + const leading_zeroes = @clz(limb); + // The most significant limb is asserted not to be all 0s (above), + // so ~limb cannot be all 1s, and ~limb + 1 cannot overflow. + sum += @popCount(~limb + carry); + sum -= leading_zeroes; // All leading zeroes were flipped and added to sum, so undo those + const remaining_ones = remaining_bits - (limb_bits - leading_zeroes); // All bits not covered by limbs + sum += remaining_ones; + } + return sum; + } + pub fn fitsInTwosComp(self: Const, signedness: Signedness, bit_count: usize) bool { if (self.eqZero()) { return true; diff --git a/lib/std/math/big/int_test.zig b/lib/std/math/big/int_test.zig index 97de06bfcc..57211ae299 100644 --- a/lib/std/math/big/int_test.zig +++ b/lib/std/math/big/int_test.zig @@ -2578,14 +2578,86 @@ test "big.int regression test for realloc with alias" { } test "big int popcount" { - var a = try Managed.initSet(testing.allocator, -1); + var a = try Managed.init(testing.allocator); defer a.deinit(); - var b = try Managed.initSet(testing.allocator, -1); + + try a.set(0); + try popCountTest(&a, 0, 0); + try popCountTest(&a, 567, 0); + try a.set(-0); + try popCountTest(&a, 0, 0); + + try a.set(1); + try popCountTest(&a, 1, 1); + try popCountTest(&a, 13, 1); + try popCountTest(&a, 432, 1); + + try a.set(255); + try popCountTest(&a, 8, 8); + try a.set(-128); + try popCountTest(&a, 8, 1); + + try a.set(-2); + try popCountTest(&a, 16, 15); + try popCountTest(&a, 15, 14); + + try a.set(-2047); + try popCountTest(&a, 12, 2); + try popCountTest(&a, 24, 14); + + try a.set(maxInt(u5000)); + try popCountTest(&a, 5000, 5000); + try a.set(minInt(i5000)); + try popCountTest(&a, 5000, 1); + + // Check -1 at various bit counts that cross Limb size multiples. + const limb_bits = @bitSizeOf(Limb); + try a.set(-1); + try popCountTest(&a, 1, 1); // i1 + try popCountTest(&a, 2, 2); + try popCountTest(&a, 16, 16); + try popCountTest(&a, 543, 543); + try popCountTest(&a, 544, 544); + try popCountTest(&a, limb_bits - 1, limb_bits - 1); + try popCountTest(&a, limb_bits, limb_bits); + try popCountTest(&a, limb_bits + 1, limb_bits + 1); + try popCountTest(&a, limb_bits * 2 - 1, limb_bits * 2 - 1); + try popCountTest(&a, limb_bits * 2, limb_bits * 2); + try popCountTest(&a, limb_bits * 2 + 1, limb_bits * 2 + 1); + + // Check very large numbers. + try a.setString(16, "ff00000100000100" ++ ("0000000000000000" ** 62)); + try popCountTest(&a, 4032, 10); + try popCountTest(&a, 6000, 10); + a.negate(); + try popCountTest(&a, 4033, 48); + try popCountTest(&a, 4133, 148); + + // Check when most significant limb is full of 1s. + const limb_size = @bitSizeOf(Limb); + try a.set(maxInt(Limb)); + try popCountTest(&a, limb_size, limb_size); + try popCountTest(&a, limb_size + 1, limb_size); + try popCountTest(&a, limb_size * 10 + 2, limb_size); + a.negate(); + try popCountTest(&a, limb_size * 2 - 2, limb_size - 1); + try popCountTest(&a, limb_size * 2 - 1, limb_size); + try popCountTest(&a, limb_size * 2, limb_size + 1); + try popCountTest(&a, limb_size * 2 + 1, limb_size + 2); + // TODO: These produce incorrect pop count for Mutable + // https://github.com/ziglang/zig/issues/13571 + // try popCountTest(&a, limb_size * 2 + 2, limb_size + 3); + // try popCountTest(&a, limb_size * 2 + 3, limb_size + 4); + // try popCountTest(&a, limb_size * 2 + 4, limb_size + 5); +} + +fn popCountTest(val: *const Managed, bit_count: usize, expected: usize) !void { + var b = try Managed.init(testing.allocator); defer b.deinit(); + try b.popCount(val, bit_count); - try a.popCount(&b, 16); - - try testing.expect(a.toConst().orderAgainstScalar(16) == .eq); + try testing.expectEqual(std.math.Order.eq, b.toConst().orderAgainstScalar(expected)); + try testing.expectEqual(expected, val.toConst().popCount(bit_count)); } test "big int conversion read/write twos complement" {