rust-hdf5 0.2.15

//! B-tree v2 (BT2) chunk index structures for HDF5.
//!
//! Implements the on-disk format for B-tree version 2 used to index chunked
//! datasets with multiple unlimited dimensions.
//!
//! Structures:
//!   - Header (BTHD): metadata about the B-tree
//!   - Internal Node (BTIN): non-leaf nodes with records, child pointers
//!   - Leaf Node (BTLF): leaf nodes containing records
//!
//! Record types:
//!   - Type 10: unfiltered chunks (scaled offsets + chunk address)
//!   - Type 11: filtered chunks (scaled offsets + chunk address + chunk_size + filter_mask)

use crate::format::bytes::{read_le_addr as read_addr, read_le_uint as read_size};
use crate::format::checksum::checksum_metadata;
use crate::format::{FormatContext, FormatError, FormatResult, UNDEF_ADDR};

/// Signature for the B-tree v2 header.
pub const BTHD_SIGNATURE: [u8; 4] = *b"BTHD";
/// Signature for the B-tree v2 internal node.
pub const BTIN_SIGNATURE: [u8; 4] = *b"BTIN";
/// Signature for the B-tree v2 leaf node.
pub const BTLF_SIGNATURE: [u8; 4] = *b"BTLF";

/// B-tree v2 version.
pub const BT2_VERSION: u8 = 0;

/// Record type: unfiltered chunks (non-filtered chunked datasets).
pub const BT2_TYPE_CHUNK_UNFILT: u8 = 10;
/// Record type: filtered chunks (filtered chunked datasets).
pub const BT2_TYPE_CHUNK_FILT: u8 = 11;

/// Bytes used to encode a filtered chunk's compressed size in a v2 B-tree
/// record (libhdf5 layout version 5 uses `sizeof_size`).
pub const BT2_FILT_CHUNK_SIZE_LEN: usize = 8;

/// A chunk record for BT2 type 10 (unfiltered).
///
/// Contains the scaled chunk coordinates and the file address.
#[derive(Debug, Clone, PartialEq)]
pub struct Bt2ChunkRecord {
    /// Scaled chunk coordinates (one per dataset dimension).
    pub scaled_offsets: Vec<u64>,
    /// File address of the chunk data.
    pub chunk_address: u64,
}

/// A filtered chunk record for BT2 type 11.
///
/// Contains scaled chunk coordinates, file address, chunk size, and filter mask.
#[derive(Debug, Clone, PartialEq)]
pub struct Bt2FilteredChunkRecord {
    /// Scaled chunk coordinates (one per dataset dimension).
    pub scaled_offsets: Vec<u64>,
    /// File address of the chunk data.
    pub chunk_address: u64,
    /// Size of the chunk after filtering (compressed size).
    pub chunk_size: u32,
    /// Filter mask (bit i set = skip filter i).
    pub filter_mask: u32,
}

/// B-tree v2 header.
///
/// On-disk layout:
/// ```text
/// "BTHD"(4) + version(1) + type(1)
/// + node_size(u32 LE) + record_size(u16 LE) + depth(u16 LE)
/// + split_percent(u8) + merge_percent(u8)
/// + root_node_addr(sizeof_addr) + num_records_in_root(u16 LE)
/// + total_num_records(sizeof_size)
/// + checksum(4)
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct Bt2Header {
    /// Record type (10=unfilt chunks, 11=filt chunks).
    pub record_type: u8,
    /// Size of each node in bytes.
    pub node_size: u32,
    /// Size of each record in bytes.
    pub record_size: u16,
    /// Depth of the B-tree (0 = root is a leaf).
    pub depth: u16,
    /// Percentage full a node must be before splitting.
    pub split_percent: u8,
    /// Percentage below which a node is merged.
    pub merge_percent: u8,
    /// Address of the root node.
    pub root_node_addr: u64,
    /// Number of records in the root node.
    pub num_records_in_root: u16,
    /// Total number of records in the entire B-tree.
    pub total_num_records: u64,
}

impl Bt2Header {
    /// Create a new B-tree v2 header for unfiltered chunk indexing.
    ///
    /// `ndims` is the number of dataset dimensions.
    pub fn new_for_chunks(ctx: &FormatContext, ndims: usize) -> Self {
        // record_size = ndims * 8 (scaled offsets) + sizeof_addr (chunk address)
        let record_size = (ndims * 8 + ctx.sizeof_addr as usize) as u16;
        Self {
            record_type: BT2_TYPE_CHUNK_UNFILT,
            node_size: 4096,
            record_size,
            depth: 0,
            split_percent: 100,
            merge_percent: 40,
            root_node_addr: UNDEF_ADDR,
            num_records_in_root: 0,
            total_num_records: 0,
        }
    }

    /// Create a new B-tree v2 header for filtered chunk indexing.
    ///
    /// `ndims` is the number of dataset dimensions.
    pub fn new_for_filtered_chunks(ctx: &FormatContext, ndims: usize) -> Self {
        // record_size = ndims * 8 + sizeof_addr + 4 (chunk_size) + 4 (filter_mask)
        let record_size = (ndims * 8 + ctx.sizeof_addr as usize + 4 + 4) as u16;
        Self {
            record_type: BT2_TYPE_CHUNK_FILT,
            node_size: 4096,
            record_size,
            depth: 0,
            split_percent: 100,
            merge_percent: 40,
            root_node_addr: UNDEF_ADDR,
            num_records_in_root: 0,
            total_num_records: 0,
        }
    }

    /// Compute the encoded size (for pre-allocation).
    pub fn encoded_size(&self, ctx: &FormatContext) -> usize {
        let sa = ctx.sizeof_addr as usize;
        let ss = ctx.sizeof_size as usize;
        // signature(4) + version(1) + type(1) + node_size(4) + record_size(2)
        // + depth(2) + split_percent(1) + merge_percent(1)
        // + root_node_addr(sa) + num_records_in_root(2) + total_num_records(ss)
        // + checksum(4)
        4 + 1 + 1 + 4 + 2 + 2 + 1 + 1 + sa + 2 + ss + 4
    }

