eth_state_diff/

balances.rs

//! Compact delta encoding and reconstruction for Ethereum validator balances.
//!
//! This module computes compact binary deltas between two validator balance
//! snapshots and applies those deltas to reconstruct the target snapshot.
//!
//! The encoding is specialized for Ethereum beacon-chain balances, where most
//! validators either retain the same balance or change by a relatively small
//! amount between consecutive states.
//!
//! ## Encoding
//!
//! For each balance in the portion shared by the base and target snapshots,
//! the delta records one of four states using a packed two-bit tag:
//!
//! - [`SET_NO_CHANGE`] — the balance is unchanged;
//! - [`SET_TO_ZERO`] — the target balance is zero;
//! - [`SET_TO_DIFF`] — the target is reconstructed by applying a signed
//!   difference to the base balance;
//! - [`SET_TO_TARGET_VALUE`] — the target balance is stored explicitly.
//!
//! Changed balances whose difference fits in an `i32` are normally encoded as
//! signed differences. The most frequently occurring difference is selected as
//! the [`BalancesDiff::mode`], and encoded differences store only the
//! difference relative to that mode.
//!
//! Signed corrected differences are encoded as zig-zag integers followed by a
//! variable-length integer encoding. This makes small and frequently occurring
//! changes inexpensive to store.
//!
//! Differences that do not fit in an `i32` are encoded as explicit target
//! values.
//!
//! Balances that exist only in the target snapshot are stored in
//! [`BalancesDiff::appended_balances`].
//!
//! ## Two-pass encoding
//!
//! [`diff_balances_iter`] cannot require its input iterators to implement
//! [`Clone`]. It therefore performs the diff in two logical passes:
//!
//! 1. the common portion of the iterators is consumed into a compact
//!    intermediate list of changes;
//! 2. the statistical mode is selected and the changes are encoded.
//!
//! Any remaining items in the target iterator are treated as newly appended
//! balances.
//!
//! ## Iterator API
//!
//! [`diff_balances`] is the convenience API for contiguous balance slices.
//! [`diff_balances_iter`] is intended for consensus clients whose balances are
//! stored in persistent lists, trees, or other non-contiguous structures.
//!
//! The iterator API avoids requiring the caller to materialize the complete
//! balance registry as a flat buffer.
//!
//! ## Reconstruction
//!
//! [`apply_balances`] and [`apply_balances_iter`] mutate the supplied balance
//! collection in place.
//!
//! The supplied collection must represent the **base snapshot** from which the
//! delta was generated. After successful application, it contains the target
//! balances.
//!
//! Existing balances are updated in place. Target balances that extend beyond
//! the base snapshot are appended from the delta.
//!
//! ## Complexity
//!
//! [`diff_balances`] and [`diff_balances_iter`] run in:
//!
//! ```text
//! O(n)
//! ```
//!
//! where `n` is the number of balances in the common portion of the snapshots.
//!
//! Delta generation additionally requires storage proportional to the number
//! of changed balances:
//!
//! ```text
//! O(k)
//! ```
//!
//! where `k` is the number of changed balances.
//!
//! [`apply_balances`] and [`apply_balances_iter`] run in:
//!
//! ```text
//! O(n + a)
//! ```
//!
//! where `n` is the number of balances represented by the tag vector and `a`
//! is the number of appended balances.
//!
//! Reconstruction operates in place and does not require allocating a second
//! balance buffer.
//!
//! ## Serialization
//!
//! [`BalancesDiff`] is designed to be serialized using `rkyv` and can then be
//! passed to a general-purpose compressor such as zstd.
//!
//! The delta representation itself is independent of the serialization and
//! compression layer.
//!
//! ## Delta validity
//!
//! Applying a delta assumes that the supplied base collection corresponds to
//! the base snapshot used to create the delta. The application functions do
//! not independently verify the original balance values.
//!
//! In particular, a `SET_TO_DIFF` entry applies its decoded difference to the
//! current value in the supplied target collection. Applying the same delta to
//! a different base snapshot therefore does not generally produce the intended
//! target snapshot.
//!
//! [`BalancesDiff`]: crate::types::BalancesDiff
//! [`SET_NO_CHANGE`]: crate::types::SET_NO_CHANGE
//! [`SET_TO_ZERO`]: crate::types::SET_TO_ZERO
//! [`SET_TO_DIFF`]: crate::types::SET_TO_DIFF
//! [`SET_TO_TARGET_VALUE`]: crate::types::SET_TO_TARGET_VALUE

