bazaar 3.5.3

//! B+Tree graph index builder.
//!
//! Port of `bzrformats.btree_index.BTreeBuilder`. See that module's docstring
//! for the wire format. This implementation supports building indexes that
//! fit in memory (no spill-to-disk) and produces byte-identical output to the
//! Python original for the common cases: empty indexes and single-leaf-row
//! trees. Multi-row trees are also supported via the propagation logic in
//! `add_key`.
//!
//! The caller can feed `(key, value, references)` tuples via [`BTreeBuilder::add_node`]
//! and then call [`BTreeBuilder::finish`] to get the serialised bytes.

use crate::chunk_writer::ChunkWriter;
use std::collections::BTreeMap;

/// Key type: an ordered sequence of byte segments.
pub type Key = Vec<Vec<u8>>;

/// One in-memory node: `(references, value)`. References have one list per
/// configured reference list.
#[derive(Debug, Clone)]
pub struct Node {
    pub references: Vec<Vec<Key>>,
    pub value: Vec<u8>,
}

#[derive(Debug)]
pub enum Error {
    DuplicateKey(Key),
    BadKey(Key, String),
    BadValue(Vec<u8>),
    BadReference(Key),
    KeyTooBig(Key),
    InconsistentKeyLength,
    InconsistentRefListCount,
}

impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Error::DuplicateKey(k) => write!(f, "duplicate key: {:?}", k),
            Error::BadKey(k, reason) => write!(f, "bad key {:?}: {}", k, reason),
            Error::BadValue(v) => write!(f, "bad value: {:?}", v),
            Error::BadReference(k) => write!(f, "bad reference: {:?}", k),
            Error::KeyTooBig(k) => write!(f, "key does not fit in one node: {:?}", k),
            Error::InconsistentKeyLength => write!(f, "inconsistent key length"),
            Error::InconsistentRefListCount => {
                write!(f, "inconsistent reference_lists count")
            }
        }
    }
}

impl std::error::Error for Error {}

const BT_SIGNATURE: &[u8] = b"B+Tree Graph Index 2\n";
pub const DEFAULT_RESERVED_HEADER_BYTES: usize = 120;
pub const DEFAULT_PAGE_SIZE: usize = 4096;
const LEAF_FLAG: &[u8] = b"type=leaf\n";
const INTERNAL_FLAG: &[u8] = b"type=internal\n";

/// Page layout parameters for [`BTreeBuilder`] and [`write_nodes`].
///
/// The Python tests override `_PAGE_SIZE` and `_RESERVED_HEADER_BYTES` to
/// force smaller trees for coverage; exposing these as parameters lets the
/// Rust port mirror those tests.
#[derive(Debug, Clone, Copy)]
pub struct Layout {
    pub page_size: usize,
    pub reserved_header_bytes: usize,
}

impl Default for Layout {
    fn default() -> Self {
        Self {
            page_size: DEFAULT_PAGE_SIZE,
            reserved_header_bytes: DEFAULT_RESERVED_HEADER_BYTES,
        }
    }
}

pub struct BTreeBuilder {
    reference_lists: usize,
    key_length: usize,
    optimize_for_size: bool,
    layout: Layout,
    nodes: BTreeMap<Key, Node>,
}

impl BTreeBuilder {
    pub fn new(reference_lists: usize, key_elements: usize) -> Self {
        Self::with_layout(reference_lists, key_elements, Layout::default())
    }

    pub fn with_layout(reference_lists: usize, key_elements: usize, layout: Layout) -> Self {
        assert!(key_elements >= 1, "key_elements must be >= 1");
        Self {
            reference_lists,
            key_length: key_elements,
            optimize_for_size: false,
            layout,
            nodes: BTreeMap::new(),
        }
    }

    pub fn set_optimize_for_size(&mut self, v: bool) {
        self.optimize_for_size = v;
    }

    pub fn key_count(&self) -> usize {
        self.nodes.len()
    }

    /// Add a single node.
    pub fn add_node(
        &mut self,
        key: Key,
        value: Vec<u8>,
        references: Vec<Vec<Key>>,
    ) -> Result<(), Error> {
        self.check_key_ref_value(&key, &references, &value)?;
        if self.nodes.contains_key(&key) {
            return Err(Error::DuplicateKey(key));
        }
        self.nodes.insert(key, Node { references, value });
        Ok(())
    }

