diff --git a/src/sections.zig b/src/sections.zig index f0a9209..59ca6c0 100644 --- a/src/sections.zig +++ b/src/sections.zig @@ -394,10 +394,16 @@ fn cmpGroup(_: void, a: Group, b: Group) bool { } // nblocks returns how many blocks a particular slice will take. -fn nblocks(arr: []const u8) u32 { - const upper = pad.roundUp(u38, section_length_bits, @intCast(u38, arr.len)); +fn nblocks(comptime T: type, arr: []const u8) T { + const B = switch (T) { + u8 => u14, + u32 => u38, + u64 => u70, + else => @compileError("only u8, u32 and u64 are supported"), + }; + const upper = pad.roundUp(B, section_length_bits, @intCast(B, arr.len)); assert(upper & (section_length - 1) == 0); - return @truncate(u32, upper >> 6); + return @truncate(T, upper >> 6); } pub const AllSections = struct { @@ -780,8 +786,12 @@ test "nblocks" { .{ 1, &[_]u8{1} ** 63 }, .{ 1, &[_]u8{1} ** 64 }, .{ 2, &[_]u8{1} ** 65 }, + .{ 255, &[_]u8{1} ** (255 * 64) }, }; - inline for (tests) |tt| - try testing.expectEqual(@as(u32, tt[0]), nblocks(tt[1])); + 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]); + } }