use rustc_hash::FxHashMap;

use crate::{
    types::{
        ArchivedBalancesDiff, BalancesDiff, BitTagVec, SET_NO_CHANGE, SET_TO_DIFF,
        SET_TO_TARGET_VALUE, SET_TO_ZERO,
    },
    Error,
};

/// Computes a compact balance delta between two contiguous balance slices.
///
/// The returned [`BalancesDiff`] contains the information required to
/// reconstruct `target` from `base`.
///
/// The common portion of the two slices is encoded using packed two-bit tags,
/// statistical mode correction, zig-zag encoded signed differences, and
/// explicit target values where necessary. If `target` contains more balances
/// than `base`, the additional balances are stored in
/// [`BalancesDiff::appended_balances`].
///
/// This is a convenience wrapper around [`diff_balances_iter`].
///
/// # Complexity
///
/// `O(n)` time and `O(k)` additional space, where `n` is the number of balances
/// in the common portion and `k` is the number of changed balances.
///
/// # Examples
///
/// ```ignore
/// let delta = diff_balances(&base, &target);
/// ```
///
/// [`BalancesDiff`]: crate::types::BalancesDiff
/// [`diff_balances_iter`]: crate::balances::diff_balances_iter
pub fn diff_balances(base: &[u64], target: &[u64]) -> BalancesDiff {
    diff_balances_iter(base.iter().copied(), target.iter().copied())
}

/// Computes a compact balance delta between two balance iterators.
///
/// This is the generic counterpart to [`diff_balances`]. It is intended for
/// consensus clients whose balance storage is not represented as a contiguous
/// `&[u64]`.
///
/// The iterators must implement [`ExactSizeIterator`], allowing the function
/// to determine the size of the common portion and distinguish existing
/// balances from balances appended to the target.
///
/// The iterators are consumed during encoding and do not need to implement
/// [`Clone`].
///
/// The common portion is first collected into a compact list of changed
/// balances so that the most frequently occurring balance difference can be
/// selected as the encoding mode. Remaining items in the target iterator are
/// stored as appended balances.
///
/// # Complexity
///
/// `O(n)` time and `O(k + a)` additional space, where:
///
/// - `n` is the number of balances in the common portion;
/// - `k` is the number of changed balances; and
/// - `a` is the number of balances appended to the target.
///
/// [`diff_balances`]: crate::balances::diff_balances
pub fn diff_balances_iter<I1, I2>(mut base: I1, mut target: I2) -> BalancesDiff
where
    I1: ExactSizeIterator<Item = u64>,
    I2: ExactSizeIterator<Item = u64>,
{
    let common_len = base.len().min(target.len());
    let mut changes = Vec::with_capacity(1024);

    for idx in 0..common_len {
        let Some(v1) = base.next() else {
            break;
        };
        let Some(v2) = target.next() else {
            break;
        };

        if v1 != v2 {
            let diff = v2
                .checked_sub(v1)
                .and_then(|value| i64::try_from(value).ok())
                .or_else(|| {
                    v1.checked_sub(v2)
                        .and_then(|value| i64::try_from(value).ok())
                        .map(|value| -value)
                });

            changes.push(Change {
                idx,
                diff,
                target: v2,
            });
        }
    }

    let mode = find_mode(&changes);
    let (tags, varint_payload, target_values) = encode(common_len, &changes, mode);

    BalancesDiff {
        tags,
        mode,
        varint_payload,
        target_values,
        appended_balances: target.collect(),
    }
}