    pub fn encode(&self, ctx: &FormatContext) -> Vec<u8> {
        let sa = ctx.sizeof_addr as usize;
        let ss = ctx.sizeof_size as usize;
        let size = self.encoded_size(ctx);
        let mut buf = Vec::with_capacity(size);

        buf.extend_from_slice(&BTHD_SIGNATURE);
        buf.push(BT2_VERSION);
        buf.push(self.record_type);
        buf.extend_from_slice(&self.node_size.to_le_bytes());
        buf.extend_from_slice(&self.record_size.to_le_bytes());
        buf.extend_from_slice(&self.depth.to_le_bytes());
        buf.push(self.split_percent);
        buf.push(self.merge_percent);
        buf.extend_from_slice(&self.root_node_addr.to_le_bytes()[..sa]);
        buf.extend_from_slice(&self.num_records_in_root.to_le_bytes());
        buf.extend_from_slice(&self.total_num_records.to_le_bytes()[..ss]);

        let cksum = checksum_metadata(&buf);
        buf.extend_from_slice(&cksum.to_le_bytes());

        debug_assert_eq!(buf.len(), size);
        buf
    }

    pub fn decode(buf: &[u8], ctx: &FormatContext) -> FormatResult<Self> {
        let sa = ctx.sizeof_addr as usize;
        let ss = ctx.sizeof_size as usize;
        let min_size = 4 + 1 + 1 + 4 + 2 + 2 + 1 + 1 + sa + 2 + ss + 4;

        if buf.len() < min_size {
            return Err(FormatError::BufferTooShort {
                needed: min_size,
                available: buf.len(),
            });
        }

        if buf[0..4] != BTHD_SIGNATURE {
            return Err(FormatError::InvalidSignature);
        }

        let version = buf[4];
        if version != BT2_VERSION {
            return Err(FormatError::InvalidVersion(version));
        }

        // Verify checksum
        let data_end = min_size - 4;
        let stored_cksum = u32::from_le_bytes([
            buf[data_end],
            buf[data_end + 1],
            buf[data_end + 2],
            buf[data_end + 3],
        ]);
        let computed_cksum = checksum_metadata(&buf[..data_end]);
        if stored_cksum != computed_cksum {
            return Err(FormatError::ChecksumMismatch {
                expected: stored_cksum,
                computed: computed_cksum,
            });
        }

        let mut pos = 5;
        let record_type = buf[pos];
        pos += 1;
        let node_size = u32::from_le_bytes([buf[pos], buf[pos + 1], buf[pos + 2], buf[pos + 3]]);
        pos += 4;
        let record_size = u16::from_le_bytes([buf[pos], buf[pos + 1]]);
        pos += 2;
        let depth = u16::from_le_bytes([buf[pos], buf[pos + 1]]);
        pos += 2;

        // Validate the geometry-driving fields, which come straight off disk.
        // node_size must hold at least the metadata prefix
        // (H5B2_METADATA_PREFIX_SIZE = 10); record_size must be non-zero; and
        // depth is capped so the node-info geometry cannot overflow and the
        // recursive node walk cannot exhaust the stack. A v2 B-tree whose
        // every node holds the minimum two records still reaches u64 capacity
        // well before depth 64, so any larger depth signals corruption.
        if (node_size as u64) < 10 {
            return Err(FormatError::InvalidData(format!(
                "v2 B-tree node_size {node_size} is smaller than the metadata prefix"
            )));
        }
        if record_size == 0 {
            return Err(FormatError::InvalidData(
                "v2 B-tree record_size must be non-zero".into(),
            ));
        }
        if depth > 64 {
            return Err(FormatError::InvalidData(format!(
                "v2 B-tree depth {depth} is implausibly large"
            )));
        }

        let split_percent = buf[pos];
        pos += 1;
        let merge_percent = buf[pos];
        pos += 1;
        let root_node_addr = read_addr(&buf[pos..], sa);
        pos += sa;
        let num_records_in_root = u16::from_le_bytes([buf[pos], buf[pos + 1]]);
        pos += 2;
        let total_num_records = read_size(&buf[pos..], ss);

        Ok(Self {
            record_type,
            node_size,
            record_size,
            depth,
            split_percent,
            merge_percent,
            root_node_addr,
            num_records_in_root,
            total_num_records,
        })
    }
}

/// B-tree v2 leaf node.
///
/// On-disk layout:
/// ```text
/// "BTLF"(4) + version(1) + type(1)
/// + records(num_records * record_size)
/// + checksum(4)
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct Bt2LeafNode {
    pub record_type: u8,
    /// Raw record data. Each record is `record_size` bytes.
    pub record_data: Vec<u8>,
    /// Number of records in this node.
    pub num_records: u16,
    /// Size of each record.
    pub record_size: u16,
}

impl Bt2LeafNode {
    /// Create a new empty leaf node.
    pub fn new(record_type: u8, record_size: u16) -> Self {
        Self {
            record_type,
            record_data: Vec::new(),
            num_records: 0,
            record_size,
        }
    }

    /// Compute the encoded size.
    pub fn encoded_size(&self) -> usize {
        // signature(4) + version(1) + type(1) + records + checksum(4)
        4 + 1 + 1 + self.record_data.len() + 4
    }

    pub fn encode(&self) -> Vec<u8> {
        let size = self.encoded_size();
        let mut buf = Vec::with_capacity(size);

        buf.extend_from_slice(&BTLF_SIGNATURE);
        buf.push(BT2_VERSION);
        buf.push(self.record_type);
        buf.extend_from_slice(&self.record_data);

        let cksum = checksum_metadata(&buf);
        buf.extend_from_slice(&cksum.to_le_bytes());

        debug_assert_eq!(buf.len(), size);
        buf
    }

