turbonss/src/DB.zig

const std = @import("std");
const os = std.os;
const mem = std.mem;
const math = std.math;
const meta = std.meta;
const sort = std.sort;
const assert = std.debug.assert;
const Allocator = std.mem.Allocator;
const ArrayList = std.ArrayList;
const AutoHashMap = std.AutoHashMap;
const BoundedArray = std.BoundedArray;

const ErrCtx = @import("ErrCtx.zig");
const Corpus = @import("Corpus.zig");
const pad = @import("padding.zig");
const compress = @import("compress.zig");
const Group = @import("Group.zig");
const CGroup = Group.CGroup;
const PackedGroup = @import("PackedGroup.zig");
const GroupStored = PackedGroup.GroupStored;
const User = @import("User.zig");
const CUser = User.CUser;
const PackedUser = @import("PackedUser.zig");
const ShellSections = @import("shell.zig").ShellWriter.ShellSections;
const ShellReader = @import("shell.zig").ShellReader;
const ShellWriter = @import("shell.zig").ShellWriter;
const InvalidHeader = @import("header.zig").Invalid;
const Header = @import("header.zig").Header;
const max_shells = @import("shell.zig").max_shells;
const section_length_bits = @import("header.zig").section_length_bits;
const section_length = @import("header.zig").section_length;
const cmph = @import("cmph.zig");
const bdz = @import("bdz.zig");

const zeroes = &[_]u8{0} ** section_length;

const DB = @This();
// All sections, as they end up in the DB. Order is important.
header: *const Header,
bdz_gid: []const u8,
bdz_groupname: []const u8,
bdz_uid: []const u8,
bdz_username: []const u8,
idx_gid2group: []const u32,
idx_groupname2group: []const u32,
idx_uid2user: []const u32,
idx_name2user: []const u32,
shell_index: []const u16,
shell_blob: []const u8,
groups: []const u8,
users: []const u8,
groupmembers: []const u8,
additional_gids: []const u8,

pub fn fromCorpus(
    allocator: Allocator,
    corpus: *const Corpus,
    err: *ErrCtx,
) error{ OutOfMemory, InvalidRecord, TooMany }!DB {
    _ = err;

    const gids = corpus.groups.items(.gid);
    const gnames = corpus.groups.items(.name);
    const uids = corpus.users.items(.uid);
    const unames = corpus.users.items(.name);

    const bdz_gid = try cmph.packU32(allocator, gids);
    errdefer allocator.free(bdz_gid);

    const bdz_groupname = try cmph.packStr(allocator, gnames);
    errdefer allocator.free(bdz_groupname);

    const bdz_uid = try cmph.packU32(allocator, uids);
    errdefer allocator.free(bdz_uid);

    const bdz_username = try cmph.packStr(allocator, unames);
    errdefer allocator.free(bdz_username);

    var shell = try shellSections(allocator, corpus);
    defer shell.deinit();

    const shell_index = try allocator.dupe(u16, shell.index.constSlice());
    errdefer allocator.free(shell_index);

    const shell_blob = try allocator.dupe(u8, shell.blob.constSlice());
    errdefer allocator.free(shell_blob);

    const additional_gids = try additionalGids(allocator, corpus);
    errdefer allocator.free(additional_gids.blob);
    defer allocator.free(additional_gids.idx2offset);

    const users = try usersSection(allocator, corpus, &additional_gids, &shell);
    errdefer allocator.free(users.blob);
    defer allocator.free(users.idx2offset);

    const groupmembers = try groupMembers(allocator, corpus, users.idx2offset);
    errdefer allocator.free(groupmembers.blob);
    defer allocator.free(groupmembers.idx2offset);

    const groups = try groupsSection(allocator, corpus, groupmembers.idx2offset);
    errdefer allocator.free(groups.blob);
    defer allocator.free(groups.idx2offset);

    const idx_gid2group = try bdzIdx(u32, allocator, bdz_gid, gids, groups.idx2offset);
    errdefer allocator.free(idx_gid2group);

    const idx_groupname2group = try bdzIdx([]const u8, allocator, bdz_groupname, gnames, groups.idx2offset);
    errdefer allocator.free(idx_groupname2group);

    const idx_uid2user = try bdzIdx(u32, allocator, bdz_uid, uids, users.idx2offset);
    errdefer allocator.free(idx_uid2user);

    const idx_name2user = try bdzIdx([]const u8, allocator, bdz_username, unames, users.idx2offset);
    errdefer allocator.free(idx_name2user);

    const header = try allocator.create(Header);
    errdefer allocator.destroy(header);

    header.* = Header{
        .nblocks_shell_blob = nblocks(u8, shell.blob.constSlice()),
        .num_shells = shell.len,
        .num_groups = groups.len,
        .num_users = users.len,
        .nblocks_bdz_gid = nblocks(u32, bdz_gid),
        .nblocks_bdz_groupname = nblocks(u32, bdz_groupname),
        .nblocks_bdz_uid = nblocks(u32, bdz_uid),
        .nblocks_bdz_username = nblocks(u32, bdz_username),
        .nblocks_groups = nblocks(u64, groups.blob),
        .nblocks_users = nblocks(u64, users.blob),
        .nblocks_groupmembers = nblocks(u64, groupmembers.blob),
        .nblocks_additional_gids = nblocks(u64, additional_gids.blob),
        .getgr_bufsize = corpus.getgr_bufsize,
        .getpw_bufsize = corpus.getpw_bufsize,
    };

    return DB{
        .header = header,
        .bdz_gid = bdz_gid,
        .bdz_groupname = bdz_groupname,
        .bdz_uid = bdz_uid,
        .bdz_username = bdz_username,
        .idx_gid2group = idx_gid2group,
        .idx_groupname2group = idx_groupname2group,
        .idx_uid2user = idx_uid2user,
        .idx_name2user = idx_name2user,
        .shell_index = shell_index,
        .shell_blob = shell_blob,
        .groups = groups.blob,
        .users = users.blob,
        .groupmembers = groupmembers.blob,
        .additional_gids = additional_gids.blob,
    };
}

