crabka-protocol 0.3.9

//! Record-batch v2 header types: `RecordBatchHeader` (zerocopy),
//! `Attributes`, `TimestampType`.

use crabka_compression::CompressionType;

/// Timestamp-type bit in the attributes word.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TimestampType {
    CreateTime,
    LogAppendTime,
}

/// Packed batch-level attributes, encoded as a 16-bit big-endian field
/// in the wire header.
///
/// - bits 0-2: compression type (matches `CompressionType::as_attribute_bits`)
/// - bit 3:    timestamp type (0 = `CreateTime`, 1 = `LogAppendTime`)
/// - bit 4:    `is_transactional`
/// - bit 5:    `is_control_batch`
/// - bit 6:    `has_delete_horizon` (KIP-534; `base_timestamp` carries the horizon)
/// - bits 7-15: reserved
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct Attributes(pub i16);

impl Attributes {
    pub const TIMESTAMP_TYPE_BIT: i16 = 1 << 3;
    pub const TRANSACTIONAL_BIT: i16 = 1 << 4;
    pub const CONTROL_BIT: i16 = 1 << 5;
    pub const DELETE_HORIZON_BIT: i16 = 1 << 6; // 0x40

    #[must_use]
    pub fn compression(self) -> CompressionType {
        // The low 3 bits are the codec id. Wider attribute bits are ignored.
        #[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)]
        let byte = (self.0 & 0x07) as u8;
        CompressionType::from_attribute_bits(byte).unwrap_or(CompressionType::None)
    }

    #[must_use]
    pub fn timestamp_type(self) -> TimestampType {
        if self.0 & Self::TIMESTAMP_TYPE_BIT != 0 {
            TimestampType::LogAppendTime
        } else {
            TimestampType::CreateTime
        }
    }

    #[must_use]
    pub fn is_transactional(self) -> bool {
        self.0 & Self::TRANSACTIONAL_BIT != 0
    }

    #[must_use]
    pub fn is_control_batch(self) -> bool {
        self.0 & Self::CONTROL_BIT != 0
    }

    #[must_use]
    pub fn has_delete_horizon(self) -> bool {
        self.0 & Self::DELETE_HORIZON_BIT != 0
    }

    #[must_use]
    pub fn with_compression(self, c: CompressionType) -> Self {
        let cleared = self.0 & !0x07;
        Self(cleared | i16::from(c.as_attribute_bits()))
    }

    #[must_use]
    pub fn with_timestamp_type(self, t: TimestampType) -> Self {
        match t {
            TimestampType::CreateTime => Self(self.0 & !Self::TIMESTAMP_TYPE_BIT),
            TimestampType::LogAppendTime => Self(self.0 | Self::TIMESTAMP_TYPE_BIT),
        }
    }

    #[must_use]
    pub fn with_transactional(self, b: bool) -> Self {
        if b {
            Self(self.0 | Self::TRANSACTIONAL_BIT)
        } else {
            Self(self.0 & !Self::TRANSACTIONAL_BIT)
        }
    }

    #[must_use]
    pub fn with_control(self, b: bool) -> Self {
        if b {
            Self(self.0 | Self::CONTROL_BIT)
        } else {
            Self(self.0 & !Self::CONTROL_BIT)
        }
    }

    #[must_use]
    pub fn with_delete_horizon(self, set: bool) -> Self {
        if set {
            Self(self.0 | Self::DELETE_HORIZON_BIT)
        } else {
            Self(self.0 & !Self::DELETE_HORIZON_BIT)
        }
    }
}

use std::mem::size_of;

use zerocopy::{
    BigEndian, FromBytes, Immutable, KnownLayout, Unaligned,
    byteorder::{I16, I32, I64, U32},
};