    pub fn decode(buf: &[u8], num_records: u16, record_size: u16) -> FormatResult<Self> {
        let records_len = (num_records as usize).saturating_mul(record_size as usize);
        let min_size = records_len.saturating_add(10);

        if buf.len() < min_size {
            return Err(FormatError::BufferTooShort {
                needed: min_size,
                available: buf.len(),
            });
        }

        if buf[0..4] != BTLF_SIGNATURE {
            return Err(FormatError::InvalidSignature);
        }

        let version = buf[4];
        if version != BT2_VERSION {
            return Err(FormatError::InvalidVersion(version));
        }

        // Verify checksum
        let data_end = min_size - 4;
        let stored_cksum = u32::from_le_bytes([
            buf[data_end],
            buf[data_end + 1],
            buf[data_end + 2],
            buf[data_end + 3],
        ]);
        let computed_cksum = checksum_metadata(&buf[..data_end]);
        if stored_cksum != computed_cksum {
            return Err(FormatError::ChecksumMismatch {
                expected: stored_cksum,
                computed: computed_cksum,
            });
        }

        let record_type = buf[5];
        let record_data = buf[6..6 + records_len].to_vec();

        Ok(Self {
            record_type,
            record_data,
            num_records,
            record_size,
        })
    }
}

/// B-tree v2 internal node.
///
/// On-disk layout:
/// ```text
/// "BTIN"(4) + version(1) + type(1)
/// + records(num_records * record_size)
/// + child_node_addrs((num_records+1) * sizeof_addr)
/// + child_nrecords((num_records+1) * nrec_size_bits, packed bytes)
/// + [if depth > 1: child_total_nrecords((num_records+1) * sizeof_size)]
/// + checksum(4)
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct Bt2InternalNode {
    pub record_type: u8,
    /// Raw record data.
    pub record_data: Vec<u8>,
    /// Number of records.
    pub num_records: u16,
    /// Size of each record.
    pub record_size: u16,
    /// Child node addresses.
    pub child_addrs: Vec<u64>,
    /// Number of records in each child.
    pub child_nrecords: Vec<u16>,
    /// Total number of records in each child's subtree (only for depth > 1).
    pub child_total_nrecords: Vec<u64>,
}

impl Bt2InternalNode {
    /// Create a new empty internal node.
    pub fn new(record_type: u8, record_size: u16) -> Self {
        Self {
            record_type,
            record_data: Vec::new(),
            num_records: 0,
            record_size,
            child_addrs: Vec::new(),
            child_nrecords: Vec::new(),
            child_total_nrecords: Vec::new(),
        }
    }

    /// Encoded size of an internal node at `depth` with `nrec` records.
    pub fn encoded_size(
        ctx: &FormatContext,
        depth: u16,
        nrec: u16,
        rrec_size: u16,
        max_nrec_size: u8,
        child_total_size: u8,
    ) -> usize {
        let sa = ctx.sizeof_addr as usize;
        let nchild = nrec as usize + 1;
        let ptr = sa
            + max_nrec_size as usize
            + if depth > 1 {
                child_total_size as usize
            } else {
                0
            };
        4 + 1 + 1 + nrec as usize * rrec_size as usize + nchild * ptr + 4
    }

    /// Encode an internal node (libhdf5 `H5B2__cache_int_serialize` layout):
    /// each child pointer is `address + node_nrec + [total_nrec if depth>1]`,
    /// where the widths come from the B-tree geometry.
    pub fn encode(
        &self,
        ctx: &FormatContext,
        depth: u16,
        max_nrec_size: u8,
        child_total_size: u8,
    ) -> Vec<u8> {
        let sa = ctx.sizeof_addr as usize;
        let nchild = self.num_records as usize + 1;
        let size = Self::encoded_size(
            ctx,
            depth,
            self.num_records,
            self.record_size,
            max_nrec_size,
            child_total_size,
        );
        debug_assert_eq!(self.child_addrs.len(), nchild);
        debug_assert_eq!(self.child_nrecords.len(), nchild);
        debug_assert!(depth <= 1 || self.child_total_nrecords.len() == nchild);
        let mut buf = Vec::with_capacity(size);
        buf.extend_from_slice(&BTIN_SIGNATURE);
        buf.push(BT2_VERSION);
        buf.push(self.record_type);
        buf.extend_from_slice(&self.record_data);
        for i in 0..nchild {
            buf.extend_from_slice(&self.child_addrs[i].to_le_bytes()[..sa]);
            buf.extend_from_slice(
                &(self.child_nrecords[i] as u64).to_le_bytes()[..max_nrec_size as usize],
            );
            if depth > 1 {
                buf.extend_from_slice(
                    &self.child_total_nrecords[i].to_le_bytes()[..child_total_size as usize],
                );
            }
        }
        let cksum = checksum_metadata(&buf);
        buf.extend_from_slice(&cksum.to_le_bytes());
        debug_assert_eq!(buf.len(), size);
        buf
    }

    /// Decode an internal node. `depth` is this node's depth (children are at
    /// `depth - 1`); `nrec` its record count; the size widths come from the
    /// B-tree geometry.
    pub fn decode(
        buf: &[u8],
        ctx: &FormatContext,
        depth: u16,
        nrec: u16,
        rrec_size: u16,
        max_nrec_size: u8,
        child_total_size: u8,
    ) -> FormatResult<Self> {
        let sa = ctx.sizeof_addr as usize;
        let nchild = nrec as usize + 1;
        let records_len = (nrec as usize).saturating_mul(rrec_size as usize);
        let ptr = sa
            + max_nrec_size as usize
            + if depth > 1 {
                child_total_size as usize
            } else {
                0
            };
        let min_size = records_len
            .saturating_add(nchild.saturating_mul(ptr))
            .saturating_add(10);
        if buf.len() < min_size {
            return Err(FormatError::BufferTooShort {
                needed: min_size,
                available: buf.len(),
            });
        }
        if buf[0..4] != BTIN_SIGNATURE {
            return Err(FormatError::InvalidSignature);
        }
        if buf[4] != BT2_VERSION {
            return Err(FormatError::InvalidVersion(buf[4]));
        }
        let data_end = min_size - 4;
        let stored = u32::from_le_bytes([
            buf[data_end],
            buf[data_end + 1],
            buf[data_end + 2],
            buf[data_end + 3],
        ]);
        let computed = checksum_metadata(&buf[..data_end]);
        if stored != computed {
            return Err(FormatError::ChecksumMismatch {
                expected: stored,
                computed,
            });
        }
        let record_type = buf[5];
        let mut pos = 6;
        let record_data = buf[pos..pos + records_len].to_vec();
        pos += records_len;
        let mut child_addrs = Vec::with_capacity(nchild);
        let mut child_nrecords = Vec::with_capacity(nchild);
        let mut child_total_nrecords = Vec::with_capacity(nchild);
        for _ in 0..nchild {
            child_addrs.push(read_addr(&buf[pos..], sa));
            pos += sa;
            child_nrecords.push(read_size(&buf[pos..], max_nrec_size as usize) as u16);
            pos += max_nrec_size as usize;
            if depth > 1 {
                child_total_nrecords.push(read_size(&buf[pos..], child_total_size as usize));
                pos += child_total_size as usize;
            }
        }
        Ok(Self {
            record_type,
            record_data,
            num_records: nrec,
            record_size: rrec_size,
            child_addrs,
            child_nrecords,
            child_total_nrecords,
        })
    }
}