    fn check_key_ref_value(
        &self,
        key: &Key,
        references: &[Vec<Key>],
        value: &[u8],
    ) -> Result<(), Error> {
        if key.len() != self.key_length {
            return Err(Error::InconsistentKeyLength);
        }
        if references.len() != self.reference_lists {
            return Err(Error::InconsistentRefListCount);
        }
        for segment in key {
            if segment.iter().any(|b| *b == b'\x00' || *b == b'\n') {
                return Err(Error::BadKey(
                    key.clone(),
                    "key segments must not contain \\x00 or \\n".to_string(),
                ));
            }
        }
        for b in value {
            if *b == 0 || *b == b'\n' {
                return Err(Error::BadValue(value.to_vec()));
            }
        }
        for ref_list in references {
            for reference in ref_list {
                if reference.len() != self.key_length {
                    return Err(Error::BadReference(reference.clone()));
                }
            }
        }
        Ok(())
    }

    /// Produce the serialised B+Tree bytes.
    pub fn finish(&self) -> Result<Vec<u8>, Error> {
        let node_iter: Vec<(Key, Node)> = self
            .nodes
            .iter()
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect();
        write_nodes(
            &node_iter,
            self.reference_lists,
            self.key_length,
            self.optimize_for_size,
            self.layout,
        )
    }
}

/// Flatten one `(key, value, references)` node into `(string_key, line_bytes)`.
///
/// Matches `_btree_serializer._flatten_node`:
///
/// * `string_key` is the key segments joined with `\0`.
/// * `line_bytes` is `string_key \0 refs_bytes \0 value \n`, where refs are
///   tab-separated lists of `\r`-separated references whose segments are
///   `\0`-joined.
pub fn flatten_node(
    key: &Key,
    value: &[u8],
    references: &[Vec<Key>],
    reference_lists: bool,
) -> (Vec<u8>, Vec<u8>) {
    let mut string_key = Vec::new();
    for (i, seg) in key.iter().enumerate() {
        if i > 0 {
            string_key.push(0);
        }
        string_key.extend_from_slice(seg);
    }

    let mut refs_bytes = Vec::new();
    if reference_lists {
        for (rl_idx, ref_list) in references.iter().enumerate() {
            if rl_idx > 0 {
                refs_bytes.push(b'\t');
            }
            for (ref_idx, reference) in ref_list.iter().enumerate() {
                if ref_idx > 0 {
                    refs_bytes.push(b'\r');
                }
                for (k_idx, seg) in reference.iter().enumerate() {
                    if k_idx > 0 {
                        refs_bytes.push(0);
                    }
                    refs_bytes.extend_from_slice(seg);
                }
            }
        }
    }

    let mut line = Vec::with_capacity(string_key.len() + refs_bytes.len() + value.len() + 3);
    line.extend_from_slice(&string_key);
    line.push(0);
    line.extend_from_slice(&refs_bytes);
    line.push(0);
    line.extend_from_slice(value);
    line.push(b'\n');
    (string_key, line)
}

struct BuilderRow {
    /// Pages already finished for this row, each exactly `PAGE_SIZE` bytes.
    /// The first PAGE_SIZE bytes of this buffer have the reserved header
    /// bytes padded at the front so the caller can patch them later.
    spool: Vec<u8>,
    /// Current open ChunkWriter for the in-progress node.
    writer: Option<ChunkWriter>,
    /// Number of nodes finished so far.
    nodes: usize,
    /// True for internal (non-leaf) rows; they must always pad.
    is_internal: bool,
}

impl BuilderRow {
    fn new(is_internal: bool) -> Self {
        Self {
            spool: Vec::new(),
            writer: None,
            nodes: 0,
            is_internal,
        }
    }