/// The fixed 61-byte v2 record-batch header, reinterpreted in place from
/// the wire bytes via `zerocopy`.
#[derive(Debug, Clone, Copy, FromBytes, KnownLayout, Immutable, Unaligned)]
#[repr(C)]
pub struct RecordBatchHeader {
    pub base_offset: I64<BigEndian>,
    pub batch_length: I32<BigEndian>,
    pub partition_leader_epoch: I32<BigEndian>,
    pub magic: i8,
    pub crc: U32<BigEndian>,
    pub attributes: I16<BigEndian>,
    pub last_offset_delta: I32<BigEndian>,
    pub base_timestamp: I64<BigEndian>,
    pub max_timestamp: I64<BigEndian>,
    pub producer_id: I64<BigEndian>,
    pub producer_epoch: I16<BigEndian>,
    pub base_sequence: I32<BigEndian>,
    pub records_count: I32<BigEndian>,
}

/// Size of the v2 record-batch header in bytes.
pub const HEADER_LEN: usize = 61;

// Compile-time assertion that the layout is exactly 61 bytes.
const _: () = assert!(size_of::<RecordBatchHeader>() == HEADER_LEN);

/// Byte offset of the `base_offset` field (i64 BE, 8 bytes) within a v2
/// record-batch header.
pub const BASE_OFFSET_RANGE: std::ops::Range<usize> = 0..8;

/// Byte offset of the `partition_leader_epoch` field (i32 BE, 4 bytes)
/// within a v2 record-batch header.
pub const LEADER_EPOCH_RANGE: std::ops::Range<usize> = 12..16;

/// First byte covered by the v2 batch CRC (`attributes` onward). Bytes
/// `0..21` — `base_offset`, `batch_length`, `partition_leader_epoch`,
/// `magic`, and the `crc` field itself — are **outside** the CRC region.
pub const CRC_COVERAGE_START: usize = 21;

/// Patch `base_offset` (bytes 0..8) and `partition_leader_epoch`
/// (bytes 12..16) in place in a writable copy of the verbatim batch
/// bytes, writing both as big-endian.
///
/// Both fields lie **before** [`CRC_COVERAGE_START`], so this never
/// invalidates the producer's CRC — the broker can stamp the assigned
/// offset and leader epoch without recomputing CRC or touching the body.
///
/// # Panics
///
/// Panics if `buf` is shorter than [`HEADER_LEN`]; callers must validate
/// the batch header (e.g. via borrowed decode) first.
pub fn patch_base_offset_and_leader_epoch(buf: &mut [u8], base_offset: i64, leader_epoch: i32) {
    assert!(
        buf.len() >= HEADER_LEN,
        "patch target shorter than v2 header"
    );
    buf[BASE_OFFSET_RANGE].copy_from_slice(&base_offset.to_be_bytes());
    buf[LEADER_EPOCH_RANGE].copy_from_slice(&leader_epoch.to_be_bytes());
}

#[cfg(test)]
mod tests {
    use assert2::{assert, check};
    use crabka_compression::CompressionType;

    use super::*;

    macro_rules! attr_case {
        ($name:ident, $bits:expr, $codec:expr, $ts:expr, $txn:expr, $ctrl:expr, $horizon:expr) => {
            #[test]
            fn $name() {
                let a = Attributes($bits);
                check!(
                    a.compression() == $codec,
                    "compression mismatch in {}",
                    stringify!($name)
                );
                check!(
                    a.timestamp_type() == $ts,
                    "timestamp_type mismatch in {}",
                    stringify!($name)
                );
                check!(
                    a.is_transactional() == $txn,
                    "is_transactional mismatch in {}",
                    stringify!($name)
                );
                check!(
                    a.is_control_batch() == $ctrl,
                    "is_control_batch mismatch in {}",
                    stringify!($name)
                );
                check!(
                    a.has_delete_horizon() == $horizon,
                    "has_delete_horizon mismatch in {}",
                    stringify!($name)
                );
            }
        };
    }