/// Per-depth v2 B-tree node geometry (`H5B2_node_info_t`).
#[derive(Debug, Clone, Copy)]
pub struct Bt2NodeInfo {
    /// Maximum records a node at this depth holds.
    pub max_nrec: u64,
    /// Maximum records in this node and all nodes below it.
    pub cum_max_nrec: u64,
    /// Bytes needed to encode `cum_max_nrec`.
    pub cum_max_nrec_size: u8,
}

/// v2 B-tree node geometry derived from the header parameters, matching
/// libhdf5 `H5B2__hdr_init`.
#[derive(Debug, Clone)]
pub struct Bt2Geometry {
    /// Bytes used to encode a child node's record count.
    pub max_nrec_size: u8,
    /// Node info indexed by depth (`0..=depth`).
    pub node_info: Vec<Bt2NodeInfo>,
}

/// libhdf5 `H5VM_limit_enc_size`: bytes to encode values in `0..=limit`.
fn limit_enc_size(limit: u64) -> u8 {
    let log2 = if limit == 0 {
        0
    } else {
        63 - limit.leading_zeros()
    };
    (log2 / 8 + 1) as u8
}

impl Bt2Geometry {
    /// v2 B-tree prefix size (`H5B2_METADATA_PREFIX_SIZE`): magic + version
    /// + type + checksum.
    const PREFIX: u64 = 10;

    pub fn new(node_size: u32, rrec_size: u16, depth: u16, sizeof_addr: u8) -> Self {
        // Saturating arithmetic throughout: `node_size` and `depth` are
        // validated by `Bt2Header::decode`, but this stays panic-free even
        // if constructed directly with degenerate parameters.
        let rrec = rrec_size.max(1) as u64;
        let leaf_max = (node_size as u64).saturating_sub(Self::PREFIX) / rrec;
        let max_nrec_size = limit_enc_size(leaf_max);
        let mut node_info = vec![Bt2NodeInfo {
            max_nrec: leaf_max,
            cum_max_nrec: leaf_max,
            cum_max_nrec_size: 0,
        }];
        for d in 1..=depth as usize {
            let ptr = sizeof_addr as u64
                + max_nrec_size as u64
                + node_info[d - 1].cum_max_nrec_size as u64;
            let max_nrec = if (node_size as u64) > Self::PREFIX + ptr {
                (node_size as u64 - Self::PREFIX - ptr) / (rrec + ptr)
            } else {
                0
            };
            let cum = max_nrec
                .saturating_add(1)
                .saturating_mul(node_info[d - 1].cum_max_nrec)
                .saturating_add(max_nrec);
            node_info.push(Bt2NodeInfo {
                max_nrec,
                cum_max_nrec: cum,
                cum_max_nrec_size: limit_enc_size(cum),
            });
        }
        Self {
            max_nrec_size,
            node_info,
        }
    }

    /// Width of a child's "total records" field for a node at `depth`
    /// (0 unless `depth > 1`).
    pub fn child_total_size(&self, depth: u16) -> u8 {
        if depth > 1 {
            self.node_info[depth as usize - 1].cum_max_nrec_size
        } else {
            0
        }
    }
}

// ==========================================================================
// In-memory BT2 chunk index (flat approach)
// ==========================================================================

/// In-memory B-tree v2 chunk index.
///
/// Keeps all records in memory. Serializes as a header + leaf node(s) for
/// small trees. For larger trees, internal nodes would be needed but we use
/// the flat approach for simplicity.
#[derive(Debug, Clone)]
pub struct Bt2ChunkIndex {
    /// Number of dataset dimensions.
    pub ndims: usize,
    /// Whether chunks are filtered.
    pub filtered: bool,
    /// Unfiltered chunk records (used when filtered == false).
    pub records: Vec<Bt2ChunkRecord>,
    /// Filtered chunk records (used when filtered == true).
    pub filtered_records: Vec<Bt2FilteredChunkRecord>,
}

impl Bt2ChunkIndex {
    /// Create a new empty B-tree v2 chunk index for unfiltered chunks.
    pub fn new_unfiltered(ndims: usize) -> Self {
        Self {
            ndims,
            filtered: false,
            records: Vec::new(),
            filtered_records: Vec::new(),
        }
    }

    /// Create a new empty B-tree v2 chunk index for filtered chunks.
    pub fn new_filtered(ndims: usize) -> Self {
        Self {
            ndims,
            filtered: true,
            records: Vec::new(),
            filtered_records: Vec::new(),
        }
    }

    /// Insert an unfiltered chunk record.
    pub fn insert(&mut self, scaled_offsets: Vec<u64>, chunk_address: u64) {
        // Check if a record with the same coordinates already exists
        if let Some(existing) = self
            .records
            .iter_mut()
            .find(|r| r.scaled_offsets == scaled_offsets)
        {
            existing.chunk_address = chunk_address;
        } else {
            self.records.push(Bt2ChunkRecord {
                scaled_offsets,
                chunk_address,
            });
        }
    }

