tidaldb/tidal/src/wal/format/batch.rs

use super::super::error::WalError;
use crate::replication::{RegionId, ShardId};

// ── Signal batch format ─────────────────────────────────────────────────────

/// Magic bytes identifying a tidalDB WAL batch frame: "TIDL" in LE byte order.
///
/// Stored as `[0x44, 0x4C, 0x49, 0x54]` which is `0x54494C44` as a u32 LE.
/// This allows `u32::from_le_bytes(magic) == 0x54494C44` to validate.
pub const MAGIC: [u8; 4] = [0x44, 0x4C, 0x49, 0x54];

/// Wire format version 1: original single-node format. Bytes 28-31 are
/// zero-padded reserved space.
pub const FORMAT_VERSION_V1: u8 = 1;

/// Wire format version 2: adds `shard_id` and `region_id`.
///
/// Bytes 28-29 carry `shard_id` (u16 LE) and bytes 30-31 carry `region_id`
/// (u16 LE) for multi-node replication. Backward compatible with v1
/// because v1 always wrote zeros at bytes 28-31, which decode as
/// `ShardId::SINGLE` and `RegionId::SINGLE`.
pub const FORMAT_VERSION_V2: u8 = 2;

/// Current wire format version. All new batches are encoded as v2.
pub const FORMAT_VERSION: u8 = FORMAT_VERSION_V2;

/// Record type discriminant for signal events.
pub const RECORD_TYPE_SIGNAL: u8 = 0x01;

/// Size of the batch header in bytes (one cache line).
pub const HEADER_SIZE: usize = 64;

/// Size of a single event record in bytes.
pub const EVENT_SIZE: usize = 21;

/// Maximum number of events in a single batch.
pub const MAX_EVENTS_PER_BATCH: u16 = 256;

/// Decoded batch header.
///
/// # Wire layout (64 bytes, one cache line)
///
/// ```text
/// Bytes  0- 3: MAGIC [0x44, 0x4C, 0x49, 0x54] ("TIDL")
/// Byte   4:    format version (u8)
/// Byte   5:    flags (u8, reserved)
/// Bytes  6- 7: event_count (u16 LE)
/// Bytes  8-15: first_seq (u64 LE)
/// Bytes 16-23: batch_timestamp (u64 LE)
/// Bytes 24-27: payload_len (u32 LE)
/// Bytes 28-29: shard_id (u16 LE) -- v2; zero in v1 and single-node
/// Bytes 30-31: region_id (u16 LE) -- v2; zero in v1 and single-node
/// Bytes 32-63: BLAKE3 checksum (32 bytes)
/// ```
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BatchHeader {
    pub version: u8,
    pub flags: u8,
    pub event_count: u16,
    pub first_seq: u64,
    pub batch_timestamp: u64,
    pub payload_len: u32,
    pub checksum: [u8; 32],
    /// Shard that produced this batch. `ShardId::SINGLE` (0) in single-node
    /// deployments and when reading v1 batches.
    pub shard_id: ShardId,
    /// Region that produced this batch. `RegionId::SINGLE` (0) in single-node
    /// deployments and when reading v1 batches.
    pub region_id: RegionId,
}

/// A single signal event record in wire format.
#[derive(Debug, Clone, PartialEq)]
pub struct EventRecord {
    pub entity_id: u64,
    pub signal_type: u8,
    pub weight: f32,
    pub timestamp_nanos: u64,
}

impl EventRecord {
    /// Serialize this event into the 21-byte wire format.
    #[must_use]
    pub fn to_bytes(&self) -> [u8; EVENT_SIZE] {
        let mut buf = [0u8; EVENT_SIZE];
        buf[0..8].copy_from_slice(&self.entity_id.to_le_bytes());
        buf[8] = self.signal_type;
        buf[9..13].copy_from_slice(&self.weight.to_le_bytes());
        buf[13..21].copy_from_slice(&self.timestamp_nanos.to_le_bytes());
        buf
    }

    /// Deserialize an event from 21 bytes of wire format.
    ///
    /// # Errors
    ///
    /// Returns `WalError::Corruption` if the slice is not exactly 21 bytes.
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, WalError> {
        if bytes.len() != EVENT_SIZE {
            return Err(WalError::Corruption {
                message: format!(
                    "event record: expected {EVENT_SIZE} bytes, got {}",
                    bytes.len()
                ),
            });
        }
        let entity_id =
            u64::from_le_bytes(bytes[0..8].try_into().map_err(|_| WalError::Corruption {
                message: "event record: invalid entity_id bytes".into(),
            })?);
        let signal_type = bytes[8];
        let weight =
            f32::from_le_bytes(bytes[9..13].try_into().map_err(|_| WalError::Corruption {
                message: "event record: invalid weight bytes".into(),
            })?);
        let timestamp_nanos =
            u64::from_le_bytes(bytes[13..21].try_into().map_err(|_| WalError::Corruption {
                message: "event record: invalid timestamp bytes".into(),
            })?);
        Ok(Self {
            entity_id,
            signal_type,
            weight,
            timestamp_nanos,
        })
    }
}

