Encode/decode variable-length integers, with optional Zig-Zag encoding (Protobuf).

varint.zig

//! Provides functions for encoding and decoding variable-length integers.

// MIT License
//
// Copyright (c) 2025 Eric Freed
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

const std = @import("std");
const Reader = std.Io.Reader;
const Writer = std.Io.Writer;

const segment_mask = 0x7F;
const continue_bit = 0x80;

/// Calculates the number of bytes required to encode the given integer value of type `T`.
pub fn calculateLength(comptime T: type, value: T, comptime zigzag: bool) usize {
    const Int = std.meta.Int(.unsigned, @typeInfo(T).int.bits);
    var unsigned: Int = if (zigzag and @typeInfo(T).int.signedness == .signed) zigzag_encoded: {
        const shift_amount: std.math.Log2Int(Int) = comptime @intCast(@typeInfo(T).int.bits - 1);
        break :zigzag_encoded @bitCast((value << 1) ^ (value >> shift_amount));
    } else @bitCast(value);

    var len: usize = 1;
    while (unsigned >= continue_bit) {
        unsigned >>= 7;
        len += 1;
    }
    return len;
}

/// Decodes a variable-length integer to a fixed-size integer of type `T` from the given buffer.
/// The `length` argument will be assigned the number of bytes used to define the value.
pub fn decode(buffer: []const u8, comptime T: type, comptime zigzag: bool, length: *usize) error{ Overflow, BufferUnderrun }!T {
    var reader = Reader.fixed(buffer);
    const result = read(&reader, T, zigzag) catch |err| return switch (err) {
        error.Overflow => error.Overflow,
        else => error.BufferUnderrun,
    };
    length.* = reader.seek;
    return result;
}

/// Encodes an integer of type `T` as a variable-length integer into the given `buffer`.
/// For signed values, Zig-Zag encoding can optionally be used.
/// Returns the number of bytes stored in `buffer`.
pub fn encode(buffer: []u8, comptime T: type, value: T, comptime zigzag: bool) error{BufferUnderrun}!usize {
    var writer = Writer.fixed(buffer);
    return write(&writer, T, value, zigzag) catch error.BufferUnderrun;
}

/// Reads a variable-length integer from the current position in the stream.
/// For signed values, Zig-Zag encoding can optionally be used.
pub fn read(reader: *Reader, comptime T: type, comptime zigzag: bool) (Reader.Error || error{Overflow})!T {
    const Int = std.meta.Int(.unsigned, @typeInfo(T).int.bits);

    var result: Int = 0;
    var shift: std.math.Log2Int(Int) = 0;
    while (true) {
        const b: u8 = try reader.takeByte();
        result |= @shlExact(@as(Int, @intCast(b & segment_mask)), shift);
        if ((b & continue_bit) == 0) break;
        shift += 7;
        if (shift >= @typeInfo(T).int.bits) return error.Overflow;
    }

    if (zigzag and @typeInfo(T).int.signedness == .signed) {
        const signed: T = @bitCast(result);
        return (signed >> 1) ^ -(signed & 1);
    } else {
        return @bitCast(result);
    }
}

/// Writes a variable-length integer to the current position in the stream.
/// For signed values, Zig-Zag encoding can optionally be used.
pub fn write(writer: *Writer, comptime T: type, value: T, comptime zigzag: bool) Writer.Error!usize {
    const Int = std.meta.Int(.unsigned, @typeInfo(T).int.bits);
    var len: usize = 0;

    var v: Int = if (zigzag and @typeInfo(T).int.signedness == .signed) zigzag_encoded: {
        const shift_amount: std.math.Log2Int(Int) = comptime @intCast(@typeInfo(T).int.bits - 1);
        break :zigzag_encoded @bitCast((value << 1) ^ (value >> shift_amount));
    } else @bitCast(value);

    while (v >= continue_bit) {
        try writer.writeByte(@intCast((v & segment_mask) | continue_bit));
        v >>= 7;
        len += 1;
    }
    try writer.writeByte(@intCast(v));
    return len + 1;
}

test "varint encode/decode" {
    var buffer: [16]u8 = undefined;
    var encode_len = try encode(&buffer, i32, 1337, false);
    var decode_len: usize = 0;
    var decoded = try decode(buffer[0..], i32, false, &decode_len);
    try std.testing.expectEqual(1337, decoded);
    try std.testing.expectEqual(encode_len, decode_len);

    encode_len = try encode(&buffer, i32, std.math.maxInt(i32), false);
    decoded = try decode(buffer[0..], i32, false, &decode_len);
    try std.testing.expectEqual(std.math.maxInt(i32), decoded);
    try std.testing.expectEqual(encode_len, decode_len);
    try std.testing.expectEqual(5, encode_len);

    // Zig-zag
    encode_len = try encode(&buffer, i32, -123987, true);
    decoded = try decode(buffer[0..], i32, true, &decode_len);
    try std.testing.expectEqual(-123987, decoded);
    try std.testing.expectEqual(encode_len, decode_len);
}