/// Applies a balance delta to a mutable balance collection in place.
///
/// `target` must initially contain the **base balance snapshot** used to
/// generate `delta`.
///
/// Existing balances are reconstructed according to the packed tags in the
/// delta. Changed balances encoded as differences are updated relative to
/// their current base value, while explicit and zero-valued entries replace
/// the corresponding balance directly.
///
/// Any balances stored in [`BalancesDiff::appended_balances`] are appended to
/// the collection after the common portion has been reconstructed.
///
/// This API accepts [`crate::ListMutTarget`] so that consensus clients can
/// apply the delta directly to tree-backed or persistent balance collections
/// without first materializing them as a `Vec<u64>`.
///
/// # Errors
///
/// Returns [`Error::MalformedDelta`] if the delta is structurally invalid,
/// contains incomplete payload data, or cannot be safely applied to the
/// supplied target collection.
pub fn apply_balances_iter<T: crate::ListMutTarget<u64>>(
    target: &mut T,
    delta: &ArchivedBalancesDiff,
) -> Result<(), Error> {
    let mode = delta.mode.to_native();

    let tag_len = usize::try_from(delta.tags.len.to_native())
        .map_err(|_| Error::MalformedDelta("delta tag length does not fit in usize".into()))?;

    if target.len() != tag_len {
        return Err(Error::MalformedDelta(format!(
            "target length {} does not match delta tag length {tag_len}",
            target.len()
        )));
    }

    let mut target_iter = delta.target_values.iter();
    let payload = delta.varint_payload.as_slice();
    let mut payload_cursor = 0usize;
    let mut base_idx = 0usize;

    for &tag_byte in delta.tags.data.iter() {
        if base_idx >= tag_len {
            break;
        }

        // Fast path: four consecutive SET_NO_CHANGE entries.
        if tag_byte == 0 {
            base_idx = (base_idx + 4).min(tag_len);
            continue;
        }

        for bit in 0..4 {
            if base_idx >= tag_len {
                break;
            }

            let tag = (tag_byte >> (bit * 2)) & 0b11;

            match tag {
                SET_NO_CHANGE => {}

                SET_TO_ZERO => {
                    let Some(value) = target.get_mut(base_idx) else {
                        return Err(Error::MalformedDelta(format!(
                            "target collection is missing balance at index {base_idx}"
                        )));
                    };

                    *value = 0;
                }

                SET_TO_TARGET_VALUE => {
                    let Some(target_value) = target_iter.next() else {
                        return Err(Error::MalformedDelta(
                            "target-value payload is shorter than the encoded target-value tags"
                                .into(),
                        ));
                    };

                    let Some(value) = target.get_mut(base_idx) else {
                        return Err(Error::MalformedDelta(format!(
                            "target collection is missing balance at index {base_idx}"
                        )));
                    };

                    *value = target_value.to_native();
                }

                SET_TO_DIFF => {
                    let encoded = read_varint(payload, &mut payload_cursor)?;
                    let corrected = zigzag_decode(encoded);

                    let Some(diff) = corrected.checked_add(mode) else {
                        return Err(Error::MalformedDelta(
                            "decoded balance difference overflows i64".into(),
                        ));
                    };

                    let Some(value) = target.get_mut(base_idx) else {
                        return Err(Error::MalformedDelta(format!(
                            "target collection is missing balance at index {base_idx}"
                        )));
                    };

                    let base_value = i64::try_from(*value).map_err(|_| {
                        Error::MalformedDelta(format!(
                            "base balance at index {base_idx} exceeds i64 range"
                        ))
                    })?;

                    let updated = base_value
                        .checked_add(diff)
                        .and_then(|value| u64::try_from(value).ok())
                        .ok_or_else(|| {
                            Error::MalformedDelta(format!(
                                "decoded balance difference produces an invalid value at index {base_idx}"
                            ))
                        })?;

                    *value = updated;
                }

                _ => {
                    return Err(Error::MalformedDelta(format!(
                        "invalid balance tag {tag} at index {base_idx}"
                    )));
                }
            }

            base_idx += 1;
        }
    }

    if !delta.appended_balances.is_empty() {
        for value in delta.appended_balances.iter() {
            target.push(value.to_native());
        }
    }

    Ok(())
}