pub fn getgrBufsize(self: *const DB) usize {
    return self.header.getgr_bufsize;
}

pub fn getpwBufsize(self: *const DB) usize {
    return self.header.getpw_bufsize;
}

pub fn deinit(self: *DB, allocator: Allocator) void {
    allocator.destroy(self.header);
    allocator.free(self.bdz_gid);
    allocator.free(self.bdz_groupname);
    allocator.free(self.bdz_uid);
    allocator.free(self.bdz_username);
    allocator.free(self.idx_gid2group);
    allocator.free(self.idx_groupname2group);
    allocator.free(self.idx_uid2user);
    allocator.free(self.idx_name2user);
    allocator.free(self.shell_index);
    allocator.free(self.shell_blob);
    allocator.free(self.groups);
    allocator.free(self.users);
    allocator.free(self.groupmembers);
    allocator.free(self.additional_gids);
    self.* = undefined;
}

const DB_fields = meta.fields(DB);
pub fn iov(self: *const DB) BoundedArray(os.iovec_const, DB_fields.len * 2) {
    var result = BoundedArray(os.iovec_const, DB_fields.len * 2).init(0) catch unreachable;
    inline for (DB_fields) |field| {
        comptime assertDefinedLayout(field.field_type);
        const value = @field(self, field.name);
        const bytes: []const u8 = switch (@TypeOf(value)) {
            *const Header => mem.asBytes(value),
            else => mem.sliceAsBytes(value),
        };
        result.appendAssumeCapacity(os.iovec_const{
            .iov_base = bytes.ptr,
            .iov_len = bytes.len,
        });
        const padding = pad.until(usize, section_length_bits, bytes.len);
        if (padding != 0)
            result.appendAssumeCapacity(.{
                .iov_base = zeroes,
                .iov_len = padding,
            });
    }

    return result;
}

pub const DBNumbers = struct {
    header: u64,
    bdz_gid: u64,
    bdz_groupname: u64,
    bdz_uid: u64,
    bdz_username: u64,
    idx_gid2group: u64,
    idx_groupname2group: u64,
    idx_uid2user: u64,
    idx_name2user: u64,
    shell_index: u64,
    shell_blob: u64,
    groups: u64,
    users: u64,
    groupmembers: u64,
    additional_gids: u64,
};

