Add modular data structures, parallel pipelines, and system profiling
Establishes foundational infrastructure across multiple crates by introducing unified persistent bit matrix storage with columnar, packed, and implicit variants, alongside De Bruijn graph node encoding and unitig iteration logic. Adds a macro-driven parallel pipeline scheduler featuring NUMA-aware runners, bounded channels, and memory budgets to enforce concurrency limits. Implements streaming nucleotide parsers with pooled page buffers for FASTA, FASTQ, and Genbank formats, complemented by system resource monitoring, progress tracking, and stage profiling utilities. Collectively, these changes provide the core data models, execution frameworks, and I/O pipelines required for downstream k-mer indexing and analysis workloads.
This commit is contained in:
Eric Coissac
@@ -1,733 +0,0 @@
|
||||
use std::io::{self, Read};
|
||||
use std::mem::ManuallyDrop;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use crate::mimetype::MimeTypeGuesser;
|
||||
use crate::xopen::open_raw;
|
||||
|
||||
pub const MAX_K: usize = 31;
|
||||
const PAGE_SIZE: usize = 65536;
|
||||
// overlap (MAX_K - 1) + page data (PAGE_SIZE) + 1 byte for the end-of-page terminating 0
|
||||
const BUF_SIZE: usize = MAX_K + PAGE_SIZE;
|
||||
|
||||
// ─── OverlapState ─────────────────────────────────────────────────────────────
|
||||
|
||||
pub(crate) struct OverlapState {
|
||||
data: [u8; MAX_K],
|
||||
len: usize,
|
||||
k: usize,
|
||||
}
|
||||
|
||||
impl OverlapState {
|
||||
pub(crate) fn new(k: usize) -> Self {
|
||||
assert!(k > 0 && k <= MAX_K);
|
||||
Self {
|
||||
data: [0u8; MAX_K],
|
||||
len: 0,
|
||||
k,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── NucParser trait ──────────────────────────────────────────────────────────
|
||||
|
||||
// Transforms a raw page into a compacted nucleotide stream in-place.
|
||||
//
|
||||
// Buffer layout on each call:
|
||||
// buf[0..overlap_len()] — overlap bytes copied by write_overlap()
|
||||
// buf[overlap_len()..overlap_len()+n] — raw bytes just read from the source
|
||||
//
|
||||
// Returns the number of output bytes in buf[0..returned].
|
||||
pub(crate) trait NucParser {
|
||||
// required: format-specific
|
||||
fn new(k: usize) -> Self
|
||||
where
|
||||
Self: Sized;
|
||||
fn overlap_state(&self) -> &OverlapState;
|
||||
fn overlap_state_mut(&mut self) -> &mut OverlapState;
|
||||
fn is_in_seq(&self) -> bool;
|
||||
fn parse_inplace(&mut self, buf: &mut [u8], n: usize) -> usize;
|
||||
|
||||
// provided: format-independent overlap management
|
||||
fn overlap_len(&self) -> usize {
|
||||
self.overlap_state().len
|
||||
}
|
||||
|
||||
fn write_overlap(&self, buf: &mut [u8]) {
|
||||
let ol = &self.overlap_state();
|
||||
buf[..ol.len].copy_from_slice(&ol.data[..ol.len]);
|
||||
}
|
||||
|
||||
// Called at end of parse_inplace: saves overlap state and returns adjusted j.
|
||||
// seq_start is the j-position where the last sequence started in this call's output.
|
||||
fn save_overlap(&mut self, buf: &mut [u8], j: usize, seq_start: usize) -> usize {
|
||||
if !self.is_in_seq() {
|
||||
self.overlap_state_mut().len = 0;
|
||||
return j;
|
||||
}
|
||||
let seq_len = j - seq_start;
|
||||
let k = self.overlap_state().k;
|
||||
if seq_len >= k {
|
||||
// Sequence long enough: save last k-1 nucleotides, terminate with 0.
|
||||
let ol = k - 1;
|
||||
self.overlap_state_mut().data[..ol].copy_from_slice(&buf[j - ol..j]);
|
||||
self.overlap_state_mut().len = ol;
|
||||
// SAFETY: j <= total - 1 < BUF_SIZE = buf.len()
|
||||
// (total = overlap_len + n <= (MAX_K-1) + PAGE_SIZE = BUF_SIZE - 1)
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = 0;
|
||||
}
|
||||
j + 1
|
||||
} else if seq_len > 0 {
|
||||
// Short sequence (< k): save whole fragment, strip from output.
|
||||
self.overlap_state_mut().data[..seq_len].copy_from_slice(&buf[seq_start..j]);
|
||||
self.overlap_state_mut().len = seq_len;
|
||||
seq_start
|
||||
} else {
|
||||
self.overlap_state_mut().len = 0;
|
||||
j
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── FASTA parser ─────────────────────────────────────────────────────────────
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum FastaState {
|
||||
OutSeq,
|
||||
InTitle,
|
||||
InSeq,
|
||||
InAmbiguous,
|
||||
}
|
||||
|
||||
pub(crate) struct FastaParser {
|
||||
state: FastaState,
|
||||
overlap: OverlapState,
|
||||
}
|
||||
|
||||
impl NucParser for FastaParser {
|
||||
fn new(k: usize) -> Self {
|
||||
Self {
|
||||
state: FastaState::OutSeq,
|
||||
overlap: OverlapState::new(k),
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn overlap_state(&self) -> &OverlapState {
|
||||
&self.overlap
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn overlap_state_mut(&mut self) -> &mut OverlapState {
|
||||
&mut self.overlap
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn is_in_seq(&self) -> bool {
|
||||
matches!(self.state, FastaState::InSeq)
|
||||
}
|
||||
|
||||
fn parse_inplace(&mut self, buf: &mut [u8], n: usize) -> usize {
|
||||
let total = self.overlap.len + n;
|
||||
let mut i = 0; // read index
|
||||
let mut j = 0; // write index (invariant: j <= i always)
|
||||
// j-position where the current sequence started in this call's output;
|
||||
// meaningful only when state is InSeq.
|
||||
let mut seq_start: usize = 0;
|
||||
|
||||
while i < total {
|
||||
// SAFETY: i < total <= BUF_SIZE = buf.len()
|
||||
let byte = unsafe { *buf.get_unchecked(i) };
|
||||
|
||||
match self.state {
|
||||
FastaState::OutSeq => {
|
||||
if byte == b'>' {
|
||||
self.state = FastaState::InTitle;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
FastaState::InTitle => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
self.state = FastaState::InSeq;
|
||||
seq_start = j;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
FastaState::InSeq => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
let nuc = byte & 0xDF; // to uppercase
|
||||
if nuc == b'A' || nuc == b'C' || nuc == b'G' || nuc == b'T' {
|
||||
// SAFETY: j <= i < total <= BUF_SIZE = buf.len()
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = nuc;
|
||||
}
|
||||
j += 1;
|
||||
i += 1;
|
||||
} else if byte == b'>' {
|
||||
if j > seq_start {
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = 0;
|
||||
}
|
||||
j += 1;
|
||||
}
|
||||
self.state = FastaState::InTitle;
|
||||
i += 1;
|
||||
} else {
|
||||
// first ambiguous base: end current sequence if non-empty
|
||||
if j > seq_start {
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = 0;
|
||||
}
|
||||
j += 1;
|
||||
}
|
||||
self.state = FastaState::InAmbiguous;
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
FastaState::InAmbiguous => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
if byte == b'>' {
|
||||
self.state = FastaState::InTitle;
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
let nuc = byte & 0xDF;
|
||||
if nuc == b'A' || nuc == b'C' || nuc == b'G' || nuc == b'T' {
|
||||
seq_start = j;
|
||||
// SAFETY: j <= i < total <= BUF_SIZE = buf.len()
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = nuc;
|
||||
}
|
||||
j += 1;
|
||||
self.state = FastaState::InSeq;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
self.save_overlap(buf, j, seq_start)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── FASTQ parser ─────────────────────────────────────────────────────────────
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum FastqState {
|
||||
OutSeq,
|
||||
InTitle,
|
||||
InSeq,
|
||||
InAmbiguous,
|
||||
InQualTitle,
|
||||
InQual,
|
||||
}
|
||||
|
||||
pub(crate) struct FastqParser {
|
||||
state: FastqState,
|
||||
overlap: OverlapState,
|
||||
}
|
||||
|
||||
impl NucParser for FastqParser {
|
||||
fn new(k: usize) -> Self {
|
||||
Self {
|
||||
state: FastqState::OutSeq,
|
||||
overlap: OverlapState::new(k),
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn overlap_state(&self) -> &OverlapState {
|
||||
&self.overlap
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn overlap_state_mut(&mut self) -> &mut OverlapState {
|
||||
&mut self.overlap
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn is_in_seq(&self) -> bool {
|
||||
matches!(self.state, FastqState::InSeq)
|
||||
}
|
||||
|
||||
fn parse_inplace(&mut self, buf: &mut [u8], n: usize) -> usize {
|
||||
let total = self.overlap.len + n;
|
||||
let mut i = 0;
|
||||
let mut j = 0;
|
||||
let mut seq_start: usize = 0;
|
||||
|
||||
while i < total {
|
||||
// SAFETY: i < total <= BUF_SIZE = buf.len()
|
||||
let byte = unsafe { *buf.get_unchecked(i) };
|
||||
|
||||
match self.state {
|
||||
FastqState::OutSeq => {
|
||||
if byte == b'@' {
|
||||
self.state = FastqState::InTitle;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
FastqState::InTitle => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
self.state = FastqState::InSeq;
|
||||
seq_start = j;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
FastqState::InSeq => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
if j > seq_start {
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = 0;
|
||||
}
|
||||
j += 1;
|
||||
}
|
||||
self.state = FastqState::InQualTitle;
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
let nuc = byte & 0xDF;
|
||||
if nuc == b'A' || nuc == b'C' || nuc == b'G' || nuc == b'T' {
|
||||
// SAFETY: j <= i < total <= BUF_SIZE = buf.len()
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = nuc;
|
||||
}
|
||||
j += 1;
|
||||
} else {
|
||||
if j > seq_start {
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = 0;
|
||||
}
|
||||
j += 1;
|
||||
}
|
||||
self.state = FastqState::InAmbiguous;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
FastqState::InAmbiguous => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
self.state = FastqState::InQualTitle;
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
let nuc = byte & 0xDF;
|
||||
if nuc == b'A' || nuc == b'C' || nuc == b'G' || nuc == b'T' {
|
||||
seq_start = j;
|
||||
// SAFETY: j <= i < total <= BUF_SIZE = buf.len()
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = nuc;
|
||||
}
|
||||
j += 1;
|
||||
self.state = FastqState::InSeq;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
FastqState::InQualTitle => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
self.state = FastqState::InQual;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
FastqState::InQual => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
self.state = FastqState::OutSeq;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
self.save_overlap(buf, j, seq_start)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── GenBank parser ───────────────────────────────────────────────────────────
|
||||
|
||||
const ORIGIN_TAIL: &[u8] = b"RIGIN";
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum GenbankState {
|
||||
OutSeq,
|
||||
MatchOrigin,
|
||||
SkipOriginLine,
|
||||
InSeq,
|
||||
InSlash,
|
||||
InAmbiguous,
|
||||
}
|
||||
|
||||
pub(crate) struct GenbankParser {
|
||||
state: GenbankState,
|
||||
overlap: OverlapState,
|
||||
keyword_pos: usize,
|
||||
at_line_start: bool,
|
||||
}
|
||||
|
||||
impl NucParser for GenbankParser {
|
||||
fn new(k: usize) -> Self {
|
||||
Self {
|
||||
state: GenbankState::OutSeq,
|
||||
overlap: OverlapState::new(k),
|
||||
keyword_pos: 0,
|
||||
at_line_start: true,
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn overlap_state(&self) -> &OverlapState {
|
||||
&self.overlap
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn overlap_state_mut(&mut self) -> &mut OverlapState {
|
||||
&mut self.overlap
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn is_in_seq(&self) -> bool {
|
||||
matches!(self.state, GenbankState::InSeq)
|
||||
}
|
||||
|
||||
fn parse_inplace(&mut self, buf: &mut [u8], n: usize) -> usize {
|
||||
let total = self.overlap.len + n;
|
||||
let mut i = 0;
|
||||
let mut j = 0;
|
||||
let mut seq_start: usize = 0;
|
||||
|
||||
while i < total {
|
||||
// SAFETY: i < total <= BUF_SIZE = buf.len()
|
||||
let byte = unsafe { *buf.get_unchecked(i) };
|
||||
|
||||
match self.state {
|
||||
GenbankState::OutSeq => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
self.at_line_start = true;
|
||||
} else if self.at_line_start && byte == b'O' {
|
||||
self.state = GenbankState::MatchOrigin;
|
||||
self.keyword_pos = 1;
|
||||
self.at_line_start = false;
|
||||
} else {
|
||||
self.at_line_start = false;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
GenbankState::MatchOrigin => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
self.state = GenbankState::OutSeq;
|
||||
self.at_line_start = true;
|
||||
} else if byte == ORIGIN_TAIL[self.keyword_pos - 1] {
|
||||
self.keyword_pos += 1;
|
||||
if self.keyword_pos == 6 {
|
||||
self.state = GenbankState::SkipOriginLine;
|
||||
}
|
||||
} else {
|
||||
self.state = GenbankState::OutSeq;
|
||||
self.at_line_start = false;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
GenbankState::SkipOriginLine => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
self.state = GenbankState::InSeq;
|
||||
seq_start = j;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
GenbankState::InSeq => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
self.at_line_start = true;
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
if self.at_line_start && byte == b'/' {
|
||||
self.state = GenbankState::InSlash;
|
||||
self.at_line_start = false;
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
self.at_line_start = false;
|
||||
let nuc = byte & 0xDF;
|
||||
if nuc == b'A' || nuc == b'C' || nuc == b'G' || nuc == b'T' {
|
||||
// SAFETY: j <= i < total <= BUF_SIZE = buf.len()
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = nuc;
|
||||
}
|
||||
j += 1;
|
||||
} else if byte.is_ascii_digit() || byte == b' ' {
|
||||
// position numbers and spacing between groups: skip
|
||||
} else {
|
||||
// ambiguous base: end current sequence if non-empty
|
||||
if j > seq_start {
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = 0;
|
||||
}
|
||||
j += 1;
|
||||
}
|
||||
self.state = GenbankState::InAmbiguous;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
GenbankState::InSlash => {
|
||||
if byte == b'/' {
|
||||
// confirmed "//": end of sequence record
|
||||
if j > seq_start {
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = 0;
|
||||
}
|
||||
j += 1;
|
||||
}
|
||||
self.state = GenbankState::OutSeq;
|
||||
self.at_line_start = false;
|
||||
} else if byte == b'\n' || byte == b'\r' {
|
||||
// single '/' line: back to sequence
|
||||
self.state = GenbankState::InSeq;
|
||||
self.at_line_start = true;
|
||||
} else {
|
||||
// false positive: single '/' mid-line, resume sequence
|
||||
self.state = GenbankState::InSeq;
|
||||
self.at_line_start = false;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
GenbankState::InAmbiguous => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
self.at_line_start = true;
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
if self.at_line_start && byte == b'/' {
|
||||
self.state = GenbankState::InSlash;
|
||||
self.at_line_start = false;
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
self.at_line_start = false;
|
||||
let nuc = byte & 0xDF;
|
||||
if nuc == b'A' || nuc == b'C' || nuc == b'G' || nuc == b'T' {
|
||||
seq_start = j;
|
||||
// SAFETY: j <= i < total <= BUF_SIZE = buf.len()
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = nuc;
|
||||
}
|
||||
j += 1;
|
||||
self.state = GenbankState::InSeq;
|
||||
}
|
||||
// digits, spaces, other ambiguous codes: skip
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
self.save_overlap(buf, j, seq_start)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── NucPage ──────────────────────────────────────────────────────────────────
|
||||
|
||||
/// Owned page of compacted nucleotides: uppercase A/C/G/T bytes separated by `0`
|
||||
/// at sequence boundaries. Automatically returns its buffer to the pool on drop.
|
||||
pub struct NucPage {
|
||||
data: ManuallyDrop<Vec<u8>>,
|
||||
len: usize,
|
||||
pool: Arc<Mutex<Vec<Vec<u8>>>>,
|
||||
}
|
||||
|
||||
impl std::ops::Deref for NucPage {
|
||||
type Target = [u8];
|
||||
fn deref(&self) -> &[u8] {
|
||||
&self.data[..self.len]
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for NucPage {
|
||||
fn drop(&mut self) {
|
||||
// SAFETY: data is never accessed after this point
|
||||
let buf = unsafe { ManuallyDrop::take(&mut self.data) };
|
||||
self.pool.lock().unwrap().push(buf);
|
||||
}
|
||||
}
|
||||
|
||||
// ─── NucPageCursor ────────────────────────────────────────────────────────────
|
||||
|
||||
/// A forward cursor over the normalised bytes of a [`NucPage`].
|
||||
///
|
||||
/// Provides the `next_byte` / `rewind` interface consumed by
|
||||
/// [`obiskbuilder::SuperKmerStreamIter`].
|
||||
pub struct NucPageCursor<'a> {
|
||||
data: &'a [u8],
|
||||
pos: usize,
|
||||
}
|
||||
|
||||
impl NucPageCursor<'_> {
|
||||
/// Returns the next byte in the page, or `None` at end.
|
||||
#[inline]
|
||||
pub fn next_byte(&mut self) -> Option<u8> {
|
||||
if self.pos < self.data.len() {
|
||||
let b = self.data[self.pos];
|
||||
self.pos += 1;
|
||||
Some(b)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Steps the cursor back by `n` bytes.
|
||||
///
|
||||
/// The caller guarantees that the last `n` bytes were all `ACGT`
|
||||
/// (no `0x00` separators), so they are still in the page buffer.
|
||||
#[inline]
|
||||
pub fn rewind(&mut self, n: usize) {
|
||||
self.pos -= n;
|
||||
}
|
||||
|
||||
/// Total number of bytes in the underlying page.
|
||||
#[inline]
|
||||
pub fn len(&self) -> usize {
|
||||
self.data.len()
|
||||
}
|
||||
|
||||
/// Returns `true` if the page contains no bytes.
|
||||
#[inline]
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.data.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
impl NucPage {
|
||||
/// Creates a forward cursor positioned at the start of this page.
|
||||
pub fn cursor(&self) -> NucPageCursor<'_> {
|
||||
NucPageCursor { data: self, pos: 0 }
|
||||
}
|
||||
}
|
||||
|
||||
// ─── NucStream ────────────────────────────────────────────────────────────────
|
||||
|
||||
pub(crate) struct NucStream<R: Read, P: NucParser> {
|
||||
reader: R,
|
||||
parser: P,
|
||||
pool: Arc<Mutex<Vec<Vec<u8>>>>,
|
||||
eof: bool,
|
||||
}
|
||||
|
||||
impl<R: Read, P: NucParser> NucStream<R, P> {
|
||||
pub(crate) fn new(reader: R, k: usize) -> Self {
|
||||
Self {
|
||||
reader,
|
||||
parser: P::new(k),
|
||||
pool: Arc::new(Mutex::new(Vec::new())),
|
||||
eof: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn read_page(&mut self) -> Option<NucPage> {
|
||||
loop {
|
||||
if self.eof {
|
||||
return None;
|
||||
}
|
||||
// take a buffer from the pool, or allocate fresh if all are in-flight
|
||||
let mut buf = self
|
||||
.pool
|
||||
.lock()
|
||||
.unwrap()
|
||||
.pop()
|
||||
.unwrap_or_else(|| vec![0u8; BUF_SIZE]);
|
||||
|
||||
let ol = self.parser.overlap_len();
|
||||
self.parser.write_overlap(&mut buf[..ol]);
|
||||
let n = self.reader.read(&mut buf[ol..ol + PAGE_SIZE]).unwrap_or(0);
|
||||
if n == 0 {
|
||||
self.eof = true;
|
||||
if ol == 0 {
|
||||
self.pool.lock().unwrap().push(buf);
|
||||
return None;
|
||||
}
|
||||
}
|
||||
let out_len = self.parser.parse_inplace(&mut buf, n);
|
||||
if out_len > 0 {
|
||||
return Some(NucPage {
|
||||
data: ManuallyDrop::new(buf),
|
||||
len: out_len,
|
||||
pool: Arc::clone(&self.pool),
|
||||
});
|
||||
}
|
||||
// empty page (all headers/ambiguous): return buf to pool and loop
|
||||
self.pool.lock().unwrap().push(buf);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<R: Read, P: NucParser> Iterator for NucStream<R, P> {
|
||||
type Item = NucPage;
|
||||
fn next(&mut self) -> Option<NucPage> {
|
||||
self.read_page()
|
||||
}
|
||||
}
|
||||
|
||||
// ─── FastaNucStream ───────────────────────────────────────────────────────────
|
||||
|
||||
pub(crate) type FastaNucStream<R> = NucStream<R, FastaParser>;
|
||||
pub(crate) type FastqNucStream<R> = NucStream<R, FastqParser>;
|
||||
pub(crate) type GenbankNucStream<R> = NucStream<R, GenbankParser>;
|
||||
|
||||
// ─── AnyNucStream ─────────────────────────────────────────────────────────────
|
||||
|
||||
pub(crate) enum AnyNucStream<R: Read> {
|
||||
Fasta(FastaNucStream<R>),
|
||||
Fastq(FastqNucStream<R>),
|
||||
Genbank(GenbankNucStream<R>),
|
||||
}
|
||||
|
||||
impl<R: Read> Iterator for AnyNucStream<R> {
|
||||
type Item = NucPage;
|
||||
fn next(&mut self) -> Option<NucPage> {
|
||||
match self {
|
||||
AnyNucStream::Fasta(s) => s.next(),
|
||||
AnyNucStream::Fastq(s) => s.next(),
|
||||
AnyNucStream::Genbank(s) => s.next(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn dispatch<R: Read>(
|
||||
mut guesser: MimeTypeGuesser<R>,
|
||||
k: usize,
|
||||
) -> Option<AnyNucStream<MimeTypeGuesser<R>>> {
|
||||
match guesser.mime_type() {
|
||||
Some("text/fasta") => Some(AnyNucStream::Fasta(NucStream::new(guesser, k))),
|
||||
Some("text/fastq") => Some(AnyNucStream::Fastq(NucStream::new(guesser, k))),
|
||||
Some("text/gbff") => Some(AnyNucStream::Genbank(NucStream::new(guesser, k))),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Wraps an already-open reader in a nucleotide stream, detecting its format.
|
||||
/// Returns `None` if the format is not recognised.
|
||||
pub(crate) fn nuc_stream<R: Read>(reader: R, k: usize) -> Option<AnyNucStream<MimeTypeGuesser<R>>> {
|
||||
dispatch(MimeTypeGuesser::new(reader), k)
|
||||
}
|
||||
|
||||
/// Opens a nucleotide stream from any source (file path, URL, or `-` for stdin),
|
||||
/// with transparent decompression and automatic format detection.
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns an `io::Error` if the source cannot be opened, decompression fails,
|
||||
/// or the format is not recognised.
|
||||
pub fn open_nuc_stream(
|
||||
source: &str,
|
||||
k: usize,
|
||||
) -> io::Result<Box<dyn Iterator<Item = NucPage> + Send>> {
|
||||
let reader = open_raw(source)?;
|
||||
nuc_stream(reader, k)
|
||||
.map(|s| Box::new(s) as Box<dyn Iterator<Item = NucPage> + Send>)
|
||||
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "unknown sequence format"))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "tests/nucstream.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,128 @@
|
||||
use super::overlap::{NucParser, OverlapState};
|
||||
|
||||
// ─── FASTA parser ─────────────────────────────────────────────────────────────
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum FastaState {
|
||||
OutSeq,
|
||||
InTitle,
|
||||
InSeq,
|
||||
InAmbiguous,
|
||||
}
|
||||
|
||||
pub(crate) struct FastaParser {
|
||||
state: FastaState,
|
||||
overlap: OverlapState,
|
||||
}
|
||||
|
||||
impl NucParser for FastaParser {
|
||||
fn new(k: usize) -> Self {
|
||||
Self {
|
||||
state: FastaState::OutSeq,
|
||||
overlap: OverlapState::new(k),
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn overlap_state(&self) -> &OverlapState {
|
||||
&self.overlap
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn overlap_state_mut(&mut self) -> &mut OverlapState {
|
||||
&mut self.overlap
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn is_in_seq(&self) -> bool {
|
||||
matches!(self.state, FastaState::InSeq)
|
||||
}
|
||||
|
||||
fn parse_inplace(&mut self, buf: &mut [u8], n: usize) -> usize {
|
||||
let total = self.overlap.len + n;
|
||||
let mut i = 0; // read index
|
||||
let mut j = 0; // write index (invariant: j <= i always)
|
||||
// j-position where the current sequence started in this call's output;
|
||||
// meaningful only when state is InSeq.
|
||||
let mut seq_start: usize = 0;
|
||||
|
||||
while i < total {
|
||||
// SAFETY: i < total <= BUF_SIZE = buf.len()
|
||||
let byte = unsafe { *buf.get_unchecked(i) };
|
||||
|
||||
match self.state {
|
||||
FastaState::OutSeq => {
|
||||
if byte == b'>' {
|
||||
self.state = FastaState::InTitle;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
FastaState::InTitle => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
self.state = FastaState::InSeq;
|
||||
seq_start = j;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
FastaState::InSeq => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
let nuc = byte & 0xDF; // to uppercase
|
||||
if nuc == b'A' || nuc == b'C' || nuc == b'G' || nuc == b'T' {
|
||||
// SAFETY: j <= i < total <= BUF_SIZE = buf.len()
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = nuc;
|
||||
}
|
||||
j += 1;
|
||||
i += 1;
|
||||
} else if byte == b'>' {
|
||||
if j > seq_start {
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = 0;
|
||||
}
|
||||
j += 1;
|
||||
}
|
||||
self.state = FastaState::InTitle;
|
||||
i += 1;
|
||||
} else {
|
||||
// first ambiguous base: end current sequence if non-empty
|
||||
if j > seq_start {
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = 0;
|
||||
}
|
||||
j += 1;
|
||||
}
|
||||
self.state = FastaState::InAmbiguous;
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
FastaState::InAmbiguous => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
if byte == b'>' {
|
||||
self.state = FastaState::InTitle;
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
let nuc = byte & 0xDF;
|
||||
if nuc == b'A' || nuc == b'C' || nuc == b'G' || nuc == b'T' {
|
||||
seq_start = j;
|
||||
// SAFETY: j <= i < total <= BUF_SIZE = buf.len()
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = nuc;
|
||||
}
|
||||
j += 1;
|
||||
self.state = FastaState::InSeq;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
self.save_overlap(buf, j, seq_start)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,132 @@
|
||||
use super::overlap::{NucParser, OverlapState};
|
||||
|
||||
// ─── FASTQ parser ─────────────────────────────────────────────────────────────
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum FastqState {
|
||||
OutSeq,
|
||||
InTitle,
|
||||
InSeq,
|
||||
InAmbiguous,
|
||||
InQualTitle,
|
||||
InQual,
|
||||
}
|
||||
|
||||
pub(crate) struct FastqParser {
|
||||
state: FastqState,
|
||||
overlap: OverlapState,
|
||||
}
|
||||
|
||||
impl NucParser for FastqParser {
|
||||
fn new(k: usize) -> Self {
|
||||
Self {
|
||||
state: FastqState::OutSeq,
|
||||
overlap: OverlapState::new(k),
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn overlap_state(&self) -> &OverlapState {
|
||||
&self.overlap
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn overlap_state_mut(&mut self) -> &mut OverlapState {
|
||||
&mut self.overlap
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn is_in_seq(&self) -> bool {
|
||||
matches!(self.state, FastqState::InSeq)
|
||||
}
|
||||
|
||||
fn parse_inplace(&mut self, buf: &mut [u8], n: usize) -> usize {
|
||||
let total = self.overlap.len + n;
|
||||
let mut i = 0;
|
||||
let mut j = 0;
|
||||
let mut seq_start: usize = 0;
|
||||
|
||||
while i < total {
|
||||
// SAFETY: i < total <= BUF_SIZE = buf.len()
|
||||
let byte = unsafe { *buf.get_unchecked(i) };
|
||||
|
||||
match self.state {
|
||||
FastqState::OutSeq => {
|
||||
if byte == b'@' {
|
||||
self.state = FastqState::InTitle;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
FastqState::InTitle => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
self.state = FastqState::InSeq;
|
||||
seq_start = j;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
FastqState::InSeq => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
if j > seq_start {
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = 0;
|
||||
}
|
||||
j += 1;
|
||||
}
|
||||
self.state = FastqState::InQualTitle;
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
let nuc = byte & 0xDF;
|
||||
if nuc == b'A' || nuc == b'C' || nuc == b'G' || nuc == b'T' {
|
||||
// SAFETY: j <= i < total <= BUF_SIZE = buf.len()
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = nuc;
|
||||
}
|
||||
j += 1;
|
||||
} else {
|
||||
if j > seq_start {
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = 0;
|
||||
}
|
||||
j += 1;
|
||||
}
|
||||
self.state = FastqState::InAmbiguous;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
FastqState::InAmbiguous => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
self.state = FastqState::InQualTitle;
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
let nuc = byte & 0xDF;
|
||||
if nuc == b'A' || nuc == b'C' || nuc == b'G' || nuc == b'T' {
|
||||
seq_start = j;
|
||||
// SAFETY: j <= i < total <= BUF_SIZE = buf.len()
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = nuc;
|
||||
}
|
||||
j += 1;
|
||||
self.state = FastqState::InSeq;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
FastqState::InQualTitle => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
self.state = FastqState::InQual;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
FastqState::InQual => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
self.state = FastqState::OutSeq;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
self.save_overlap(buf, j, seq_start)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
use super::overlap::{NucParser, OverlapState};
|
||||
|
||||
// ─── GenBank parser ───────────────────────────────────────────────────────────
|
||||
|
||||
const ORIGIN_TAIL: &[u8] = b"RIGIN";
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum GenbankState {
|
||||
OutSeq,
|
||||
MatchOrigin,
|
||||
SkipOriginLine,
|
||||
InSeq,
|
||||
InSlash,
|
||||
InAmbiguous,
|
||||
}
|
||||
|
||||
pub(crate) struct GenbankParser {
|
||||
state: GenbankState,
|
||||
overlap: OverlapState,
|
||||
keyword_pos: usize,
|
||||
at_line_start: bool,
|
||||
}
|
||||
|
||||
impl NucParser for GenbankParser {
|
||||
fn new(k: usize) -> Self {
|
||||
Self {
|
||||
state: GenbankState::OutSeq,
|
||||
overlap: OverlapState::new(k),
|
||||
keyword_pos: 0,
|
||||
at_line_start: true,
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn overlap_state(&self) -> &OverlapState {
|
||||
&self.overlap
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn overlap_state_mut(&mut self) -> &mut OverlapState {
|
||||
&mut self.overlap
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn is_in_seq(&self) -> bool {
|
||||
matches!(self.state, GenbankState::InSeq)
|
||||
}
|
||||
|
||||
fn parse_inplace(&mut self, buf: &mut [u8], n: usize) -> usize {
|
||||
let total = self.overlap.len + n;
|
||||
let mut i = 0;
|
||||
let mut j = 0;
|
||||
let mut seq_start: usize = 0;
|
||||
|
||||
while i < total {
|
||||
// SAFETY: i < total <= BUF_SIZE = buf.len()
|
||||
let byte = unsafe { *buf.get_unchecked(i) };
|
||||
|
||||
match self.state {
|
||||
GenbankState::OutSeq => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
self.at_line_start = true;
|
||||
} else if self.at_line_start && byte == b'O' {
|
||||
self.state = GenbankState::MatchOrigin;
|
||||
self.keyword_pos = 1;
|
||||
self.at_line_start = false;
|
||||
} else {
|
||||
self.at_line_start = false;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
GenbankState::MatchOrigin => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
self.state = GenbankState::OutSeq;
|
||||
self.at_line_start = true;
|
||||
} else if byte == ORIGIN_TAIL[self.keyword_pos - 1] {
|
||||
self.keyword_pos += 1;
|
||||
if self.keyword_pos == 6 {
|
||||
self.state = GenbankState::SkipOriginLine;
|
||||
}
|
||||
} else {
|
||||
self.state = GenbankState::OutSeq;
|
||||
self.at_line_start = false;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
GenbankState::SkipOriginLine => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
self.state = GenbankState::InSeq;
|
||||
seq_start = j;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
GenbankState::InSeq => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
self.at_line_start = true;
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
if self.at_line_start && byte == b'/' {
|
||||
self.state = GenbankState::InSlash;
|
||||
self.at_line_start = false;
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
self.at_line_start = false;
|
||||
let nuc = byte & 0xDF;
|
||||
if nuc == b'A' || nuc == b'C' || nuc == b'G' || nuc == b'T' {
|
||||
// SAFETY: j <= i < total <= BUF_SIZE = buf.len()
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = nuc;
|
||||
}
|
||||
j += 1;
|
||||
} else if byte.is_ascii_digit() || byte == b' ' {
|
||||
// position numbers and spacing between groups: skip
|
||||
} else {
|
||||
// ambiguous base: end current sequence if non-empty
|
||||
if j > seq_start {
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = 0;
|
||||
}
|
||||
j += 1;
|
||||
}
|
||||
self.state = GenbankState::InAmbiguous;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
GenbankState::InSlash => {
|
||||
if byte == b'/' {
|
||||
// confirmed "//": end of sequence record
|
||||
if j > seq_start {
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = 0;
|
||||
}
|
||||
j += 1;
|
||||
}
|
||||
self.state = GenbankState::OutSeq;
|
||||
self.at_line_start = false;
|
||||
} else if byte == b'\n' || byte == b'\r' {
|
||||
// single '/' line: back to sequence
|
||||
self.state = GenbankState::InSeq;
|
||||
self.at_line_start = true;
|
||||
} else {
|
||||
// false positive: single '/' mid-line, resume sequence
|
||||
self.state = GenbankState::InSeq;
|
||||
self.at_line_start = false;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
GenbankState::InAmbiguous => {
|
||||
if byte == b'\n' || byte == b'\r' {
|
||||
self.at_line_start = true;
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
if self.at_line_start && byte == b'/' {
|
||||
self.state = GenbankState::InSlash;
|
||||
self.at_line_start = false;
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
self.at_line_start = false;
|
||||
let nuc = byte & 0xDF;
|
||||
if nuc == b'A' || nuc == b'C' || nuc == b'G' || nuc == b'T' {
|
||||
seq_start = j;
|
||||
// SAFETY: j <= i < total <= BUF_SIZE = buf.len()
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = nuc;
|
||||
}
|
||||
j += 1;
|
||||
self.state = GenbankState::InSeq;
|
||||
}
|
||||
// digits, spaces, other ambiguous codes: skip
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
self.save_overlap(buf, j, seq_start)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
//! Streaming, in-place normalisation of raw sequence bytes into compacted
|
||||
//! nucleotide pages: uppercase A/C/G/T separated by `0` at sequence
|
||||
//! boundaries, ready for k-mer extraction without re-scanning for case or
|
||||
//! ambiguity codes.
|
||||
//!
|
||||
//! Submodules: [`overlap`] (format-independent k-1 overlap bookkeeping
|
||||
//! shared by every parser), [`fasta`]/[`fastq`]/[`genbank`] (one
|
||||
//! format-specific in-place state machine each), [`page`] (the pooled
|
||||
//! output buffer, [`NucPage`]), [`stream`] (format dispatch and the public
|
||||
//! [`open_nuc_stream`] entry point).
|
||||
|
||||
mod fasta;
|
||||
mod fastq;
|
||||
mod genbank;
|
||||
mod overlap;
|
||||
mod page;
|
||||
mod stream;
|
||||
|
||||
pub use page::{NucPage, NucPageCursor};
|
||||
pub use stream::open_nuc_stream;
|
||||
|
||||
// Only used by `tests.rs` (`use super::*`) below — the crate itself always
|
||||
// reaches these through their defining submodule directly.
|
||||
#[cfg(test)]
|
||||
pub(crate) use fasta::FastaParser;
|
||||
#[cfg(test)]
|
||||
pub(crate) use fastq::FastqParser;
|
||||
#[cfg(test)]
|
||||
pub(crate) use genbank::GenbankParser;
|
||||
#[cfg(test)]
|
||||
pub(crate) use stream::NucStream;
|
||||
|
||||
pub(crate) const MAX_K: usize = 31;
|
||||
pub(crate) const PAGE_SIZE: usize = 65536;
|
||||
// overlap (MAX_K - 1) + page data (PAGE_SIZE) + 1 byte for the end-of-page terminating 0
|
||||
pub(crate) const BUF_SIZE: usize = MAX_K + PAGE_SIZE;
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "../tests/nucstream.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,81 @@
|
||||
use super::MAX_K;
|
||||
|
||||
// ─── OverlapState ─────────────────────────────────────────────────────────────
|
||||
|
||||
pub(crate) struct OverlapState {
|
||||
data: [u8; MAX_K],
|
||||
pub(super) len: usize,
|
||||
k: usize,
|
||||
}
|
||||
|
||||
impl OverlapState {
|
||||
pub(crate) fn new(k: usize) -> Self {
|
||||
assert!(k > 0 && k <= MAX_K);
|
||||
Self {
|
||||
data: [0u8; MAX_K],
|
||||
len: 0,
|
||||
k,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── NucParser trait ──────────────────────────────────────────────────────────
|
||||
|
||||
// Transforms a raw page into a compacted nucleotide stream in-place.
|
||||
//
|
||||
// Buffer layout on each call:
|
||||
// buf[0..overlap_len()] — overlap bytes copied by write_overlap()
|
||||
// buf[overlap_len()..overlap_len()+n] — raw bytes just read from the source
|
||||
//
|
||||
// Returns the number of output bytes in buf[0..returned].
|
||||
pub(crate) trait NucParser {
|
||||
// required: format-specific
|
||||
fn new(k: usize) -> Self
|
||||
where
|
||||
Self: Sized;
|
||||
fn overlap_state(&self) -> &OverlapState;
|
||||
fn overlap_state_mut(&mut self) -> &mut OverlapState;
|
||||
fn is_in_seq(&self) -> bool;
|
||||
fn parse_inplace(&mut self, buf: &mut [u8], n: usize) -> usize;
|
||||
|
||||
// provided: format-independent overlap management
|
||||
fn overlap_len(&self) -> usize {
|
||||
self.overlap_state().len
|
||||
}
|
||||
|
||||
fn write_overlap(&self, buf: &mut [u8]) {
|
||||
let ol = &self.overlap_state();
|
||||
buf[..ol.len].copy_from_slice(&ol.data[..ol.len]);
|
||||
}
|
||||
|
||||
// Called at end of parse_inplace: saves overlap state and returns adjusted j.
|
||||
// seq_start is the j-position where the last sequence started in this call's output.
|
||||
fn save_overlap(&mut self, buf: &mut [u8], j: usize, seq_start: usize) -> usize {
|
||||
if !self.is_in_seq() {
|
||||
self.overlap_state_mut().len = 0;
|
||||
return j;
|
||||
}
|
||||
let seq_len = j - seq_start;
|
||||
let k = self.overlap_state().k;
|
||||
if seq_len >= k {
|
||||
// Sequence long enough: save last k-1 nucleotides, terminate with 0.
|
||||
let ol = k - 1;
|
||||
self.overlap_state_mut().data[..ol].copy_from_slice(&buf[j - ol..j]);
|
||||
self.overlap_state_mut().len = ol;
|
||||
// SAFETY: j <= total - 1 < BUF_SIZE = buf.len()
|
||||
// (total = overlap_len + n <= (MAX_K-1) + PAGE_SIZE = BUF_SIZE - 1)
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(j) = 0;
|
||||
}
|
||||
j + 1
|
||||
} else if seq_len > 0 {
|
||||
// Short sequence (< k): save whole fragment, strip from output.
|
||||
self.overlap_state_mut().data[..seq_len].copy_from_slice(&buf[seq_start..j]);
|
||||
self.overlap_state_mut().len = seq_len;
|
||||
seq_start
|
||||
} else {
|
||||
self.overlap_state_mut().len = 0;
|
||||
j
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
use std::mem::ManuallyDrop;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
// ─── NucPage ──────────────────────────────────────────────────────────────────
|
||||
|
||||
/// Owned page of compacted nucleotides: uppercase A/C/G/T bytes separated by `0`
|
||||
/// at sequence boundaries. Automatically returns its buffer to the pool on drop.
|
||||
pub struct NucPage {
|
||||
pub(super) data: ManuallyDrop<Vec<u8>>,
|
||||
pub(super) len: usize,
|
||||
pub(super) pool: Arc<Mutex<Vec<Vec<u8>>>>,
|
||||
}
|
||||
|
||||
impl std::ops::Deref for NucPage {
|
||||
type Target = [u8];
|
||||
fn deref(&self) -> &[u8] {
|
||||
&self.data[..self.len]
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for NucPage {
|
||||
fn drop(&mut self) {
|
||||
// SAFETY: data is never accessed after this point
|
||||
let buf = unsafe { ManuallyDrop::take(&mut self.data) };
|
||||
self.pool.lock().unwrap().push(buf);
|
||||
}
|
||||
}
|
||||
|
||||
// ─── NucPageCursor ────────────────────────────────────────────────────────────
|
||||
|
||||
/// A forward cursor over the normalised bytes of a [`NucPage`].
|
||||
///
|
||||
/// Provides the `next_byte` / `rewind` interface consumed by
|
||||
/// [`obiskbuilder::SuperKmerStreamIter`].
|
||||
pub struct NucPageCursor<'a> {
|
||||
data: &'a [u8],
|
||||
pos: usize,
|
||||
}
|
||||
|
||||
impl NucPageCursor<'_> {
|
||||
/// Returns the next byte in the page, or `None` at end.
|
||||
#[inline]
|
||||
pub fn next_byte(&mut self) -> Option<u8> {
|
||||
if self.pos < self.data.len() {
|
||||
let b = self.data[self.pos];
|
||||
self.pos += 1;
|
||||
Some(b)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Steps the cursor back by `n` bytes.
|
||||
///
|
||||
/// The caller guarantees that the last `n` bytes were all `ACGT`
|
||||
/// (no `0x00` separators), so they are still in the page buffer.
|
||||
#[inline]
|
||||
pub fn rewind(&mut self, n: usize) {
|
||||
self.pos -= n;
|
||||
}
|
||||
|
||||
/// Total number of bytes in the underlying page.
|
||||
#[inline]
|
||||
pub fn len(&self) -> usize {
|
||||
self.data.len()
|
||||
}
|
||||
|
||||
/// Returns `true` if the page contains no bytes.
|
||||
#[inline]
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.data.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
impl NucPage {
|
||||
/// Creates a forward cursor positioned at the start of this page.
|
||||
pub fn cursor(&self) -> NucPageCursor<'_> {
|
||||
NucPageCursor { data: self, pos: 0 }
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,135 @@
|
||||
use std::io::{self, Read};
|
||||
use std::mem::ManuallyDrop;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use crate::mimetype::MimeTypeGuesser;
|
||||
use crate::xopen::open_raw;
|
||||
|
||||
use super::fasta::FastaParser;
|
||||
use super::fastq::FastqParser;
|
||||
use super::genbank::GenbankParser;
|
||||
use super::overlap::NucParser;
|
||||
use super::page::NucPage;
|
||||
use super::{BUF_SIZE, PAGE_SIZE};
|
||||
|
||||
// ─── NucStream ────────────────────────────────────────────────────────────────
|
||||
|
||||
pub(crate) struct NucStream<R: Read, P: NucParser> {
|
||||
reader: R,
|
||||
parser: P,
|
||||
pool: Arc<Mutex<Vec<Vec<u8>>>>,
|
||||
eof: bool,
|
||||
}
|
||||
|
||||
impl<R: Read, P: NucParser> NucStream<R, P> {
|
||||
pub(crate) fn new(reader: R, k: usize) -> Self {
|
||||
Self {
|
||||
reader,
|
||||
parser: P::new(k),
|
||||
pool: Arc::new(Mutex::new(Vec::new())),
|
||||
eof: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn read_page(&mut self) -> Option<NucPage> {
|
||||
loop {
|
||||
if self.eof {
|
||||
return None;
|
||||
}
|
||||
// take a buffer from the pool, or allocate fresh if all are in-flight
|
||||
let mut buf = self
|
||||
.pool
|
||||
.lock()
|
||||
.unwrap()
|
||||
.pop()
|
||||
.unwrap_or_else(|| vec![0u8; BUF_SIZE]);
|
||||
|
||||
let ol = self.parser.overlap_len();
|
||||
self.parser.write_overlap(&mut buf[..ol]);
|
||||
let n = self.reader.read(&mut buf[ol..ol + PAGE_SIZE]).unwrap_or(0);
|
||||
if n == 0 {
|
||||
self.eof = true;
|
||||
if ol == 0 {
|
||||
self.pool.lock().unwrap().push(buf);
|
||||
return None;
|
||||
}
|
||||
}
|
||||
let out_len = self.parser.parse_inplace(&mut buf, n);
|
||||
if out_len > 0 {
|
||||
return Some(NucPage {
|
||||
data: ManuallyDrop::new(buf),
|
||||
len: out_len,
|
||||
pool: Arc::clone(&self.pool),
|
||||
});
|
||||
}
|
||||
// empty page (all headers/ambiguous): return buf to pool and loop
|
||||
self.pool.lock().unwrap().push(buf);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<R: Read, P: NucParser> Iterator for NucStream<R, P> {
|
||||
type Item = NucPage;
|
||||
fn next(&mut self) -> Option<NucPage> {
|
||||
self.read_page()
|
||||
}
|
||||
}
|
||||
|
||||
// ─── FastaNucStream ───────────────────────────────────────────────────────────
|
||||
|
||||
pub(crate) type FastaNucStream<R> = NucStream<R, FastaParser>;
|
||||
pub(crate) type FastqNucStream<R> = NucStream<R, FastqParser>;
|
||||
pub(crate) type GenbankNucStream<R> = NucStream<R, GenbankParser>;
|
||||
|
||||
// ─── AnyNucStream ─────────────────────────────────────────────────────────────
|
||||
|
||||
pub(crate) enum AnyNucStream<R: Read> {
|
||||
Fasta(FastaNucStream<R>),
|
||||
Fastq(FastqNucStream<R>),
|
||||
Genbank(GenbankNucStream<R>),
|
||||
}
|
||||
|
||||
impl<R: Read> Iterator for AnyNucStream<R> {
|
||||
type Item = NucPage;
|
||||
fn next(&mut self) -> Option<NucPage> {
|
||||
match self {
|
||||
AnyNucStream::Fasta(s) => s.next(),
|
||||
AnyNucStream::Fastq(s) => s.next(),
|
||||
AnyNucStream::Genbank(s) => s.next(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn dispatch<R: Read>(
|
||||
mut guesser: MimeTypeGuesser<R>,
|
||||
k: usize,
|
||||
) -> Option<AnyNucStream<MimeTypeGuesser<R>>> {
|
||||
match guesser.mime_type() {
|
||||
Some("text/fasta") => Some(AnyNucStream::Fasta(NucStream::new(guesser, k))),
|
||||
Some("text/fastq") => Some(AnyNucStream::Fastq(NucStream::new(guesser, k))),
|
||||
Some("text/gbff") => Some(AnyNucStream::Genbank(NucStream::new(guesser, k))),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Wraps an already-open reader in a nucleotide stream, detecting its format.
|
||||
/// Returns `None` if the format is not recognised.
|
||||
pub(crate) fn nuc_stream<R: Read>(reader: R, k: usize) -> Option<AnyNucStream<MimeTypeGuesser<R>>> {
|
||||
dispatch(MimeTypeGuesser::new(reader), k)
|
||||
}
|
||||
|
||||
/// Opens a nucleotide stream from any source (file path, URL, or `-` for stdin),
|
||||
/// with transparent decompression and automatic format detection.
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns an `io::Error` if the source cannot be opened, decompression fails,
|
||||
/// or the format is not recognised.
|
||||
pub fn open_nuc_stream(
|
||||
source: &str,
|
||||
k: usize,
|
||||
) -> io::Result<Box<dyn Iterator<Item = NucPage> + Send>> {
|
||||
let reader = open_raw(source)?;
|
||||
nuc_stream(reader, k)
|
||||
.map(|s| Box::new(s) as Box<dyn Iterator<Item = NucPage> + Send>)
|
||||
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "unknown sequence format"))
|
||||
}
|
||||
Reference in New Issue
Block a user