/// Applies a balance delta to a contiguous balance vector in place.
///
/// `base` must initially contain the balance snapshot from which `delta` was
/// generated. After successful application, `base` contains the reconstructed
/// target snapshot.
///
/// This is a convenience wrapper around [`apply_balances_iter`] for clients
/// that store balances contiguously.
///
/// # Errors
///
/// Returns an error if `delta` is malformed or if its represented base length
/// does not match `base`.
///
/// [`apply_balances_iter`]: crate::balances::apply_balances_iter
pub fn apply_balances(base: &mut Vec<u64>, delta: &ArchivedBalancesDiff) -> Result<(), Error> {
    apply_balances_iter(base, delta)
}

/// Intermediate representation of a changed balance.
///
/// This is retained between the comparison pass and encoding pass so that
/// `diff_balances_iter` can determine the statistical mode without requiring
/// the source iterators to be cloned.
struct Change {
    idx: usize,
    diff: Option<i64>,
    target: u64,
}

/// Finds the most frequently occurring representable balance difference.
///
/// Only differences that fit within `i32` participate in mode selection.
/// Larger differences are always encoded as explicit target values and
/// therefore do not benefit from mode correction.
///
/// Returns `0` when no representable balance difference exists.
fn find_mode(changes: &[Change]) -> i64 {
    let mut freq_map = FxHashMap::default();
    freq_map.reserve(256);

    for change in changes {
        let Some(diff) = change.diff else {
            continue;
        };

        if i32::try_from(diff).is_ok() {
            *freq_map.entry(diff).or_insert(0usize) += 1;
        }
    }

    freq_map
        .into_iter()
        .max_by_key(|&(_, count)| count)
        .map(|(value, _)| value)
        .unwrap_or(0)
}

/// Encodes changed balances using the selected statistical mode.
///
/// Each changed balance is assigned a two-bit tag. Differences that fit in
/// `i32` are encoded as zig-zag varints after subtracting `mode`; values that
/// cannot use the difference representation are stored as absolute target
/// values.
fn encode(common_len: usize, changes: &[Change], mode: i64) -> (BitTagVec, Vec<u8>, Vec<u64>) {
    let mut tags = BitTagVec::new(common_len);
    let mut varint_payload = Vec::with_capacity(changes.len());
    let mut target_values = Vec::new();

    for change in changes {
        let Change { idx, diff, target } = *change;

        if target == 0 {
            tags.set(idx, SET_TO_ZERO);
            continue;
        }

        let Some(diff) = diff else {
            tags.set(idx, SET_TO_TARGET_VALUE);
            target_values.push(target);
            continue;
        };

        if i32::try_from(diff).is_ok() {
            tags.set(idx, SET_TO_DIFF);

            // `diff` and `mode` are both representable i32 values, so their
            // difference is within the i64 range.
            let corrected = diff - mode;
            write_varint(zigzag_encode(corrected), &mut varint_payload);
        } else {
            tags.set(idx, SET_TO_TARGET_VALUE);
            target_values.push(target);
        }
    }

    (tags, varint_payload, target_values)
}

/// Encodes a signed integer using zig-zag encoding.
///
/// Negative and positive values of similar magnitude are mapped to nearby
/// unsigned values, making small signed differences efficient to represent
/// with the subsequent variable-length encoding.
#[inline]
fn zigzag_encode(n: i64) -> u64 {
    ((n as u64) << 1) ^ ((n >> 63) as u64)
}

/// Decodes a value previously encoded with [`zigzag_encode`].
#[inline]
fn zigzag_decode(n: u64) -> i64 {
    ((n >> 1) as i64) ^ -((n & 1) as i64)
}

/// Encodes an unsigned integer using the variable-length format used by the
/// balance delta representation.
///
/// Each output byte stores seven payload bits. The high bit indicates whether
/// another byte follows.
#[inline]
pub(super) fn write_varint(mut val: u64, buf: &mut Vec<u8>) {
    loop {
        if val < 0x80 {
            buf.push(val as u8);
            break;
        }

        buf.push((val as u8) | 0x80);
        val >>= 7;
    }
}