    attr_case!(
        zero,
        0,
        CompressionType::None,
        TimestampType::CreateTime,
        false,
        false,
        false
    );
    attr_case!(
        gzip_only,
        0b0000_0000_0000_0001,
        CompressionType::Gzip,
        TimestampType::CreateTime,
        false,
        false,
        false
    );
    attr_case!(
        snappy_only,
        0b0000_0000_0000_0010,
        CompressionType::Snappy,
        TimestampType::CreateTime,
        false,
        false,
        false
    );
    attr_case!(
        lz4_only,
        0b0000_0000_0000_0011,
        CompressionType::Lz4,
        TimestampType::CreateTime,
        false,
        false,
        false
    );
    attr_case!(
        zstd_only,
        0b0000_0000_0000_0100,
        CompressionType::Zstd,
        TimestampType::CreateTime,
        false,
        false,
        false
    );
    attr_case!(
        log_append,
        0b0000_0000_0000_1000,
        CompressionType::None,
        TimestampType::LogAppendTime,
        false,
        false,
        false
    );
    attr_case!(
        transactional,
        0b0000_0000_0001_0000,
        CompressionType::None,
        TimestampType::CreateTime,
        true,
        false,
        false
    );
    attr_case!(
        control,
        0b0000_0000_0010_0000,
        CompressionType::None,
        TimestampType::CreateTime,
        false,
        true,
        false
    );
    attr_case!(
        delete_horizon_only,
        0b0000_0000_0100_0000,
        CompressionType::None,
        TimestampType::CreateTime,
        false,
        false,
        true
    );
    attr_case!(
        all_set,
        0b0000_0000_0111_1100,
        CompressionType::Zstd,
        TimestampType::LogAppendTime,
        true,
        true,
        true
    );

    #[test]
    fn builder_round_trip() {
        let a = Attributes::default()
            .with_compression(CompressionType::Snappy)
            .with_timestamp_type(TimestampType::LogAppendTime)
            .with_transactional(true)
            .with_control(false);

        // Snappy = bits 0-2 = 010, LogAppendTime = bit 3, transactional = bit 4.
        assert!(a == Attributes(0b0000_0000_0001_1010));
    }

    #[test]
    fn replacing_compression_clears_old_bits() {
        // Starting with Lz4 (bits 0-2 = 011), switching to Gzip (= 001)
        // must clear bit 1, not OR over it.
        let a = Attributes::default().with_compression(CompressionType::Lz4);
        let b = a.with_compression(CompressionType::Gzip);
        assert!(b.compression() == CompressionType::Gzip);
        assert!(b.0 & 0x07 == 1);
    }

    #[test]
    fn delete_horizon_bit_round_trips() {
        // Default has no delete horizon.
        let base = Attributes::default();
        assert!(!base.has_delete_horizon());

        // Setting it flips exactly bit 6 (mask 0x40).
        let set = base.with_delete_horizon(true);
        assert!(set.has_delete_horizon());
        assert!(set.0 & Attributes::DELETE_HORIZON_BIT == 0x40);

        // Orthogonal to control / transactional: setting those does not touch
        // bit 6, and bit 6 does not touch them.
        let combo = Attributes::default()
            .with_control(true)
            .with_transactional(true)
            .with_delete_horizon(true);
        // control = bit 5, transactional = bit 4, delete horizon = bit 6.
        assert!(combo == Attributes(0b0000_0000_0111_0000));

        // Clearing bit 6 leaves the others intact.
        let cleared = combo.with_delete_horizon(false);
        assert!(cleared == Attributes(0b0000_0000_0011_0000));
    }