// in blocks
pub fn fieldLengths(header: *const Header) DBNumbers {
    return DBNumbers{
        .header = comptime DB.nblocks_n(u64, @sizeOf(Header)),
        .bdz_gid = header.nblocks_bdz_gid,
        .bdz_groupname = header.nblocks_bdz_groupname,
        .bdz_uid = header.nblocks_bdz_uid,
        .bdz_username = header.nblocks_bdz_username,
        .idx_gid2group = nblocks_n(u32, header.num_groups * 4),
        .idx_groupname2group = nblocks_n(u32, header.num_groups * 4),
        .idx_uid2user = nblocks_n(u32, header.num_users * 4),
        .idx_name2user = nblocks_n(u32, header.num_users * 4),
        .shell_index = nblocks_n(u16, header.num_shells * 2),
        .shell_blob = header.nblocks_shell_blob,
        .groups = header.nblocks_groups,
        .users = header.nblocks_users,
        .groupmembers = header.nblocks_groupmembers,
        .additional_gids = header.nblocks_additional_gids,
    };
}

// in blocks
pub fn fieldOffsets(lengths: DBNumbers) DBNumbers {
    var result: DBNumbers = undefined;
    result.header = 0;
    var offset = comptime nblocks_n(u64, @sizeOf(Header));
    comptime assert(mem.eql(u8, DB_fields[0].name, "header"));
    inline for (DB_fields[1..]) |field, i| {
        comptime assert(mem.eql(u8, field.name, meta.fields(DBNumbers)[i + 1].name));
        @field(result, field.name) = offset;
        offset += @field(lengths, field.name);
    }
    return result;
}

pub fn fromBytes(buf: []align(8) const u8) InvalidHeader!DB {
    const header = try Header.fromBytes(buf[0..@sizeOf(Header)]);
    const lengths = fieldLengths(header);
    const offsets = fieldOffsets(lengths);

    var result: DB = undefined;
    result.header = header;
    inline for (DB_fields[1..]) |field| {
        const start_block = @field(offsets, field.name);
        const end = (start_block + @field(lengths, field.name)) << section_length_bits;
        const start = start_block << section_length_bits;
        const slice_type = meta.Child(field.field_type);
        const value = mem.bytesAsSlice(slice_type, buf[start..end]);
        @field(result, field.name) = value;
    }

    return result;
}

// dumps PackedGroup to []u8 and returns a CGroup.
fn getGroup(self: *const DB, group: PackedGroup, buf: []u8) error{OutOfMemory}!CGroup {
    const members_slice = self.groupmembers[group.members_offset..];
    var vit = compress.VarintSliceIteratorMust(members_slice);
    const num_members = vit.remaining;

    const ptr_end = @sizeOf(?[*:0]const u8) * (num_members + 1);
    if (ptr_end > buf.len) return error.OutOfMemory;
    var member_ptrs = mem.bytesAsSlice(?[*:0]const u8, buf[0..ptr_end]);
    member_ptrs[member_ptrs.len - 1] = null;
    var buf_offset: usize = ptr_end;

    var it = compress.DeltaDecompressionIterator(&vit);
    var i: usize = 0;
    while (it.nextMust()) |member_offset| : (i += 1) {
        const entry = PackedUser.fromBytes(self.users[member_offset << 3 ..]);
        const start = buf_offset;
        const name = entry.user.name();
        if (buf_offset + name.len + 1 > buf.len) return error.OutOfMemory;
        mem.copy(u8, buf[buf_offset..], name);
        buf_offset += name.len;
        buf[buf_offset] = 0;
        buf_offset += 1;

        // TODO: arr[i] = buf[...] triggers a bug in zig pre-0.10
        const terminated = buf[start .. buf_offset - 1 :0];
        member_ptrs[i] = terminated;
    }

    const name = group.name();
    if (buf_offset + name.len + 1 > buf.len) return error.OutOfMemory;
    mem.copy(u8, buf[buf_offset..], name);
    buf[buf_offset + name.len] = 0;

    return CGroup{
        .gid = group.gid(),
        .name = buf[buf_offset .. buf_offset + name.len].ptr,
        .members = member_ptrs.ptr,
    };
}