    fn finish_node(&mut self, pad: bool, layout: Layout) {
        if self.is_internal {
            assert!(pad, "internal rows must be padded");
        }
        let writer = self
            .writer
            .take()
            .expect("finish_node called with no open writer");
        let finished = writer.finish();
        if self.nodes == 0 {
            // Reserve the header bytes at the very start of the first page.
            self.spool
                .extend_from_slice(&vec![0u8; layout.reserved_header_bytes]);
        }
        let mut byte_lines = finished.bytes_list;
        let mut skipped_bytes = 0usize;
        if !pad && finished.nulls_needed > 0 {
            byte_lines.pop();
            skipped_bytes = finished.nulls_needed;
        }
        for b in &byte_lines {
            self.spool.extend_from_slice(b);
        }
        let remainder = (self.spool.len() + skipped_bytes) % layout.page_size;
        assert_eq!(
            remainder,
            0,
            "incorrect node length: spool={}, remainder={}",
            self.spool.len(),
            remainder
        );
        self.nodes += 1;
    }
}

pub fn write_nodes(
    node_iter: &[(Key, Node)],
    reference_lists: usize,
    key_length: usize,
    optimize_for_size: bool,
    layout: Layout,
) -> Result<Vec<u8>, Error> {
    let mut rows: Vec<BuilderRow> = Vec::new();
    let mut key_count = 0usize;
    for (key, node) in node_iter {
        if key_count == 0 {
            rows.push(BuilderRow::new(false));
        }
        key_count += 1;
        let (string_key, line) =
            flatten_node(key, &node.value, &node.references, reference_lists > 0);
        add_key(
            &string_key,
            &line,
            &mut rows,
            optimize_for_size,
            /*allow_optimize=*/ true,
            layout,
        )?;
    }
    // Finish every row that still has an open writer, in reverse so the leaf
    // finishes before its internal rows.
    let rows_len = rows.len();
    for (idx, row) in rows.iter_mut().enumerate().rev() {
        let pad = idx < rows_len - 1 || row.is_internal;
        if row.writer.is_some() {
            row.finish_node(pad, layout);
        }
    }

    // Header lines.
    let mut header = Vec::new();
    header.extend_from_slice(BT_SIGNATURE);
    header.extend_from_slice(format!("node_ref_lists={}\n", reference_lists).as_bytes());
    header.extend_from_slice(format!("key_elements={}\n", key_length).as_bytes());
    header.extend_from_slice(format!("len={}\n", key_count).as_bytes());
    let row_lengths: Vec<usize> = rows.iter().map(|r| r.nodes).collect();
    let row_lengths_str = row_lengths
        .iter()
        .map(|n| n.to_string())
        .collect::<Vec<_>>()
        .join(",");
    header.extend_from_slice(b"row_lengths=");
    header.extend_from_slice(row_lengths_str.as_bytes());
    header.push(b'\n');
    assert!(
        header.len() <= layout.reserved_header_bytes,
        "Could not fit the header in the reserved space: {} > {}",
        header.len(),
        layout.reserved_header_bytes
    );

    let mut result = header;
    let header_len = result.len();
    // Now write each row. For the first row, the first page's header
    // placeholder is overlaid by the header bytes we just wrote, and the
    // remainder of the reserved region is zero-padded. For subsequent rows,
    // the first page is realigned to a page boundary by emitting
    // `reserved_header_bytes` zeros (since `position` resets to 0 after the
    // first row).
    let mut position = header_len;
    for row in &rows {
        if row.spool.is_empty() {
            continue;
        }
        // The first page of this row: copy `spool[reserved..min(page_size, spool.len())]`
        // (skipping the reserved header placeholder).
        let first_page_end = std::cmp::min(layout.page_size, row.spool.len());
        result.extend_from_slice(&row.spool[layout.reserved_header_bytes..first_page_end]);
        if row.spool.len() >= layout.page_size {
            assert!(position <= layout.reserved_header_bytes);
            let pad = layout.reserved_header_bytes - position;
            result.extend_from_slice(&vec![0u8; pad]);
        }
        // Remaining pages of this row, each exactly page_size.
        if row.spool.len() > layout.page_size {
            result.extend_from_slice(&row.spool[layout.page_size..]);
        }
        position = 0;
    }
    Ok(result)
}

