gwseq-io 0.2.0

//! The bbi chromosome B+ tree.
//!
//! Keys are fixed-width and NUL-padded, and the width the file declares is
//! carried out of here — not to match against (ids are compared whole; see
//! [`crate::genomic::chr`]) but so that a failed lookup against a file whose
//! key field is narrower than the query can say so.

use crate::bytes::LeCursor;
use crate::error::{Error, Result};
use crate::genomic::ChrMap;
use crate::source::ByteSource;

use super::header::{ChrTreeHeader, CHR_TREE_HEADER_SIZE, CHR_TREE_MAGIC, CHR_TREE_MAGIC_SWAPPED};

/// A corrupt child offset pointing back up the tree would otherwise descend
/// until memory runs out. A real tree over even a scaffolded assembly is two or
/// three levels.
const MAX_DEPTH: usize = 64;

/// Keys are read into a `String` each; a file declaring a huge key size and a
/// huge item count would otherwise allocate before anything checked it.
const MAX_ITEMS: u64 = 1 << 24;

pub fn read(source: &dyn ByteSource, offset: u64) -> Result<(ChrMap, ChrTreeHeader)> {
    let path = source.path();
    let buf = source.read_exact_at(offset, CHR_TREE_HEADER_SIZE as usize)?;
    let mut c = LeCursor::new(&buf, offset, path);

    let magic = c.read_u32()?;
    if magic != CHR_TREE_MAGIC {
        return Err(Error::format(
            path,
            if magic == CHR_TREE_MAGIC_SWAPPED {
                "incompatible endianness (chromosome tree)".to_string()
            } else {
                "invalid chr tree magic number".to_string()
            },
        ));
    }
    let header = ChrTreeHeader {
        block_size: c.read_u32()?,
        key_size: c.read_u32()?,
        val_size: c.read_u32()?,
        item_count: c.read_u64()?,
    };
    if header.key_size == 0 {
        return Err(Error::corrupt(
            path,
            offset,
            "chromosome tree declares a key size of 0",
        ));
    }
    if header.item_count > MAX_ITEMS {
        return Err(Error::corrupt(
            path,
            offset,
            format!(
                "chromosome tree declares {} chromosomes, more than the {MAX_ITEMS} \
                 this reader will allocate for",
                header.item_count
            ),
        ));
    }

    let mut entries = Vec::with_capacity(header.item_count.min(4096) as usize);
    walk(
        source,
        offset + CHR_TREE_HEADER_SIZE,
        header.key_size as usize,
        header.item_count as usize,
        0,
        &mut entries,
    )?;

    Ok((
        ChrMap::from_indexed_entries(entries).with_key_size(header.key_size as usize),
        header,
    ))
}

/// Depth-first over the node at `offset`, appending every leaf item.
fn walk(
    source: &dyn ByteSource,
    offset: u64,
    key_size: usize,
    // The index a leaf item may not reach. Readers turn these indices into a
    // dense table — `BbiReader::chr_names` is `vec![String::new(); max + 1]` —
    // so one leaf carrying `0xFFFF_FFFF` would ask for 4.3e9 `String`s. The
    // tree's own `item_count` is how many chromosomes it holds and therefore
    // how many indices it can name; it is already capped at [`MAX_ITEMS`].
    item_count: usize,
    depth: usize,
    out: &mut Vec<(String, i64, usize)>,
) -> Result<()> {
    if depth > MAX_DEPTH {
        return Err(Error::corrupt(
            source.path(),
            offset,
            format!("chromosome tree is deeper than {MAX_DEPTH} levels"),
        ));
    }

    let head = source.read_exact_at(offset, 4)?;
    let is_leaf = head[0] != 0;
    let count = u16::from_le_bytes([head[2], head[3]]) as usize;
    if count == 0 {
        return Ok(());
    }

    // A leaf item is key + u32 index + u32 size; an internal one is key + u64
    // child offset. Both are key_size + 8.
    let item_size = key_size + 8;
    let body_offset = offset + 4;
    let buf = source.read_exact_at(body_offset, count * item_size)?;
    let mut c = LeCursor::new(&buf, body_offset, source.path());

    if is_leaf {
        for _ in 0..count {
            let id = c.take_padded_str(key_size)?.to_string();
            let index = c.read_u32()? as usize;
            let size = c.read_u32()? as i64;
            if index >= item_count {
                return Err(Error::corrupt(
                    source.path(),
                    offset,
                    format!(
                        "chromosome {id} carries index {index}, and the tree declares only                          {item_count} chromosome(s)"
                    ),
                ));
            }
            out.push((id, size, index));
        }
        return Ok(());
    }

    // Children are collected before descending: `c` borrows `buf`, and the
    // recursion reads more of the file through the same source.
    let mut children = Vec::with_capacity(count);
    for _ in 0..count {
        c.skip(key_size)?;
        children.push(c.read_u64()?);
    }
    for child in children {
        walk(source, child, key_size, item_count, depth + 1, out)?;
    }
    Ok(())
}