// get a CGroup entry by name.
fn getgrnam(self: *const DB, name: []const u8, buf: []u8) error{OutOfMemory}!?CGroup {
    const idx = bdz.search(self.bdz_groupname, name);
    if (idx >= self.header.num_groups) return null;
    const offset = self.idx_groupname2group[idx];
    const nbits = PackedGroup.alignment_bits;
    const group = PackedGroup.fromBytes(self.groups[offset << nbits ..]).group;
    if (!mem.eql(u8, name, group.name())) return null;
    return try self.getGroup(group, buf);
}

// get a CGroup entry by it's gid.
fn getgrgid(self: *const DB, gid: u32, buf: []u8) error{OutOfMemory}!?CGroup {
    const idx = bdz.search_u32(self.bdz_gid, gid);
    if (idx >= self.header.num_groups) return null;
    const offset = self.idx_gid2group[idx];
    const nbits = PackedGroup.alignment_bits;
    const group = PackedGroup.fromBytes(self.groups[offset << nbits ..]).group;
    if (gid != group.gid()) return null;
    return try self.getGroup(group, buf);
}

fn pushStr(str: []const u8, buf: []u8, offset: *usize) [*]const u8 {
    const start = offset.*;
    mem.copy(u8, buf[offset.*..], str);
    buf[offset.* + str.len] = 0;
    offset.* += str.len + 1;
    return @ptrCast([*]const u8, &buf[start]);
}

fn getUser(self: *const DB, user: PackedUser, buf: []u8) error{OutOfMemory}!CUser {
    const shell_reader = ShellReader{
        .index = self.shell_index,
        .blob = self.shell_blob,
    };
    const name = user.name();
    const gecos = user.gecos();
    const home = user.home();
    const shell = user.shell(shell_reader);

    const strlen =
        name.len + 1 +
        gecos.len + 1 +
        home.len + 1 +
        shell.len + 1;
    if (strlen > buf.len) return error.OutOfMemory;

    var offset: usize = 0;
    const pw_name = pushStr(name, buf, &offset);
    const pw_gecos = pushStr(gecos, buf, &offset);
    const pw_dir = pushStr(home, buf, &offset);
    const pw_shell = pushStr(shell, buf, &offset);

    return CUser{
        .pw_name = pw_name,
        .pw_uid = user.uid(),
        .pw_gid = user.gid(),
        .pw_gecos = pw_gecos,
        .pw_dir = pw_dir,
        .pw_shell = pw_shell,
    };
}

// get a CUser entry by name.
pub fn getpwnam(self: *const DB, name: []const u8, buf: []u8) error{OutOfMemory}!?CUser {
    const idx = bdz.search(self.bdz_username, name);
    // bdz may return a hash that's bigger than the number of users
    if (idx >= self.header.num_users) return null;
    const offset = self.idx_name2user[idx];
    const nbits = PackedUser.alignment_bits;
    const user = PackedUser.fromBytes(self.users[offset << nbits ..]).user;
    if (!mem.eql(u8, name, user.name())) return null;
    return try self.getUser(user, buf);
}

// get a CUser entry by uid.
pub fn getpwuid(self: *const DB, uid: u32, buf: []u8) error{OutOfMemory}!?CUser {
    const idx = bdz.search_u32(self.bdz_uid, uid);
    if (idx >= self.header.num_users) return null;
    const offset = self.idx_uid2user[idx];
    const nbits = PackedUser.alignment_bits;
    const user = PackedUser.fromBytes(self.users[offset << nbits ..]).user;
    if (uid != user.uid()) return null;
    return try self.getUser(user, buf);
}

fn shellSections(
    allocator: Allocator,
    corpus: *const Corpus,
) error{OutOfMemory}!ShellSections {
    var popcon = ShellWriter.init(allocator);
    errdefer popcon.deinit();
    for (corpus.users.items(.shell)) |shell|
        try popcon.put(shell);
    return popcon.toOwnedSections(max_shells);
}

const AdditionalGids = struct {
    // user index -> offset in blob
    idx2offset: []const u64,
    // compressed user gids blob. A blob contains N <= users.len items,
    // an item is:
    //   len: varint
    //   gid: [varint]varint,
    // ... and the gid list is delta-compressed.
    blob: []const u8,
};

