std.math.big.int: implement popCount() for Const

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;

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" {