fn add_key(
    string_key: &[u8],
    line: &[u8],
    rows: &mut Vec<BuilderRow>,
    optimize_for_size: bool,
    allow_optimize: bool,
    layout: Layout,
) -> Result<(), Error> {
    let mut new_leaf = false;
    // Ensure the leaf (and any internal rows above with no writer) have an
    // open writer.
    if rows.last().unwrap().writer.is_none() {
        new_leaf = true;
        let rows_len = rows.len();
        for pos in 0..rows_len - 1 {
            if rows[pos].writer.is_none() {
                let length = if rows[pos].nodes == 0 {
                    layout.page_size - layout.reserved_header_bytes
                } else {
                    layout.page_size
                };
                let opt = if allow_optimize {
                    optimize_for_size
                } else {
                    false
                };
                let mut writer = ChunkWriter::new(length, 0, opt);
                let _ = writer.write(INTERNAL_FLAG, false);
                let offset_line = format!("offset={}\n", rows[pos + 1].nodes);
                let _ = writer.write(offset_line.as_bytes(), false);
                rows[pos].writer = Some(writer);
            }
        }
        let leaf_idx = rows_len - 1;
        let length = if rows[leaf_idx].nodes == 0 {
            layout.page_size - layout.reserved_header_bytes
        } else {
            layout.page_size
        };
        let mut writer = ChunkWriter::new(length, 0, optimize_for_size);
        let _ = writer.write(LEAF_FLAG, false);
        rows[leaf_idx].writer = Some(writer);
    }
    let overflow = rows
        .last_mut()
        .unwrap()
        .writer
        .as_mut()
        .unwrap()
        .write(line, false);
    if overflow {
        if new_leaf {
            return Err(Error::KeyTooBig(
                string_key.split(|b| *b == 0).map(|s| s.to_vec()).collect(),
            ));
        }
        // The leaf is full; finish it and propagate the divider key upwards.
        // Intermediate leaf pages are padded to the full page size — only
        // the very last leaf page (flushed by the top-level write_nodes
        // loop) is unpadded.
        let leaf_last = rows.len() - 1;
        rows[leaf_last].finish_node(true, layout);
        let mut key_line = string_key.to_vec();
        key_line.push(b'\n');
        let mut new_row_needed = true;
        for pos in (0..rows.len() - 1).rev() {
            let writer = rows[pos].writer.as_mut().unwrap();
            let overflow = writer.write(&key_line, false);
            if overflow {
                rows[pos].finish_node(true, layout);
            } else {
                new_row_needed = false;
                break;
            }
        }
        if new_row_needed {
            // Insert a new root.
            let mut new_row = BuilderRow::new(true);
            let mut writer = ChunkWriter::new(
                layout.page_size - layout.reserved_header_bytes,
                0,
                optimize_for_size,
            );
            let _ = writer.write(INTERNAL_FLAG, false);
            let offset_line = format!("offset={}\n", rows[0].nodes - 1);
            let _ = writer.write(offset_line.as_bytes(), false);
            let _ = writer.write(&key_line, false);
            new_row.writer = Some(writer);
            rows.insert(0, new_row);
        }
        return add_key(
            string_key,
            line,
            rows,
            optimize_for_size,
            allow_optimize,
            layout,
        );
    }
    Ok(())
}

/// Decide where in `backing_indices` a spilled merge should land, and
/// how many leading slots get merged into it.
///
/// Mirrors the power-of-2 strategy implemented in the Python
/// `BTreeBuilder._spill_mem_keys_and_combine`: scan `backing_indices`
/// from slot 0; the first `None` slot is the landing position, and
/// every non-`None` slot before it is included in the merge.
///
/// Returns the landing slot index. Caller obligations after the merge:
///   * extend `backing_indices` with `None` if `slot == len`;
///   * write the new index into `backing_indices[slot]`;
///   * null out `backing_indices[0..slot]`.
///
/// `occupancy[i]` is `true` iff `backing_indices[i]` is `Some(_)`.
pub fn spill_landing_slot(occupancy: &[bool]) -> usize {
    occupancy
        .iter()
        .position(|&filled| !filled)
        .unwrap_or(occupancy.len())
}

#[cfg(test)]
mod tests {
    use super::*;
    use flate2::read::ZlibDecoder;
    use std::io::Read;

    #[test]
    fn empty_1_0_matches_python_header() {
        let builder = BTreeBuilder::new(0, 1);
        let content = builder.finish().unwrap();
        assert_eq!(
            content,
            b"B+Tree Graph Index 2\nnode_ref_lists=0\nkey_elements=1\nlen=0\nrow_lengths=\n"
                .to_vec()
        );
    }