/// Decodes one unsigned variable-length integer from `buf`.
///
/// `cursor` is advanced past the decoded integer.
///
/// Returns an error if the payload is truncated or contains more than ten
/// bytes, or if the final byte contains bits outside the `u64` range.
#[inline]
pub(super) fn read_varint(buf: &[u8], cursor: &mut usize) -> Result<u64, Error> {
    let mut value = 0u64;

    for shift in (0..=63).step_by(7) {
        let Some(&byte) = buf.get(*cursor) else {
            return Err(Error::MalformedDelta(format!(
                "truncated varint payload at byte offset {}",
                *cursor
            )));
        };

        *cursor += 1;

        let payload = (byte & 0x7f) as u64;

        if shift == 63 && payload > 1 {
            return Err(Error::MalformedDelta(format!(
                "varint overflows u64 at byte offset {}",
                *cursor - 1
            )));
        }

        value |= payload << shift;

        if byte & 0x80 == 0 {
            return Ok(value);
        }
    }

    Err(Error::MalformedDelta(
        "varint exceeds the maximum length for a u64".into(),
    ))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::ArchivedBalancesDiff;
    use crate::ListMutTarget;

    /// A minimal mock target for testing `apply_balances_iter` directly.
    struct MockTarget {
        inner: Vec<u64>,
    }

    impl ListMutTarget<u64> for MockTarget {
        fn len(&self) -> usize {
            self.inner.len()
        }

        fn get_mut(&mut self, index: usize) -> Option<&mut u64> {
            self.inner.get_mut(index)
        }

        fn push(&mut self, value: u64) {
            self.inner.push(value);
        }
    }

    /// Helper to perform a full roundtrip: diff -> rkyv serialize -> rkyv access -> apply
    fn assert_roundtrip(base: &[u64], target: &[u64]) {
        let delta = diff_balances(base, target);

        let bytes = rkyv::to_bytes::<rkyv::rancor::Error>(&delta)
            .expect("test setup: failed to serialize delta");

        let archived = rkyv::access::<ArchivedBalancesDiff, rkyv::rancor::Error>(&bytes)
            .expect("test setup: failed to access archived delta");

        let mut reconstructed = base.to_vec();
        apply_balances(&mut reconstructed, archived)
            .expect("test setup: failed to apply valid delta");

        assert_eq!(
            reconstructed, target,
            "Roundtrip failed: base={base:?}, target={target:?}, delta={delta:?}"
        );
    }

    #[test]
    fn test_no_changes() {
        let state = vec![1000; 10];
        assert_roundtrip(&state, &state);
    }

    #[test]
    fn test_single_balance_change() {
        let base = vec![100, 200, 300];
        let target = vec![100, 205, 300];
        assert_roundtrip(&base, &target);
    }

    #[test]
    fn test_all_set_to_zero() {
        let base = vec![100, 200, 300];
        let target = vec![0, 0, 0];
        assert_roundtrip(&base, &target);
    }

    #[test]
    fn test_appended_balances() {
        let base = vec![100, 200];
        let target = vec![100, 200, 300, 400];
        assert_roundtrip(&base, &target);
    }

    #[test]
    fn test_only_appended_balances() {
        let base = vec![];
        let target = vec![10, 20, 30];
        assert_roundtrip(&base, &target);
    }

    #[test]
    fn test_mode_selection_frequent_diff() {
        // 1000 balances. 500 of them change by exactly +1000.
        // Mode should be 1000, making those diffs very cheap to encode.
        let base: Vec<u64> = (0..1000).map(|i| (i * 100) as u64).collect();
        let mut target = base.clone();

        for i in (0..1000).step_by(2) {
            target[i] += 1000;
        }

        let delta = diff_balances(&base, &target);
        assert_eq!(
            delta.mode, 1000,
            "Mode should select the most frequent difference"
        );
        assert_roundtrip(&base, &target);
    }

    #[test]
    fn test_i32_boundary_max_diff() {
        let base = vec![100];
        // diff = i32::MAX = 2147483647
        let target = vec![100 + i32::MAX as u64];
        assert_roundtrip(&base, &target);
    }

    #[test]
    fn test_i32_boundary_min_diff() {
        // diff = -2147483647 (which is i32::MIN + 1, fits in i32)
        let base = vec![100 + 2147483647];
        let target = vec![100];
        assert_roundtrip(&base, &target);
    }

    #[test]
    fn test_i32_overflow_uses_target_value() {
        // diff = i32::MAX + 1 = 2147483648 (does NOT fit in i32)
        let base = vec![100];
        let target = vec![100 + i32::MAX as u64 + 1];

        let delta = diff_balances(&base, &target);

        // It must fall back to SET_TO_TARGET_VALUE, which means no varint payload
        // should be generated for this diff.
        assert_eq!(
            delta.target_values.len(),
            1,
            "u64 diff exceeding i32::MAX must use SET_TO_TARGET_VALUE"
        );
        assert_eq!(
            delta.varint_payload.len(),
            0,
            "Should not use varint payload for unrepresentable diff"
        );

        assert_roundtrip(&base, &target);
    }

    #[test]
    fn test_u64_extreme_values_use_target_value() {
        // base = 0, target = u64::MAX. Cannot be encoded as a diff.
        let base = vec![0];
        let target = vec![u64::MAX];
        let delta = diff_balances(&base, &target);
        assert_eq!(
            delta.target_values.len(),
            1,
            "u64::MAX diff must use SET_TO_TARGET_VALUE"
        );
        assert_eq!(
            delta.varint_payload.len(),
            0,
            "Should not use varint payload for unrepresentable diff"
        );

        // base = u64::MAX, target = 0. Diff is unrepresentable, but the encoder
        // optimizes this to SET_TO_ZERO before it even considers the diff.
        // We just verify it roundtrips correctly.
        assert_roundtrip(&[0], &[u64::MAX]);
        assert_roundtrip(&[u64::MAX], &[0]);
    }

    #[test]
    fn test_apply_length_mismatch_returns_error() {
        let base = vec![100, 200];
        let target = vec![100, 200];

        let delta = diff_balances(&base, &target);
        let bytes = rkyv::to_bytes::<rkyv::rancor::Error>(&delta)
            .expect("test setup: failed to serialize delta");
        let archived = rkyv::access::<ArchivedBalancesDiff, rkyv::rancor::Error>(&bytes)
            .expect("test setup: failed to access archived delta");

        let mut wrong_base = MockTarget { inner: vec![100] }; // Length 1 instead of 2
        let result = apply_balances_iter(&mut wrong_base, archived);

        assert!(result.is_err(), "Should error on length mismatch");
        let err_str = format!("{}", result.expect_err("test setup: expected error"));
        assert!(
            err_str.contains("does not match delta tag length"),
            "Error message should mention length mismatch"
        );
    }

    #[test]
    fn test_apply_i64_overflow_returns_error() {
        // Setup a valid delta where mode = 1, and it applies a diff of 1.
        let base = vec![100, 200];
        let target = vec![101, 201];
        let delta = diff_balances(&base, &target);
        assert_eq!(delta.mode, 1);

        let mut bytes = rkyv::to_bytes::<rkyv::rancor::Error>(&delta)
            .expect("test setup: failed to serialize delta");

        // Find the mode (i64 = 1) and maliciously patch it to i64::MAX
        let mode_bytes = 1i64.to_le_bytes();
        let max_bytes = i64::MAX.to_le_bytes();

        let pos = bytes
            .windows(8)
            .position(|w| w == mode_bytes)
            .expect("test setup: could not find mode in serialized delta");

        bytes
            .get_mut(pos..pos + 8)
            .expect("test setup: `pos` is valid and derived from a window of length 8")
            .copy_from_slice(&max_bytes);

        let archived = rkyv::access::<ArchivedBalancesDiff, rkyv::rancor::Error>(&bytes)
            .expect("test setup: failed to access archived delta");
        let mut state = MockTarget { inner: base };

        // Patching mode to i64::MAX means the decoded diff becomes i64::MAX.
        // Applying i64::MAX to a base of 100 overflows during the final u64 conversion.
        let result = apply_balances_iter(&mut state, archived);
        assert!(result.is_err(), "Should error on overflow during apply");
        let err_str = format!("{}", result.expect_err("test setup"));
        assert!(
            err_str.contains("produces an invalid value"),
            "Error message should mention the invalid resulting value"
        );
    }

    #[test]
    fn test_apply_base_exceeds_i64_range_returns_error() {
        // Base balance is i64::MAX + 1. Even a diff of 0 (no change) wouldn't fail,
        // but applying a diff of 1 requires reading base as i64, which will fail.
        let base = vec![i64::MAX as u64 + 1];
        let target = vec![i64::MAX as u64 + 2]; // diff is 1

        let delta = diff_balances(&base, &target);

        let bytes = rkyv::to_bytes::<rkyv::rancor::Error>(&delta)
            .expect("test setup: failed to serialize delta");
        let archived = rkyv::access::<ArchivedBalancesDiff, rkyv::rancor::Error>(&bytes)
            .expect("test setup: failed to access archived delta");

        let mut state = MockTarget { inner: base };

        let result = apply_balances_iter(&mut state, archived);
        assert!(result.is_err());
        let err_str = format!("{}", result.expect_err("test setup"));
        assert!(
            err_str.contains("exceeds i64 range"),
            "Error message should mention base balance exceeding i64"
        );
    }

    #[test]
    fn test_zigzag_roundtrip() {
        let values = [
            0i64,
            1,
            -1,
            2,
            -2,
            i32::MAX as i64,
            i32::MIN as i64,
            i64::MAX,
            i64::MIN,
        ];
        for &v in &values {
            assert_eq!(zigzag_decode(zigzag_encode(v)), v);
        }
    }

    #[test]
    fn test_write_and_read_varint_roundtrip() {
        let values = [0u64, 1, 127, 128, 255, 16383, 16384, u64::MAX];
        for &v in &values {
            let mut buf = Vec::new();
            write_varint(v, &mut buf);
            let mut cursor = 0;
            let decoded = read_varint(&buf, &mut cursor).expect("valid write");
            assert_eq!(decoded, v);
            assert_eq!(cursor, buf.len());
        }
    }

    #[test]
    fn test_read_varint_zero() {
        let buf = [0u8];
        let mut cursor = 0;
        assert_eq!(read_varint(&buf, &mut cursor).expect("valid zero"), 0);
        assert_eq!(cursor, 1);
    }

    #[test]
    fn test_read_varint_two_bytes() {
        // 128 = 0x80. Encoded as: 0x80 (low 7 bits = 0, continuation), 0x01 (low 7 bits = 1).
        let buf = [0x80, 0x01];
        let mut cursor = 0;
        assert_eq!(read_varint(&buf, &mut cursor).expect("valid 128"), 128);
        assert_eq!(cursor, 2);
    }

    #[test]
    fn test_read_varint_max_u64() {
        // u64::MAX encoded as 10 bytes
        let buf: [u8; 10] = [0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x01];
        let mut cursor = 0;
        assert_eq!(read_varint(&buf, &mut cursor).expect("valid max"), u64::MAX);
        assert_eq!(cursor, 10);
    }

    #[test]
    fn test_read_varint_overflow_u64() {
        // Attempt to encode u64::MAX + 1 (bit 64 set).
        // Same as max, but last byte has payload 2 instead of 1.
        let buf: [u8; 10] = [0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x02];
        let mut cursor = 0;
        let result = read_varint(&buf, &mut cursor);
        assert!(result.is_err());
        assert!(format!("{}", result.expect_err("test setup")).contains("overflows u64"));
    }

    #[test]
    fn test_read_varint_truncated() {
        let buf = [0xFF]; // Indicates continuation, but buffer ends
        let mut cursor = 0;
        let result = read_varint(&buf, &mut cursor);
        assert!(result.is_err());
        assert!(format!("{}", result.expect_err("test setup")).contains("truncated"));
    }

    #[test]
    fn test_read_varint_exceeds_max_length() {
        // 9 bytes of 0xFF, followed by 0x81.
        // 0x81 has a valid payload (1), but the continuation bit is set,
        // forcing the loop to attempt an 11th byte which exceeds u64 max length.
        let buf: [u8; 10] = [0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x81];
        let mut cursor = 0;
        let result = read_varint(&buf, &mut cursor);
        assert!(result.is_err());
        assert!(
            format!("{}", result.expect_err("test setup")).contains("exceeds the maximum length")
        );
    }
}