use std::collections::HashMap;
use std::fs::File;
use std::io::{self, BufRead, BufReader, Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use crate::error::{Error, Result};
use crate::record::Sequence;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FaiRecord {
pub name: String,
pub length: u64,
pub offset: u64,
pub line_bases: u64,
pub line_width: u64,
}
impl FaiRecord {
fn byte_offset(&self, base: u64) -> u64 {
let full_lines = base / self.line_bases;
let within = base % self.line_bases;
self.offset + full_lines * self.line_width + within
}
}
#[derive(Debug, Clone, Default)]
pub struct FastaIndex {
records: Vec<FaiRecord>,
by_name: HashMap<String, usize>,
}
impl FastaIndex {
pub fn build<R: Read>(reader: R) -> Result<FastaIndex> {
let mut reader = BufReader::with_capacity(128 * 1024, reader);
let mut index = FastaIndex::default();
let mut line = Vec::new();
let mut offset: u64 = 0;
let mut current: Option<(FaiRecord, bool)> = None;
loop {
line.clear();
let read = reader.read_until(b'\n', &mut line)?;
if read == 0 {
break;
}
let content = trim_newline(&line);
if line.first() == Some(&b'>') {
if let Some((record, _)) = current.take() {
index.push(record)?;
}
let name =
String::from_utf8_lossy(crate::record::header_id(&content[1..])).into_owned();
if name.is_empty() {
return Err(Error::Index(format!(
"record header at byte {offset} has no name"
)));
}
current = Some((
FaiRecord {
name,
length: 0,
offset: offset + read as u64,
line_bases: 0,
line_width: 0,
},
false,
));
} else if let Some((record, saw_short_line)) = current.as_mut() {
let bases = content.len() as u64;
if record.line_bases == 0 && bases == 0 {
record.offset = offset + read as u64;
} else if record.line_bases == 0 {
record.line_bases = bases;
record.line_width = read as u64;
} else if *saw_short_line || bases > record.line_bases {
return Err(Error::Index(format!(
"sequence {:?} has different line lengths",
record.name
)));
} else if bases < record.line_bases {
*saw_short_line = true;
}
record.length += bases;
} else if !content.is_empty() {
return Err(Error::Index(format!(
"sequence data at byte {offset} precedes any header"
)));
}
offset += read as u64;
}
if let Some((record, _)) = current.take() {
index.push(record)?;
}
Ok(index)
}
pub fn build_from_path<P: AsRef<Path>>(path: P) -> Result<FastaIndex> {
let path = path.as_ref();
FastaIndex::build(open_for_scanning(path)?)
}
pub fn parse<R: Read>(reader: R) -> Result<FastaIndex> {
let reader = BufReader::new(reader);
let mut index = FastaIndex::default();
for (lineno, line) in reader.lines().enumerate() {
let line = line?;
if line.trim().is_empty() {
continue;
}
let mut fields = line.split('\t');
let mut next = |what: &str| -> Result<String> {
fields.next().map(str::to_string).ok_or_else(|| {
Error::Index(format!("line {}: missing {what} column", lineno + 1))
})
};
let name = next("name")?;
let numbers: Vec<u64> = ["length", "offset", "line_bases", "line_width"]
.into_iter()
.map(|what| -> Result<u64> {
next(what)?.trim().parse::<u64>().map_err(|e| {
Error::Index(format!("line {}: bad {what} column: {e}", lineno + 1))
})
})
.collect::<Result<Vec<u64>>>()?;
index.push(FaiRecord {
name,
length: numbers[0],
offset: numbers[1],
line_bases: numbers[2],
line_width: numbers[3],
})?;
}
Ok(index)
}
pub fn from_path<P: AsRef<Path>>(path: P) -> Result<FastaIndex> {
let path = path.as_ref();
FastaIndex::parse(
File::open(path).map_err(|e| {
Error::Io(io::Error::new(e.kind(), format!("{}: {e}", path.display())))
})?,
)
}
pub fn write<W: Write>(&self, out: &mut W) -> Result<()> {
for r in &self.records {
writeln!(
out,
"{}\t{}\t{}\t{}\t{}",
r.name, r.length, r.offset, r.line_bases, r.line_width
)?;
}
Ok(())
}
pub fn write_to_path<P: AsRef<Path>>(&self, fasta_path: P) -> Result<PathBuf> {
let target = fai_path(fasta_path.as_ref());
let mut file = File::create(&target)?;
self.write(&mut file)?;
file.flush()?;
Ok(target)
}
fn push(&mut self, record: FaiRecord) -> Result<()> {
if self.by_name.contains_key(&record.name) {
return Err(Error::Index(format!(
"duplicate sequence name {:?}",
record.name
)));
}
if record.length > 0 && record.line_bases == 0 {
return Err(Error::Index(format!(
"sequence {:?} has a zero line length",
record.name
)));
}
self.by_name.insert(record.name.clone(), self.records.len());
self.records.push(record);
Ok(())
}
pub fn get(&self, name: &str) -> Option<&FaiRecord> {
self.by_name.get(name).map(|&i| &self.records[i])
}
pub fn records(&self) -> &[FaiRecord] {
&self.records
}
pub fn len(&self) -> usize {
self.records.len()
}
pub fn is_empty(&self) -> bool {
self.records.is_empty()
}
pub fn total_length(&self) -> u64 {
self.records.iter().map(|r| r.length).sum()
}
}
pub fn fai_path(fasta: &Path) -> PathBuf {
let mut name = fasta.as_os_str().to_os_string();
name.push(".fai");
PathBuf::from(name)
}
pub trait ReadSeek: Read + Seek {}
impl<T: Read + Seek> ReadSeek for T {}
pub type BoxedSource = Box<dyn ReadSeek + Send>;
pub struct IndexedFasta<S: ReadSeek = BoxedSource> {
source: S,
index: FastaIndex,
path: PathBuf,
}
impl IndexedFasta<BoxedSource> {
pub fn open<P: AsRef<Path>>(path: P) -> Result<IndexedFasta> {
let path = path.as_ref();
let index = match FastaIndex::from_path(fai_path(path)) {
Ok(index) => index,
Err(Error::Io(e)) if e.kind() == io::ErrorKind::NotFound => {
FastaIndex::build_from_path(path)?
}
Err(e) => return Err(e),
};
IndexedFasta::with_index(path, index)
}
pub fn with_index<P: AsRef<Path>>(path: P, index: FastaIndex) -> Result<IndexedFasta> {
let path = path.as_ref().to_path_buf();
let source = open_for_seeking(&path)?;
Ok(IndexedFasta {
source,
index,
path,
})
}
}
impl<S: ReadSeek> IndexedFasta<S> {
pub fn from_source(source: S, index: FastaIndex) -> IndexedFasta<S> {
IndexedFasta {
source,
index,
path: PathBuf::from("<memory>"),
}
}
pub fn index(&self) -> &FastaIndex {
&self.index
}
pub fn path(&self) -> &Path {
&self.path
}
pub fn fetch(&mut self, name: &str) -> Result<Sequence> {
let length = self
.index
.get(name)
.ok_or_else(|| Error::UnknownSequence(name.to_string()))?
.length;
self.fetch_region(name, 0, length)
}
pub fn fetch_region(&mut self, name: &str, start: u64, end: u64) -> Result<Sequence> {
let record = self
.index
.get(name)
.ok_or_else(|| Error::UnknownSequence(name.to_string()))?
.clone();
if start > end || end > record.length {
return Err(Error::OutOfBounds {
id: name.to_string(),
start,
end,
length: record.length,
});
}
let bases = (end - start) as usize;
let mut seq = Vec::with_capacity(bases);
if bases > 0 {
let from = record.byte_offset(start);
let to = record.byte_offset(end - 1) + 1;
self.source.seek(SeekFrom::Start(from))?;
let mut raw = vec![0u8; (to - from) as usize];
self.source.read_exact(&mut raw).map_err(|e| {
if e.kind() == io::ErrorKind::UnexpectedEof {
Error::Index(format!(
"{}: index does not match the FASTA file (truncated at {name})",
self.path.display()
))
} else {
Error::Io(e)
}
})?;
seq.extend(raw.into_iter().filter(|b| *b != b'\n' && *b != b'\r'));
}
if seq.len() != bases {
return Err(Error::Index(format!(
"{}: index does not match the FASTA file (got {} of {bases} bases for {name})",
self.path.display(),
seq.len()
)));
}
let id = if start == 0 && end == record.length {
name.to_string()
} else {
format!("{name}:{}-{end}", start + 1)
};
Ok(Sequence::fasta(id, seq))
}
pub fn fetch_locus(&mut self, locus: &str) -> Result<Sequence> {
let (name, range) = match locus.rsplit_once(':') {
None => return self.fetch(locus),
Some((name, range)) => (name, range),
};
let range = range.replace("..", "-").replace(',', "");
let (start, end) = match range.split_once('-') {
Some((s, e)) => (s, e),
None => (range.as_str(), range.as_str()),
};
let parse = |text: &str, what: &str| -> Result<u64> {
text.trim()
.parse::<u64>()
.map_err(|e| Error::Index(format!("bad {what} coordinate in {locus:?}: {e}")))
};
let start = parse(start, "start")?;
let end = parse(end, "end")?;
if start == 0 {
return Err(Error::Index(format!("{locus:?}: coordinates are 1-based")));
}
self.fetch_region(name, start - 1, end)
}
}
fn open_for_scanning(path: &Path) -> Result<Box<dyn Read>> {
let file = open_file(path)?;
match probe(path)? {
Container::Plain => Ok(Box::new(file)),
#[cfg(feature = "gzip")]
Container::Bgzf => Ok(Box::new(crate::bgzf::BgzfReader::new(file)?)),
#[cfg(feature = "gzip")]
Container::PlainGzip => Err(plain_gzip_error(path)),
}
}
fn open_for_seeking(path: &Path) -> Result<BoxedSource> {
let file = open_file(path)?;
match probe(path)? {
Container::Plain => Ok(Box::new(file)),
#[cfg(feature = "gzip")]
Container::Bgzf => {
let index = match crate::bgzf::GziIndex::from_path(crate::bgzf::gzi_path(path)) {
Ok(index) => index,
Err(Error::Io(e)) if e.kind() == io::ErrorKind::NotFound => {
crate::bgzf::GziIndex::build(open_file(path)?)?
}
Err(e) => return Err(e),
};
Ok(Box::new(
crate::bgzf::BgzfReader::new(file)?.with_index(index),
))
}
#[cfg(feature = "gzip")]
Container::PlainGzip => Err(plain_gzip_error(path)),
}
}
enum Container {
Plain,
#[cfg(feature = "gzip")]
Bgzf,
#[cfg(feature = "gzip")]
PlainGzip,
}
fn probe(path: &Path) -> Result<Container> {
let mut file = open_file(path)?;
let mut head = [0u8; 128];
let mut filled = 0;
while filled < head.len() {
match file.read(&mut head[filled..]) {
Ok(0) => break,
Ok(n) => filled += n,
Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
Err(e) => return Err(Error::Io(e)),
}
}
let head = &head[..filled];
if crate::format::Compression::from_magic(head) == crate::format::Compression::None {
return Ok(Container::Plain);
}
#[cfg(feature = "gzip")]
{
if crate::bgzf::is_bgzf(head) {
Ok(Container::Bgzf)
} else {
Ok(Container::PlainGzip)
}
}
#[cfg(not(feature = "gzip"))]
Err(Error::FeatureDisabled("gzip"))
}
#[cfg(feature = "gzip")]
fn plain_gzip_error(path: &Path) -> Error {
Error::Index(format!(
"{}: this is plain gzip, which cannot be randomly accessed. \
Recompress it as BGZF (`bgzip`, or fastx's own gzip output) to index it.",
path.display()
))
}
fn open_file(path: &Path) -> Result<File> {
File::open(path)
.map_err(|e| Error::Io(io::Error::new(e.kind(), format!("{}: {e}", path.display()))))
}
fn trim_newline(line: &[u8]) -> &[u8] {
let mut end = line.len();
if end > 0 && line[end - 1] == b'\n' {
end -= 1;
}
if end > 0 && line[end - 1] == b'\r' {
end -= 1;
}
&line[..end]
}
#[cfg(test)]
mod tests {
use super::*;
const FASTA: &[u8] = b">chr1 first\nACGTACGTAC\nGGGG\n>chr2\nTTTTTTTTTT\nTTTTTTTTTT\n";
#[test]
fn builds_index_matching_samtools_layout() {
let index = FastaIndex::build(FASTA).unwrap();
assert_eq!(index.len(), 2);
let chr1 = index.get("chr1").unwrap();
assert_eq!(chr1.name, "chr1");
assert_eq!(chr1.length, 14);
assert_eq!(chr1.offset, 12); assert_eq!(chr1.line_bases, 10);
assert_eq!(chr1.line_width, 11);
let chr2 = index.get("chr2").unwrap();
assert_eq!(chr2.length, 20);
assert_eq!(chr2.offset, 12 + 11 + 5 + 6);
assert_eq!(index.total_length(), 34);
}
#[test]
fn round_trips_fai_text() {
let index = FastaIndex::build(FASTA).unwrap();
let mut text = Vec::new();
index.write(&mut text).unwrap();
assert_eq!(text, b"chr1\t14\t12\t10\t11\nchr2\t20\t34\t10\t11\n");
let reparsed = FastaIndex::parse(&text[..]).unwrap();
assert_eq!(reparsed.records(), index.records());
}
#[test]
fn rejects_ragged_lines() {
let ragged = b">a\nACGT\nAC\nACGT\n";
assert!(matches!(
FastaIndex::build(&ragged[..]),
Err(Error::Index(_))
));
assert!(FastaIndex::build(&b">a\nACGT\nAC\n"[..]).is_ok());
}
#[test]
fn blank_line_before_the_sequence_does_not_shift_the_offset() {
let index = FastaIndex::build(&b">a\n\nACGT\n>b\nAC\n"[..]).unwrap();
let a = index.get("a").unwrap();
assert_eq!((a.offset, a.length, a.line_bases), (4, 4, 4));
let index = FastaIndex::build(&b">a\n\n\n>b\nAC\n"[..]).unwrap();
assert_eq!(index.get("a").unwrap().length, 0);
assert_eq!(index.get("b").unwrap().length, 2);
}
#[test]
fn rejects_duplicate_names() {
assert!(matches!(
FastaIndex::build(&b">a\nAC\n>a\nGT\n"[..]),
Err(Error::Index(_))
));
}
#[test]
fn names_match_what_the_reader_produces() {
let fasta = ">chr\u{a0}1 description\nACGT\n";
let index = FastaIndex::build(fasta.as_bytes()).unwrap();
let records = crate::read_all_from(fasta.as_bytes()).unwrap();
assert_eq!(records[0].id, "chr\u{a0}1");
assert!(index.get(&records[0].id).is_some());
assert!(matches!(
FastaIndex::build(&b"> a\nACGT\n"[..]),
Err(Error::Index(_))
));
assert!(crate::read_all_from(&b"> a\nACGT\n"[..]).is_err());
}
#[test]
fn rejects_data_before_header() {
assert!(matches!(
FastaIndex::build(&b"ACGT\n>a\nAC\n"[..]),
Err(Error::Index(_))
));
}
#[test]
fn fetches_regions() {
let dir = std::env::temp_dir().join(format!("fastx-index-{}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
let path = dir.join("ref.fa");
std::fs::write(&path, FASTA).unwrap();
let index = FastaIndex::build_from_path(&path).unwrap();
index.write_to_path(&path).unwrap();
assert!(fai_path(&path).exists());
let mut fasta = IndexedFasta::open(&path).unwrap();
assert_eq!(fasta.fetch("chr1").unwrap().seq, b"ACGTACGTACGGGG");
assert_eq!(fasta.fetch("chr2").unwrap().seq, b"T".repeat(20));
assert_eq!(fasta.fetch_region("chr1", 0, 4).unwrap().seq, b"ACGT");
assert_eq!(fasta.fetch_region("chr1", 8, 12).unwrap().seq, b"ACGG");
assert_eq!(fasta.fetch_region("chr1", 10, 14).unwrap().seq, b"GGGG");
assert!(fasta.fetch_region("chr1", 5, 5).unwrap().seq.is_empty());
assert_eq!(fasta.fetch_locus("chr1:1-4").unwrap().seq, b"ACGT");
assert_eq!(fasta.fetch_locus("chr1:9..12").unwrap().seq, b"ACGG");
assert_eq!(fasta.fetch_locus("chr2").unwrap().seq.len(), 20);
assert_eq!(fasta.fetch_locus("chr1:1-4").unwrap().id, "chr1:1-4");
assert!(matches!(
fasta.fetch("nope"),
Err(Error::UnknownSequence(_))
));
assert!(matches!(
fasta.fetch_region("chr1", 0, 99),
Err(Error::OutOfBounds { .. })
));
assert!(fasta.fetch_locus("chr1:0-4").is_err());
assert!(fasta.fetch_locus("chr1:x-4").is_err());
std::fs::remove_dir_all(&dir).ok();
}
#[cfg(feature = "gzip")]
#[test]
fn indexes_and_fetches_from_bgzf() {
use std::io::Write;
let dir = std::env::temp_dir().join(format!("fastx-bgzf-idx-{}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
let mut fasta = Vec::new();
for chromosome in 0..3 {
fasta.extend_from_slice(format!(">chr{chromosome} test\n").as_bytes());
for line in 0..1_500 {
let base = b"ACGT"[(chromosome + line) % 4];
fasta.extend(std::iter::repeat(base).take(60));
fasta.push(b'\n');
}
}
let bgzf_path = dir.join("ref.fa.gz");
let mut writer = crate::bgzf::BgzfWriter::create(&bgzf_path).unwrap();
writer.write_all(&fasta).unwrap();
let (_, gzi) = writer.finish_with_index().unwrap();
gzi.write_to_path(&bgzf_path).unwrap();
let index = FastaIndex::build_from_path(&bgzf_path).unwrap();
let plain_index = FastaIndex::build(&fasta[..]).unwrap();
assert_eq!(index.records(), plain_index.records());
index.write_to_path(&bgzf_path).unwrap();
let plain_path = dir.join("ref.fa");
std::fs::write(&plain_path, &fasta).unwrap();
let mut compressed = IndexedFasta::open(&bgzf_path).unwrap();
let mut plain = IndexedFasta::open(&plain_path).unwrap();
for name in ["chr0", "chr1", "chr2"] {
assert_eq!(compressed.fetch(name).unwrap(), plain.fetch(name).unwrap());
for (start, end) in [(0, 1), (59, 61), (1_000, 1_100), (89_000, 90_000)] {
assert_eq!(
compressed.fetch_region(name, start, end).unwrap(),
plain.fetch_region(name, start, end).unwrap(),
"{name}:{start}-{end}"
);
}
}
std::fs::remove_file(crate::bgzf::gzi_path(&bgzf_path)).unwrap();
let mut rescanned = IndexedFasta::open(&bgzf_path).unwrap();
assert_eq!(
rescanned.fetch_region("chr1", 1_000, 1_100).unwrap(),
plain.fetch_region("chr1", 1_000, 1_100).unwrap()
);
std::fs::remove_dir_all(&dir).ok();
}
#[cfg(feature = "gzip")]
#[test]
fn refuses_to_index_plain_gzip() {
use std::io::Write;
let dir = std::env::temp_dir().join(format!("fastx-plain-gz-{}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
let path = dir.join("ref.fa.gz");
let file = std::fs::File::create(&path).unwrap();
let mut encoder = flate2::write::GzEncoder::new(file, flate2::Compression::fast());
encoder.write_all(FASTA).unwrap();
encoder.finish().unwrap();
let error = FastaIndex::build_from_path(&path).unwrap_err();
assert!(matches!(error, Error::Index(_)), "{error}");
assert!(error.to_string().contains("BGZF"), "{error}");
assert_eq!(crate::read_all(&path).unwrap().len(), 2);
std::fs::remove_dir_all(&dir).ok();
}
}