/// Encode a batch of events into the WAL wire format (single-node convenience).
///
/// Equivalent to `encode_batch_with_shard(events, first_seq, batch_ts,
/// ShardId::SINGLE, RegionId::SINGLE)`.
///
/// # Errors
///
/// Returns `WalError::Corruption` if `events` is empty or exceeds
/// `MAX_EVENTS_PER_BATCH`.
pub fn encode_batch(
    events: &[EventRecord],
    first_seq: u64,
    batch_ts: u64,
) -> Result<Vec<u8>, WalError> {
    encode_batch_with_shard(
        events,
        first_seq,
        batch_ts,
        ShardId::SINGLE,
        RegionId::SINGLE,
    )
}

/// Encode a batch of events with explicit shard and region identity.
///
/// Produces a byte vector containing the 64-byte header followed by
/// tightly packed 21-byte event records. The BLAKE3 checksum covers
/// `header[0..32] || event_bytes`, which includes the shard/region bytes.
///
/// # Errors
///
/// Returns `WalError::Corruption` if `events` is empty or exceeds
/// `MAX_EVENTS_PER_BATCH`.
pub fn encode_batch_with_shard(
    events: &[EventRecord],
    first_seq: u64,
    batch_ts: u64,
    shard_id: ShardId,
    region_id: RegionId,
) -> Result<Vec<u8>, WalError> {
    let event_count = events.len();
    if event_count == 0 || event_count > usize::from(MAX_EVENTS_PER_BATCH) {
        return Err(WalError::Corruption {
            message: format!(
                "batch event count {event_count} out of range [1, {MAX_EVENTS_PER_BATCH}]"
            ),
        });
    }

    let payload_len = event_count * EVENT_SIZE;
    let total_len = HEADER_SIZE + payload_len;
    let mut buf = vec![0u8; total_len];

    // Write header fields [0..32]
    buf[0..4].copy_from_slice(&MAGIC);
    buf[4] = FORMAT_VERSION;
    buf[5] = 0; // flags: reserved
    #[allow(clippy::cast_possible_truncation)]
    let count_u16 = event_count as u16;
    buf[6..8].copy_from_slice(&count_u16.to_le_bytes());
    buf[8..16].copy_from_slice(&first_seq.to_le_bytes());
    buf[16..24].copy_from_slice(&batch_ts.to_le_bytes());
    #[allow(clippy::cast_possible_truncation)]
    let payload_len_u32 = payload_len as u32;
    buf[24..28].copy_from_slice(&payload_len_u32.to_le_bytes());
    // v2: shard_id and region_id at bytes 28-31.
    buf[28..30].copy_from_slice(&shard_id.0.to_le_bytes());
    buf[30..32].copy_from_slice(&region_id.0.to_le_bytes());

    // Write event records starting at offset 64
    for (i, event) in events.iter().enumerate() {
        let offset = HEADER_SIZE + i * EVENT_SIZE;
        buf[offset..offset + EVENT_SIZE].copy_from_slice(&event.to_bytes());
    }

    // Compute BLAKE3 over header[0..32] || event_bytes
    let checksum = compute_checksum(&buf[0..32], &buf[HEADER_SIZE..]);
    buf[32..64].copy_from_slice(checksum.as_bytes());

    Ok(buf)
}

