stage2: implement @popCount for non-vectors

Author: andrewrk

lib/std/math/big/int.zig

@@ -733,6 +733,27 @@ pub const Mutable = struct {
         rma.truncate(rma.toConst(), signedness, bit_count);
     }
 
+    /// r = @popCount(a) with 2s-complement semantics.
+    /// r and a may be aliases.
+    ///
+    /// Assets the result fits in `r`. Upper bound on the number of limbs needed by
+    /// r is `calcTwosCompLimbCount(bit_count)`.
+    pub fn popCount(r: *Mutable, a: Const, bit_count: usize) void {
+        r.copy(a);
+
+        if (!a.positive) {
+            r.positive = true; // Negate.
+            r.bitNotWrap(r.toConst(), .unsigned, bit_count); // Bitwise NOT.
+            r.addScalar(r.toConst(), 1); // Add one.
+        }
+
+        var sum: Limb = 0;
+        for (r.limbs[0..r.len]) |limb| {
+            sum += @popCount(Limb, limb);
+        }
+        r.set(sum);
+    }
+
     /// rma = a * a
     ///
     /// `rma` may not alias with `a`.
@@ -2735,6 +2756,15 @@ pub const Managed = struct {
         m.saturate(a, signedness, bit_count);
         r.setMetadata(m.positive, m.len);
     }
+
+    /// r = @popCount(a) with 2s-complement semantics.
+    /// r and a may be aliases.
+    pub fn popCount(r: *Managed, a: Const, bit_count: usize) !void {
+        try r.ensureCapacity(calcTwosCompLimbCount(bit_count));
+        var m = r.toMutable();
+        m.popCount(a, bit_count);
+        r.setMetadata(m.positive, m.len);
+    }
 };
 
 /// Different operators which can be used in accumulation style functions

lib/std/math/big/int_test.zig

@@ -2434,3 +2434,14 @@ test "big.int regression test for realloc with alias" {
 
     try testing.expect(a.toConst().orderAgainstScalar(14691098406862188148944207245954912110548093601382197697835) == .eq);
 }
+
+test "big int popcount" {
+    var a = try Managed.initSet(testing.allocator, -1);
+    defer a.deinit();
+    var b = try Managed.initSet(testing.allocator, -1);
+    defer b.deinit();
+
+    try a.popCount(b.toConst(), 16);
+
+    try testing.expect(a.toConst().orderAgainstScalar(16) == .eq);
+}

src/Air.zig