    /// Insert a filtered chunk record.
    pub fn insert_filtered(
        &mut self,
        scaled_offsets: Vec<u64>,
        chunk_address: u64,
        chunk_size: u32,
        filter_mask: u32,
    ) {
        if let Some(existing) = self
            .filtered_records
            .iter_mut()
            .find(|r| r.scaled_offsets == scaled_offsets)
        {
            existing.chunk_address = chunk_address;
            existing.chunk_size = chunk_size;
            existing.filter_mask = filter_mask;
        } else {
            self.filtered_records.push(Bt2FilteredChunkRecord {
                scaled_offsets,
                chunk_address,
                chunk_size,
                filter_mask,
            });
        }
    }

    /// Look up a chunk by its scaled coordinates. Returns the record if found.
    pub fn lookup(&self, scaled_offsets: &[u64]) -> Option<&Bt2ChunkRecord> {
        self.records
            .iter()
            .find(|r| r.scaled_offsets == scaled_offsets)
    }

    /// Look up a filtered chunk by its scaled coordinates.
    pub fn lookup_filtered(&self, scaled_offsets: &[u64]) -> Option<&Bt2FilteredChunkRecord> {
        self.filtered_records
            .iter()
            .find(|r| r.scaled_offsets == scaled_offsets)
    }

    /// Iterate all unfiltered records.
    pub fn iter(&self) -> impl Iterator<Item = &Bt2ChunkRecord> {
        self.records.iter()
    }

    /// Iterate all filtered records.
    pub fn iter_filtered(&self) -> impl Iterator<Item = &Bt2FilteredChunkRecord> {
        self.filtered_records.iter()
    }

    /// Total number of records.
    pub fn num_records(&self) -> usize {
        if self.filtered {
            self.filtered_records.len()
        } else {
            self.records.len()
        }
    }

    /// Compute the record size in bytes.
    ///
    /// Filtered records encode the compressed size in a fixed
    /// [`BT2_FILT_CHUNK_SIZE_LEN`]-byte field, matching libhdf5's layout
    /// version 5.
    pub fn record_size(&self, ctx: &FormatContext) -> u16 {
        let sa = ctx.sizeof_addr as usize;
        if self.filtered {
            (sa + BT2_FILT_CHUNK_SIZE_LEN + 4 + self.ndims * 8) as u16
        } else {
            (self.ndims * 8 + sa) as u16
        }
    }

    /// Encode all records as raw bytes.
    fn encode_records(&self, ctx: &FormatContext) -> Vec<u8> {
        let sa = ctx.sizeof_addr as usize;
        let rec_size = self.record_size(ctx) as usize;
        let num = self.num_records();
        let mut buf = Vec::with_capacity(num * rec_size);

        if self.filtered {
            for rec in &self.filtered_records {
                // libhdf5 H5D__bt2_filt_encode: address, chunk size,
                // filter mask, then the scaled offsets.
                buf.extend_from_slice(&rec.chunk_address.to_le_bytes()[..sa]);
                buf.extend_from_slice(
                    &(rec.chunk_size as u64).to_le_bytes()[..BT2_FILT_CHUNK_SIZE_LEN],
                );
                buf.extend_from_slice(&rec.filter_mask.to_le_bytes());
                for &offset in &rec.scaled_offsets {
                    buf.extend_from_slice(&offset.to_le_bytes());
                }
            }
        } else {
            for rec in &self.records {
                // libhdf5 layout: chunk address first, then scaled offsets.
                buf.extend_from_slice(&rec.chunk_address.to_le_bytes()[..sa]);
                for &offset in &rec.scaled_offsets {
                    buf.extend_from_slice(&offset.to_le_bytes());
                }
            }
        }

        buf
    }

    /// Encode the B-tree as a header + leaf node.
    ///
    /// Returns `(header_bytes, leaf_bytes)`.
    pub fn encode(&self, ctx: &FormatContext) -> (Vec<u8>, Vec<u8>) {
        let rec_size = self.record_size(ctx);
        let num = self.num_records() as u16;

        let record_data = self.encode_records(ctx);

        let leaf = Bt2LeafNode {
            record_type: if self.filtered {
                BT2_TYPE_CHUNK_FILT
            } else {
                BT2_TYPE_CHUNK_UNFILT
            },
            record_data,
            num_records: num,
            record_size: rec_size,
        };
        let leaf_encoded = leaf.encode();

        // We'll set root_node_addr to UNDEF_ADDR; the caller sets it to the
        // actual leaf address after allocating.
        let header = Bt2Header {
            record_type: if self.filtered {
                BT2_TYPE_CHUNK_FILT
            } else {
                BT2_TYPE_CHUNK_UNFILT
            },
            node_size: leaf_encoded.len() as u32,
            record_size: rec_size,
            depth: 0,
            split_percent: 100,
            merge_percent: 40,
            root_node_addr: UNDEF_ADDR,
            num_records_in_root: num,
            total_num_records: num as u64,
        };
        let header_encoded = header.encode(ctx);

        (header_encoded, leaf_encoded)
    }

    /// Decode unfiltered records from a leaf node's raw record data.
    pub fn decode_unfiltered_records(
        record_data: &[u8],
        num_records: usize,
        ndims: usize,
        ctx: &FormatContext,
    ) -> FormatResult<Vec<Bt2ChunkRecord>> {
        let sa = ctx.sizeof_addr as usize;
        let rec_size = ndims * 8 + sa;
        if record_data.len() < num_records * rec_size {
            return Err(FormatError::BufferTooShort {
                needed: num_records * rec_size,
                available: record_data.len(),
            });
        }

        let mut records = Vec::with_capacity(num_records);
        let mut pos = 0;
        for _ in 0..num_records {
            // libhdf5 record layout: chunk address first, then scaled offsets.
            let chunk_address = read_addr(&record_data[pos..], sa);
            pos += sa;
            let mut scaled_offsets = Vec::with_capacity(ndims);
            for _ in 0..ndims {
                let offset = u64::from_le_bytes([
                    record_data[pos],
                    record_data[pos + 1],
                    record_data[pos + 2],
                    record_data[pos + 3],
                    record_data[pos + 4],
                    record_data[pos + 5],
                    record_data[pos + 6],
                    record_data[pos + 7],
                ]);
                scaled_offsets.push(offset);
                pos += 8;
            }
            records.push(Bt2ChunkRecord {
                scaled_offsets,
                chunk_address,
            });
        }

        Ok(records)
    }