fn additionalGids(
    allocator: Allocator,
    corpus: *const Corpus,
) error{OutOfMemory}!AdditionalGids {
    var blob = ArrayList(u8).init(allocator);
    errdefer blob.deinit();
    var idx2offset = try allocator.alloc(u64, corpus.users.len);
    errdefer allocator.free(idx2offset);

    // zero'th entry is empty, so groupless users can refer to it.
    try compress.appendUvarint(&blob, 0);

    var scratch = try allocator.alloc(u32, 256);
    var scratch_allocated: bool = true;
    defer if (scratch_allocated) allocator.free(scratch);
    for (corpus.user2groups) |usergroups, user_idx| {
        if (usergroups.len == 0) {
            idx2offset[user_idx] = 0;
            continue;
        }
        idx2offset[user_idx] = blob.items.len;
        if (scratch.len < usergroups.len) {
            allocator.free(scratch);
            scratch_allocated = false;
            scratch = try allocator.alloc(u32, usergroups.len);
            scratch_allocated = true;
        }
        scratch.len = usergroups.len;
        const corpusGids = corpus.groups.items(.gid);
        for (usergroups) |group_idx, i|
            scratch[i] = corpusGids[group_idx];
        compress.deltaCompress(u32, scratch) catch |err| switch (err) {
            error.NotSorted => unreachable,
        };
        try compress.appendUvarint(&blob, usergroups.len);
        for (scratch) |gid|
            try compress.appendUvarint(&blob, gid);
    }

    return AdditionalGids{
        .idx2offset = idx2offset,
        .blob = blob.toOwnedSlice(),
    };
}

const UsersSection = struct {
    // number of users in this section
    len: u32,
    // user index -> offset in blob
    idx2offset: []const u32,
    blob: []const u8,
};

fn usersSection(
    allocator: Allocator,
    corpus: *const Corpus,
    gids: *const AdditionalGids,
    shells: *const ShellSections,
) error{ OutOfMemory, InvalidRecord, TooMany }!UsersSection {
    var idx2offset = try allocator.alloc(u32, corpus.users.len);
    errdefer allocator.free(idx2offset);
    // as of writing each user takes 12 bytes + blobs + padding, padded to
    // 8 bytes. 24 is an optimistic lower bound for an average record size.
    var blob = try ArrayList(u8).initCapacity(allocator, 24 * corpus.users.len);
    errdefer blob.deinit();
    var i: usize = 0;
    while (i < corpus.users.len) : (i += 1) {
        // TODO: this is inefficient by calling `.slice()` on every iteration
        const user = corpus.users.get(i);
        const user_offset = math.cast(u35, blob.items.len) orelse return error.TooMany;
        assert(user_offset & 7 == 0);
        idx2offset[i] = @truncate(u32, user_offset >> 3);
        try PackedUser.packTo(
            &blob,
            user,
            gids.idx2offset[i],
            shells.shell2idx,
        );
        try pad.arrayList(&blob, PackedUser.alignment_bits);
    }
    return UsersSection{
        .len = @intCast(u32, corpus.users.len),
        .idx2offset = idx2offset,
        .blob = blob.toOwnedSlice(),
    };
}

const GroupMembers = struct {
    // group index to it's offset in blob
    idx2offset: []const u64,
    // members are delta-varint encoded byte-offsets to the user struct
    blob: []const u8,
};

fn groupMembers(
    allocator: Allocator,
    corpus: *const Corpus,
    user2offset: []const u32,
) error{OutOfMemory}!GroupMembers {
    var idx2offset = try allocator.alloc(u64, corpus.groups.len);
    errdefer allocator.free(idx2offset);
    var blob = ArrayList(u8).init(allocator);
    errdefer blob.deinit();
    // zero'th entry is empty, so empty groups can refer to it
    try compress.appendUvarint(&blob, 0);

    var scratch = try ArrayList(u32).initCapacity(allocator, 1024);
    defer scratch.deinit();

    for (corpus.group2users) |members, group_idx| {
        if (members.len == 0) {
            idx2offset[group_idx] = 0;
            continue;
        }

        idx2offset[group_idx] = blob.items.len;
        try scratch.ensureTotalCapacity(members.len);
        scratch.items.len = members.len;
        for (members) |user_idx, i|
            scratch.items[i] = user2offset[user_idx];

        compress.deltaCompress(u32, scratch.items) catch |err| switch (err) {
            error.NotSorted => unreachable,
        };
        try compress.appendUvarint(&blob, members.len);
        for (scratch.items) |elem|
            try compress.appendUvarint(&blob, elem);
    }
    return GroupMembers{
        .idx2offset = idx2offset,
        .blob = blob.toOwnedSlice(),
    };
}