    #[test]
    fn empty_2_1_matches_python_header() {
        let builder = BTreeBuilder::new(1, 2);
        let content = builder.finish().unwrap();
        assert_eq!(
            content,
            b"B+Tree Graph Index 2\nnode_ref_lists=1\nkey_elements=2\nlen=0\nrow_lengths=\n"
                .to_vec()
        );
    }

    fn pos_to_key(pos: usize, lead: &[u8]) -> Vec<u8> {
        let mut out = Vec::new();
        out.extend_from_slice(lead);
        let digit = format!("{}", pos).into_bytes();
        for _ in 0..40 {
            out.extend_from_slice(&digit);
        }
        out
    }

    #[test]
    fn root_leaf_1_0_round_trips_five_keys() {
        // Mirrors test_btree_index.test_root_leaf_1_0 except we check the
        // serialised node content (not the Python tempfile plumbing).
        let mut builder = BTreeBuilder::new(0, 1);
        for i in 0..5 {
            let key = vec![pos_to_key(i, b"")];
            let value = format!("value:{}", i).into_bytes();
            builder.add_node(key, value, vec![]).unwrap();
        }
        let content = builder.finish().unwrap();
        let header =
            b"B+Tree Graph Index 2\nnode_ref_lists=0\nkey_elements=1\nlen=5\nrow_lengths=1\n";
        assert_eq!(&content[..header.len()], header);
        // The compressed leaf follows the header.
        let node_content = &content[header.len()..];
        let mut decoder = ZlibDecoder::new(node_content);
        let mut node_bytes = Vec::new();
        decoder.read_to_end(&mut node_bytes).unwrap();
        // Decompressed payload should have the leaf flag and five entries.
        assert!(node_bytes.starts_with(b"type=leaf\n"));
        for i in 0..5 {
            let digit = format!("{}", i);
            let mut line = Vec::new();
            for _ in 0..40 {
                line.extend_from_slice(digit.as_bytes());
            }
            line.extend_from_slice(format!("\x00\x00value:{}\n", i).as_bytes());
            assert!(
                node_bytes.windows(line.len()).any(|w| w == line.as_slice()),
                "missing line for index {}",
                i
            );
        }
    }

    #[test]
    fn root_leaf_1_0_decompresses_byte_exact() {
        // The expected decompressed leaf contents should match Python's
        // exactly since the input nodes are deterministic.
        let mut builder = BTreeBuilder::new(0, 1);
        for i in 0..5 {
            let key = vec![pos_to_key(i, b"")];
            let value = format!("value:{}", i).into_bytes();
            builder.add_node(key, value, vec![]).unwrap();
        }
        let content = builder.finish().unwrap();
        let header =
            b"B+Tree Graph Index 2\nnode_ref_lists=0\nkey_elements=1\nlen=5\nrow_lengths=1\n";
        let node_content = &content[header.len()..];
        let mut decoder = ZlibDecoder::new(node_content);
        let mut node_bytes = Vec::new();
        decoder.read_to_end(&mut node_bytes).unwrap();
        let expected = concat_expected(&[
            b"type=leaf\n",
            b"0000000000000000000000000000000000000000\x00\x00value:0\n",
            b"1111111111111111111111111111111111111111\x00\x00value:1\n",
            b"2222222222222222222222222222222222222222\x00\x00value:2\n",
            b"3333333333333333333333333333333333333333\x00\x00value:3\n",
            b"4444444444444444444444444444444444444444\x00\x00value:4\n",
        ]);
        assert_eq!(node_bytes, expected);
    }

    fn concat_expected(parts: &[&[u8]]) -> Vec<u8> {
        let mut out = Vec::new();
        for p in parts {
            out.extend_from_slice(p);
        }
        out
    }