/// A chromosome as the writer records it: the name it is stored under, the size
/// it turned out to be, and the index its data records carry.
#[derive(Debug, Clone)]
pub struct WriteEntry {
    pub id: String,
    pub size: u32,
    pub index: u32,
}

/// A leaf value: the index and the size, eight bytes.
const CHR_TREE_VALUE_SIZE: usize = 8;

/// Write the chromosome B+ tree over `entries` (Supp. Tables 8–11).
///
/// `entries` must be **sorted by id**: a B+ tree is searched by key, so the
/// order of the leaves is the order of the names, not of the indices they
/// carry. The format has never required the two to agree, and they generally do
/// not — the indices follow the order the chromosomes were written in, which is
/// what keeps the data sorted by (chromosome, start).
///
/// `key_size` is the longest name rather than the shortest distinguishing
/// prefix UCSC computes. It costs a few bytes per chromosome and spares the
/// whole question of whether a truncated key is still unique.
///
/// Leaf and branch nodes come out the same size here, since a value is eight
/// bytes and so is a child offset, which is why one node size serves both.
pub fn write_tree(
    entries: &[WriteEntry],
    tree_offset: u64,
    block_size: u32,
    sink: &mut dyn FnMut(&[u8]),
) -> Result<()> {
    let item_count = entries.len() as u64;
    let mut key_size = 1usize;
    for (i, entry) in entries.iter().enumerate() {
        key_size = key_size.max(entry.id.len());
        if i > 0 && entries[i - 1].id >= entry.id {
            return Err(Error::invalid(format!(
                "chromosome {} does not sort after {}",
                entry.id,
                entries[i - 1].id
            )));
        }
    }
    // A node holds at most as many children as there are chromosomes, matching
    // UCSC: a tree of three contigs has no use for 256 slots of padding per
    // node, and every file carries at least a root.
    let node_block_size = block_size
        .min(item_count.min(u32::MAX as u64) as u32)
        .max(2);
    let shape = super::rtree::TreeShape::new(item_count, node_block_size)?;
    let levels = shape.levels();
    let block = node_block_size as u64;
    let item_size = key_size + CHR_TREE_VALUE_SIZE;
    let node_size = super::rtree::TREE_NODE_HEADER_SIZE as u64 + block * item_size as u64;

    let mut header = Vec::with_capacity(CHR_TREE_HEADER_SIZE as usize);
    header.extend_from_slice(&CHR_TREE_MAGIC.to_le_bytes());
    header.extend_from_slice(&node_block_size.to_le_bytes());
    header.extend_from_slice(&(key_size as u32).to_le_bytes());
    header.extend_from_slice(&(CHR_TREE_VALUE_SIZE as u32).to_le_bytes());
    header.extend_from_slice(&item_count.to_le_bytes());
    header.extend_from_slice(&[0u8; 8]); // reserved
    debug_assert_eq!(header.len(), CHR_TREE_HEADER_SIZE as usize);
    sink(&header);

    // Every node of this tree is the same size, so a level starts wherever the
    // one above it left off.
    let mut level_offsets = vec![0u64; levels];
    let mut offset = tree_offset + CHR_TREE_HEADER_SIZE;
    for (level, slot) in level_offsets.iter_mut().enumerate() {
        *slot = offset;
        offset += shape.counts[level] * node_size;
    }

    let mut node = Vec::new();
    for level in 0..levels {
        let is_leaf = shape.is_leaf_level(level);
        let child_level_offset = if is_leaf { 0 } else { level_offsets[level + 1] };
        for index in 0..shape.counts[level] {
            let total = if is_leaf {
                item_count
            } else {
                shape.counts[level + 1]
            };
            let count = block.min(total.saturating_sub(index * block));
            node.clear();
            node.reserve(super::rtree::TREE_NODE_HEADER_SIZE + block as usize * item_size);
            node.push(u8::from(is_leaf));
            node.push(0); // reserved
            node.extend_from_slice(&(count as u16).to_le_bytes());
            for i in 0..count {
                if is_leaf {
                    let entry = &entries[(index * block + i) as usize];
                    push_key(&mut node, &entry.id, key_size);
                    node.extend_from_slice(&entry.index.to_le_bytes());
                    node.extend_from_slice(&entry.size.to_le_bytes());
                } else {
                    // The key of a branch is the first key of the subtree it
                    // points at, which is what a search compares against to
                    // choose a child.
                    let child = index * block + i;
                    let first = (child * shape.spans[level + 1]) as usize;
                    push_key(&mut node, &entries[first].id, key_size);
                    node.extend_from_slice(&(child_level_offset + child * node_size).to_le_bytes());
                }
            }
            node.resize(
                super::rtree::TREE_NODE_HEADER_SIZE + block as usize * item_size,
                0,
            );
            sink(&node);
        }
    }
    Ok(())
}