const GroupsSection = struct {
    // number of groups in this section
    len: u32,
    // group index -> offset in blob
    idx2offset: []const u32,
    blob: []const u8,
};

fn groupsSection(
    allocator: Allocator,
    corpus: *const Corpus,
    members_offset: []const u64,
) error{ OutOfMemory, InvalidRecord }!GroupsSection {
    var idx2offset = try allocator.alloc(u32, corpus.groups.len);
    errdefer allocator.free(idx2offset);

    var blob = try ArrayList(u8).initCapacity(allocator, 8 * corpus.groups.len);
    errdefer blob.deinit();

    var i: usize = 0;
    while (i < corpus.groups.len) : (i += 1) {
        // TODO: this is inefficient; it's calling `.slice()` on every iteration
        const group = corpus.groups.get(i);
        const group_offset = @intCast(u32, blob.items.len);
        assert(group_offset & 7 == 0);
        idx2offset[i] = @truncate(u32, group_offset >> 3);
        const group_stored = GroupStored{
            .gid = group.gid,
            .name = group.name,
            .members_offset = members_offset[i],
        };
        try PackedGroup.packTo(&blob, group_stored);
        try pad.arrayList(&blob, PackedGroup.alignment_bits);
    }

    return GroupsSection{
        .len = @intCast(u32, corpus.groups.len),
        .idx2offset = idx2offset,
        .blob = blob.toOwnedSlice(),
    };
}

// creates a bdz index using packed_mphf.
//   hash = bdz_search(packed_mphf, keys[i]);
//   result[hash] = idx2offset[i];
fn bdzIdx(
    comptime T: type,
    allocator: Allocator,
    packed_mphf: []const u8,
    keys: []const T,
    idx2offset: []const u32,
) error{OutOfMemory}![]const u32 {
    const search_fn = switch (T) {
        u32 => bdz.search_u32,
        []const u8 => bdz.search,
        else => unreachable,
    };
    assert(keys.len <= math.maxInt(u32));
    var result = try allocator.alloc(u32, keys.len);
    errdefer allocator.free(result);
    for (keys) |key, i|
        result[search_fn(packed_mphf, key)] = idx2offset[i];
    return result;
}

// nblocks_n returns how many blocks a given number of bytes will take
pub fn nblocks_n(comptime T: type, nbytes: usize) T {
    const B = switch (T) {
        u8 => u14,
        u16 => u22,
        u32 => u38,
        u64 => u70,
        else => @compileError("got " ++ @typeName(T) ++ ", only u8, u32 and u64 are supported"),
    };
    const upper = pad.roundUp(B, section_length_bits, @intCast(B, nbytes));
    assert(upper & (section_length - 1) == 0);
    return @truncate(T, upper >> section_length_bits);
}

// nblocks returns how many blocks a particular slice will take.
fn nblocks(comptime T: type, arr: []const u8) T {
    return nblocks_n(T, arr.len);
}

fn assertDefinedLayout(comptime T: type) void {
    return switch (T) {
        u4, u8, u16, u32, u64 => {},
        else => switch (@typeInfo(T)) {
            .Array => assertDefinedLayout(meta.Elem(T)),
            .Pointer => |info| assertDefinedLayout(info.child),
            .Enum => assertDefinedLayout(meta.Tag(T)),
            .Struct => {
                if (meta.containerLayout(T) == .Auto)
                    @compileError("layout of " ++ @typeName(T) ++ " is undefined");
                for (meta.fields(T)) |field|
                    assertDefinedLayout(field.field_type);
            },
            else => @compileError("unexpected type " ++ @typeName(T)),
        },
    };
}

const testing = std.testing;

