tiny.simd.thread.ranges

1const std = @import("std");23pub const IndexRange = struct {4    begin: usize = 0,5    end: usize = 0,67    pub fn init(begin: usize, end: usize) error{InvalidRange}!IndexRange {8        if (begin >= end) return error.InvalidRange;9        return .{ .begin = begin, .end = end };10    }1112    pub fn count(self: IndexRange) usize {13        std.debug.assert(self.begin <= self.end);14        return self.end - self.begin;15    }1617    pub fn contains(self: IndexRange, index: usize) bool {18        return self.begin <= index and index < self.end;19    }2021    pub fn containsRange(self: IndexRange, other: IndexRange) bool {22        return other.begin >= self.begin and other.end <= self.end;23    }24};2526pub fn makeIndexRange(27    begin: usize,28    end: usize,29    max_size: usize,30) error{ InvalidRange, EmptyTask }!IndexRange {31    if (begin >= end) return error.InvalidRange;32    if (max_size == 0) return error.EmptyTask;33    return .{34        .begin = begin,35        .end = begin + @min(max_size, end - begin),36    };37}3839pub const IndexRangePartition = struct {40    range: IndexRange,41    task_size: u32,42    task_count: u32,4344    pub fn single(45        task_size: usize,46    ) error{ EmptyTask, TaskTooLarge }!IndexRangePartition {47        if (task_size == 0) return error.EmptyTask;48        if (task_size > std.math.maxInt(u32)) return error.TaskTooLarge;49        return .{50            .range = .{ .begin = 0, .end = task_size },51            .task_size = @intCast(task_size),52            .task_count = 1,53        };54    }5556    pub fn init(57        range: IndexRange,58        task_size: usize,59    ) error{ InvalidRange, EmptyTask, RangeTooLarge, TaskTooLarge }!IndexRangePartition {60        if (range.begin >= range.end) return error.InvalidRange;61        if (task_size == 0) return error.EmptyTask;62        if (range.count() > std.math.maxInt(u32)) return error.RangeTooLarge;63        if (task_size > std.math.maxInt(u32)) return error.TaskTooLarge;64        const task_count = std.math.divCeil(65            u32,66            @intCast(range.count()),67            @intCast(task_size),68        ) catch unreachable;69        std.debug.assert(task_count != 0);70        return .{71            .range = range,72            .task_size = @intCast(task_size),73            .task_count = task_count,74        };75    }7677    pub fn taskSize(self: IndexRangePartition) usize {78        return self.task_size;79    }8081    pub fn taskCount(self: IndexRangePartition) usize {82        return self.task_count;83    }8485    pub fn rangeAt(86        self: IndexRangePartition,87        task_index: usize,88    ) error{TaskOutOfBounds}!IndexRange {89        if (task_index >= self.task_count) return error.TaskOutOfBounds;90        const begin = self.range.begin + task_index * self.task_size;91        return .{92            .begin = begin,93            .end = begin + @min(self.task_size, self.range.end - begin),94        };95    }9697    pub fn visitAll(98        self: IndexRangePartition,99        context: anytype,100        comptime visit: fn (@TypeOf(context), IndexRange) void,101    ) void {102        for (0..self.task_count) |task_index| {103            visit(context, self.rangeAt(task_index) catch unreachable);104        }105    }106107    pub fn visitFirst(108        self: IndexRangePartition,109        context: anytype,110        comptime visit: fn (@TypeOf(context), IndexRange) void,111    ) void {112        visit(context, self.rangeAt(0) catch unreachable);113    }114115    pub fn visitRemaining(116        self: IndexRangePartition,117        context: anytype,118        comptime visit: fn (@TypeOf(context), IndexRange) void,119    ) void {120        for (1..self.task_count) |task_index| {121            visit(context, self.rangeAt(task_index) catch unreachable);122        }123    }124};125126const RangeVisits = struct {127    ranges: [8]IndexRange = @splat(.{}),128    count: usize = 0,129130    fn append(self: *RangeVisits, range: IndexRange) void {131        self.ranges[self.count] = range;132        self.count += 1;133    }134};135136test "Highway index ranges preserve half-open containment" {137    const range = try IndexRange.init(7, 19);138    try std.testing.expectEqual(@as(usize, 12), range.count());139    try std.testing.expect(range.contains(7));140    try std.testing.expect(range.contains(18));141    try std.testing.expect(!range.contains(19));142    try std.testing.expect(range.containsRange(try IndexRange.init(9, 13)));143    try std.testing.expect(!range.containsRange(try IndexRange.init(6, 13)));144    try std.testing.expectEqual(145        IndexRange{ .begin = 17, .end = 19 },146        try makeIndexRange(17, 19, std.math.maxInt(usize)),147    );148}149150test "Highway index range partitions cover the remainder once" {151    const partition = try IndexRangePartition.init(152        try IndexRange.init(11, 40),153        6,154    );155    try std.testing.expectEqual(@as(usize, 6), partition.taskSize());156    try std.testing.expectEqual(@as(usize, 5), partition.taskCount());157    var cursor: usize = 11;158    for (0..partition.taskCount()) |task_index| {159        const range = try partition.rangeAt(task_index);160        try std.testing.expectEqual(cursor, range.begin);161        try std.testing.expect(range.count() >= 1);162        try std.testing.expect(range.count() <= partition.taskSize());163        cursor = range.end;164    }165    try std.testing.expectEqual(@as(usize, 40), cursor);166    try std.testing.expectError(167        error.TaskOutOfBounds,168        partition.rangeAt(partition.taskCount()),169    );170    var all = RangeVisits{};171    partition.visitAll(&all, RangeVisits.append);172    try std.testing.expectEqual(partition.taskCount(), all.count);173    for (all.ranges[0..all.count], 0..) |range, task| {174        try std.testing.expectEqual(try partition.rangeAt(task), range);175    }176    var first = RangeVisits{};177    partition.visitFirst(&first, RangeVisits.append);178    try std.testing.expectEqual(@as(usize, 1), first.count);179    try std.testing.expectEqual(try partition.rangeAt(0), first.ranges[0]);180    var remaining = RangeVisits{};181    partition.visitRemaining(&remaining, RangeVisits.append);182    try std.testing.expectEqual(partition.taskCount() - 1, remaining.count);183    for (remaining.ranges[0..remaining.count], 1..) |range, task| {184        try std.testing.expectEqual(try partition.rangeAt(task), range);185    }186}187188test "Highway one-task partition covers the requested prefix" {189    const partition = try IndexRangePartition.single(17);190    try std.testing.expectEqual(@as(usize, 17), partition.taskSize());191    try std.testing.expectEqual(@as(usize, 1), partition.taskCount());192    try std.testing.expectEqual(193        IndexRange{ .begin = 0, .end = 17 },194        try partition.rangeAt(0),195    );196    try std.testing.expectError(197        error.EmptyTask,198        IndexRangePartition.single(0),199    );200}201202test "Highway index range partition rejects unrepresentable limits" {203    try std.testing.expectError(204        error.EmptyTask,205        IndexRangePartition.init(try IndexRange.init(0, 1), 0),206    );207    if (@sizeOf(usize) > @sizeOf(u32)) {208        try std.testing.expectError(209            error.RangeTooLarge,210            IndexRangePartition.init(211                .{ .begin = 0, .end = @as(usize, std.math.maxInt(u32)) + 2 },212                1,213            ),214        );215        try std.testing.expectError(216            error.TaskTooLarge,217            IndexRangePartition.init(218                try IndexRange.init(0, 1),219                @as(usize, std.math.maxInt(u32)) + 1,220            ),221        );222    }223}

1pub const ranges = @import("range.zig");