/// A key, NUL-padded to the tree's fixed key width.
fn push_key(out: &mut Vec<u8>, id: &str, key_size: usize) {
    let bytes = id.as_bytes();
    out.extend_from_slice(&bytes[..bytes.len().min(key_size)]);
    out.resize(out.len() + key_size.saturating_sub(bytes.len()), 0);
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::source::testing::MemorySource;

    /// A one-node leaf tree at offset 0.
    fn flat_tree(key_size: usize, items: &[(&str, u32, u32)]) -> Vec<u8> {
        let mut b = Vec::new();
        b.extend_from_slice(&CHR_TREE_MAGIC.to_le_bytes());
        b.extend_from_slice(&256u32.to_le_bytes()); // blockSize
        b.extend_from_slice(&(key_size as u32).to_le_bytes());
        b.extend_from_slice(&8u32.to_le_bytes()); // valSize
        b.extend_from_slice(&(items.len() as u64).to_le_bytes());
        b.extend_from_slice(&0u64.to_le_bytes()); // reserved
        assert_eq!(b.len(), CHR_TREE_HEADER_SIZE as usize);
        b.push(1); // isLeaf
        b.push(0); // reserved
        b.extend_from_slice(&(items.len() as u16).to_le_bytes());
        for (name, index, size) in items {
            let mut key = name.as_bytes().to_vec();
            key.resize(key_size, 0);
            b.extend_from_slice(&key);
            b.extend_from_slice(&index.to_le_bytes());
            b.extend_from_slice(&size.to_le_bytes());
        }
        b
    }

    #[test]
    fn reads_a_flat_tree_and_keeps_the_files_indices() {
        // Deliberately not in index order in the file.
        let bytes = flat_tree(6, &[("chr2", 1, 200), ("chr1", 0, 100), ("chrX", 2, 300)]);
        let source = MemorySource::new(bytes);
        let (map, header) = read(&source, 0).unwrap();
        assert_eq!(header.key_size, 6);
        assert_eq!(header.item_count, 3);
        // Sorted by the file's index, which is the reference order.
        assert_eq!(map.names(), ["chr1", "chr2", "chrX"]);
        assert_eq!(map.resolve("chr2").unwrap().size, 200);
        assert_eq!(map.resolve("chr2").unwrap().index, 1);
        assert_eq!(map.by_index(2).unwrap().id, "chrX");
        assert_eq!(map.declared_key_size(), Some(6));
    }

    #[test]
    fn a_name_exactly_as_long_as_the_key_field_keeps_every_character() {
        let bytes = flat_tree(4, &[("chr1", 0, 100)]);
        let source = MemorySource::new(bytes);
        let (map, _) = read(&source, 0).unwrap();
        assert_eq!(map.names(), ["chr1"]);
    }

    #[test]
    fn walks_an_internal_node_to_its_leaves() {
        // Root at 32 points at two leaves placed after it.
        let key_size = 4usize;
        let item = key_size + 8;
        let mut b = Vec::new();
        b.extend_from_slice(&CHR_TREE_MAGIC.to_le_bytes());
        b.extend_from_slice(&2u32.to_le_bytes());
        b.extend_from_slice(&(key_size as u32).to_le_bytes());
        b.extend_from_slice(&8u32.to_le_bytes());
        b.extend_from_slice(&2u64.to_le_bytes());
        b.extend_from_slice(&0u64.to_le_bytes());
        let root_at = b.len() as u64;
        let root_len = 4 + 2 * item;
        let leaf0_at = root_at + root_len as u64;
        let leaf_len = 4 + item;
        let leaf1_at = leaf0_at + leaf_len as u64;

        b.push(0); // internal
        b.push(0);
        b.extend_from_slice(&2u16.to_le_bytes());
        for (name, child) in [("chr1", leaf0_at), ("chr2", leaf1_at)] {
            let mut key = name.as_bytes().to_vec();
            key.resize(key_size, 0);
            b.extend_from_slice(&key);
            b.extend_from_slice(&child.to_le_bytes());
        }
        for (name, index, size) in [("chr1", 0u32, 111u32), ("chr2", 1, 222)] {
            b.push(1); // leaf
            b.push(0);
            b.extend_from_slice(&1u16.to_le_bytes());
            let mut key = name.as_bytes().to_vec();
            key.resize(key_size, 0);
            b.extend_from_slice(&key);
            b.extend_from_slice(&index.to_le_bytes());
            b.extend_from_slice(&size.to_le_bytes());
        }

        let source = MemorySource::new(b);
        let (map, _) = read(&source, 0).unwrap();
        assert_eq!(map.names(), ["chr1", "chr2"]);
        assert_eq!(map.resolve("chr1").unwrap().size, 111);
        assert_eq!(map.resolve("chr2").unwrap().size, 222);
    }

    #[test]
    fn a_wrong_magic_is_refused_both_ways() {
        let mut bytes = flat_tree(4, &[("chr1", 0, 1)]);
        bytes[..4].copy_from_slice(&CHR_TREE_MAGIC_SWAPPED.to_le_bytes());
        let err = read(&MemorySource::new(bytes.clone()), 0)
            .unwrap_err()
            .to_string();
        assert!(err.contains("incompatible endianness"), "{err}");
        bytes[..4].copy_from_slice(&0xDEAD_BEEFu32.to_le_bytes());
        let err = read(&MemorySource::new(bytes), 0).unwrap_err().to_string();
        assert!(err.contains("invalid chr tree magic"), "{err}");
    }

    #[test]
    fn a_truncated_node_is_corrupt_not_a_panic() {
        let mut bytes = flat_tree(6, &[("chr1", 0, 100), ("chr2", 1, 200)]);
        bytes.truncate(bytes.len() - 5);
        assert!(matches!(
            read(&MemorySource::new(bytes), 0),
            Err(Error::Corrupt { .. })
        ));
    }

    #[test]
    fn an_absurd_item_count_is_refused_before_allocating() {
        let mut bytes = flat_tree(4, &[("chr1", 0, 1)]);
        bytes[16..24].copy_from_slice(&u64::MAX.to_le_bytes());
        let err = read(&MemorySource::new(bytes), 0).unwrap_err().to_string();
        assert!(err.contains("more than the"), "{err}");
    }

    /// The other half of "nothing is allocated from a number a file names":
    /// `BbiReader` turns these indices into `vec![String::new(); max + 1]`, so
    /// a leaf naming 0xFFFF_FFFF asks for 4.3e9 `String`s. The tree says how
    /// many chromosomes it holds, and that is the bound.
    #[test]
    fn an_index_past_the_declared_chromosome_count_is_refused() {
        let bytes = flat_tree(4, &[("chr1", u32::MAX, 100)]);
        let err = read(&MemorySource::new(bytes), 0).unwrap_err().to_string();
        assert!(err.contains("carries index 4294967295"), "{err}");

        // The boundary: two chromosomes may be 0 and 1, not 0 and 2.
        let ok = flat_tree(4, &[("chr1", 0, 100), ("chr2", 1, 200)]);
        assert!(read(&MemorySource::new(ok), 0).is_ok());
        let bad = flat_tree(4, &[("chr1", 0, 100), ("chr2", 2, 200)]);
        let err = read(&MemorySource::new(bad), 0).unwrap_err().to_string();
        assert!(err.contains("carries index 2"), "{err}");
    }

    #[test]
    fn a_zero_key_size_is_refused() {
        let mut bytes = flat_tree(4, &[("chr1", 0, 1)]);
        bytes[8..12].copy_from_slice(&0u32.to_le_bytes());
        let err = read(&MemorySource::new(bytes), 0).unwrap_err().to_string();
        assert!(err.contains("key size of 0"), "{err}");
    }

    #[test]
    fn a_child_offset_pointing_at_its_own_parent_stops_at_the_depth_limit() {
        // Root at 32 points back at itself: without the limit this recurses
        // until the stack goes.
        let key_size = 4usize;
        let mut b = Vec::new();
        b.extend_from_slice(&CHR_TREE_MAGIC.to_le_bytes());
        b.extend_from_slice(&2u32.to_le_bytes());
        b.extend_from_slice(&(key_size as u32).to_le_bytes());
        b.extend_from_slice(&8u32.to_le_bytes());
        b.extend_from_slice(&1u64.to_le_bytes());
        b.extend_from_slice(&0u64.to_le_bytes());
        b.push(0);
        b.push(0);
        b.extend_from_slice(&1u16.to_le_bytes());
        b.extend_from_slice(b"chr1");
        b.extend_from_slice(&32u64.to_le_bytes()); // itself
        let err = read(&MemorySource::new(b), 0).unwrap_err().to_string();
        assert!(err.contains("deeper than 64 levels"), "{err}");
    }
}