    /// Decode filtered records from a node's raw record data.
    ///
    /// `record_size` is the on-disk record width from the B-tree header; the
    /// compressed-size field width is derived from it (libhdf5 encodes the
    /// chunk size in `record_size - sizeof_addr - 4 - ndims*8` bytes).
    pub fn decode_filtered_records(
        record_data: &[u8],
        num_records: usize,
        ndims: usize,
        record_size: u16,
        ctx: &FormatContext,
    ) -> FormatResult<Vec<Bt2FilteredChunkRecord>> {
        let sa = ctx.sizeof_addr as usize;
        let rec_size = record_size as usize;
        // chunk-size field width = record - address - filter_mask - offsets.
        let chunk_size_len = rec_size.checked_sub(sa + 4 + ndims * 8).ok_or_else(|| {
            FormatError::InvalidData(format!(
                "filtered v2 B-tree record size {} too small",
                rec_size
            ))
        })?;
        if record_data.len() < num_records * rec_size {
            return Err(FormatError::BufferTooShort {
                needed: num_records * rec_size,
                available: record_data.len(),
            });
        }

        let mut records = Vec::with_capacity(num_records);
        let mut pos = 0;
        for _ in 0..num_records {
            // libhdf5 H5D__bt2_filt_encode: address, chunk size, filter
            // mask, then the scaled offsets.
            let chunk_address = read_addr(&record_data[pos..], sa);
            pos += sa;
            let chunk_size = read_size(&record_data[pos..], chunk_size_len) as u32;
            pos += chunk_size_len;
            let filter_mask = u32::from_le_bytes([
                record_data[pos],
                record_data[pos + 1],
                record_data[pos + 2],
                record_data[pos + 3],
            ]);
            pos += 4;
            let mut scaled_offsets = Vec::with_capacity(ndims);
            for _ in 0..ndims {
                scaled_offsets.push(read_size(&record_data[pos..], 8));
                pos += 8;
            }
            records.push(Bt2FilteredChunkRecord {
                scaled_offsets,
                chunk_address,
                chunk_size,
                filter_mask,
            });
        }

        Ok(records)
    }
}

// ========================================================================= helpers

// ======================================================================= tests

#[cfg(test)]
mod tests {
    use super::*;

    fn ctx8() -> FormatContext {
        FormatContext {
            sizeof_addr: 8,
            sizeof_size: 8,
        }
    }

    fn ctx4() -> FormatContext {
        FormatContext {
            sizeof_addr: 4,
            sizeof_size: 4,
        }
    }

    // ---- Header tests ----

    #[test]
    fn header_roundtrip() {
        let mut hdr = Bt2Header::new_for_chunks(&ctx8(), 2);
        hdr.root_node_addr = 0x3000;
        hdr.num_records_in_root = 5;
        hdr.total_num_records = 5;

        let encoded = hdr.encode(&ctx8());
        assert_eq!(encoded.len(), hdr.encoded_size(&ctx8()));
        assert_eq!(&encoded[..4], b"BTHD");

        let decoded = Bt2Header::decode(&encoded, &ctx8()).unwrap();
        assert_eq!(decoded, hdr);
    }

    /// Re-encode a header buffer's checksum after patching a field, so the
    /// decoder reaches field validation instead of failing on the checksum.
    fn rechecksum(buf: &mut [u8]) {
        let data_end = buf.len() - 4;
        let cksum = checksum_metadata(&buf[..data_end]);
        buf[data_end..].copy_from_slice(&cksum.to_le_bytes());
    }

    #[test]
    fn header_decode_rejects_malformed_geometry_fields() {
        // node_size at offset 6..10, record_size at 10..12, depth at 12..14.
        let make = || {
            let mut hdr = Bt2Header::new_for_chunks(&ctx8(), 2);
            hdr.root_node_addr = 0x3000;
            hdr.encode(&ctx8())
        };

        // node_size smaller than the metadata prefix.
        let mut buf = make();
        buf[6..10].copy_from_slice(&4u32.to_le_bytes());
        rechecksum(&mut buf);
        assert!(matches!(
            Bt2Header::decode(&buf, &ctx8()),
            Err(FormatError::InvalidData(_))
        ));

        // record_size zero.
        let mut buf = make();
        buf[10..12].copy_from_slice(&0u16.to_le_bytes());
        rechecksum(&mut buf);
        assert!(matches!(
            Bt2Header::decode(&buf, &ctx8()),
            Err(FormatError::InvalidData(_))
        ));

        // Implausibly large depth.
        let mut buf = make();
        buf[12..14].copy_from_slice(&5000u16.to_le_bytes());
        rechecksum(&mut buf);
        assert!(matches!(
            Bt2Header::decode(&buf, &ctx8()),
            Err(FormatError::InvalidData(_))
        ));

        // A well-formed header still decodes.
        assert!(Bt2Header::decode(&make(), &ctx8()).is_ok());
    }

    #[test]
    fn bt2_geometry_is_panic_free_on_degenerate_params() {
        // node_size below the prefix, zero record size, oversized depth.
        let _ = Bt2Geometry::new(4, 24, 0, 8);
        let _ = Bt2Geometry::new(0, 0, 0, 8);
        let _ = Bt2Geometry::new(4096, 24, 600, 8);
        let _ = Bt2Geometry::new(u32::MAX, 1, 64, 8);
    }

    #[test]
    fn header_roundtrip_ctx4() {
        let hdr = Bt2Header::new_for_chunks(&ctx4(), 3);
        let encoded = hdr.encode(&ctx4());
        let decoded = Bt2Header::decode(&encoded, &ctx4()).unwrap();
        assert_eq!(decoded, hdr);
    }