@@ -202,6 +202,10 @@ pub const Inst = struct {
         /// Result type will always be an unsigned integer big enough to fit the answer.
         /// Uses the `ty_op` field.
         ctz,
+        /// Count number of 1 bits in an integer according to its representation in twos complement.
+        /// Result type will always be an unsigned integer big enough to fit the answer.
+        /// Uses the `ty_op` field.
+        popcount,
 
         /// `<`. Result type is always bool.
         /// Uses the `bin_op` field.
@@ -744,6 +748,7 @@ pub fn typeOfIndex(air: Air, inst: Air.Inst.Index) Type {
         .get_union_tag,
         .clz,
         .ctz,
+        .popcount,
         => return air.getRefType(datas[inst].ty_op.ty),
 
         .loop,

src/Liveness.zig

@@ -313,6 +313,7 @@ fn analyzeInst(
         .get_union_tag,
         .clz,
         .ctz,
+        .popcount,
         => {
             const o = inst_datas[inst].ty_op;
             return trackOperands(a, new_set, inst, main_tomb, .{ o.operand, .none, .none });

src/Sema.zig

@@ -9904,8 +9904,30 @@ fn zirCtz(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air.Ins
 
 fn zirPopCount(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air.Inst.Ref {
     const inst_data = sema.code.instructions.items(.data)[inst].un_node;
-    const src = inst_data.src();
-    return sema.fail(block, src, "TODO: Sema.zirPopCount", .{});
+    const ty_src: LazySrcLoc = .{ .node_offset_builtin_call_arg0 = inst_data.src_node };
+    const operand_src: LazySrcLoc = .{ .node_offset_builtin_call_arg1 = inst_data.src_node };
+    const operand = sema.resolveInst(inst_data.operand);
+    const operand_ty = sema.typeOf(operand);
+    // TODO implement support for vectors
+    if (operand_ty.zigTypeTag() != .Int) {
+        return sema.fail(block, ty_src, "expected integer type, found '{}'", .{
+            operand_ty,
+        });
+    }
+    const target = sema.mod.getTarget();
+    const bits = operand_ty.intInfo(target).bits;
+    if (bits == 0) return Air.Inst.Ref.zero;
+
+    const result_ty = try Type.smallestUnsignedInt(sema.arena, bits);
+
+    const runtime_src = if (try sema.resolveMaybeUndefVal(block, operand_src, operand)) |val| {
+        if (val.isUndef()) return sema.addConstUndef(result_ty);
+        const result_val = try val.popCount(operand_ty, target, sema.arena);
+        return sema.addConstant(result_ty, result_val);
+    } else operand_src;
+
+    try sema.requireRuntimeBlock(block, runtime_src);
+    return block.addTyOp(.popcount, result_ty, operand);
 }
 
 fn zirByteSwap(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air.Inst.Ref {

src/arch/aarch64/CodeGen.zig

@@ -481,6 +481,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
                     .get_union_tag   => try self.airGetUnionTag(inst),
                     .clz             => try self.airClz(inst),
                     .ctz             => try self.airCtz(inst),
+                    .popcount        => try self.airPopcount(inst),
 
                     .atomic_store_unordered => try self.airAtomicStore(inst, .Unordered),
                     .atomic_store_monotonic => try self.airAtomicStore(inst, .Monotonic),
@@ -1138,6 +1139,12 @@ fn airCtz(self: *Self, inst: Air.Inst.Index) !void {
     return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
 }
 
+fn airPopcount(self: *Self, inst: Air.Inst.Index) !void {
+    const ty_op = self.air.instructions.items(.data)[inst].ty_op;
+    const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airPopcount for {}", .{self.target.cpu.arch});
+    return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
+}
+
 fn reuseOperand(self: *Self, inst: Air.Inst.Index, operand: Air.Inst.Ref, op_index: Liveness.OperandInt, mcv: MCValue) bool {
     if (!self.liveness.operandDies(inst, op_index))
         return false;

src/codegen.zig

@@ -836,6 +836,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
                     .get_union_tag   => try self.airGetUnionTag(inst),
                     .clz             => try self.airClz(inst),
                     .ctz             => try self.airCtz(inst),
+                    .popcount        => try self.airPopcount(inst),
 
                     .atomic_store_unordered => try self.airAtomicStore(inst, .Unordered),
                     .atomic_store_monotonic => try self.airAtomicStore(inst, .Monotonic),
@@ -1610,6 +1611,14 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
             return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
         }
 
+        fn airPopcount(self: *Self, inst: Air.Inst.Index) !void {
+            const ty_op = self.air.instructions.items(.data)[inst].ty_op;
+            const result: MCValue = if (self.liveness.isUnused(inst)) .dead else switch (arch) {
+                else => return self.fail("TODO implement airPopcount for {}", .{self.target.cpu.arch}),
+            };
+            return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
+        }
+
         fn reuseOperand(self: *Self, inst: Air.Inst.Index, operand: Air.Inst.Ref, op_index: Liveness.OperandInt, mcv: MCValue) bool {
             if (!self.liveness.operandDies(inst, op_index))
                 return false;

src/codegen/c.zig

@@ -1107,6 +1107,7 @@ fn genBody(f: *Function, body: []const Air.Inst.Index) error{ AnalysisFail, OutO
             .get_union_tag    => try airGetUnionTag(f, inst),
             .clz              => try airBuiltinCall(f, inst, "clz"),
             .ctz              => try airBuiltinCall(f, inst, "ctz"),
+            .popcount         => try airBuiltinCall(f, inst, "popcount"),
 
             .int_to_float,
             .float_to_int,

src/codegen/llvm.zig

@@ -1777,6 +1777,7 @@ pub const FuncGen = struct {
                 .get_union_tag  => try self.airGetUnionTag(inst),
                 .clz            => try self.airClzCtz(inst, "ctlz"),
                 .ctz            => try self.airClzCtz(inst, "cttz"),
+                .popcount       => try self.airPopCount(inst, "ctpop"),
 
                 .atomic_store_unordered => try self.airAtomicStore(inst, .Unordered),
                 .atomic_store_monotonic => try self.airAtomicStore(inst, .Monotonic),
@@ -3679,6 +3680,40 @@ pub const FuncGen = struct {
         }
     }
 
+    fn airPopCount(self: *FuncGen, inst: Air.Inst.Index, prefix: [*:0]const u8) !?*const llvm.Value {
+        if (self.liveness.isUnused(inst)) return null;
+
+        const ty_op = self.air.instructions.items(.data)[inst].ty_op;
+        const operand_ty = self.air.typeOf(ty_op.operand);
+        const operand = try self.resolveInst(ty_op.operand);
+        const target = self.dg.module.getTarget();
+        const bits = operand_ty.intInfo(target).bits;
+
+        var fn_name_buf: [100]u8 = undefined;
+        const llvm_fn_name = std.fmt.bufPrintZ(&fn_name_buf, "llvm.{s}.i{d}", .{
+            prefix, bits,
+        }) catch unreachable;
+        const fn_val = self.dg.object.llvm_module.getNamedFunction(llvm_fn_name) orelse blk: {
+            const operand_llvm_ty = try self.dg.llvmType(operand_ty);
+            const param_types = [_]*const llvm.Type{operand_llvm_ty};
+            const fn_type = llvm.functionType(operand_llvm_ty, &param_types, param_types.len, .False);
+            break :blk self.dg.object.llvm_module.addFunction(llvm_fn_name, fn_type);
+        };
+
+        const params = [_]*const llvm.Value{operand};
+        const wrong_size_result = self.builder.buildCall(fn_val, &params, params.len, .C, .Auto, "");
+        const result_ty = self.air.typeOfIndex(inst);
+        const result_llvm_ty = try self.dg.llvmType(result_ty);
+        const result_bits = result_ty.intInfo(target).bits;
+        if (bits > result_bits) {
+            return self.builder.buildTrunc(wrong_size_result, result_llvm_ty, "");
+        } else if (bits < result_bits) {
+            return self.builder.buildZExt(wrong_size_result, result_llvm_ty, "");
+        } else {
+            return wrong_size_result;
+        }
+    }
+
     fn callFloor(self: *FuncGen, arg: *const llvm.Value, ty: Type) !*const llvm.Value {
         return self.callFloatUnary(arg, ty, "floor");
     }

src/print_air.zig

@@ -196,6 +196,7 @@ const Writer = struct {
             .get_union_tag,
             .clz,
             .ctz,
+            .popcount,
             => try w.writeTyOp(s, inst),
 
             .block,