    /// Build a sample 61-byte header with known values. Reused across the
    /// header table tests below.
    fn sample_header_bytes() -> [u8; HEADER_LEN] {
        let mut buf = [0u8; HEADER_LEN];
        buf[0..8].copy_from_slice(&100i64.to_be_bytes()); // base_offset
        buf[8..12].copy_from_slice(&77i32.to_be_bytes()); // batch_length
        buf[12..16].copy_from_slice(&1i32.to_be_bytes()); // partition_leader_epoch
        buf[16] = 2; // magic
        buf[17..21].copy_from_slice(&0x1234_5678u32.to_be_bytes()); // crc
        buf[21..23].copy_from_slice(&0i16.to_be_bytes()); // attributes
        buf[23..27].copy_from_slice(&3i32.to_be_bytes()); // last_offset_delta
        buf[27..35].copy_from_slice(&111i64.to_be_bytes()); // base_timestamp
        buf[35..43].copy_from_slice(&222i64.to_be_bytes()); // max_timestamp
        buf[43..51].copy_from_slice(&(-1i64).to_be_bytes()); // producer_id
        buf[51..53].copy_from_slice(&7i16.to_be_bytes()); // producer_epoch
        buf[53..57].copy_from_slice(&(-1i32).to_be_bytes()); // base_sequence
        buf[57..61].copy_from_slice(&4i32.to_be_bytes()); // records_count
        buf
    }

    macro_rules! header_field {
        ($name:ident, $field:ident, $expected:expr) => {
            #[test]
            fn $name() {
                let buf = sample_header_bytes();
                let h = RecordBatchHeader::ref_from_bytes(&buf[..]).expect("header reinterpret");
                assert!(h.$field.get() == $expected);
            }
        };
    }

    header_field!(reads_base_offset, base_offset, 100i64);
    header_field!(reads_batch_length, batch_length, 77i32);
    header_field!(reads_partition_leader_epoch, partition_leader_epoch, 1i32);
    header_field!(reads_crc, crc, 0x1234_5678u32);
    header_field!(reads_last_offset_delta, last_offset_delta, 3i32);
    header_field!(reads_base_timestamp, base_timestamp, 111i64);
    header_field!(reads_max_timestamp, max_timestamp, 222i64);
    header_field!(reads_producer_id, producer_id, -1i64);
    header_field!(reads_producer_epoch, producer_epoch, 7i16);
    header_field!(reads_base_sequence, base_sequence, -1i32);
    header_field!(reads_records_count, records_count, 4i32);

    #[test]
    fn reads_magic_directly() {
        let buf = sample_header_bytes();
        let h = RecordBatchHeader::ref_from_bytes(&buf[..]).unwrap();
        assert!(h.magic == 2);
    }

    #[test]
    fn header_is_exactly_61_bytes() {
        assert!(std::mem::size_of::<RecordBatchHeader>() == HEADER_LEN);
    }

    #[test]
    fn too_short_buffer_errors() {
        let buf = [0u8; HEADER_LEN - 1];
        assert!(RecordBatchHeader::ref_from_bytes(&buf[..]).is_err());
    }

    #[test]
    fn patch_writes_only_pre_crc_fields() {
        let mut buf = sample_header_bytes().to_vec();
        let crc_region_before = buf[CRC_COVERAGE_START..].to_vec();
        let crc_field_before = buf[17..21].to_vec();

        patch_base_offset_and_leader_epoch(&mut buf, 9_001, 42);

        // The two stamped fields changed to the expected big-endian values.
        let h = RecordBatchHeader::ref_from_bytes(&buf[..]).unwrap();
        check!(h.base_offset.get() == 9_001);
        check!(h.partition_leader_epoch.get() == 42);

        // The CRC field itself is untouched (no recompute).
        check!(&buf[17..21] == &crc_field_before[..]);
        // Everything in the CRC-covered region is byte-identical.
        check!(&buf[CRC_COVERAGE_START..] == &crc_region_before[..]);
    }

    #[test]
    fn crc_coverage_constants_match_field_layout() {
        // base_offset and leader_epoch are entirely below the CRC start;
        // attributes (the first CRC-covered field) begins exactly at 21.
        check!(BASE_OFFSET_RANGE.end <= CRC_COVERAGE_START);
        check!(LEADER_EPOCH_RANGE.end <= CRC_COVERAGE_START);
        check!(CRC_COVERAGE_START == 21);
    }
}