    #[test]
    fn header_filtered_roundtrip() {
        let hdr = Bt2Header::new_for_filtered_chunks(&ctx8(), 2);
        assert_eq!(hdr.record_type, BT2_TYPE_CHUNK_FILT);
        // record_size = 2*8 + 8 + 4 + 4 = 32
        assert_eq!(hdr.record_size, 32);

        let encoded = hdr.encode(&ctx8());
        let decoded = Bt2Header::decode(&encoded, &ctx8()).unwrap();
        assert_eq!(decoded, hdr);
    }

    #[test]
    fn header_bad_signature() {
        let hdr = Bt2Header::new_for_chunks(&ctx8(), 2);
        let mut encoded = hdr.encode(&ctx8());
        encoded[0] = b'X';
        let err = Bt2Header::decode(&encoded, &ctx8()).unwrap_err();
        assert!(matches!(err, FormatError::InvalidSignature));
    }

    #[test]
    fn header_checksum_mismatch() {
        let hdr = Bt2Header::new_for_chunks(&ctx8(), 2);
        let mut encoded = hdr.encode(&ctx8());
        encoded[6] ^= 0xFF;
        let err = Bt2Header::decode(&encoded, &ctx8()).unwrap_err();
        assert!(matches!(err, FormatError::ChecksumMismatch { .. }));
    }

    // ---- Leaf node tests ----

    #[test]
    fn leaf_node_roundtrip() {
        let mut leaf = Bt2LeafNode::new(BT2_TYPE_CHUNK_UNFILT, 24);
        // 2 records, each 24 bytes (2 dims * 8 + 8 addr)
        let rec1 = [0u8; 24];
        let mut rec2 = [0u8; 24];
        rec2[0] = 1; // scaled_offset[0] = 1
        leaf.record_data.extend_from_slice(&rec1);
        leaf.record_data.extend_from_slice(&rec2);
        leaf.num_records = 2;

        let encoded = leaf.encode();
        assert_eq!(&encoded[..4], b"BTLF");

        let decoded = Bt2LeafNode::decode(&encoded, 2, 24).unwrap();
        assert_eq!(decoded.record_data, leaf.record_data);
        assert_eq!(decoded.record_type, BT2_TYPE_CHUNK_UNFILT);
    }

    #[test]
    fn leaf_node_empty() {
        let leaf = Bt2LeafNode::new(BT2_TYPE_CHUNK_UNFILT, 24);
        let encoded = leaf.encode();
        let decoded = Bt2LeafNode::decode(&encoded, 0, 24).unwrap();
        assert!(decoded.record_data.is_empty());
    }

    #[test]
    fn leaf_node_bad_checksum() {
        let mut leaf = Bt2LeafNode::new(BT2_TYPE_CHUNK_UNFILT, 8);
        leaf.record_data = vec![0u8; 8];
        leaf.num_records = 1;
        let mut encoded = leaf.encode();
        encoded[6] ^= 0xFF;
        let err = Bt2LeafNode::decode(&encoded, 1, 8).unwrap_err();
        assert!(matches!(err, FormatError::ChecksumMismatch { .. }));
    }

    // ---- Internal node tests ----

    #[test]
    fn internal_node_roundtrip() {
        let rec_size = 24u16;
        let mut node = Bt2InternalNode::new(BT2_TYPE_CHUNK_UNFILT, rec_size);
        node.record_data = vec![0xAA; rec_size as usize]; // 1 record
        node.num_records = 1;
        node.child_addrs = vec![0x1000, 0x2000]; // 2 children
        node.child_nrecords = vec![3, 5];

        // depth 1: children are leaves, no total-nrec field.
        let encoded = node.encode(&ctx8(), 1, 1, 0);
        assert_eq!(&encoded[..4], b"BTIN");

        let decoded = Bt2InternalNode::decode(&encoded, &ctx8(), 1, 1, rec_size, 1, 0).unwrap();
        assert_eq!(decoded.record_data, node.record_data);
        assert_eq!(decoded.child_addrs, node.child_addrs);
        assert_eq!(decoded.child_nrecords, node.child_nrecords);
    }

    #[test]
    fn internal_node_depth2_roundtrip() {
        let rec_size = 16u16;
        let mut node = Bt2InternalNode::new(BT2_TYPE_CHUNK_UNFILT, rec_size);
        node.record_data = vec![0xBB; rec_size as usize * 2]; // 2 records
        node.num_records = 2;
        node.child_addrs = vec![0x1000, 0x2000, 0x3000]; // 3 children
        node.child_nrecords = vec![4, 6, 2];
        node.child_total_nrecords = vec![100, 200, 50];

        // depth 2: children carry a 2-byte total-nrec field.
        let encoded = node.encode(&ctx8(), 2, 1, 2);
        let decoded = Bt2InternalNode::decode(&encoded, &ctx8(), 2, 2, rec_size, 1, 2).unwrap();
        assert_eq!(decoded.child_total_nrecords, vec![100, 200, 50]);
        assert_eq!(decoded.child_nrecords, vec![4, 6, 2]);
    }

    #[test]
    fn bt2_geometry_matches_libhdf5() {
        // node_size 2048, record_size 24, depth 1: leaf holds (2048-10)/24 = 84.
        let g = Bt2Geometry::new(2048, 24, 1, 8);
        assert_eq!(g.node_info[0].max_nrec, 84);
        assert_eq!(g.max_nrec_size, 1);
        assert_eq!(g.child_total_size(1), 0);
    }

    // ---- In-memory index tests ----

    #[test]
    fn chunk_index_insert_and_lookup() {
        let mut idx = Bt2ChunkIndex::new_unfiltered(2);
        idx.insert(vec![0, 0], 0x1000);
        idx.insert(vec![0, 1], 0x2000);
        idx.insert(vec![1, 0], 0x3000);

        assert_eq!(idx.num_records(), 3);

        let r = idx.lookup(&[0, 1]).unwrap();
        assert_eq!(r.chunk_address, 0x2000);

        assert!(idx.lookup(&[2, 2]).is_none());
    }