    /// Port of the test suite's `make_nodes(count, key_elements,
    /// reference_lists)` helper. Generates `count * key_elements` sample
    /// nodes with the same deterministic key/value/reference layout.
    fn make_nodes(
        count: usize,
        key_elements: usize,
        reference_lists: usize,
    ) -> Vec<(Key, Vec<u8>, Vec<Vec<Key>>)> {
        let mut nodes = Vec::new();
        for prefix_pos in 0..key_elements {
            let prefix: Vec<Vec<u8>> = if key_elements > 1 {
                vec![pos_to_key(prefix_pos, b"")]
            } else {
                vec![]
            };
            for pos in 0..count {
                let mut key = prefix.clone();
                key.push(pos_to_key(pos, b""));
                let value = format!("value:{pos}").into_bytes();
                let mut refs: Vec<Vec<Key>> = Vec::new();
                for list_pos in 0..reference_lists {
                    let mut list: Vec<Key> = Vec::new();
                    for ref_pos in 0..(list_pos + pos % 2) {
                        let mut r = prefix.clone();
                        if pos % 2 == 1 {
                            r.push(pos_to_key(pos - 1, b"ref"));
                        } else {
                            r.push(pos_to_key(ref_pos, b"ref"));
                        }
                        list.push(r);
                    }
                    refs.push(list);
                }
                nodes.push((key, value, refs));
            }
        }
        nodes
    }

    fn build_with_nodes(
        reference_lists: usize,
        key_elements: usize,
        nodes: &[(Key, Vec<u8>, Vec<Vec<Key>>)],
    ) -> Vec<u8> {
        let layout = Layout {
            page_size: 4096,
            reserved_header_bytes: 100,
        };
        let mut builder = BTreeBuilder::with_layout(reference_lists, key_elements, layout);
        for (key, value, refs) in nodes {
            builder
                .add_node(key.clone(), value.clone(), refs.clone())
                .unwrap();
        }
        builder.finish().unwrap()
    }

    fn decompress(data: &[u8]) -> Vec<u8> {
        let mut decoder = ZlibDecoder::new(data);
        let mut out = Vec::new();
        decoder.read_to_end(&mut out).unwrap();
        out
    }

    fn sorted_keys(nodes: &[(Key, Vec<u8>, Vec<Vec<Key>>)]) -> Vec<Key> {
        let mut keys: Vec<Key> = nodes.iter().map(|(k, _, _)| k.clone()).collect();
        keys.sort();
        keys
    }

    // Mirror of test_btree_index.test_root_leaf_2_2: a single-leaf index
    // with two key elements and two reference lists. This is the byte-exact
    // check that reference lists serialise correctly through write_nodes.
    #[test]
    fn root_leaf_2_2_byte_exact() {
        let nodes = make_nodes(5, 2, 2);
        let content = build_with_nodes(2, 2, &nodes);
        assert_eq!(content.len(), 238);
        assert_eq!(
            &content[..74],
            b"B+Tree Graph Index 2\nnode_ref_lists=2\nkey_elements=2\nlen=10\nrow_lengths=1\n"
        );
        let node_bytes = decompress(&content[74..]);
        let expected = concat_expected(&[
            b"type=leaf\n",
            b"0000000000000000000000000000000000000000\x000000000000000000000000000000000000000000\x00\t0000000000000000000000000000000000000000\x00ref0000000000000000000000000000000000000000\x00value:0\n",
            b"0000000000000000000000000000000000000000\x001111111111111111111111111111111111111111\x000000000000000000000000000000000000000000\x00ref0000000000000000000000000000000000000000\t0000000000000000000000000000000000000000\x00ref0000000000000000000000000000000000000000\r0000000000000000000000000000000000000000\x00ref0000000000000000000000000000000000000000\x00value:1\n",
            b"0000000000000000000000000000000000000000\x002222222222222222222222222222222222222222\x00\t0000000000000000000000000000000000000000\x00ref0000000000000000000000000000000000000000\x00value:2\n",
            b"0000000000000000000000000000000000000000\x003333333333333333333333333333333333333333\x000000000000000000000000000000000000000000\x00ref2222222222222222222222222222222222222222\t0000000000000000000000000000000000000000\x00ref2222222222222222222222222222222222222222\r0000000000000000000000000000000000000000\x00ref2222222222222222222222222222222222222222\x00value:3\n",
            b"0000000000000000000000000000000000000000\x004444444444444444444444444444444444444444\x00\t0000000000000000000000000000000000000000\x00ref0000000000000000000000000000000000000000\x00value:4\n",
            b"1111111111111111111111111111111111111111\x000000000000000000000000000000000000000000\x00\t1111111111111111111111111111111111111111\x00ref0000000000000000000000000000000000000000\x00value:0\n",
            b"1111111111111111111111111111111111111111\x001111111111111111111111111111111111111111\x001111111111111111111111111111111111111111\x00ref0000000000000000000000000000000000000000\t1111111111111111111111111111111111111111\x00ref0000000000000000000000000000000000000000\r1111111111111111111111111111111111111111\x00ref0000000000000000000000000000000000000000\x00value:1\n",
            b"1111111111111111111111111111111111111111\x002222222222222222222222222222222222222222\x00\t1111111111111111111111111111111111111111\x00ref0000000000000000000000000000000000000000\x00value:2\n",
            b"1111111111111111111111111111111111111111\x003333333333333333333333333333333333333333\x001111111111111111111111111111111111111111\x00ref2222222222222222222222222222222222222222\t1111111111111111111111111111111111111111\x00ref2222222222222222222222222222222222222222\r1111111111111111111111111111111111111111\x00ref2222222222222222222222222222222222222222\x00value:3\n",
            b"1111111111111111111111111111111111111111\x004444444444444444444444444444444444444444\x00\t1111111111111111111111111111111111111111\x00ref0000000000000000000000000000000000000000\x00value:4\n",
        ]);
        assert_eq!(node_bytes, expected);
    }

