std.mem: Add readPackedInt, writePackedInt, etc.

These utility functions allow reading from (stage2) packed memory at
runtime-known offsets.
This commit is contained in:
Cody Tapscott
2022-10-18 11:33:24 -07:00
parent f28e4e03ee
commit c639c22544

View File

@@ -1298,6 +1298,76 @@ pub fn readVarInt(comptime ReturnType: type, bytes: []const u8, endian: Endian)
return result;
}
/// Loads an integer from packed memory with provided bit_count, bit_offset, and signedness.
/// Asserts that T is large enough to store the read value.
///
/// Example:
/// const T = packed struct(u16){ a: u3, b: u7, c: u6 };
/// var st = T{ .a = 1, .b = 2, .c = 4 };
/// const b_field = readVarPackedInt(u64, std.mem.asBytes(&st), @bitOffsetOf(T, "b"), 7, builtin.cpu.arch.endian(), .unsigned);
///
pub fn readVarPackedInt(
comptime T: type,
bytes: []const u8,
bit_offset: usize,
bit_count: usize,
endian: std.builtin.Endian,
signedness: std.builtin.Signedness,
) T {
const uN = std.meta.Int(.unsigned, @bitSizeOf(T));
const iN = std.meta.Int(.signed, @bitSizeOf(T));
const Log2N = std.math.Log2Int(T);
const read_size = (bit_count + (bit_offset % 8) + 7) / 8;
const bit_shift = @intCast(u3, bit_offset % 8);
const pad = @intCast(Log2N, @bitSizeOf(T) - bit_count);
const lowest_byte = switch (endian) {
.Big => bytes.len - (bit_offset / 8) - read_size,
.Little => bit_offset / 8,
};
const read_bytes = bytes[lowest_byte..][0..read_size];
if (@bitSizeOf(T) <= 8) {
// These are the same shifts/masks we perform below, but adds `@truncate`/`@intCast`
// where needed since int is smaller than a byte.
const value = if (read_size == 1) b: {
break :b @truncate(uN, read_bytes[0] >> bit_shift);
} else b: {
const i: u1 = @boolToInt(endian == .Big);
const head = @truncate(uN, read_bytes[i] >> bit_shift);
const tail_shift = @intCast(Log2N, @as(u4, 8) - bit_shift);
const tail = @truncate(uN, read_bytes[1 - i]);
break :b (tail << tail_shift) | head;
};
switch (signedness) {
.signed => return @intCast(T, (@bitCast(iN, value) << pad) >> pad),
.unsigned => return @intCast(T, (@bitCast(uN, value) << pad) >> pad),
}
}
// Copy the value out (respecting endianness), accounting for bit_shift
var int: uN = 0;
switch (endian) {
.Big => {
for (read_bytes[0 .. read_size - 1]) |elem| {
int = elem | (int << 8);
}
int = (read_bytes[read_size - 1] >> bit_shift) | (int << (@as(u4, 8) - bit_shift));
},
.Little => {
int = read_bytes[0] >> bit_shift;
for (read_bytes[1..]) |elem, i| {
int |= (@as(uN, elem) << @intCast(Log2N, (8 * (i + 1) - bit_shift)));
}
},
}
switch (signedness) {
.signed => return @intCast(T, (@bitCast(iN, int) << pad) >> pad),
.unsigned => return @intCast(T, (@bitCast(uN, int) << pad) >> pad),
}
}
/// Reads an integer from memory with bit count specified by T.
/// The bit count of T must be evenly divisible by 8.
/// This function cannot fail and cannot cause undefined behavior.
@@ -1365,6 +1435,84 @@ pub fn readInt(comptime T: type, bytes: *const [@divExact(@typeInfo(T).Int.bits,
}
}
fn readPackedIntLittle(comptime T: type, bytes: []const u8, bit_offset: usize) T {
const uN = std.meta.Int(.unsigned, @bitSizeOf(T));
const Log2N = std.math.Log2Int(T);
const bit_count = @as(usize, @bitSizeOf(T));
const bit_shift = @intCast(u3, bit_offset % 8);
const load_size = (bit_count + 7) / 8;
const load_tail_bits = @intCast(u3, (load_size * 8) - bit_count);
const LoadInt = std.meta.Int(.unsigned, load_size * 8);
if (bit_count == 0)
return 0;
// Read by loading a LoadInt, and then follow it up with a 1-byte read
// of the tail if bit_offset pushed us over a byte boundary.
const read_bytes = bytes[bit_offset / 8 ..];
const val = @truncate(uN, readIntLittle(LoadInt, read_bytes[0..load_size]) >> bit_shift);
if (bit_shift > load_tail_bits) {
const tail_bits = @intCast(Log2N, bit_shift - load_tail_bits);
const tail_byte = read_bytes[load_size];
const tail_truncated = if (bit_count < 8) @truncate(uN, tail_byte) else @as(uN, tail_byte);
return @bitCast(T, val | (tail_truncated << (@truncate(Log2N, bit_count) -% tail_bits)));
} else return @bitCast(T, val);
}
fn readPackedIntBig(comptime T: type, bytes: []const u8, bit_offset: usize) T {
const uN = std.meta.Int(.unsigned, @bitSizeOf(T));
const Log2N = std.math.Log2Int(T);
const bit_count = @as(usize, @bitSizeOf(T));
const bit_shift = @intCast(u3, bit_offset % 8);
const byte_count = (@as(usize, bit_shift) + bit_count + 7) / 8;
const load_size = (bit_count + 7) / 8;
const load_tail_bits = @intCast(u3, (load_size * 8) - bit_count);
const LoadInt = std.meta.Int(.unsigned, load_size * 8);
if (bit_count == 0)
return 0;
// Read by loading a LoadInt, and then follow it up with a 1-byte read
// of the tail if bit_offset pushed us over a byte boundary.
const end = bytes.len - (bit_offset / 8);
const read_bytes = bytes[(end - byte_count)..end];
const val = @truncate(uN, readIntBig(LoadInt, bytes[(end - load_size)..end][0..load_size]) >> bit_shift);
if (bit_shift > load_tail_bits) {
const tail_bits = @intCast(Log2N, bit_shift - load_tail_bits);
const tail_byte = if (bit_count < 8) @truncate(uN, read_bytes[0]) else @as(uN, read_bytes[0]);
return @bitCast(T, val | (tail_byte << (@truncate(Log2N, bit_count) -% tail_bits)));
} else return @bitCast(T, val);
}
pub const readPackedIntNative = switch (native_endian) {
.Little => readPackedIntLittle,
.Big => readPackedIntBig,
};
pub const readPackedIntForeign = switch (native_endian) {
.Little => readPackedIntBig,
.Big => readPackedIntLittle,
};
/// Loads an integer from packed memory.
/// Asserts that buffer contains at least bit_offset + @bitSizeOf(T) bits.
///
/// Example:
/// const T = packed struct(u16){ a: u3, b: u7, c: u6 };
/// var st = T{ .a = 1, .b = 2, .c = 4 };
/// const b_field = readPackedInt(u7, std.mem.asBytes(&st), @bitOffsetOf(T, "b"), builtin.cpu.arch.endian());
///
pub fn readPackedInt(comptime T: type, bytes: []const u8, bit_offset: usize, endian: Endian) T {
switch (endian) {
.Little => return readPackedIntLittle(T, bytes, bit_offset),
.Big => return readPackedIntBig(T, bytes, bit_offset),
}
}
/// Asserts that bytes.len >= @typeInfo(T).Int.bits / 8. Reads the integer starting from index 0
/// and ignores extra bytes.
/// The bit count of T must be evenly divisible by 8.
@@ -1447,6 +1595,100 @@ pub fn writeInt(comptime T: type, buffer: *[@divExact(@typeInfo(T).Int.bits, 8)]
}
}
pub fn writePackedIntLittle(comptime T: type, bytes: []u8, bit_offset: usize, value: T) void {
const uN = std.meta.Int(.unsigned, @bitSizeOf(T));
const Log2N = std.math.Log2Int(T);
const bit_count = @as(usize, @bitSizeOf(T));
const bit_shift = @intCast(u3, bit_offset % 8);
const store_size = (@bitSizeOf(T) + 7) / 8;
const store_tail_bits = @intCast(u3, (store_size * 8) - bit_count);
const StoreInt = std.meta.Int(.unsigned, store_size * 8);
if (bit_count == 0)
return;
// Write by storing a StoreInt, and then follow it up with a 1-byte tail
// if bit_offset pushed us over a byte boundary.
const write_bytes = bytes[bit_offset / 8 ..];
const head = write_bytes[0] & ((@as(u8, 1) << bit_shift) - 1);
var write_value = (@as(StoreInt, @bitCast(uN, value)) << bit_shift) | @intCast(StoreInt, head);
if (bit_shift > store_tail_bits) {
const tail_len = @intCast(Log2N, bit_shift - store_tail_bits);
write_bytes[store_size] &= ~((@as(u8, 1) << @intCast(u3, tail_len)) - 1);
write_bytes[store_size] |= @intCast(u8, (@bitCast(uN, value) >> (@truncate(Log2N, bit_count) -% tail_len)));
} else if (bit_shift < store_tail_bits) {
const tail_len = store_tail_bits - bit_shift;
const tail = write_bytes[store_size - 1] & (@as(u8, 0xfe) << (7 - tail_len));
write_value |= @as(StoreInt, tail) << (8 * (store_size - 1));
}
writeIntLittle(StoreInt, write_bytes[0..store_size], write_value);
}
pub fn writePackedIntBig(comptime T: type, bytes: []u8, bit_offset: usize, value: T) void {
const uN = std.meta.Int(.unsigned, @bitSizeOf(T));
const Log2N = std.math.Log2Int(T);
const bit_count = @as(usize, @bitSizeOf(T));
const bit_shift = @intCast(u3, bit_offset % 8);
const byte_count = (bit_shift + bit_count + 7) / 8;
const store_size = (@bitSizeOf(T) + 7) / 8;
const store_tail_bits = @intCast(u3, (store_size * 8) - bit_count);
const StoreInt = std.meta.Int(.unsigned, store_size * 8);
if (bit_count == 0)
return;
// Write by storing a StoreInt, and then follow it up with a 1-byte tail
// if bit_offset pushed us over a byte boundary.
const end = bytes.len - (bit_offset / 8);
const write_bytes = bytes[(end - byte_count)..end];
const head = write_bytes[byte_count - 1] & ((@as(u8, 1) << bit_shift) - 1);
var write_value = (@as(StoreInt, @bitCast(uN, value)) << bit_shift) | @intCast(StoreInt, head);
if (bit_shift > store_tail_bits) {
const tail_len = @intCast(Log2N, bit_shift - store_tail_bits);
write_bytes[0] &= ~((@as(u8, 1) << @intCast(u3, tail_len)) - 1);
write_bytes[0] |= @intCast(u8, (@bitCast(uN, value) >> (@truncate(Log2N, bit_count) -% tail_len)));
} else if (bit_shift < store_tail_bits) {
const tail_len = store_tail_bits - bit_shift;
const tail = write_bytes[0] & (@as(u8, 0xfe) << (7 - tail_len));
write_value |= @as(StoreInt, tail) << (8 * (store_size - 1));
}
writeIntBig(StoreInt, write_bytes[(byte_count - store_size)..][0..store_size], write_value);
}
pub const writePackedIntNative = switch (native_endian) {
.Little => writePackedIntLittle,
.Big => writePackedIntBig,
};
pub const writePackedIntForeign = switch (native_endian) {
.Little => writePackedIntBig,
.Big => writePackedIntLittle,
};
/// Stores an integer to packed memory.
/// Asserts that buffer contains at least bit_offset + @bitSizeOf(T) bits.
///
/// Example:
/// const T = packed struct(u16){ a: u3, b: u7, c: u6 };
/// var st = T{ .a = 1, .b = 2, .c = 4 };
/// // st.b = 0x7f;
/// writePackedInt(u7, std.mem.asBytes(&st), @bitOffsetOf(T, "b"), 0x7f, builtin.cpu.arch.endian());
///
pub fn writePackedInt(comptime T: type, bytes: []u8, bit_offset: usize, value: T, endian: Endian) void {
switch (endian) {
.Little => writePackedIntLittle(T, bytes, bit_offset, value),
.Big => writePackedIntBig(T, bytes, bit_offset, value),
}
}
/// Writes a twos-complement little-endian integer to memory.
/// Asserts that buf.len >= @typeInfo(T).Int.bits / 8.
/// The bit count of T must be divisible by 8.
@@ -1523,6 +1765,69 @@ pub fn writeIntSlice(comptime T: type, buffer: []u8, value: T, endian: Endian) v
};
}
/// Stores an integer to packed memory with provided bit_count, bit_offset, and signedness.
/// If negative, the written value is sign-extended.
///
/// Example:
/// const T = packed struct(u16){ a: u3, b: u7, c: u6 };
/// var st = T{ .a = 1, .b = 2, .c = 4 };
/// // st.b = 0x7f;
/// var value: u64 = 0x7f;
/// writeVarPackedInt(std.mem.asBytes(&st), @bitOffsetOf(T, "b"), 7, value, builtin.cpu.arch.endian());
///
pub fn writeVarPackedInt(bytes: []u8, bit_offset: usize, bit_count: usize, value: anytype, endian: std.builtin.Endian) void {
const T = @TypeOf(value);
const uN = std.meta.Int(.unsigned, @bitSizeOf(T));
const Log2N = std.math.Log2Int(T);
const bit_shift = @intCast(u3, bit_offset % 8);
const write_size = (bit_count + bit_shift + 7) / 8;
const lowest_byte = switch (endian) {
.Big => bytes.len - (bit_offset / 8) - write_size,
.Little => bit_offset / 8,
};
const write_bytes = bytes[lowest_byte..][0..write_size];
if (write_size == 1) {
// Single byte writes are handled specially, since we need to mask bits
// on both ends of the byte.
const mask = (@as(u8, 0xff) >> @intCast(u3, 8 - bit_count));
const new_bits = @intCast(u8, @bitCast(uN, value) & mask) << bit_shift;
write_bytes[0] = (write_bytes[0] & ~(mask << bit_shift)) | new_bits;
return;
}
var remaining: T = value;
// Iterate bytes forward for Little-endian, backward for Big-endian
const delta: i2 = if (endian == .Big) -1 else 1;
const start = if (endian == .Big) @intCast(isize, write_bytes.len - 1) else 0;
var i: isize = start; // isize for signed index arithmetic
// Write first byte, using a mask to protects bits preceding bit_offset
const head_mask = @as(u8, 0xff) >> bit_shift;
write_bytes[@intCast(usize, i)] &= ~(head_mask << bit_shift);
write_bytes[@intCast(usize, i)] |= @intCast(u8, @bitCast(uN, remaining) & head_mask) << bit_shift;
remaining >>= @intCast(Log2N, @as(u4, 8) - bit_shift);
i += delta;
// Write bytes[1..bytes.len - 1]
if (@bitSizeOf(T) > 8) {
const loop_end = start + delta * (@intCast(isize, write_size) - 1);
while (i != loop_end) : (i += delta) {
write_bytes[@intCast(usize, i)] = @truncate(u8, @bitCast(uN, remaining));
remaining >>= 8;
}
}
// Write last byte, using a mask to protect bits following bit_offset + bit_count
const following_bits = -%@truncate(u3, bit_shift + bit_count);
const tail_mask = (@as(u8, 0xff) << following_bits) >> following_bits;
write_bytes[@intCast(usize, i)] &= ~tail_mask;
write_bytes[@intCast(usize, i)] |= @intCast(u8, @bitCast(uN, remaining) & tail_mask);
}
test "writeIntBig and writeIntLittle" {
var buf0: [0]u8 = undefined;
var buf1: [1]u8 = undefined;
@@ -3393,3 +3698,158 @@ pub fn alignInSlice(slice: anytype, comptime new_alignment: usize) ?AlignedSlice
const aligned_slice = bytesAsSlice(Element, aligned_bytes[0..slice_length_bytes]);
return @alignCast(new_alignment, aligned_slice);
}
test "read/write(Var)PackedInt" {
const foreign_endian: Endian = if (native_endian == .Big) .Little else .Big;
const expect = std.testing.expect;
var prng = std.rand.DefaultPrng.init(1234);
const random = prng.random();
@setEvalBranchQuota(10_000);
inline for ([_]type{ u8, u16, u32, u128 }) |BackingType| {
for ([_]BackingType{
@as(BackingType, 0), // all zeros
-%@as(BackingType, 1), // all ones
random.int(BackingType), // random
random.int(BackingType), // random
random.int(BackingType), // random
}) |init_value| {
const uTs = [_]type{ u1, u3, u7, u8, u9, u10, u15, u16, u86 };
const iTs = [_]type{ i1, i3, i7, i8, i9, i10, i15, i16, i86 };
inline for (uTs ++ iTs) |PackedType| {
if (@bitSizeOf(PackedType) > @bitSizeOf(BackingType))
continue;
const iPackedType = std.meta.Int(.signed, @bitSizeOf(PackedType));
const uPackedType = std.meta.Int(.unsigned, @bitSizeOf(PackedType));
const Log2T = std.math.Log2Int(BackingType);
const offset_at_end = @bitSizeOf(BackingType) - @bitSizeOf(PackedType);
for ([_]usize{ 0, 1, 7, 8, 9, 10, 15, 16, 86, offset_at_end }) |offset| {
if (offset > offset_at_end or offset == @bitSizeOf(BackingType))
continue;
for ([_]PackedType{
~@as(PackedType, 0), // all ones: -1 iN / maxInt uN
@as(PackedType, 0), // all zeros: 0 iN / 0 uN
@bitCast(PackedType, @as(iPackedType, math.maxInt(iPackedType))), // maxInt iN
@bitCast(PackedType, @as(iPackedType, math.minInt(iPackedType))), // maxInt iN
random.int(PackedType), // random
random.int(PackedType), // random
}) |write_value| {
{ // Fixed-size Read/Write (Native-endian)
// Initialize Value
var value: BackingType = init_value;
// Read
const read_value1 = readPackedInt(PackedType, asBytes(&value), offset, native_endian);
try expect(read_value1 == @bitCast(PackedType, @truncate(uPackedType, value >> @intCast(Log2T, offset))));
// Write
writePackedInt(PackedType, asBytes(&value), offset, write_value, native_endian);
try expect(write_value == @bitCast(PackedType, @truncate(uPackedType, value >> @intCast(Log2T, offset))));
// Read again
const read_value2 = readPackedInt(PackedType, asBytes(&value), offset, native_endian);
try expect(read_value2 == write_value);
// Verify bits outside of the target integer are unmodified
const diff_bits = init_value ^ value;
if (offset != offset_at_end)
try expect(diff_bits >> @intCast(Log2T, offset + @bitSizeOf(PackedType)) == 0);
if (offset != 0)
try expect(diff_bits << @intCast(Log2T, @bitSizeOf(BackingType) - offset) == 0);
}
{ // Fixed-size Read/Write (Foreign-endian)
// Initialize Value
var value: BackingType = @byteSwap(init_value);
// Read
const read_value1 = readPackedInt(PackedType, asBytes(&value), offset, foreign_endian);
try expect(read_value1 == @bitCast(PackedType, @truncate(uPackedType, @byteSwap(value) >> @intCast(Log2T, offset))));
// Write
writePackedInt(PackedType, asBytes(&value), offset, write_value, foreign_endian);
try expect(write_value == @bitCast(PackedType, @truncate(uPackedType, @byteSwap(value) >> @intCast(Log2T, offset))));
// Read again
const read_value2 = readPackedInt(PackedType, asBytes(&value), offset, foreign_endian);
try expect(read_value2 == write_value);
// Verify bits outside of the target integer are unmodified
const diff_bits = init_value ^ @byteSwap(value);
if (offset != offset_at_end)
try expect(diff_bits >> @intCast(Log2T, offset + @bitSizeOf(PackedType)) == 0);
if (offset != 0)
try expect(diff_bits << @intCast(Log2T, @bitSizeOf(BackingType) - offset) == 0);
}
const signedness = @typeInfo(PackedType).Int.signedness;
const NextPowerOfTwoInt = std.meta.Int(signedness, comptime try std.math.ceilPowerOfTwo(u16, @bitSizeOf(PackedType)));
const ui64 = std.meta.Int(signedness, 64);
inline for ([_]type{ PackedType, NextPowerOfTwoInt, ui64 }) |U| {
{ // Variable-size Read/Write (Native-endian)
if (@bitSizeOf(U) < @bitSizeOf(PackedType))
continue;
// Initialize Value
var value: BackingType = init_value;
// Read
const read_value1 = readVarPackedInt(U, asBytes(&value), offset, @bitSizeOf(PackedType), native_endian, signedness);
try expect(read_value1 == @bitCast(PackedType, @truncate(uPackedType, value >> @intCast(Log2T, offset))));
// Write
writeVarPackedInt(asBytes(&value), offset, @bitSizeOf(PackedType), @as(U, write_value), native_endian);
try expect(write_value == @bitCast(PackedType, @truncate(uPackedType, value >> @intCast(Log2T, offset))));
// Read again
const read_value2 = readVarPackedInt(U, asBytes(&value), offset, @bitSizeOf(PackedType), native_endian, signedness);
try expect(read_value2 == write_value);
// Verify bits outside of the target integer are unmodified
const diff_bits = init_value ^ value;
if (offset != offset_at_end)
try expect(diff_bits >> @intCast(Log2T, offset + @bitSizeOf(PackedType)) == 0);
if (offset != 0)
try expect(diff_bits << @intCast(Log2T, @bitSizeOf(BackingType) - offset) == 0);
}
{ // Variable-size Read/Write (Foreign-endian)
if (@bitSizeOf(U) < @bitSizeOf(PackedType))
continue;
// Initialize Value
var value: BackingType = @byteSwap(init_value);
// Read
const read_value1 = readVarPackedInt(U, asBytes(&value), offset, @bitSizeOf(PackedType), foreign_endian, signedness);
try expect(read_value1 == @bitCast(PackedType, @truncate(uPackedType, @byteSwap(value) >> @intCast(Log2T, offset))));
// Write
writeVarPackedInt(asBytes(&value), offset, @bitSizeOf(PackedType), @as(U, write_value), foreign_endian);
try expect(write_value == @bitCast(PackedType, @truncate(uPackedType, @byteSwap(value) >> @intCast(Log2T, offset))));
// Read again
const read_value2 = readVarPackedInt(U, asBytes(&value), offset, @bitSizeOf(PackedType), foreign_endian, signedness);
try expect(read_value2 == write_value);
// Verify bits outside of the target integer are unmodified
const diff_bits = init_value ^ @byteSwap(value);
if (offset != offset_at_end)
try expect(diff_bits >> @intCast(Log2T, offset + @bitSizeOf(PackedType)) == 0);
if (offset != 0)
try expect(diff_bits << @intCast(Log2T, @bitSizeOf(BackingType) - offset) == 0);
}
}
}
}
}
}
}
}