std.equalRange: Compute lower and upper bounds simultaneously.

[LucasSantos91]

The current implementation of equalRange just calls lowerRange and upperRange, but a lot of
the work done by these two functions can be shared. Specifically, each iteration gives information about whether the lower bound or the upper bound can be tightened. This leads to fewer iterations and, since there is one comparison per iteration, fewer comparisons.
Implementation adapted from GCC.
This sample demonstrates the difference between the current implementation and mine:

fn S(comptime T: type) type {
    return struct {
        needle: T,
        count: *usize,

        pub fn order(context: @This(), item: T) std.math.Order {
            context.count.* += 1;
            return std.math.order(item, context.needle);
        }
        pub fn orderLength(context: @This(), item: []const u8) std.math.Order {
            context.count.* += 1;
            return std.math.order(item.len, context.needle);
        }
    };
}
pub fn main() !void {
    var count: usize = 0;

    try std.testing.expectEqual(.{ 0, 0 }, equalRange(i32, &[_]i32{}, S(i32){ .needle = 0, .count = &count }, S(i32).order));
    try std.testing.expectEqual(.{ 0, 0 }, equalRange(i32, &[_]i32{ 2, 4, 8, 16, 32, 64 }, S(i32){ .needle = 0, .count = &count }, S(i32).order));
    try std.testing.expectEqual(.{ 0, 1 }, equalRange(i32, &[_]i32{ 2, 4, 8, 16, 32, 64 }, S(i32){ .needle = 2, .count = &count }, S(i32).order));
    try std.testing.expectEqual(.{ 2, 2 }, equalRange(i32, &[_]i32{ 2, 4, 8, 16, 32, 64 }, S(i32){ .needle = 5, .count = &count }, S(i32).order));
    try std.testing.expectEqual(.{ 2, 3 }, equalRange(i32, &[_]i32{ 2, 4, 8, 16, 32, 64 }, S(i32){ .needle = 8, .count = &count }, S(i32).order));
    try std.testing.expectEqual(.{ 5, 6 }, equalRange(i32, &[_]i32{ 2, 4, 8, 16, 32, 64 }, S(i32){ .needle = 64, .count = &count }, S(i32).order));
    try std.testing.expectEqual(.{ 6, 6 }, equalRange(i32, &[_]i32{ 2, 4, 8, 16, 32, 64 }, S(i32){ .needle = 100, .count = &count }, S(i32).order));
    try std.testing.expectEqual(.{ 2, 6 }, equalRange(i32, &[_]i32{ 2, 4, 8, 8, 8, 8, 15, 22 }, S(i32){ .needle = 8, .count = &count }, S(i32).order));
    try std.testing.expectEqual(.{ 2, 2 }, equalRange(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, S(u32){ .needle = 5, .count = &count }, S(u32).order));
    try std.testing.expectEqual(.{ 1, 1 }, equalRange(f32, &[_]f32{ -54.2, -26.7, 0.0, 56.55, 100.1, 322.0 }, S(f32){ .needle = -33.4, .count = &count }, S(f32).order));
    try std.testing.expectEqual(.{ 3, 5 }, equalRange(
        []const u8,
        &[_][]const u8{ "Mars", "Venus", "Earth", "Saturn", "Uranus", "Mercury", "Jupiter", "Neptune" },
        S(usize){ .needle = 6, .count = &count },
        S(usize).orderLength,
    ));

    std.debug.print("Count: {}\n", .{count});
}

For each comparison, we bump the count. With the current implementation, we get 57 comparisons. With mine, we get 43.

This optimization is orthogonal to left-bias proposed by 21278

[Olvilock]

After the equal element is found, you discard information about bounds already computed in low and high, essentially precomputing binarySearch, and using lowerBound and upperBound on both halves. This can still be improved with little effort:

pub fn equalRange(
    comptime T: type,
    items: []const T,
    context: anytype,
    comptime compareFn: fn (@TypeOf(context), T) std.math.Order,
) struct { usize, usize } {
    var low: usize = 0;
    var high: usize = items.len;

    while (low < high) {
        const mid = low + (high - low) / 2;
        switch (compareFn(context, items[mid])) {
            .lt => low = mid + 1,
            .gt => high = mid,
            .eq => return .{
                low + lowerBound(T, items[low..mid], context, compareFn),
                mid + upperBound(T, items[mid..high], context, compareFn),
            },
        }
    }
    return .{ low, low };
}

[LucasSantos91]

Thanks, @Olvilock.

[Olvilock]

Example test in this gist (run on zig master).
Seek equal range of element 5 in array [_]u32{ 2, 3, 4, 5, 5 }

The problem is that current implementation of lowerBound and upperBound does not match the contract in the docs (the assumed order is opposite), the tests in std.sort for those functions use precisely the order defined in struct S.

[LucasSantos91]

Ops, messed up the rebase. I will fix the bug mentioned by @Olvilock and reopen.