use std::collections::{HashMap, HashSet};
use std::io::{BufRead, BufReader, Read, Seek, SeekFrom};
use std::ops::Range;
use crate::error::{FastqError, FastqPosition, Result};
#[cfg(feature = "bgzf")]
use crate::{BgzfDecodedBlockReader, BgzfIndex, BgzfIndexEntry, BgzfSeekReader, BgzfVirtualOffset};
use memchr::memchr;
const DEFAULT_BATCH_RECORDS: usize = 1024;
const DEFAULT_READER_BUFFER_SIZE: usize = 64 * 1024;
const REFERENCE_BATCH_RECORDS: usize = 16;
const REFERENCE_READER_BUFFER_SIZE: usize = 256 * 1024;
const REFERENCE_EXPECTED_SEQ_LEN: usize = 1024 * 1024;
const TWO_LINE_STREAM_BUFFER_SIZE: usize = 64 * 1024;
const INDEXED_FETCH_SCRATCH_SIZE: usize = 64 * 1024;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct ByteRange {
start: u32,
end: u32,
}
impl ByteRange {
fn to_usize(self) -> Range<usize> {
self.start as usize..self.end as usize
}
}
#[derive(Debug, Clone)]
pub struct FastaConfig {
pub batch_records: usize,
pub buffer_size: usize,
pub expected_seq_len: usize,
}
impl Default for FastaConfig {
fn default() -> Self {
Self {
batch_records: DEFAULT_BATCH_RECORDS,
buffer_size: DEFAULT_READER_BUFFER_SIZE,
expected_seq_len: 0,
}
}
}
impl FastaConfig {
pub fn reference() -> Self {
Self {
batch_records: REFERENCE_BATCH_RECORDS,
buffer_size: REFERENCE_READER_BUFFER_SIZE,
expected_seq_len: REFERENCE_EXPECTED_SEQ_LEN,
}
}
pub fn two_line(mut self) -> Self {
self.buffer_size = TWO_LINE_STREAM_BUFFER_SIZE;
self
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FastaRecordRef {
pub name: Range<u32>,
pub seq: Range<u32>,
}
#[derive(Debug, Clone, Copy)]
pub struct FastaRecord<'a> {
bytes: &'a [u8],
record: &'a FastaRecordRef,
id_token: ByteRange,
}
impl<'a> FastaRecord<'a> {
pub fn name(self) -> &'a [u8] {
&self.bytes[to_usize(self.record.name.clone())]
}
pub fn name_without_gt(self) -> &'a [u8] {
let name = self.name();
name.strip_prefix(b">").unwrap_or(name)
}
pub fn id_token(self) -> &'a [u8] {
&self.bytes[self.id_token.to_usize()]
}
pub fn seq(self) -> &'a [u8] {
&self.bytes[to_usize(self.record.seq.clone())]
}
}
#[derive(Debug, Clone, Copy)]
pub struct FastaVisitRecord<'a> {
name: &'a [u8],
seq: &'a [u8],
}
impl<'a> FastaVisitRecord<'a> {
pub fn name(self) -> &'a [u8] {
self.name
}
pub fn name_without_gt(self) -> &'a [u8] {
self.name.strip_prefix(b">").unwrap_or(self.name)
}
pub fn id_token(self) -> &'a [u8] {
let name = self.name_without_gt();
let end = name
.iter()
.position(u8::is_ascii_whitespace)
.unwrap_or(name.len());
&name[..end]
}
pub fn seq(self) -> &'a [u8] {
self.seq
}
}
const FASTA_STATS_CHECKSUM_INIT: u64 = 0xcbf2_9ce4_8422_2325;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FastaStats {
pub records: u64,
pub bases: u64,
pub checksum: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FastaShape {
Empty,
TwoLine,
Multiline,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FastaIndexEntry {
pub name: Vec<u8>,
pub len: u64,
pub offset: u64,
pub line_bases: u64,
pub line_width: u64,
#[cfg(feature = "bgzf")]
pub virtual_offset: Option<BgzfVirtualOffset>,
}
impl FastaIndexEntry {
pub fn sequence_offset(&self, pos: u64) -> Result<u64> {
if pos > self.len {
return Err(FastqError::Format(
"FASTA sequence position exceeds reference length".into(),
));
}
if pos == self.len {
return self.sequence_end_offset();
}
if self.line_bases == 0 {
return Err(FastqError::Format(
"FASTA index entry has zero line_bases for non-empty sequence".into(),
));
}
let line = pos / self.line_bases;
let in_line = pos % self.line_bases;
let line_offset = line.checked_mul(self.line_width).ok_or_else(|| {
FastqError::Format("FASTA index offset calculation overflowed".into())
})?;
self.offset
.checked_add(line_offset)
.and_then(|offset| offset.checked_add(in_line))
.ok_or_else(|| FastqError::Format("FASTA index offset calculation overflowed".into()))
}
pub fn sequence_spans(&self, range: Range<u64>) -> Result<Vec<Range<u64>>> {
self.validate_range(range.clone())?;
let mut spans = Vec::new();
let mut pos = range.start;
while pos < range.end {
if self.line_bases == 0 {
return Err(FastqError::Format(
"FASTA index entry has zero line_bases for non-empty range".into(),
));
}
let in_line = pos % self.line_bases;
let take = (range.end - pos).min(self.line_bases - in_line);
let start = self.sequence_offset(pos)?;
let end = start.checked_add(take).ok_or_else(|| {
FastqError::Format("FASTA index span calculation overflowed".into())
})?;
spans.push(start..end);
pos += take;
}
Ok(spans)
}
fn sequence_end_offset(&self) -> Result<u64> {
if self.len == 0 {
return Ok(self.offset);
}
self.sequence_offset(self.len - 1).map(|offset| offset + 1)
}
fn validate_range(&self, range: Range<u64>) -> Result<()> {
if range.start > range.end {
return Err(FastqError::Format(
"FASTA range start must be <= end".into(),
));
}
if range.end > self.len {
return Err(FastqError::Format(
"FASTA range end exceeds reference length".into(),
));
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FastaPartitionConfig {
pub target_partitions: usize,
pub overlap_bases: u64,
}
impl FastaPartitionConfig {
pub fn new(target_partitions: usize, overlap_bases: u64) -> Self {
Self {
target_partitions,
overlap_bases,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FastaPartition {
pub partition_index: usize,
pub name: Vec<u8>,
pub core: Range<u64>,
pub fetch: Range<u64>,
}
impl FastaPartition {
pub fn core_len(&self) -> u64 {
self.core.end - self.core.start
}
pub fn core_offset_in_fetch(&self) -> u64 {
self.core.start - self.fetch.start
}
}
pub fn plan_fasta_partitions(
index: &FastaIndex,
config: FastaPartitionConfig,
) -> Result<Vec<FastaPartition>> {
let target_partitions = config.target_partitions.max(1);
let total_bases = index.entries().iter().try_fold(0_u64, |acc, entry| {
acc.checked_add(entry.len)
.ok_or_else(|| FastqError::Format("FASTA partition total length overflowed".into()))
})?;
if total_bases == 0 {
return Ok(Vec::new());
}
let target_partitions_u64 = u64::try_from(target_partitions)
.map_err(|_| FastqError::Format("FASTA partition count exceeds u64 range".into()))?;
let target_bases = total_bases.div_ceil(target_partitions_u64).max(1);
let mut partitions = Vec::new();
for entry in index.entries() {
let mut start = 0_u64;
while start < entry.len {
let remaining = entry.len - start;
let take = remaining.min(target_bases);
let core = start..start + take;
let fetch_start = core.start.saturating_sub(config.overlap_bases);
let fetch_end = core.end.saturating_add(config.overlap_bases).min(entry.len);
partitions.push(FastaPartition {
partition_index: partitions.len(),
name: entry.name.clone(),
core: core.clone(),
fetch: fetch_start..fetch_end,
});
start = core.end;
}
}
Ok(partitions)
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FastaReferenceChunk {
pub name: Vec<u8>,
pub global_offset: u64,
pub seq: Vec<u8>,
}
#[derive(Debug, Clone, Copy)]
pub struct FastaReferenceChunkRef<'a> {
pub name: &'a [u8],
pub global_offset: u64,
pub seq: &'a [u8],
}
pub trait FastaReferenceChunkSink {
fn chunk(&mut self, chunk: FastaReferenceChunkRef<'_>) -> Result<()>;
}
impl<F> FastaReferenceChunkSink for F
where
F: FnMut(FastaReferenceChunkRef<'_>) -> Result<()>,
{
fn chunk(&mut self, chunk: FastaReferenceChunkRef<'_>) -> Result<()> {
self(chunk)
}
}
pub struct FastaReferenceChunks<'a, R> {
reader: &'a mut IndexedFastaReader<R>,
name: Vec<u8>,
entry: FastaIndexEntry,
next_offset: u64,
end: u64,
chunk_bases: u64,
initialized: bool,
}
#[cfg(feature = "bgzf")]
pub struct BgzfFastaReferenceChunks<'a, R> {
reader: &'a mut BgzfIndexedFastaReader<R>,
name: Vec<u8>,
entry: FastaIndexEntry,
next_offset: u64,
end: u64,
chunk_bases: u64,
initialized: bool,
}
impl<R: Read + Seek> Iterator for FastaReferenceChunks<'_, R> {
type Item = Result<FastaReferenceChunk>;
fn next(&mut self) -> Option<Self::Item> {
if self.next_offset >= self.end {
return None;
}
let start = self.next_offset;
if !self.initialized {
if let Err(err) =
seek_indexed_sequence_start(&mut self.reader.inner, &self.entry, start)
{
self.next_offset = self.end;
return Some(Err(err));
}
self.initialized = true;
}
let capacity = indexed_chunk_capacity(start, self.end, self.chunk_bases);
let mut seq = match capacity {
Ok(capacity) => Vec::with_capacity(capacity),
Err(err) => {
self.next_offset = self.end;
return Some(Err(err));
}
};
match read_next_indexed_sequence_chunk(
&mut self.reader.inner,
&self.entry,
start,
self.end,
self.chunk_bases,
&mut seq,
&mut self.reader.scratch,
) {
Ok(next_offset) => {
self.next_offset = next_offset;
Some(Ok(FastaReferenceChunk {
name: self.name.clone(),
global_offset: start,
seq,
}))
}
Err(err) => {
self.next_offset = self.end;
Some(Err(err))
}
}
}
}
#[cfg(feature = "bgzf")]
impl<R: Read + Seek> Iterator for BgzfFastaReferenceChunks<'_, R> {
type Item = Result<FastaReferenceChunk>;
fn next(&mut self) -> Option<Self::Item> {
if self.next_offset >= self.end {
return None;
}
let start = self.next_offset;
if !self.initialized {
if let Err(err) = seek_bgzf_indexed_sequence_start(
&mut self.reader.inner,
&self.reader.bgzf_index,
&self.entry,
start,
) {
self.next_offset = self.end;
return Some(Err(err));
}
self.initialized = true;
}
let capacity = indexed_chunk_capacity(start, self.end, self.chunk_bases);
let mut seq = match capacity {
Ok(capacity) => Vec::with_capacity(capacity),
Err(err) => {
self.next_offset = self.end;
return Some(Err(err));
}
};
match read_next_indexed_sequence_chunk(
&mut self.reader.inner,
&self.entry,
start,
self.end,
self.chunk_bases,
&mut seq,
&mut self.reader.scratch,
) {
Ok(next_offset) => {
self.next_offset = next_offset;
Some(Ok(FastaReferenceChunk {
name: self.name.clone(),
global_offset: start,
seq,
}))
}
Err(err) => {
self.next_offset = self.end;
Some(Err(err))
}
}
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct FastaIndex {
entries: Vec<FastaIndexEntry>,
name_to_index: HashMap<Vec<u8>, usize>,
}
impl FastaIndex {
pub fn entries(&self) -> &[FastaIndexEntry] {
&self.entries
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn get(&self, name: &[u8]) -> Option<&FastaIndexEntry> {
self.name_to_index
.get(name)
.and_then(|&idx| self.entries.get(idx))
}
pub fn to_fai_string(&self) -> String {
let mut out = String::new();
for entry in &self.entries {
out.push_str(&String::from_utf8_lossy(&entry.name));
out.push('\t');
out.push_str(&entry.len.to_string());
out.push('\t');
out.push_str(&entry.offset.to_string());
out.push('\t');
out.push_str(&entry.line_bases.to_string());
out.push('\t');
out.push_str(&entry.line_width.to_string());
out.push('\n');
}
out
}
pub fn from_fai_bytes(bytes: &[u8]) -> Result<Self> {
let mut entries = Vec::new();
let mut seen = HashSet::new();
for (line_idx, line) in bytes.split(|&b| b == b'\n').enumerate() {
let line = trim_line(line);
if line.is_empty() {
continue;
}
let fields = line.split(|&b| b == b'\t').collect::<Vec<_>>();
if fields.len() != 5 {
return Err(FastqError::Format(format!(
"invalid .fai line {}: expected 5 tab-delimited fields",
line_idx + 1
)));
}
if fields[0].is_empty() {
return Err(FastqError::Format(format!(
"invalid .fai line {}: empty reference name",
line_idx + 1
)));
}
if !seen.insert(fields[0].to_vec()) {
return Err(FastqError::Format(format!(
"invalid .fai line {}: duplicate reference name",
line_idx + 1
)));
}
let len = parse_fai_u64(fields[1], line_idx + 1, "length")?;
let offset = parse_fai_u64(fields[2], line_idx + 1, "offset")?;
let line_bases = parse_fai_u64(fields[3], line_idx + 1, "line_bases")?;
let line_width = parse_fai_u64(fields[4], line_idx + 1, "line_width")?;
if len > 0 && line_bases == 0 {
return Err(FastqError::Format(format!(
"invalid .fai line {}: non-empty reference has zero line_bases",
line_idx + 1
)));
}
if line_width < line_bases {
return Err(FastqError::Format(format!(
"invalid .fai line {}: line_width is smaller than line_bases",
line_idx + 1
)));
}
entries.push(FastaIndexEntry {
name: fields[0].to_vec(),
len,
offset,
line_bases,
line_width,
#[cfg(feature = "bgzf")]
virtual_offset: None,
});
}
Ok(Self::from_entries(entries))
}
pub fn from_fai_str(text: &str) -> Result<Self> {
Self::from_fai_bytes(text.as_bytes())
}
pub fn from_fai_read<R: Read>(mut reader: R) -> Result<Self> {
let mut bytes = Vec::new();
reader.read_to_end(&mut bytes)?;
Self::from_fai_bytes(&bytes)
}
fn from_entries(entries: Vec<FastaIndexEntry>) -> Self {
let name_to_index = entries
.iter()
.enumerate()
.map(|(idx, entry)| (entry.name.clone(), idx))
.collect();
Self {
entries,
name_to_index,
}
}
}
fn parse_fai_u64(value: &[u8], line: usize, field: &str) -> Result<u64> {
let text = std::str::from_utf8(value).map_err(|_| {
FastqError::Format(format!(
"invalid .fai line {line}: {field} is not valid UTF-8"
))
})?;
text.parse::<u64>().map_err(|_| {
FastqError::Format(format!(
"invalid .fai line {line}: {field} is not an unsigned integer"
))
})
}
fn indexed_physical_window(entry: &FastaIndexEntry, range: Range<u64>) -> Result<Range<u64>> {
entry.validate_range(range.clone())?;
let start = entry.sequence_offset(range.start)?;
let end = if range.start == range.end {
start
} else {
entry
.sequence_offset(range.end - 1)?
.checked_add(1)
.ok_or_else(|| FastqError::Format("FASTA index span calculation overflowed".into()))?
};
if end < start {
return Err(FastqError::Format(
"FASTA index physical range is invalid".into(),
));
}
Ok(start..end)
}
fn copy_indexed_sequence_window<R: Read>(
reader: &mut R,
entry: &FastaIndexEntry,
mut physical_offset: u64,
mut len: u64,
expected_bases: usize,
out: &mut Vec<u8>,
scratch: &mut Vec<u8>,
) -> Result<()> {
if len == 0 {
return Ok(());
}
if entry.line_width == 0 {
return Err(FastqError::Format(
"FASTA index entry has zero line_width for non-empty range".into(),
));
}
if scratch.len() < INDEXED_FETCH_SCRATCH_SIZE {
scratch.resize(INDEXED_FETCH_SCRATCH_SIZE, 0);
}
while len > 0 {
let take = usize::try_from(len.min(scratch.len() as u64))
.map_err(|_| FastqError::Format("FASTA fetch span exceeds usize range".into()))?;
reader.read_exact(&mut scratch[..take])?;
let mut cursor = 0;
while cursor < take {
let cursor_offset = physical_offset.checked_add(cursor as u64).ok_or_else(|| {
FastqError::Format("FASTA fetch physical offset overflowed".into())
})?;
let rel = cursor_offset.checked_sub(entry.offset).ok_or_else(|| {
FastqError::Format("FASTA index physical offset precedes sequence offset".into())
})?;
let in_line = rel % entry.line_width;
let remaining = take - cursor;
if in_line >= entry.line_bases {
let skip = (entry.line_width - in_line).min(remaining as u64);
cursor += usize::try_from(skip).map_err(|_| {
FastqError::Format("FASTA fetch span exceeds usize range".into())
})?;
continue;
}
let run = (entry.line_bases - in_line).min(remaining as u64);
let run = usize::try_from(run)
.map_err(|_| FastqError::Format("FASTA fetch span exceeds usize range".into()))?;
out.extend_from_slice(&scratch[cursor..cursor + run]);
if out.len() > expected_bases {
return Err(FastqError::Format(
"FASTA fetch produced more bases than expected".into(),
));
}
cursor += run;
}
physical_offset = physical_offset
.checked_add(take as u64)
.ok_or_else(|| FastqError::Format("FASTA fetch physical offset overflowed".into()))?;
len -= take as u64;
}
if out.len() != expected_bases {
return Err(FastqError::Format(
"FASTA fetch produced fewer bases than expected".into(),
));
}
Ok(())
}
fn seek_indexed_sequence_start<R: Seek>(
reader: &mut R,
entry: &FastaIndexEntry,
seq_pos: u64,
) -> Result<()> {
let offset = entry.sequence_offset(seq_pos)?;
reader.seek(SeekFrom::Start(offset))?;
Ok(())
}
#[cfg(feature = "bgzf")]
fn seek_bgzf_indexed_sequence_start<R: Read + Seek>(
reader: &mut BgzfSeekReader<R>,
bgzf_index: &BgzfIndex,
entry: &FastaIndexEntry,
seq_pos: u64,
) -> Result<()> {
let offset = entry.sequence_offset(seq_pos)?;
let virtual_offset = bgzf_index
.virtual_offset_for_uncompressed_offset(offset)?
.ok_or_else(|| FastqError::Bgzf("BGZF span offset is not indexed".into()))?;
reader.seek_virtual_offset(virtual_offset)?;
Ok(())
}
fn indexed_chunk_capacity(start: u64, end: u64, chunk_bases: u64) -> Result<usize> {
usize::try_from(end.saturating_sub(start).min(chunk_bases.max(1)))
.map_err(|_| FastqError::Format("FASTA chunk length exceeds usize range".into()))
}
fn read_next_indexed_sequence_chunk<R: Read>(
reader: &mut R,
entry: &FastaIndexEntry,
start: u64,
end: u64,
chunk_bases: u64,
out: &mut Vec<u8>,
scratch: &mut Vec<u8>,
) -> Result<u64> {
if entry.line_bases == 0 {
return Err(FastqError::Format(
"FASTA index entry has zero line_bases for non-empty range".into(),
));
}
if entry.line_width < entry.line_bases {
return Err(FastqError::Format(
"FASTA index entry has line_width smaller than line_bases".into(),
));
}
let chunk_end = start.saturating_add(chunk_bases.max(1)).min(end);
out.clear();
out.reserve(indexed_chunk_capacity(start, chunk_end, chunk_bases)?);
let mut pos = start;
while pos < chunk_end {
let in_line = pos % entry.line_bases;
let take = (chunk_end - pos).min(entry.line_bases - in_line);
read_exact_into_vec(reader, out, take)?;
pos += take;
if pos < end && pos.is_multiple_of(entry.line_bases) {
skip_indexed_separator(reader, entry.line_width - entry.line_bases, scratch)?;
}
}
Ok(pos)
}
fn read_exact_into_vec<R: Read>(reader: &mut R, out: &mut Vec<u8>, len: u64) -> Result<()> {
let len = usize::try_from(len)
.map_err(|_| FastqError::Format("FASTA sequence run exceeds usize range".into()))?;
let old_len = out.len();
out.resize(old_len + len, 0);
reader.read_exact(&mut out[old_len..])?;
Ok(())
}
fn skip_indexed_separator<R: Read>(
reader: &mut R,
mut len: u64,
scratch: &mut Vec<u8>,
) -> Result<()> {
if len == 0 {
return Ok(());
}
if scratch.len() < INDEXED_FETCH_SCRATCH_SIZE {
scratch.resize(INDEXED_FETCH_SCRATCH_SIZE, 0);
}
while len > 0 {
let take = usize::try_from(len.min(scratch.len() as u64))
.map_err(|_| FastqError::Format("FASTA separator span exceeds usize range".into()))?;
reader.read_exact(&mut scratch[..take])?;
len -= take as u64;
}
Ok(())
}
struct IndexedSequenceStream<'a, R> {
reader: &'a mut R,
scratch: &'a mut Vec<u8>,
entry: &'a FastaIndexEntry,
next_offset: u64,
end: u64,
chunk_bases: u64,
}
impl<R: Read> IndexedSequenceStream<'_, R> {
fn next_chunk_into(&mut self, out: &mut Vec<u8>) -> Result<Option<u64>> {
if self.next_offset >= self.end {
return Ok(None);
}
let start = self.next_offset;
self.next_offset = read_next_indexed_sequence_chunk(
self.reader,
self.entry,
start,
self.end,
self.chunk_bases,
out,
self.scratch,
)?;
Ok(Some(start))
}
}
fn stream_indexed_reference_chunks_into<R, S>(
stream: &mut IndexedSequenceStream<'_, R>,
name: &[u8],
out: &mut Vec<u8>,
sink: &mut S,
) -> Result<()>
where
R: Read,
S: FastaReferenceChunkSink,
{
while let Some(start) = stream.next_chunk_into(out)? {
sink.chunk(FastaReferenceChunkRef {
name,
global_offset: start,
seq: out,
})?;
}
Ok(())
}
pub struct IndexedFastaReader<R> {
inner: R,
index: FastaIndex,
scratch: Vec<u8>,
}
impl<R: Read + Seek> IndexedFastaReader<R> {
pub fn new(inner: R, index: FastaIndex) -> Self {
Self {
inner,
index,
scratch: Vec::new(),
}
}
pub fn index(&self) -> &FastaIndex {
&self.index
}
pub fn fetch_into(&mut self, name: &[u8], range: Range<u64>, out: &mut Vec<u8>) -> Result<()> {
let entry = self.index.get(name).ok_or_else(|| {
FastqError::Format(format!(
"FASTA reference not found in index: {}",
String::from_utf8_lossy(name)
))
})?;
let window = indexed_physical_window(entry, range.clone())?;
out.clear();
let expected_bases = usize::try_from(range.end - range.start).map_err(|_| {
FastqError::Format("FASTA fetch range length exceeds usize range".into())
})?;
out.reserve(expected_bases);
self.inner.seek(SeekFrom::Start(window.start))?;
copy_indexed_sequence_window(
&mut self.inner,
entry,
window.start,
window.end - window.start,
expected_bases,
out,
&mut self.scratch,
)?;
Ok(())
}
pub fn fetch(&mut self, name: &[u8], range: Range<u64>) -> Result<Vec<u8>> {
let mut out = Vec::new();
self.fetch_into(name, range, &mut out)?;
Ok(out)
}
pub fn reference_chunks(
&mut self,
name: &[u8],
range: Range<u64>,
chunk_bases: u64,
) -> Result<FastaReferenceChunks<'_, R>> {
let entry = self.index.get(name).ok_or_else(|| {
FastqError::Format(format!(
"FASTA reference not found in index: {}",
String::from_utf8_lossy(name)
))
})?;
entry.validate_range(range.clone())?;
let entry = entry.clone();
Ok(FastaReferenceChunks {
reader: self,
name: name.to_vec(),
entry,
next_offset: range.start,
end: range.end,
chunk_bases: chunk_bases.max(1),
initialized: false,
})
}
pub fn reference_chunks_into<S>(
&mut self,
name: &[u8],
range: Range<u64>,
chunk_bases: u64,
out: &mut Vec<u8>,
sink: &mut S,
) -> Result<()>
where
S: FastaReferenceChunkSink,
{
let entry = self.index.get(name).ok_or_else(|| {
FastqError::Format(format!(
"FASTA reference not found in index: {}",
String::from_utf8_lossy(name)
))
})?;
entry.validate_range(range.clone())?;
if range.start < range.end {
seek_indexed_sequence_start(&mut self.inner, entry, range.start)?;
}
let mut stream = IndexedSequenceStream {
reader: &mut self.inner,
scratch: &mut self.scratch,
entry,
next_offset: range.start,
end: range.end,
chunk_bases,
};
stream_indexed_reference_chunks_into(&mut stream, name, out, sink)
}
pub fn fetch_partition(&mut self, partition: &FastaPartition) -> Result<FastaReferenceChunk> {
self.fetch(&partition.name, partition.fetch.clone())
.map(|seq| FastaReferenceChunk {
name: partition.name.clone(),
global_offset: partition.fetch.start,
seq,
})
}
pub fn into_inner(self) -> (R, FastaIndex) {
(self.inner, self.index)
}
}
#[cfg(feature = "bgzf")]
pub struct BgzfIndexedFastaReader<R> {
inner: BgzfSeekReader<R>,
fasta_index: FastaIndex,
bgzf_index: BgzfIndex,
scratch: Vec<u8>,
}
#[cfg(feature = "bgzf")]
impl<R: Read + Seek> BgzfIndexedFastaReader<R> {
pub fn new(inner: R, fasta_index: FastaIndex, bgzf_index: BgzfIndex) -> Self {
Self {
inner: BgzfSeekReader::new(inner),
fasta_index,
bgzf_index,
scratch: Vec::new(),
}
}
pub fn fasta_index(&self) -> &FastaIndex {
&self.fasta_index
}
pub fn bgzf_index(&self) -> &BgzfIndex {
&self.bgzf_index
}
pub fn fetch_into(&mut self, name: &[u8], range: Range<u64>, out: &mut Vec<u8>) -> Result<()> {
let entry = self.fasta_index.get(name).ok_or_else(|| {
FastqError::Format(format!(
"FASTA reference not found in index: {}",
String::from_utf8_lossy(name)
))
})?;
let window = indexed_physical_window(entry, range.clone())?;
out.clear();
let expected_bases = usize::try_from(range.end - range.start).map_err(|_| {
FastqError::Format("FASTA fetch range length exceeds usize range".into())
})?;
out.reserve(expected_bases);
if expected_bases == 0 {
return Ok(());
}
let virtual_offset = self
.bgzf_index
.virtual_offset_for_uncompressed_offset(window.start)?
.ok_or_else(|| FastqError::Bgzf("BGZF span offset is not indexed".into()))?;
self.inner.seek_virtual_offset(virtual_offset)?;
copy_indexed_sequence_window(
&mut self.inner,
entry,
window.start,
window.end - window.start,
expected_bases,
out,
&mut self.scratch,
)?;
Ok(())
}
pub fn fetch(&mut self, name: &[u8], range: Range<u64>) -> Result<Vec<u8>> {
let mut out = Vec::new();
self.fetch_into(name, range, &mut out)?;
Ok(out)
}
pub fn reference_chunks(
&mut self,
name: &[u8],
range: Range<u64>,
chunk_bases: u64,
) -> Result<BgzfFastaReferenceChunks<'_, R>> {
let entry = self.fasta_index.get(name).ok_or_else(|| {
FastqError::Format(format!(
"FASTA reference not found in index: {}",
String::from_utf8_lossy(name)
))
})?;
entry.validate_range(range.clone())?;
let entry = entry.clone();
Ok(BgzfFastaReferenceChunks {
reader: self,
name: name.to_vec(),
entry,
next_offset: range.start,
end: range.end,
chunk_bases: chunk_bases.max(1),
initialized: false,
})
}
pub fn reference_chunks_into<S>(
&mut self,
name: &[u8],
range: Range<u64>,
chunk_bases: u64,
out: &mut Vec<u8>,
sink: &mut S,
) -> Result<()>
where
S: FastaReferenceChunkSink,
{
let entry = self.fasta_index.get(name).ok_or_else(|| {
FastqError::Format(format!(
"FASTA reference not found in index: {}",
String::from_utf8_lossy(name)
))
})?;
entry.validate_range(range.clone())?;
if range.start < range.end {
seek_bgzf_indexed_sequence_start(
&mut self.inner,
&self.bgzf_index,
entry,
range.start,
)?;
}
let mut stream = IndexedSequenceStream {
reader: &mut self.inner,
scratch: &mut self.scratch,
entry,
next_offset: range.start,
end: range.end,
chunk_bases,
};
stream_indexed_reference_chunks_into(&mut stream, name, out, sink)
}
pub fn fetch_partition(&mut self, partition: &FastaPartition) -> Result<FastaReferenceChunk> {
self.fetch(&partition.name, partition.fetch.clone())
.map(|seq| FastaReferenceChunk {
name: partition.name.clone(),
global_offset: partition.fetch.start,
seq,
})
}
}
impl Default for FastaStats {
fn default() -> Self {
Self {
records: 0,
bases: 0,
checksum: FASTA_STATS_CHECKSUM_INIT,
}
}
}
impl FastaStats {
pub fn observe_sequence(&mut self, seq: &[u8]) {
self.observe_sequence_parts(
seq.len() as u64,
seq.first().copied().unwrap_or_default(),
seq.last().copied().unwrap_or_default(),
);
}
fn observe_sequence_parts(&mut self, len: u64, first: u8, last: u8) {
self.records += 1;
self.bases += len;
self.checksum ^= len;
self.checksum = self
.checksum
.rotate_left(5)
.wrapping_mul(0x0000_0100_0000_01b3);
self.checksum ^= first as u64;
self.checksum ^= (last as u64) << 8;
}
}
pub trait FastaRecordSink {
fn record(&mut self, record: FastaVisitRecord<'_>) -> Result<()>;
}
impl<F> FastaRecordSink for F
where
F: FnMut(FastaVisitRecord<'_>) -> Result<()>,
{
fn record(&mut self, record: FastaVisitRecord<'_>) -> Result<()> {
self(record)
}
}
#[derive(Debug)]
pub struct FastaBatch<'a> {
bytes: &'a [u8],
records: &'a [FastaRecordRef],
id_tokens: &'a [ByteRange],
first_record_index: u64,
}
impl<'a> FastaBatch<'a> {
pub fn len(&self) -> usize {
self.records.len()
}
pub fn is_empty(&self) -> bool {
self.records.is_empty()
}
pub fn bytes(&self) -> &'a [u8] {
self.bytes
}
pub fn record_refs(&self) -> &'a [FastaRecordRef] {
self.records
}
pub fn first_record_index(&self) -> u64 {
self.first_record_index
}
pub fn records(&self) -> impl Iterator<Item = FastaRecord<'_>> {
self.records
.iter()
.zip(self.id_tokens.iter().cloned())
.map(|(record, id_token)| FastaRecord {
bytes: self.bytes,
record,
id_token,
})
}
pub fn to_owned_batch(&self) -> OwnedFastaBatch {
OwnedFastaBatch {
bytes: self.bytes.to_vec(),
records: self.records.to_vec(),
id_tokens: self.id_tokens.to_vec(),
first_record_index: self.first_record_index,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OwnedFastaBatch {
bytes: Vec<u8>,
records: Vec<FastaRecordRef>,
id_tokens: Vec<ByteRange>,
first_record_index: u64,
}
impl OwnedFastaBatch {
pub fn len(&self) -> usize {
self.records.len()
}
pub fn is_empty(&self) -> bool {
self.records.is_empty()
}
pub fn bytes(&self) -> &[u8] {
&self.bytes
}
pub fn record_refs(&self) -> &[FastaRecordRef] {
&self.records
}
pub fn first_record_index(&self) -> u64 {
self.first_record_index
}
pub fn records(&self) -> impl Iterator<Item = OwnedFastaRecord<'_>> {
self.records
.iter()
.zip(self.id_tokens.iter().cloned())
.map(|(record, id_token)| OwnedFastaRecord {
bytes: &self.bytes,
record,
id_token,
})
}
}
#[derive(Debug, Clone, Copy)]
pub struct OwnedFastaRecord<'a> {
bytes: &'a [u8],
record: &'a FastaRecordRef,
id_token: ByteRange,
}
impl<'a> OwnedFastaRecord<'a> {
pub fn name(self) -> &'a [u8] {
&self.bytes[to_usize(self.record.name.clone())]
}
pub fn name_without_gt(self) -> &'a [u8] {
let name = self.name();
name.strip_prefix(b">").unwrap_or(name)
}
pub fn id_token(self) -> &'a [u8] {
&self.bytes[self.id_token.to_usize()]
}
pub fn seq(self) -> &'a [u8] {
&self.bytes[to_usize(self.record.seq.clone())]
}
}
#[derive(Debug)]
pub struct FastaReader<R> {
reader: BufReader<R>,
config: FastaConfig,
bytes: Vec<u8>,
records: Vec<FastaRecordRef>,
id_tokens: Vec<ByteRange>,
line: Vec<u8>,
pending_header: Option<PendingHeader>,
byte_offset: u64,
record_index: u64,
eof: bool,
}
#[derive(Debug)]
struct PendingHeader {
bytes: Vec<u8>,
byte_offset: u64,
}
impl<R: Read> FastaReader<R> {
pub fn new(reader: R) -> Self {
Self::with_config(reader, FastaConfig::default())
}
pub fn with_config(reader: R, config: FastaConfig) -> Self {
let config = FastaConfig {
batch_records: config.batch_records.max(1),
buffer_size: config.buffer_size.max(1024),
expected_seq_len: config.expected_seq_len,
};
let seq_capacity = config
.batch_records
.saturating_mul(config.expected_seq_len)
.min(8 * 1024 * 1024);
Self {
reader: BufReader::with_capacity(config.buffer_size, reader),
bytes: Vec::with_capacity(seq_capacity),
records: Vec::with_capacity(config.batch_records),
id_tokens: Vec::with_capacity(config.batch_records),
config,
line: Vec::new(),
pending_header: None,
byte_offset: 0,
record_index: 0,
eof: false,
}
}
pub fn into_inner(self) -> R {
self.reader.into_inner()
}
pub fn next_batch(&mut self) -> Result<Option<FastaBatch<'_>>> {
self.bytes.clear();
self.records.clear();
self.id_tokens.clear();
let first_record_index = self.record_index;
while self.records.len() < self.config.batch_records {
let Some(header) = self.next_header()? else {
break;
};
self.push_record(header)?;
}
if self.records.is_empty() {
return Ok(None);
}
self.record_index += self.records.len() as u64;
Ok(Some(FastaBatch {
bytes: &self.bytes,
records: &self.records,
id_tokens: &self.id_tokens,
first_record_index,
}))
}
pub fn next_owned_batch(&mut self) -> Result<Option<OwnedFastaBatch>> {
self.next_batch()
.map(|batch| batch.map(|batch| batch.to_owned_batch()))
}
pub fn visit_records<F>(&mut self, mut visit: F) -> Result<()>
where
F: FnMut(FastaVisitRecord<'_>) -> Result<()>,
{
self.visit_records_with_sink(&mut visit)
}
pub fn visit_records_with_sink<S>(&mut self, sink: &mut S) -> Result<()>
where
S: FastaRecordSink,
{
while let Some(header) = self.next_header()? {
self.visit_record(header, sink)?;
self.record_index += 1;
}
Ok(())
}
pub fn stats(&mut self) -> Result<FastaStats> {
let mut stats = FastaStats::default();
while let Some(header) = self.next_header()? {
self.observe_record_stats(header, &mut stats)?;
self.record_index += 1;
}
Ok(stats)
}
fn observe_record_stats(
&mut self,
_header: PendingHeader,
stats: &mut FastaStats,
) -> Result<()> {
let record_index = self.record_index;
let mut bases = 0_u64;
let mut first = 0_u8;
let mut last = 0_u8;
loop {
let line_start = self.byte_offset;
let n = self.read_line()?;
if n == 0 {
self.eof = true;
break;
}
let trimmed = trim_line(&self.line);
if trimmed.starts_with(b">") {
validate_header(trimmed, line_start, record_index + 1)?;
self.pending_header = Some(PendingHeader {
bytes: trimmed.to_vec(),
byte_offset: line_start,
});
break;
}
if trimmed.is_empty() {
continue;
}
if bases == 0 {
first = trimmed[0];
}
last = *trimmed.last().unwrap_or(&0);
bases += trimmed.len() as u64;
}
stats.observe_sequence_parts(bases, first, last);
Ok(())
}
fn visit_record<S>(&mut self, header: PendingHeader, sink: &mut S) -> Result<()>
where
S: FastaRecordSink,
{
let record_index = self.record_index;
self.bytes.clear();
loop {
let line_start = self.byte_offset;
let n = self.read_line()?;
if n == 0 {
self.eof = true;
break;
}
let trimmed = trim_line(&self.line);
if trimmed.starts_with(b">") {
validate_header(trimmed, line_start, record_index + 1)?;
self.pending_header = Some(PendingHeader {
bytes: trimmed.to_vec(),
byte_offset: line_start,
});
break;
}
if trimmed.is_empty() {
continue;
}
self.bytes.extend_from_slice(trimmed);
}
let _ = header.byte_offset;
sink.record(FastaVisitRecord {
name: &header.bytes,
seq: &self.bytes,
})
}
fn next_header(&mut self) -> Result<Option<PendingHeader>> {
if let Some(header) = self.pending_header.take() {
return Ok(Some(header));
}
if self.eof {
return Ok(None);
}
loop {
let line_start = self.byte_offset;
let n = self.read_line()?;
if n == 0 {
self.eof = true;
return Ok(None);
}
let trimmed = trim_line(&self.line);
if trimmed.is_empty() {
continue;
}
if !trimmed.starts_with(b">") {
return Err(format_at(
"FASTA record header must start with `>`",
line_start,
self.record_index,
));
}
validate_header(trimmed, line_start, self.record_index)?;
return Ok(Some(PendingHeader {
bytes: trimmed.to_vec(),
byte_offset: line_start,
}));
}
}
fn push_record(&mut self, header: PendingHeader) -> Result<()> {
let record_index = self.record_index + self.records.len() as u64;
let name_start = checked_u32(self.bytes.len())?;
self.bytes.extend_from_slice(&header.bytes);
let name_end = checked_u32(self.bytes.len())?;
let id_token = id_token_range(&self.bytes, name_start..name_end)?;
let seq_start = checked_u32(self.bytes.len())?;
loop {
let line_start = self.byte_offset;
let n = self.read_line()?;
if n == 0 {
self.eof = true;
break;
}
let trimmed = trim_line(&self.line);
if trimmed.starts_with(b">") {
validate_header(trimmed, line_start, record_index + 1)?;
self.pending_header = Some(PendingHeader {
bytes: trimmed.to_vec(),
byte_offset: line_start,
});
break;
}
if trimmed.is_empty() {
continue;
}
self.bytes.extend_from_slice(trimmed);
}
let seq_end = checked_u32(self.bytes.len())?;
self.records.push(FastaRecordRef {
name: name_start..name_end,
seq: seq_start..seq_end,
});
self.id_tokens.push(id_token);
let _ = header.byte_offset;
Ok(())
}
fn read_line(&mut self) -> Result<usize> {
self.line.clear();
let n = self.reader.read_until(b'\n', &mut self.line)?;
self.byte_offset += n as u64;
Ok(n)
}
}
pub fn count_fasta_read<R: Read>(reader: R) -> Result<FastaStats> {
let mut reader = BufReader::with_capacity(DEFAULT_READER_BUFFER_SIZE, reader);
count_fasta_bufread(&mut reader)
}
fn count_fasta_bufread<R: BufRead>(reader: &mut R) -> Result<FastaStats> {
let mut scanner = FastaCountScanner::default();
let mut carry = Vec::new();
let mut carry_start = 0_u64;
let mut byte_offset = 0_u64;
loop {
let available = reader.fill_buf()?;
if available.is_empty() {
break;
}
let mut consumed = 0;
while consumed < available.len() {
let line_start = byte_offset + consumed as u64;
let Some(relative_newline) = memchr(b'\n', &available[consumed..]) else {
if carry.is_empty() {
carry_start = line_start;
}
carry.extend_from_slice(&available[consumed..]);
consumed = available.len();
break;
};
let line_end = consumed + relative_newline;
if carry.is_empty() {
scanner.observe_line(trim_line(&available[consumed..line_end]), line_start)?;
} else {
carry.extend_from_slice(&available[consumed..line_end]);
scanner.observe_line(trim_line(&carry), carry_start)?;
carry.clear();
}
consumed = line_end + 1;
}
reader.consume(consumed);
byte_offset += consumed as u64;
}
if !carry.is_empty() {
scanner.observe_line(trim_line(&carry), carry_start)?;
}
scanner.finish()
}
#[derive(Debug, Default)]
struct FastaCountScanner {
stats: FastaStats,
in_record: bool,
record_bases: u64,
first_base: u8,
last_base: u8,
record_index: u64,
}
impl FastaCountScanner {
fn observe_line(&mut self, line: &[u8], line_start: u64) -> Result<()> {
if line.is_empty() {
return Ok(());
}
if line.starts_with(b">") {
self.finish_record();
validate_header(line, line_start, self.record_index)?;
self.in_record = true;
return Ok(());
}
if !self.in_record {
return Err(format_at(
"FASTA record header must start with `>`",
line_start,
self.record_index,
));
}
if self.record_bases == 0 {
self.first_base = line[0];
}
self.last_base = line.last().copied().unwrap_or_default();
self.record_bases += line.len() as u64;
Ok(())
}
fn finish(mut self) -> Result<FastaStats> {
self.finish_record();
Ok(self.stats)
}
fn finish_record(&mut self) {
if !self.in_record {
return;
}
self.stats
.observe_sequence_parts(self.record_bases, self.first_base, self.last_base);
self.record_index += 1;
self.in_record = false;
self.record_bases = 0;
self.first_base = 0;
self.last_base = 0;
}
}
pub fn count_fasta_bytes(bytes: &[u8]) -> Result<FastaStats> {
let mut scanner = FastaCountScanner::default();
let mut cursor = 0;
while cursor < bytes.len() {
let line_start = cursor as u64;
let line = next_trimmed_line(bytes, &mut cursor);
scanner.observe_line(line, line_start)?;
}
scanner.finish()
}
pub fn build_fasta_index<R: Read>(reader: R) -> Result<FastaIndex> {
let mut builder = FastaIndexBuilder::default();
let mut reader = BufReader::new(reader);
build_fasta_index_bufread(&mut reader, &mut builder)?;
Ok(builder.finish())
}
#[cfg(feature = "bgzf")]
pub fn build_fasta_index_bgzf<R: Read>(reader: R) -> Result<FastaIndex> {
let mut block_reader = BgzfDecodedBlockReader::new(reader);
let mut builder = FastaIndexBuilder::default();
let mut bgzf_entries = Vec::new();
let mut line = Vec::new();
while let Some(block) = block_reader.next_block()? {
bgzf_entries.push(block.index_entry());
observe_fasta_index_chunk(block.bytes(), &mut line, &mut builder)?;
}
if !line.is_empty() {
builder.observe_physical_line(&line)?;
}
builder.finish_current()?;
let mut index = builder.finish();
for entry in &mut index.entries {
entry.virtual_offset = bgzf_virtual_offset_for(&bgzf_entries, entry.offset)?;
}
Ok(index)
}
#[cfg(feature = "bgzf")]
fn bgzf_virtual_offset_for(
entries: &[BgzfIndexEntry],
offset: u64,
) -> Result<Option<BgzfVirtualOffset>> {
let idx = entries.partition_point(|entry| entry.uncompressed_offset <= offset);
let Some(entry) = idx.checked_sub(1).and_then(|idx| entries.get(idx)) else {
return Ok(None);
};
entry.virtual_offset_for(offset)
}
#[derive(Debug, Default)]
struct FastaIndexBuilder {
entries: Vec<FastaIndexEntry>,
seen_names: HashSet<Vec<u8>>,
current: Option<FastaIndexRecord>,
byte_offset: u64,
}
#[derive(Debug)]
struct FastaIndexRecord {
name: Vec<u8>,
len: u64,
offset: Option<u64>,
line_bases: Option<u64>,
line_width: Option<u64>,
last_line_bases: Option<u64>,
last_line_width: Option<u64>,
record_index: u64,
}
fn build_fasta_index_bufread<R: BufRead>(
reader: &mut R,
builder: &mut FastaIndexBuilder,
) -> Result<()> {
let mut line = Vec::new();
loop {
let available = reader.fill_buf()?;
if available.is_empty() {
break;
}
let consumed = observe_fasta_index_chunk(available, &mut line, builder)?;
reader.consume(consumed);
}
if !line.is_empty() {
builder.observe_physical_line(&line)?;
}
builder.finish_current()?;
Ok(())
}
fn observe_fasta_index_chunk(
bytes: &[u8],
carry: &mut Vec<u8>,
builder: &mut FastaIndexBuilder,
) -> Result<usize> {
let mut consumed = 0;
while consumed < bytes.len() {
let Some(relative_newline) = memchr(b'\n', &bytes[consumed..]) else {
carry.extend_from_slice(&bytes[consumed..]);
return Ok(bytes.len());
};
let line_end = consumed + relative_newline + 1;
if carry.is_empty() {
builder.observe_physical_line(&bytes[consumed..line_end])?;
} else {
carry.extend_from_slice(&bytes[consumed..line_end]);
builder.observe_physical_line(carry)?;
carry.clear();
}
consumed = line_end;
}
Ok(consumed)
}
impl FastaIndexBuilder {
fn observe_physical_line(&mut self, line: &[u8]) -> Result<()> {
let line_start = self.byte_offset;
self.byte_offset += line.len() as u64;
let trimmed = trim_line(line);
if trimmed.is_empty() {
self.observe_blank(line_start)
} else if trimmed.starts_with(b">") {
self.start_record(trimmed, line_start)
} else {
self.observe_sequence_line(trimmed.len() as u64, line.len() as u64, line_start)
}
}
fn start_record(&mut self, header: &[u8], byte_offset: u64) -> Result<()> {
self.finish_current()?;
validate_header(header, byte_offset, self.entries.len() as u64)?;
let name = fasta_index_name(header);
if self.seen_names.contains(name) {
return Err(format_at(
"duplicate FASTA index reference name",
byte_offset,
self.entries.len() as u64,
));
}
self.seen_names.insert(name.to_vec());
self.current = Some(FastaIndexRecord {
name: name.to_vec(),
len: 0,
offset: None,
line_bases: None,
line_width: None,
last_line_bases: None,
last_line_width: None,
record_index: self.entries.len() as u64,
});
Ok(())
}
fn observe_blank(&mut self, byte_offset: u64) -> Result<()> {
if self
.current
.as_ref()
.and_then(|record| record.offset)
.is_some()
{
return Err(format_at(
"FASTA index does not support blank sequence lines",
byte_offset,
self.current
.as_ref()
.map_or(0, |record| record.record_index),
));
}
Ok(())
}
fn observe_sequence_line(&mut self, bases: u64, width: u64, byte_offset: u64) -> Result<()> {
let Some(record) = self.current.as_mut() else {
return Err(format_at(
"FASTA record header must start with `>`",
byte_offset,
self.entries.len() as u64,
));
};
if record.offset.is_none() {
record.offset = Some(byte_offset);
record.line_bases = Some(bases);
record.line_width = Some(width);
} else if let (Some(last_bases), Some(last_width)) =
(record.last_line_bases, record.last_line_width)
{
let expected_bases = record.line_bases.unwrap_or(last_bases);
let expected_width = record.line_width.unwrap_or(last_width);
if bases > expected_bases {
return Err(format_at(
"FASTA final sequence line is longer than the first sequence line",
byte_offset,
record.record_index,
));
}
if last_bases != expected_bases || last_width != expected_width {
return Err(format_at(
"non-final FASTA sequence line has inconsistent wrapping",
byte_offset,
record.record_index,
));
}
}
record.len += bases;
record.last_line_bases = Some(bases);
record.last_line_width = Some(width);
Ok(())
}
fn finish_current(&mut self) -> Result<()> {
let Some(record) = self.current.take() else {
return Ok(());
};
let offset = record.offset.unwrap_or(self.byte_offset);
let line_bases = record.line_bases.unwrap_or(0);
let line_width = record.line_width.unwrap_or(0);
self.entries.push(FastaIndexEntry {
name: record.name,
len: record.len,
offset,
line_bases,
line_width,
#[cfg(feature = "bgzf")]
virtual_offset: None,
});
Ok(())
}
fn finish(mut self) -> FastaIndex {
let name_to_index = self
.entries
.iter()
.enumerate()
.map(|(idx, entry)| (entry.name.clone(), idx))
.collect();
FastaIndex {
entries: std::mem::take(&mut self.entries),
name_to_index,
}
}
}
fn fasta_index_name(header: &[u8]) -> &[u8] {
let name = header.strip_prefix(b">").unwrap_or(header);
let end = name
.iter()
.position(u8::is_ascii_whitespace)
.unwrap_or(name.len());
&name[..end]
}
pub fn visit_fasta_bytes<F>(bytes: &[u8], mut visit: F) -> Result<u64>
where
F: FnMut(FastaVisitRecord<'_>) -> Result<()>,
{
let mut cursor = 0;
let mut pending_header = None;
let mut record_index = 0;
let mut folded = Vec::new();
while let Some((header_offset, header)) =
take_next_header(bytes, &mut cursor, &mut pending_header, record_index)?
{
validate_header(header, header_offset, record_index)?;
folded.clear();
let mut first_seq = None;
loop {
if cursor >= bytes.len() {
break;
}
let line_start = cursor as u64;
let line = next_trimmed_line(bytes, &mut cursor);
if line.starts_with(b">") {
validate_header(line, line_start, record_index + 1)?;
pending_header = Some((line_start, line));
break;
}
if line.is_empty() {
continue;
}
first_seq = Some(line);
break;
}
let Some(first_seq) = first_seq else {
visit(FastaVisitRecord {
name: header,
seq: b"",
})?;
record_index += 1;
continue;
};
if cursor >= bytes.len() || bytes[cursor] == b'>' {
visit(FastaVisitRecord {
name: header,
seq: first_seq,
})?;
record_index += 1;
continue;
}
let mut seq_line_count = 1_usize;
loop {
if cursor >= bytes.len() {
break;
}
let line_start = cursor as u64;
let line = next_trimmed_line(bytes, &mut cursor);
if line.starts_with(b">") {
validate_header(line, line_start, record_index + 1)?;
pending_header = Some((line_start, line));
break;
}
if line.is_empty() {
continue;
}
seq_line_count += 1;
if seq_line_count == 2 {
folded.extend_from_slice(first_seq);
}
folded.extend_from_slice(line);
}
let seq = if seq_line_count == 1 {
first_seq
} else {
&folded
};
visit(FastaVisitRecord { name: header, seq })?;
record_index += 1;
}
Ok(record_index)
}
pub fn detect_fasta_shape(bytes: &[u8]) -> Result<FastaShape> {
let mut cursor = 0;
let mut saw_record = false;
while cursor < bytes.len() {
let header_offset = cursor as u64;
let header = next_trimmed_line(bytes, &mut cursor);
if header.is_empty() {
if saw_record {
return Ok(FastaShape::Multiline);
}
continue;
}
if !header.starts_with(b">") {
return Err(format_at(
"FASTA record header must start with `>`",
header_offset,
0,
));
}
validate_header(header, header_offset, 0)?;
saw_record = true;
let mut seq_lines = 0_usize;
while cursor < bytes.len() {
let checkpoint = cursor;
let line = next_trimmed_line(bytes, &mut cursor);
if line.starts_with(b">") {
cursor = checkpoint;
break;
}
if line.is_empty() {
return Ok(FastaShape::Multiline);
}
seq_lines += 1;
if seq_lines > 1 {
return Ok(FastaShape::Multiline);
}
}
if seq_lines != 1 {
return Ok(FastaShape::Multiline);
}
}
if saw_record {
Ok(FastaShape::TwoLine)
} else {
Ok(FastaShape::Empty)
}
}
pub fn visit_fasta_bytes_auto<F>(bytes: &[u8], visit: F) -> Result<u64>
where
F: FnMut(FastaVisitRecord<'_>) -> Result<()>,
{
match detect_fasta_shape(bytes)? {
FastaShape::TwoLine => visit_two_line_fasta_bytes(bytes, visit),
FastaShape::Empty | FastaShape::Multiline => visit_fasta_bytes(bytes, visit),
}
}
pub fn count_two_line_fasta_bytes(bytes: &[u8]) -> Result<FastaStats> {
let mut cursor = 0;
let mut record_index = 0;
let mut stats = FastaStats::default();
while cursor < bytes.len() {
let header_offset = cursor as u64;
if bytes[cursor] != b'>' {
return Err(format_at(
"two-line FASTA record header must start with `>`",
header_offset,
record_index,
));
}
let header = next_trimmed_line(bytes, &mut cursor);
validate_header(header, header_offset, record_index)?;
if cursor >= bytes.len() {
return Err(format_at(
"two-line FASTA record is missing a sequence line",
cursor as u64,
record_index,
));
}
let seq_offset = cursor as u64;
let seq = next_trimmed_line(bytes, &mut cursor);
if seq.is_empty() || seq.starts_with(b">") {
return Err(format_at(
"two-line FASTA record is missing a sequence line",
seq_offset,
record_index,
));
}
if cursor < bytes.len() && bytes[cursor] != b'>' {
return Err(format_at(
"two-line FASTA sequence must be followed by a header",
cursor as u64,
record_index + 1,
));
}
stats.observe_sequence(seq);
record_index += 1;
}
Ok(stats)
}
pub fn visit_two_line_fasta_bytes<F>(bytes: &[u8], mut visit: F) -> Result<u64>
where
F: FnMut(FastaVisitRecord<'_>) -> Result<()>,
{
let mut cursor = 0;
let mut record_index = 0;
while cursor < bytes.len() {
let header_offset = cursor as u64;
if bytes[cursor] != b'>' {
return Err(format_at(
"two-line FASTA record header must start with `>`",
header_offset,
record_index,
));
}
let header = next_trimmed_line(bytes, &mut cursor);
validate_header(header, header_offset, record_index)?;
if cursor >= bytes.len() {
return Err(format_at(
"two-line FASTA record is missing a sequence line",
cursor as u64,
record_index,
));
}
let seq_offset = cursor as u64;
let seq = next_trimmed_line(bytes, &mut cursor);
if seq.is_empty() || seq.starts_with(b">") {
return Err(format_at(
"two-line FASTA record is missing a sequence line",
seq_offset,
record_index,
));
}
if cursor < bytes.len() && bytes[cursor] != b'>' {
return Err(format_at(
"two-line FASTA sequence must be followed by a header",
cursor as u64,
record_index + 1,
));
}
visit(FastaVisitRecord { name: header, seq })?;
record_index += 1;
}
Ok(record_index)
}
pub fn visit_two_line_fasta_read<R, F>(reader: R, visit: F) -> Result<u64>
where
R: Read,
F: FnMut(FastaVisitRecord<'_>) -> Result<()>,
{
let mut reader = BufReader::with_capacity(TWO_LINE_STREAM_BUFFER_SIZE, reader);
visit_two_line_fasta_bufread(&mut reader, visit)
}
pub fn count_two_line_fasta_read<R: Read>(reader: R) -> Result<FastaStats> {
let mut reader = BufReader::with_capacity(TWO_LINE_STREAM_BUFFER_SIZE, reader);
count_two_line_fasta_bufread(&mut reader)
}
fn visit_two_line_fasta_bufread<R, F>(reader: &mut R, mut visit: F) -> Result<u64>
where
R: BufRead,
F: FnMut(FastaVisitRecord<'_>) -> Result<()>,
{
let mut state = TwoLineStreamState::Header;
let mut record_index = 0;
let mut header = Vec::new();
let mut carry = Vec::new();
loop {
let available = reader.fill_buf()?;
if available.is_empty() {
break;
}
let mut consumed = 0;
while consumed < available.len() {
let Some(relative_newline) = memchr(b'\n', &available[consumed..]) else {
carry.extend_from_slice(&available[consumed..]);
consumed = available.len();
break;
};
let line_end = consumed + relative_newline;
let line = &available[consumed..line_end];
process_two_line_stream_line(
line,
&mut carry,
&mut state,
&mut header,
&mut record_index,
&mut visit,
)?;
consumed = line_end + 1;
}
reader.consume(consumed);
}
if !carry.is_empty() {
process_two_line_stream_line(
b"",
&mut carry,
&mut state,
&mut header,
&mut record_index,
&mut visit,
)?;
}
match state {
TwoLineStreamState::Header => Ok(record_index),
TwoLineStreamState::Seq => Err(format_at(
"two-line FASTA record is missing a sequence line",
0,
record_index,
)),
}
}
fn count_two_line_fasta_bufread<R: BufRead>(reader: &mut R) -> Result<FastaStats> {
let mut stats = FastaStats::default();
let mut state = TwoLineStreamState::Header;
let mut record_index = 0;
let mut carry = Vec::new();
loop {
let available = reader.fill_buf()?;
if available.is_empty() {
break;
}
let mut consumed = 0;
while consumed < available.len() {
let Some(relative_newline) = memchr(b'\n', &available[consumed..]) else {
carry.extend_from_slice(&available[consumed..]);
consumed = available.len();
break;
};
let line_end = consumed + relative_newline;
let line = &available[consumed..line_end];
process_two_line_count_line(
line,
&mut carry,
&mut state,
&mut record_index,
&mut stats,
)?;
consumed = line_end + 1;
}
reader.consume(consumed);
}
if !carry.is_empty() {
process_two_line_count_line(b"", &mut carry, &mut state, &mut record_index, &mut stats)?;
}
if state == TwoLineStreamState::Seq {
return Err(format_at(
"two-line FASTA record is missing a sequence line",
0,
record_index,
));
}
Ok(stats)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum TwoLineStreamState {
Header,
Seq,
}
fn process_two_line_stream_line<F>(
line: &[u8],
carry: &mut Vec<u8>,
state: &mut TwoLineStreamState,
header: &mut Vec<u8>,
record_index: &mut u64,
visit: &mut F,
) -> Result<()>
where
F: FnMut(FastaVisitRecord<'_>) -> Result<()>,
{
if carry.is_empty() {
process_complete_two_line_stream_line(line, state, header, record_index, visit)
} else {
carry.extend_from_slice(line);
let owned_line = trim_line(carry);
process_complete_two_line_stream_line(owned_line, state, header, record_index, visit)?;
carry.clear();
Ok(())
}
}
fn process_complete_two_line_stream_line<F>(
line: &[u8],
state: &mut TwoLineStreamState,
header: &mut Vec<u8>,
record_index: &mut u64,
visit: &mut F,
) -> Result<()>
where
F: FnMut(FastaVisitRecord<'_>) -> Result<()>,
{
let line = trim_line(line);
if *state == TwoLineStreamState::Header {
if !line.starts_with(b">") {
return Err(format_at(
"two-line FASTA record header must start with `>`",
0,
*record_index,
));
}
validate_header(line, 0, *record_index)?;
header.clear();
header.extend_from_slice(line);
*state = TwoLineStreamState::Seq;
return Ok(());
}
if line.is_empty() || line.starts_with(b">") {
return Err(format_at(
"two-line FASTA record is missing a sequence line",
0,
*record_index,
));
}
visit(FastaVisitRecord {
name: header,
seq: line,
})?;
*record_index += 1;
*state = TwoLineStreamState::Header;
Ok(())
}
fn process_two_line_count_line(
line: &[u8],
carry: &mut Vec<u8>,
state: &mut TwoLineStreamState,
record_index: &mut u64,
stats: &mut FastaStats,
) -> Result<()> {
if carry.is_empty() {
process_complete_two_line_count_line(line, state, record_index, stats)
} else {
carry.extend_from_slice(line);
let owned_line = trim_line(carry);
process_complete_two_line_count_line(owned_line, state, record_index, stats)?;
carry.clear();
Ok(())
}
}
fn process_complete_two_line_count_line(
line: &[u8],
state: &mut TwoLineStreamState,
record_index: &mut u64,
stats: &mut FastaStats,
) -> Result<()> {
let line = trim_line(line);
if *state == TwoLineStreamState::Header {
if !line.starts_with(b">") {
return Err(format_at(
"two-line FASTA record header must start with `>`",
0,
*record_index,
));
}
validate_header(line, 0, *record_index)?;
*state = TwoLineStreamState::Seq;
return Ok(());
}
if line.is_empty() || line.starts_with(b">") {
return Err(format_at(
"two-line FASTA record is missing a sequence line",
0,
*record_index,
));
}
stats.observe_sequence(line);
*record_index += 1;
*state = TwoLineStreamState::Header;
Ok(())
}
fn take_next_header<'a>(
bytes: &'a [u8],
cursor: &mut usize,
pending_header: &mut Option<(u64, &'a [u8])>,
record_index: u64,
) -> Result<Option<(u64, &'a [u8])>> {
if let Some(header) = pending_header.take() {
return Ok(Some(header));
}
while *cursor < bytes.len() {
let line_start = *cursor as u64;
let line = next_trimmed_line(bytes, cursor);
if line.is_empty() {
continue;
}
if !line.starts_with(b">") {
return Err(format_at(
"FASTA record header must start with `>`",
line_start,
record_index,
));
}
return Ok(Some((line_start, line)));
}
Ok(None)
}
fn next_trimmed_line<'a>(bytes: &'a [u8], cursor: &mut usize) -> &'a [u8] {
let start = *cursor;
match memchr(b'\n', &bytes[start..]) {
Some(relative) => {
let end = start + relative;
*cursor = end + 1;
trim_line(&bytes[start..end])
}
None => {
*cursor = bytes.len();
trim_line(&bytes[start..])
}
}
}
fn trim_line(line: &[u8]) -> &[u8] {
let line = line.strip_suffix(b"\n").unwrap_or(line);
line.strip_suffix(b"\r").unwrap_or(line)
}
fn validate_header(header: &[u8], byte_offset: u64, record_index: u64) -> Result<()> {
if header.len() == 1 || fasta_index_name(header).is_empty() {
return Err(format_at("empty FASTA id", byte_offset, record_index));
}
Ok(())
}
fn id_token_range(bytes: &[u8], name: Range<u32>) -> Result<ByteRange> {
let name_range = to_usize(name.clone());
let header = bytes
.get(name_range)
.ok_or_else(|| FastqError::Format("FASTA header byte range exceeds batch buffer".into()))?;
let without_gt = header.strip_prefix(b">").unwrap_or(header);
let token_end = without_gt
.iter()
.position(u8::is_ascii_whitespace)
.unwrap_or(without_gt.len());
let start = name.start + u32::from(header.starts_with(b">"));
let token_end_u32 = u32::try_from(token_end)
.map_err(|_| FastqError::Format("FASTA id token range exceeds u32 range".into()))?;
let end = start
.checked_add(token_end_u32)
.ok_or_else(|| FastqError::Format("FASTA id token range overflowed".into()))?;
Ok(ByteRange { start, end })
}
fn checked_u32(value: usize) -> Result<u32> {
u32::try_from(value)
.map_err(|_| FastqError::Format("FASTA batch byte offsets exceed u32 range".into()))
}
fn format_at(message: impl Into<String>, byte_offset: u64, record_index: u64) -> FastqError {
FastqError::FormatAt {
message: message.into(),
position: FastqPosition::new(byte_offset, record_index, 0),
}
}
fn to_usize(range: Range<u32>) -> Range<usize> {
range.start as usize..range.end as usize
}
#[cfg(test)]
mod tests;