test "read/write via iovec" {
    const allocator = testing.allocator;
    var corpus = try Corpus.testCorpus(allocator);
    defer corpus.deinit();

    var errc = ErrCtx{};
    var db = try DB.fromCorpus(allocator, &corpus, &errc);
    defer db.deinit(allocator);

    const fd = try os.memfd_create("test_turbonss_db", 0);
    defer os.close(fd);

    const len = try os.writev(fd, db.iov().constSlice());
    const buf = try os.mmap(null, len, os.PROT.READ, os.MAP.SHARED, fd, 0);
    const db2 = try fromBytes(buf);
    try testing.expectEqual(corpus.groups.len, db.header.num_groups);
    try testing.expectEqual(corpus.users.len, db.header.num_users);
    try testing.expectEqual(db.header.num_groups, db2.header.num_groups);
    try testing.expectEqual(db.header.num_users, db2.header.num_users);
    const num_groups = db2.header.num_groups;
    const num_users = db2.header.num_users;

    try testing.expectEqualSlices(u32, db.idx_gid2group, db2.idx_gid2group[0..num_groups]);
    try testing.expectEqualSlices(u32, db.idx_uid2user, db2.idx_uid2user[0..num_users]);
    try testing.expectEqualStrings("", errc.unwrap().constSlice());
}

test "getgrnam/getgrgid" {
    var corpus = try Corpus.testCorpus(testing.allocator);
    defer corpus.deinit();
    var errc = ErrCtx{};
    var db = try DB.fromCorpus(testing.allocator, &corpus, &errc);
    defer db.deinit(testing.allocator);
    var buf = try testing.allocator.alloc(u8, db.getgrBufsize());
    defer testing.allocator.free(buf);

    {
        try testing.expectEqual(try db.getgrnam("doesnotexist", buf), null);
        const all = (try db.getgrnam("all", buf)).?;
        try testing.expectEqual(all.gid, 9999);
        try testing.expectEqualStrings(all.name[0..4], "all\x00");
        const members = all.members;
        try testing.expectEqualStrings(mem.sliceTo(members[0].?, 0), "Name" ** 8);
        try testing.expectEqualStrings(mem.sliceTo(members[1].?, 0), "root");
        try testing.expectEqualStrings(mem.sliceTo(members[2].?, 0), "svc-bar");
        try testing.expectEqualStrings(mem.sliceTo(members[3].?, 0), "vidmantas");
        try testing.expectEqual(members[4], null);
    }

    {
        try testing.expectEqual(try db.getgrgid(42, buf), null);
        const all = (try db.getgrgid(9999, buf)).?;
        try testing.expectEqual(all.gid, 9999);
        try testing.expectEqualStrings(all.name[0..3], "all");
    }
    try testing.expectEqualStrings("", errc.unwrap().constSlice());

    _ = try db.getgrnam("all", buf);
    buf.len -= 1;
    try testing.expectError(error.OutOfMemory, db.getgrnam("all", buf));
}

test "getpwnam/getpwuid" {
    var corpus = try Corpus.testCorpus(testing.allocator);
    defer corpus.deinit();
    var errc = ErrCtx{};
    var db = try DB.fromCorpus(testing.allocator, &corpus, &errc);
    defer db.deinit(testing.allocator);
    var buf = try testing.allocator.alloc(u8, db.getpwBufsize());
    defer testing.allocator.free(buf);

    {
        try testing.expectEqual(try db.getpwnam("doesnotexist", buf), null);
        const vidmantas = (try db.getpwnam("vidmantas", buf)).?;
        try testing.expectEqual(vidmantas.pw_uid, 128);
        try testing.expectEqual(vidmantas.pw_gid, 128);
        try testing.expectEqualStrings(vidmantas.pw_name[0..10], "vidmantas\x00");
        try testing.expectEqualStrings(vidmantas.pw_gecos[0..20], "Vidmantas Kaminskas\x00");
        try testing.expectEqualStrings(vidmantas.pw_dir[0..16], "/home/vidmantas\x00");
    }

    {
        try testing.expectEqual(try db.getpwuid(123456, buf), null);
        const vidmantas = (try db.getpwuid(128, buf)).?;
        try testing.expectEqual(vidmantas.pw_uid, 128);
        try testing.expectEqual(vidmantas.pw_gid, 128);
        try testing.expectEqualStrings(vidmantas.pw_name[0..10], "vidmantas\x00");
    }
    try testing.expectEqualStrings("", errc.unwrap().constSlice());

    const long = try db.getpwnam("Name" ** 8, buf);
    try testing.expectEqualStrings(long.?.pw_name[0..33], "Name" ** 8 ++ "\x00");
    buf.len -= 1;
    try testing.expectError(error.OutOfMemory, db.getpwnam("Name" ** 8, buf));
}