    // Mirror of test_btree_index.test_last_page_rounded_1_layer: 10 nodes
    // fit on a single leaf page.
    #[test]
    fn last_page_rounded_1_layer() {
        let nodes = make_nodes(10, 1, 0);
        let content = build_with_nodes(0, 1, &nodes);
        assert_eq!(
            &content[..74],
            b"B+Tree Graph Index 2\nnode_ref_lists=0\nkey_elements=1\nlen=10\nrow_lengths=1\n"
        );
        // The single leaf holds all 10 keys, in sorted order.
        let leaf = decompress(&content[74..]);
        let parsed = crate::btree_index::parse_leaf_lines(&leaf, 1, 0).expect("leaf parses");
        let mut got: Vec<Key> = parsed.into_iter().map(|(k, _, _)| k).collect();
        got.sort();
        assert_eq!(got, sorted_keys(&nodes));
    }

    // Mirror of test_btree_index.test_last_page_not_rounded_2_layer and the
    // leaf-split portion of test_2_leaves_1_0: 400 nodes split across two
    // leaves under a single internal root, 231 keys landing on the first leaf.
    #[test]
    fn last_page_not_rounded_2_layer() {
        let nodes = make_nodes(400, 1, 0);
        let content = build_with_nodes(0, 1, &nodes);
        assert_eq!(
            &content[..77],
            b"B+Tree Graph Index 2\nnode_ref_lists=0\nkey_elements=1\nlen=400\nrow_lengths=1,2\n"
        );
        let sorted = sorted_keys(&nodes);
        // Root internal node points at the split key.
        let root = decompress(&content[77..4096]);
        let expected_root =
            concat_expected(&[b"type=internal\noffset=0\n", &pos_to_key(307, b""), b"\n"]);
        assert_eq!(root, expected_root);
        // First leaf holds the first 231 keys, the second the remainder.
        let leaf1 = decompress(&content[4096..8192]);
        let mut keys1: Vec<Key> = crate::btree_index::parse_leaf_lines(&leaf1, 1, 0)
            .unwrap()
            .into_iter()
            .map(|(k, _, _)| k)
            .collect();
        keys1.sort();
        assert_eq!(keys1, sorted[..231]);
        let leaf2 = decompress(&content[8192..]);
        let mut keys2: Vec<Key> = crate::btree_index::parse_leaf_lines(&leaf2, 1, 0)
            .unwrap()
            .into_iter()
            .map(|(k, _, _)| k)
            .collect();
        keys2.sort();
        assert_eq!(keys2, sorted[231..]);
    }

    #[test]
    fn flatten_node_without_references() {
        let key = vec![b"file-id".to_vec()];
        let value = b"val";
        let (string_key, line) = flatten_node(&key, value, &[], false);
        assert_eq!(string_key, b"file-id");
        assert_eq!(line, b"file-id\x00\x00val\n");
    }