#[cfg(test)]
mod write_tests {
    use super::*;
    use crate::source::testing::MemorySource;

    fn entries(names: &[(&str, u32, u32)]) -> Vec<WriteEntry> {
        names
            .iter()
            .map(|(id, size, index)| WriteEntry {
                id: (*id).to_string(),
                size: *size,
                index: *index,
            })
            .collect()
    }

    /// Write a tree at offset 0 and read it straight back. The round trip is
    /// the whole test: the writer's node offsets are computed from the shape
    /// before any node exists, so a mistake in them is a tree that only fails
    /// when something walks it.
    fn round_trip(items: &[(&str, u32, u32)], block_size: u32) -> ChrMap {
        let mut bytes = Vec::new();
        write_tree(&entries(items), 0, block_size, &mut |b| {
            bytes.extend_from_slice(b)
        })
        .unwrap();
        let (map, header) = read(&MemorySource::new(bytes), 0).unwrap();
        assert_eq!(header.item_count, items.len() as u64);
        map
    }

    #[test]
    fn a_flat_tree_reads_back_with_every_chromosome() {
        let map = round_trip(
            &[("chr1", 1000, 0), ("chr2", 2000, 1), ("chrX", 500, 2)],
            256,
        );
        assert_eq!(map.len(), 3);
        assert_eq!(map.resolve("chr2").unwrap().size, 2000);
        assert_eq!(map.resolve("chrX").unwrap().index, 2);
    }