test "additionalGids" {
    const allocator = testing.allocator;
    var corpus = try Corpus.testCorpus(allocator);
    defer corpus.deinit();

    var additional_gids = try additionalGids(allocator, &corpus);
    defer allocator.free(additional_gids.idx2offset);
    defer allocator.free(additional_gids.blob);

    var user_idx: usize = 0;
    while (user_idx < corpus.users.len) : (user_idx += 1) {
        const groups = corpus.user2groups[user_idx];
        const offset = additional_gids.idx2offset[user_idx];
        if (groups.len == 0) {
            try testing.expect(offset == 0);
            continue;
        }
        var vit = try compress.VarintSliceIterator(additional_gids.blob[offset..]);
        var it = compress.DeltaDecompressionIterator(&vit);
        try testing.expectEqual(it.remaining(), groups.len);
        var i: u64 = 0;
        const corpusGids = corpus.groups.items(.gid);
        while (try it.next()) |gid| : (i += 1) {
            try testing.expectEqual(gid, corpusGids[groups[i]]);
        }
        try testing.expectEqual(i, groups.len);
    }
}

test "pack gids" {
    const allocator = testing.allocator;
    var corpus = try Corpus.testCorpus(allocator);
    defer corpus.deinit();

    const cmph_gid = try cmph.packU32(allocator, corpus.groups.items(.gid));
    defer allocator.free(cmph_gid);

    const k1 = bdz.search_u32(cmph_gid, 0);
    const k2 = bdz.search_u32(cmph_gid, 128);
    const k3 = bdz.search_u32(cmph_gid, 9999);
    const k4 = bdz.search_u32(cmph_gid, 100000);
    var hashes = &[_]u32{ k1, k2, k3, k4 };
    sort.sort(u32, hashes, {}, comptime sort.asc(u32));
    for (hashes) |hash, i|
        try testing.expectEqual(i, hash);
}

const hash_offsets = &[_]u32{ 0, 10, 20, 30 };

fn expectUsedHashes(allocator: Allocator, arr: []const u32) !void {
    var used = AutoHashMap(u32, void).init(allocator);
    defer used.deinit();

    for (arr) |elem|
        try used.putNoClobber(elem, {});
    for (hash_offsets) |item|
        try testing.expect(used.get(item) != null);
}

test "bdzIdx on u32" {
    const keys = [_]u32{ 42, 1, 2, 3 };
    const mphf = try cmph.packU32(testing.allocator, keys[0..]);
    defer testing.allocator.free(mphf);
    var result = try bdzIdx(u32, testing.allocator, mphf, keys[0..], hash_offsets);
    defer testing.allocator.free(result);
    try expectUsedHashes(testing.allocator, result);
}

test "bdzIdx on str" {
    const keys = [_][]const u8{ "42", "1", "2", "3" };
    const mphf = try cmph.packStr(testing.allocator, keys[0..]);
    defer testing.allocator.free(mphf);
    var result = try bdzIdx([]const u8, testing.allocator, mphf, keys[0..], hash_offsets);
    defer testing.allocator.free(result);
    try expectUsedHashes(testing.allocator, result);
}

test "nblocks" {
    const tests = .{
        .{ 0, &[_]u8{} },
        .{ 1, &[_]u8{ 1, 2, 42 } },
        .{ 1, &[_]u8{1} ** 63 },
        .{ 1, &[_]u8{1} ** 64 },
        .{ 2, &[_]u8{1} ** 65 },
        .{ 255, &[_]u8{1} ** (255 * 64) },
    };

    inline for (tests) |tt| {
        try testing.expectEqual(nblocks(u8, tt[1]), tt[0]);
        try testing.expectEqual(nblocks(u32, tt[1]), tt[0]);
        try testing.expectEqual(nblocks(u64, tt[1]), tt[0]);
    }
}