gwseq_io/

bytes.rs

//! Little-endian field reads and writes over a byte buffer.
//!
//! Buffer ownership is [`bytes::Bytes`] — refcounted, cheaply sliced, which
//! is what lets a BAM record borrow its decompressed BGZF block and a bbi item
//! its decompressed data block without a copy. What is left is the typed-field
//! half, and that is this file.
//!
//! Every format here is little-endian and its readers refuse a byte-swapped
//! file rather than swapping it, so there is no big-endian path. Reads copy
//! rather than transmute: a record offset in a binary format is under no
//! obligation to be aligned, and `from_le_bytes` over a copied array is both
//! correct and, at these widths, free.
//!
//! Bounds checking is not optional. Every read is checked and returns
//! [`Error::Corrupt`] naming the offset, which is the same check the language
//! emits anyway.

use bytes::Bytes;

use crate::error::{Error, Result};

/// A forward cursor over a buffer, reading little-endian fields.
///
/// Carries the path it came from so that a short read names the file rather
/// than reporting a bare offset the caller has to attribute.
pub struct LeCursor<'a> {
    buf: &'a [u8],
    pos: usize,
    /// Where `buf[0]` sits in the file, so errors report file offsets.
    base: u64,
    path: &'a str,
}

impl<'a> LeCursor<'a> {
    pub fn new(buf: &'a [u8], base: u64, path: &'a str) -> Self {
        Self {
            buf,
            pos: 0,
            base,
            path,
        }
    }

    pub fn position(&self) -> usize {
        self.pos
    }

    pub fn file_offset(&self) -> u64 {
        self.base + self.pos as u64
    }

    pub fn remaining(&self) -> usize {
        self.buf.len() - self.pos
    }

    pub fn is_empty(&self) -> bool {
        self.remaining() == 0
    }

    pub fn seek(&mut self, pos: usize) -> Result<()> {
        if pos > self.buf.len() {
            return Err(self.short(pos - self.buf.len(), "seek"));
        }
        self.pos = pos;
        Ok(())
    }

    pub fn skip(&mut self, n: usize) -> Result<()> {
        // Checked, as `take` is. Every caller today passes a count derived from
        // a `u32`, so on 64-bit this cannot wrap — but a primitive that is safe
        // only because of what its callers happen to be is a primitive waiting
        // for a new caller.
        let pos = self
            .pos
            .checked_add(n)
            .ok_or_else(|| self.short(n, "skip"))?;
        self.seek(pos)
    }

    pub fn take(&mut self, n: usize) -> Result<&'a [u8]> {
        let end = self
            .pos
            .checked_add(n)
            .ok_or_else(|| self.short(n, "take"))?;
        if end > self.buf.len() {
            return Err(self.short(n, "take"));
        }
        let out = &self.buf[self.pos..end];
        self.pos = end;
        Ok(out)
    }

    /// A NUL-terminated string, cursor left after the NUL. BAM read names and
    /// SAM header text are stored this way.
    pub fn take_cstr(&mut self) -> Result<&'a str> {
        let rest = &self.buf[self.pos..];
        let nul = memchr::memchr(0, rest).ok_or_else(|| self.corrupt("unterminated string"))?;
        let out = std::str::from_utf8(&rest[..nul])
            .map_err(|_| self.corrupt("string is not valid UTF-8"))?;
        self.pos += nul + 1;
        Ok(out)
    }

    /// A fixed-width, NUL-padded string, as the bbi chromosome B+ tree stores
    /// its keys. The whole field is consumed whatever the string's length.
    pub fn take_padded_str(&mut self, width: usize) -> Result<&'a str> {
        let field = self.take(width)?;
        let end = memchr::memchr(0, field).unwrap_or(width);
        std::str::from_utf8(&field[..end]).map_err(|_| self.corrupt("string is not valid UTF-8"))
    }

    fn short(&self, wanted: usize, what: &str) -> Error {
        Error::corrupt(
            self.path,
            self.file_offset(),
            format!(
                "{what} of {wanted} bytes with {} left in a {}-byte buffer",
                self.remaining(),
                self.buf.len()
            ),
        )
    }

    fn corrupt(&self, what: &str) -> Error {
        Error::corrupt(self.path, self.file_offset(), what)
    }
}

