eth-state-diff 0.2.2

//! 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(),
    ))
}