    #[test]
    fn a_deep_tree_reads_back_with_every_chromosome() {
        // block_size 2 over 9 leaves is four levels, so every branch offset
        // this computes has to be right for the walk to reach the leaves.
        let names: Vec<String> = (0..9).map(|i| format!("chr{i}")).collect();
        let items: Vec<(&str, u32, u32)> = names
            .iter()
            .enumerate()
            .map(|(i, n)| (n.as_str(), (i as u32 + 1) * 100, i as u32))
            .collect();
        let map = round_trip(&items, 2);
        assert_eq!(map.len(), 9);
        for (i, name) in names.iter().enumerate() {
            let entry = map.resolve(name).unwrap();
            assert_eq!(entry.index, i, "{name}");
            assert_eq!(entry.size, (i as i64 + 1) * 100, "{name}");
        }
    }

    #[test]
    fn the_leaf_order_is_the_name_order_and_the_indices_need_not_follow() {
        // Sorted by name, numbered by write order, which is the ordinary case:
        // chr10 sorts before chr2 and is written after it.
        let map = round_trip(&[("chr1", 10, 0), ("chr10", 30, 2), ("chr2", 20, 1)], 256);
        assert_eq!(map.resolve("chr10").unwrap().index, 2);
        assert_eq!(map.resolve("chr2").unwrap().index, 1);
        // And the map comes back in the file's *reference* order, which the
        // reader restores from the indices — not in the leaf order the tree
        // stores, which is by name.
        assert_eq!(map.names(), ["chr1", "chr2", "chr10"]);
    }

    #[test]
    fn a_name_shorter_than_the_key_is_padded_and_reads_back_untrimmed() {
        let map = round_trip(&[("chr1", 10, 0), ("scaffold_1234", 20, 1)], 256);
        assert_eq!(map.names(), ["chr1", "scaffold_1234"]);
        assert_eq!(map.resolve("chr1").unwrap().size, 10);
    }

    #[test]
    fn entries_out_of_name_order_are_refused() {
        let err = write_tree(
            &entries(&[("chr2", 1, 0), ("chr1", 1, 1)]),
            0,
            256,
            &mut |_| {},
        )
        .unwrap_err()
        .to_string();
        assert!(err.contains("does not sort after"), "{err}");
        // And so is a duplicate, which `>=` is what catches.
        let err = write_tree(
            &entries(&[("chr1", 1, 0), ("chr1", 1, 1)]),
            0,
            256,
            &mut |_| {},
        )
        .unwrap_err()
        .to_string();
        assert!(err.contains("does not sort after"), "{err}");
    }

    #[test]
    fn a_single_chromosome_still_gets_a_root() {
        let map = round_trip(&[("chr1", 4096, 0)], 256);
        assert_eq!(map.len(), 1);
        assert_eq!(map.resolve("chr1").unwrap().size, 4096);
    }
}