/// Generates `read_u32`, `read_i64`, `read_f32`, … on [`LeCursor`], plus a
/// `peek_` form that does not advance.
macro_rules! le_readers {
    ($($name:ident, $peek:ident => $t:ty),* $(,)?) => {
        impl LeCursor<'_> {
            $(
                #[inline]
                pub fn $name(&mut self) -> Result<$t> {
                    const N: usize = std::mem::size_of::<$t>();
                    let bytes = self.take(N)?;
                    let mut arr = [0u8; N];
                    arr.copy_from_slice(bytes);
                    Ok(<$t>::from_le_bytes(arr))
                }

                #[inline]
                pub fn $peek(&self) -> Result<$t> {
                    const N: usize = std::mem::size_of::<$t>();
                    if self.remaining() < N {
                        return Err(self.short(N, stringify!($peek)));
                    }
                    let mut arr = [0u8; N];
                    arr.copy_from_slice(&self.buf[self.pos..self.pos + N]);
                    Ok(<$t>::from_le_bytes(arr))
                }
            )*
        }
    };
}

le_readers! {
    read_u8,  peek_u8  => u8,
    read_i8,  peek_i8  => i8,
    read_u16, peek_u16 => u16,
    read_i16, peek_i16 => i16,
    read_u32, peek_u32 => u32,
    read_i32, peek_i32 => i32,
    read_u64, peek_u64 => u64,
    read_i64, peek_i64 => i64,
    read_f32, peek_f32 => f32,
    read_f64, peek_f64 => f64,
}

/// The write half: an append-only buffer with the same field vocabulary, plus
/// the patch-back-in-place the bbi writer needs for counts and offsets whose
/// values are only known once what follows them has been built.
#[derive(Default)]
pub struct LeBuf {
    buf: Vec<u8>,
}

impl LeBuf {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_capacity(n: usize) -> Self {
        Self {
            buf: Vec::with_capacity(n),
        }
    }

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

    pub fn is_empty(&self) -> bool {
        self.buf.is_empty()
    }

    pub fn clear(&mut self) {
        self.buf.clear();
    }

    pub fn as_slice(&self) -> &[u8] {
        &self.buf
    }

    pub fn into_bytes(self) -> Bytes {
        Bytes::from(self.buf)
    }

    pub fn put_bytes(&mut self, data: &[u8]) {
        self.buf.extend_from_slice(data);
    }

    pub fn put_zeros(&mut self, n: usize) {
        self.buf.resize(self.buf.len() + n, 0);
    }

    /// Zero-padded to a fixed width, as the chromosome B+ tree stores keys.
    /// A name longer than `width` is an error rather than a truncation: a
    /// silently truncated key breaks chromosome resolution in a way that only
    /// surfaces at read time.
    pub fn put_padded_str(&mut self, s: &str, width: usize) -> Result<()> {
        if s.len() > width {
            return Err(Error::invalid(format!(
                "chromosome name {s:?} is longer than the {width}-byte key field"
            )));
        }
        self.buf.extend_from_slice(s.as_bytes());
        self.put_zeros(width - s.len());
        Ok(())
    }
}

macro_rules! le_writers {
    ($($name:ident, $set:ident => $t:ty),* $(,)?) => {
        impl LeBuf {
            $(
                #[inline]
                pub fn $name(&mut self, value: $t) {
                    self.buf.extend_from_slice(&value.to_le_bytes());
                }

                /// Overwrite a field already placed. Panics if it does not fit,
                /// which would be a bug in the writer rather than bad input.
                #[inline]
                pub fn $set(&mut self, offset: usize, value: $t) {
                    let bytes = value.to_le_bytes();
                    self.buf[offset..offset + bytes.len()].copy_from_slice(&bytes);
                }
            )*
        }
    };
}

le_writers! {
    put_u8,  set_u8  => u8,
    put_u16, set_u16 => u16,
    put_u32, set_u32 => u32,
    put_i32, set_i32 => i32,
    put_u64, set_u64 => u64,
    put_i64, set_i64 => i64,
    put_f32, set_f32 => f32,
    put_f64, set_f64 => f64,
}