    #[test]
    fn two_leaves_with_reserved_100() {
        // Mirrors test_btree_index.test_2_leaves_1_0: 400 nodes with the
        // Python _RESERVED_HEADER_BYTES override set to 100.
        let layout = Layout {
            page_size: 4096,
            reserved_header_bytes: 100,
        };
        let mut builder = BTreeBuilder::with_layout(0, 1, layout);
        for i in 0..400 {
            let key = vec![pos_to_key(i, b"")];
            let value = format!("value:{}", i).into_bytes();
            builder.add_node(key, value, vec![]).unwrap();
        }
        let content = builder.finish().unwrap();
        // The header should start with the signature and claim 400 keys and
        // some multi-row row_lengths (e.g. "1,2").
        assert!(content.starts_with(
            b"B+Tree Graph Index 2\nnode_ref_lists=0\nkey_elements=1\nlen=400\nrow_lengths="
        ));
    }

    #[test]
    fn flatten_node_with_references() {
        let key = vec![b"f".to_vec(), b"r".to_vec()];
        let value = b"value:0";
        let references = vec![vec![
            vec![b"f".to_vec(), b"p1".to_vec()],
            vec![b"f".to_vec(), b"p2".to_vec()],
        ]];
        let (string_key, line) = flatten_node(&key, value, &references, true);
        assert_eq!(string_key, b"f\x00r");
        assert_eq!(line, b"f\x00r\x00f\x00p1\rf\x00p2\x00value:0\n");
    }

    #[test]
    fn spill_landing_slot_first_spill_lands_at_zero() {
        // No backings yet -> new spill lands at slot 0.
        assert_eq!(spill_landing_slot(&[]), 0);
    }

    #[test]
    fn spill_landing_slot_after_one_spill_lands_at_one() {
        // [Some] -> next spill merges slot 0 with mem, lands at slot 1.
        assert_eq!(spill_landing_slot(&[true]), 1);
    }

    #[test]
    fn spill_landing_slot_skips_filled_to_next_none() {
        // [Some, Some] -> next spill merges both slots with mem, lands at 2.
        assert_eq!(spill_landing_slot(&[true, true]), 2);
        // [None, Some] -> next spill lands at the first None (slot 0); the
        // populated slot 1 is left undisturbed.
        assert_eq!(spill_landing_slot(&[false, true]), 0);
        // [Some, None, Some] -> next spill merges slot 0 with mem, lands at 1.
        assert_eq!(spill_landing_slot(&[true, false, true]), 1);
    }

    #[test]
    fn add_node_rejects_duplicate_key() {
        let mut builder = BTreeBuilder::new(0, 1);
        builder
            .add_node(vec![b"key".to_vec()], b"value".to_vec(), vec![])
            .unwrap();
        match builder.add_node(vec![b"key".to_vec()], b"other".to_vec(), vec![]) {
            Err(Error::DuplicateKey(k)) => assert_eq!(k, vec![b"key".to_vec()]),
            other => panic!("expected DuplicateKey, got {:?}", other),
        }
    }

    #[test]
    fn add_node_rejects_value_with_newline() {
        let mut builder = BTreeBuilder::new(0, 1);
        match builder.add_node(vec![b"key".to_vec()], b"bad\nvalue".to_vec(), vec![]) {
            Err(Error::BadValue(v)) => assert_eq!(v, b"bad\nvalue"),
            other => panic!("expected BadValue, got {:?}", other),
        }
    }

    #[test]
    fn add_node_rejects_key_with_null() {
        let mut builder = BTreeBuilder::new(0, 1);
        match builder.add_node(vec![b"bad\x00key".to_vec()], b"value".to_vec(), vec![]) {
            Err(Error::BadKey(k, _)) => assert_eq!(k, vec![b"bad\x00key".to_vec()]),
            other => panic!("expected BadKey, got {:?}", other),
        }
    }

    #[test]
    fn finish_rejects_key_too_big_for_one_node() {
        // A key whose compressed size exceeds a page cannot be placed; finish
        // must report KeyTooBig. Mirrors test_btree_index.test_key_too_big:
        // the key is incompressible (concatenated distinct numbers), so a
        // simple character repeat would not do.
        let mut big_key = Vec::new();
        for n in 0..4096u32 {
            big_key.extend_from_slice(n.to_string().as_bytes());
        }
        let mut builder = BTreeBuilder::new(0, 1);
        builder
            .add_node(vec![big_key], b"value".to_vec(), vec![])
            .unwrap();
        match builder.finish() {
            Err(Error::KeyTooBig(_)) => {}
            other => panic!("expected KeyTooBig, got {:?}", other.map(|_| "Ok(bytes)")),
        }
    }
}