/// Decode a batch from raw bytes.
///
/// Two-phase validation:
/// - Phase 1: magic bytes, version, payload length bounds
/// - Phase 2: BLAKE3 checksum verification
///
/// # Errors
///
/// Returns `WalError::Corruption` on any validation failure.
pub fn decode_batch(bytes: &[u8]) -> Result<(BatchHeader, Vec<EventRecord>), WalError> {
    if bytes.len() < HEADER_SIZE {
        return Err(WalError::Corruption {
            message: format!(
                "batch too short for header: {} bytes, need {HEADER_SIZE}",
                bytes.len()
            ),
        });
    }

    // Phase 1: structural validation
    if bytes[0..4] != MAGIC {
        return Err(WalError::Corruption {
            message: "invalid magic bytes".into(),
        });
    }

    let version = bytes[4];
    if version != FORMAT_VERSION_V1 && version != FORMAT_VERSION_V2 {
        return Err(WalError::Corruption {
            message: format!("unsupported format version: {version}"),
        });
    }

    let flags = bytes[5];
    let event_count =
        u16::from_le_bytes(bytes[6..8].try_into().map_err(|_| WalError::Corruption {
            message: "invalid event_count bytes".into(),
        })?);

    if event_count == 0 || event_count > MAX_EVENTS_PER_BATCH {
        return Err(WalError::Corruption {
            message: format!("event count {event_count} out of range [1, {MAX_EVENTS_PER_BATCH}]"),
        });
    }

    let first_seq =
        u64::from_le_bytes(bytes[8..16].try_into().map_err(|_| WalError::Corruption {
            message: "invalid first_seq bytes".into(),
        })?);

    let batch_timestamp =
        u64::from_le_bytes(bytes[16..24].try_into().map_err(|_| WalError::Corruption {
            message: "invalid batch_timestamp bytes".into(),
        })?);

    let payload_len =
        u32::from_le_bytes(bytes[24..28].try_into().map_err(|_| WalError::Corruption {
            message: "invalid payload_len bytes".into(),
        })?);

    // v2 fields at bytes 28-31. For v1 batches these bytes are always zero,
    // which correctly decodes as ShardId::SINGLE and RegionId::SINGLE.
    let shard_id = ShardId(u16::from_le_bytes(bytes[28..30].try_into().map_err(
        |_| WalError::Corruption {
            message: "invalid shard_id bytes".into(),
        },
    )?));
    let region_id = RegionId(u16::from_le_bytes(bytes[30..32].try_into().map_err(
        |_| WalError::Corruption {
            message: "invalid region_id bytes".into(),
        },
    )?));

    let expected_payload = u32::from(event_count) * EVENT_SIZE as u32;
    if payload_len != expected_payload {
        return Err(WalError::Corruption {
            message: format!(
                "payload_len {payload_len} != event_count {event_count} * {EVENT_SIZE}"
            ),
        });
    }

    let total_len = HEADER_SIZE + payload_len as usize;
    if bytes.len() < total_len {
        return Err(WalError::Corruption {
            message: format!(
                "batch truncated: have {} bytes, need {total_len}",
                bytes.len()
            ),
        });
    }

    // Extract stored checksum
    let mut checksum = [0u8; 32];
    checksum.copy_from_slice(&bytes[32..64]);

    // Phase 2: BLAKE3 verification
    let event_bytes = &bytes[HEADER_SIZE..total_len];
    let computed = compute_checksum(&bytes[0..32], event_bytes);
    if computed.as_bytes() != &checksum {
        return Err(WalError::Corruption {
            message: "BLAKE3 checksum mismatch".into(),
        });
    }

    // Parse event records
    let mut events = Vec::with_capacity(usize::from(event_count));
    for i in 0..usize::from(event_count) {
        let offset = i * EVENT_SIZE;
        let event = EventRecord::from_bytes(&event_bytes[offset..offset + EVENT_SIZE])?;
        events.push(event);
    }

    let header = BatchHeader {
        version,
        flags,
        event_count,
        first_seq,
        batch_timestamp,
        payload_len,
        checksum,
        shard_id,
        region_id,
    };

    Ok((header, events))
}

/// Compute the BLAKE3 checksum for a batch.
///
/// Input: `header_prefix[0..32] || event_bytes`.
/// The hash field at `[32..64]` is NOT part of the hash input.
fn compute_checksum(header_prefix: &[u8], event_bytes: &[u8]) -> blake3::Hash {
    let mut hasher = blake3::Hasher::new();
    hasher.update(header_prefix);
    hasher.update(event_bytes);
    hasher.finalize()
}

/// Compute the per-event content hash used for deduplication.
///
/// Returns the first 128 bits of the BLAKE3 hash of the 21-byte event record.
///
/// # Panics
///
/// Cannot panic. The `expect` is on a `try_into` converting a 16-byte slice
/// (from a 32-byte BLAKE3 hash) into `[u8; 16]`, which is infallible.
#[must_use]
pub fn event_content_hash(event: &EventRecord) -> u128 {
    let bytes = event.to_bytes();
    let hash = blake3::hash(&bytes);
    let hash_bytes: &[u8; 32] = hash.as_bytes();
    u128::from_le_bytes(
        hash_bytes[..16]
            .try_into()
            .expect("BLAKE3 hash is always 32 bytes; first 16 is infallible"),
    )
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
#[path = "batch_tests.rs"]
mod tests;