    #[test]
    fn chunk_index_insert_replaces() {
        let mut idx = Bt2ChunkIndex::new_unfiltered(2);
        idx.insert(vec![0, 0], 0x1000);
        idx.insert(vec![0, 0], 0x2000); // replace
        assert_eq!(idx.num_records(), 1);
        assert_eq!(idx.lookup(&[0, 0]).unwrap().chunk_address, 0x2000);
    }

    #[test]
    fn chunk_index_iterate() {
        let mut idx = Bt2ChunkIndex::new_unfiltered(1);
        for i in 0..5 {
            idx.insert(vec![i], 0x1000 + i * 0x100);
        }
        let addrs: Vec<u64> = idx.iter().map(|r| r.chunk_address).collect();
        assert_eq!(addrs.len(), 5);
    }

    #[test]
    fn chunk_index_encode_decode_roundtrip() {
        let ctx = ctx8();
        let mut idx = Bt2ChunkIndex::new_unfiltered(2);
        idx.insert(vec![0, 0], 0x1000);
        idx.insert(vec![0, 1], 0x2000);
        idx.insert(vec![1, 0], 0x3000);

        let (hdr_bytes, leaf_bytes) = idx.encode(&ctx);

        // Decode header
        let hdr = Bt2Header::decode(&hdr_bytes, &ctx).unwrap();
        assert_eq!(hdr.record_type, BT2_TYPE_CHUNK_UNFILT);
        assert_eq!(hdr.depth, 0);
        assert_eq!(hdr.total_num_records, 3);
        assert_eq!(hdr.num_records_in_root, 3);
        // record_size = 2*8 + 8 = 24
        assert_eq!(hdr.record_size, 24);

        // Decode leaf
        let leaf = Bt2LeafNode::decode(&leaf_bytes, 3, hdr.record_size).unwrap();
        let records =
            Bt2ChunkIndex::decode_unfiltered_records(&leaf.record_data, 3, 2, &ctx).unwrap();

        assert_eq!(records.len(), 3);
        assert_eq!(records[0].scaled_offsets, vec![0, 0]);
        assert_eq!(records[0].chunk_address, 0x1000);
        assert_eq!(records[1].scaled_offsets, vec![0, 1]);
        assert_eq!(records[1].chunk_address, 0x2000);
        assert_eq!(records[2].scaled_offsets, vec![1, 0]);
        assert_eq!(records[2].chunk_address, 0x3000);
    }

    #[test]
    fn unfiltered_record_is_address_first() {
        // libhdf5 (H5D__bt2_unfilt_encode) writes the chunk address before
        // the scaled offsets. Lock that byte order in.
        let ctx = ctx8();
        let mut idx = Bt2ChunkIndex::new_unfiltered(2);
        idx.insert(vec![3, 7], 0xABCD);
        let bytes = idx.encode_records(&ctx);
        // First 8 bytes = address, then two 8-byte scaled offsets.
        assert_eq!(&bytes[0..8], &0xABCDu64.to_le_bytes());
        assert_eq!(&bytes[8..16], &3u64.to_le_bytes());
        assert_eq!(&bytes[16..24], &7u64.to_le_bytes());
    }

    #[test]
    fn filtered_chunk_index_encode_decode_roundtrip() {
        let ctx = ctx8();
        let mut idx = Bt2ChunkIndex::new_filtered(2);
        idx.insert_filtered(vec![0, 0], 0x1000, 512, 0);
        idx.insert_filtered(vec![1, 0], 0x2000, 300, 1);

        let (hdr_bytes, leaf_bytes) = idx.encode(&ctx);

        let hdr = Bt2Header::decode(&hdr_bytes, &ctx).unwrap();
        assert_eq!(hdr.record_type, BT2_TYPE_CHUNK_FILT);
        assert_eq!(hdr.total_num_records, 2);
        // record_size = sizeof_addr(8) + chunk_size_len(8) + filter_mask(4)
        //             + ndims*8(16) = 36
        assert_eq!(hdr.record_size, 36);

        let leaf = Bt2LeafNode::decode(&leaf_bytes, 2, hdr.record_size).unwrap();
        let records =
            Bt2ChunkIndex::decode_filtered_records(&leaf.record_data, 2, 2, hdr.record_size, &ctx)
                .unwrap();

        assert_eq!(records.len(), 2);
        assert_eq!(records[0].chunk_address, 0x1000);
        assert_eq!(records[0].chunk_size, 512);
        assert_eq!(records[0].filter_mask, 0);
        assert_eq!(records[1].chunk_address, 0x2000);
        assert_eq!(records[1].chunk_size, 300);
        assert_eq!(records[1].filter_mask, 1);
    }

    #[test]
    fn chunk_index_ctx4_roundtrip() {
        let ctx = ctx4();
        let mut idx = Bt2ChunkIndex::new_unfiltered(1);
        idx.insert(vec![0], 0x100);
        idx.insert(vec![1], 0x200);

        let (hdr_bytes, leaf_bytes) = idx.encode(&ctx);
        let hdr = Bt2Header::decode(&hdr_bytes, &ctx).unwrap();
        // record_size = 1*8 + 4 = 12
        assert_eq!(hdr.record_size, 12);

        let leaf = Bt2LeafNode::decode(&leaf_bytes, 2, hdr.record_size).unwrap();
        let records =
            Bt2ChunkIndex::decode_unfiltered_records(&leaf.record_data, 2, 1, &ctx).unwrap();
        assert_eq!(records[0].chunk_address, 0x100);
        assert_eq!(records[1].chunk_address, 0x200);
    }

    #[test]
    fn empty_chunk_index() {
        let ctx = ctx8();
        let idx = Bt2ChunkIndex::new_unfiltered(3);
        assert_eq!(idx.num_records(), 0);

        let (hdr_bytes, leaf_bytes) = idx.encode(&ctx);
        let hdr = Bt2Header::decode(&hdr_bytes, &ctx).unwrap();
        assert_eq!(hdr.total_num_records, 0);

        let leaf = Bt2LeafNode::decode(&leaf_bytes, 0, hdr.record_size).unwrap();
        assert!(leaf.record_data.is_empty());
    }
}