1use std::collections::HashMap;
19use std::fs::File;
20use std::io::{self, BufRead, BufReader, Read, Seek, SeekFrom, Write};
21use std::path::{Path, PathBuf};
22
23use crate::error::{Error, Result};
24use crate::record::Sequence;
25
26#[derive(Debug, Clone, PartialEq, Eq)]
28pub struct FaiRecord {
29 pub name: String,
31 pub length: u64,
33 pub offset: u64,
35 pub line_bases: u64,
37 pub line_width: u64,
39}
40
41impl FaiRecord {
42 fn byte_offset(&self, base: u64) -> u64 {
44 let full_lines = base / self.line_bases;
45 let within = base % self.line_bases;
46 self.offset + full_lines * self.line_width + within
47 }
48}
49
50#[derive(Debug, Clone, Default)]
52pub struct FastaIndex {
53 records: Vec<FaiRecord>,
54 by_name: HashMap<String, usize>,
55}
56
57impl FastaIndex {
58 pub fn build<R: Read>(reader: R) -> Result<FastaIndex> {
64 let mut reader = BufReader::with_capacity(128 * 1024, reader);
65 let mut index = FastaIndex::default();
66 let mut line = Vec::new();
67 let mut offset: u64 = 0;
68 let mut current: Option<(FaiRecord, bool)> = None;
70
71 loop {
72 line.clear();
73 let read = reader.read_until(b'\n', &mut line)?;
74 if read == 0 {
75 break;
76 }
77 let content = trim_newline(&line);
78 if line.first() == Some(&b'>') {
79 if let Some((record, _)) = current.take() {
80 index.push(record)?;
81 }
82 let name =
85 String::from_utf8_lossy(crate::record::header_id(&content[1..])).into_owned();
86 if name.is_empty() {
87 return Err(Error::Index(format!(
88 "record header at byte {offset} has no name"
89 )));
90 }
91 current = Some((
92 FaiRecord {
93 name,
94 length: 0,
95 offset: offset + read as u64,
96 line_bases: 0,
97 line_width: 0,
98 },
99 false,
100 ));
101 } else if let Some((record, saw_short_line)) = current.as_mut() {
102 let bases = content.len() as u64;
103 if record.line_bases == 0 && bases == 0 {
104 record.offset = offset + read as u64;
108 } else if record.line_bases == 0 {
109 record.line_bases = bases;
110 record.line_width = read as u64;
111 } else if *saw_short_line || bases > record.line_bases {
112 return Err(Error::Index(format!(
115 "sequence {:?} has different line lengths",
116 record.name
117 )));
118 } else if bases < record.line_bases {
119 *saw_short_line = true;
120 }
121 record.length += bases;
122 } else if !content.is_empty() {
123 return Err(Error::Index(format!(
124 "sequence data at byte {offset} precedes any header"
125 )));
126 }
127 offset += read as u64;
128 }
129 if let Some((record, _)) = current.take() {
130 index.push(record)?;
131 }
132 Ok(index)
133 }
134
135 pub fn build_from_path<P: AsRef<Path>>(path: P) -> Result<FastaIndex> {
145 let path = path.as_ref();
146 FastaIndex::build(open_for_scanning(path)?)
147 }
148
149 pub fn parse<R: Read>(reader: R) -> Result<FastaIndex> {
151 let reader = BufReader::new(reader);
152 let mut index = FastaIndex::default();
153 for (lineno, line) in reader.lines().enumerate() {
154 let line = line?;
155 if line.trim().is_empty() {
156 continue;
157 }
158 let mut fields = line.split('\t');
159 let mut next = |what: &str| -> Result<String> {
160 fields.next().map(str::to_string).ok_or_else(|| {
161 Error::Index(format!("line {}: missing {what} column", lineno + 1))
162 })
163 };
164 let name = next("name")?;
165 let numbers: Vec<u64> = ["length", "offset", "line_bases", "line_width"]
166 .into_iter()
167 .map(|what| -> Result<u64> {
168 next(what)?.trim().parse::<u64>().map_err(|e| {
169 Error::Index(format!("line {}: bad {what} column: {e}", lineno + 1))
170 })
171 })
172 .collect::<Result<Vec<u64>>>()?;
173 index.push(FaiRecord {
174 name,
175 length: numbers[0],
176 offset: numbers[1],
177 line_bases: numbers[2],
178 line_width: numbers[3],
179 })?;
180 }
181 Ok(index)
182 }
183
184 pub fn from_path<P: AsRef<Path>>(path: P) -> Result<FastaIndex> {
186 let path = path.as_ref();
187 FastaIndex::parse(
188 File::open(path).map_err(|e| {
189 Error::Io(io::Error::new(e.kind(), format!("{}: {e}", path.display())))
190 })?,
191 )
192 }
193
194 pub fn write<W: Write>(&self, out: &mut W) -> Result<()> {
196 for r in &self.records {
197 writeln!(
198 out,
199 "{}\t{}\t{}\t{}\t{}",
200 r.name, r.length, r.offset, r.line_bases, r.line_width
201 )?;
202 }
203 Ok(())
204 }
205
206 pub fn write_to_path<P: AsRef<Path>>(&self, fasta_path: P) -> Result<PathBuf> {
208 let target = fai_path(fasta_path.as_ref());
209 let mut file = File::create(&target)?;
210 self.write(&mut file)?;
211 file.flush()?;
212 Ok(target)
213 }
214
215 fn push(&mut self, record: FaiRecord) -> Result<()> {
216 if self.by_name.contains_key(&record.name) {
217 return Err(Error::Index(format!(
218 "duplicate sequence name {:?}",
219 record.name
220 )));
221 }
222 if record.length > 0 && record.line_bases == 0 {
223 return Err(Error::Index(format!(
224 "sequence {:?} has a zero line length",
225 record.name
226 )));
227 }
228 self.by_name.insert(record.name.clone(), self.records.len());
229 self.records.push(record);
230 Ok(())
231 }
232
233 pub fn get(&self, name: &str) -> Option<&FaiRecord> {
235 self.by_name.get(name).map(|&i| &self.records[i])
236 }
237
238 pub fn records(&self) -> &[FaiRecord] {
240 &self.records
241 }
242
243 pub fn len(&self) -> usize {
245 self.records.len()
246 }
247
248 pub fn is_empty(&self) -> bool {
250 self.records.is_empty()
251 }
252
253 pub fn total_length(&self) -> u64 {
255 self.records.iter().map(|r| r.length).sum()
256 }
257}
258
259pub fn fai_path(fasta: &Path) -> PathBuf {
261 let mut name = fasta.as_os_str().to_os_string();
262 name.push(".fai");
263 PathBuf::from(name)
264}
265
266pub trait ReadSeek: Read + Seek {}
269
270impl<T: Read + Seek> ReadSeek for T {}
271
272pub type BoxedSource = Box<dyn ReadSeek + Send>;
274
275pub struct IndexedFasta<S: ReadSeek = BoxedSource> {
281 source: S,
282 index: FastaIndex,
283 path: PathBuf,
284}
285
286impl IndexedFasta<BoxedSource> {
287 pub fn open<P: AsRef<Path>>(path: P) -> Result<IndexedFasta> {
294 let path = path.as_ref();
295 let index = match FastaIndex::from_path(fai_path(path)) {
296 Ok(index) => index,
297 Err(Error::Io(e)) if e.kind() == io::ErrorKind::NotFound => {
298 FastaIndex::build_from_path(path)?
299 }
300 Err(e) => return Err(e),
301 };
302 IndexedFasta::with_index(path, index)
303 }
304
305 pub fn with_index<P: AsRef<Path>>(path: P, index: FastaIndex) -> Result<IndexedFasta> {
307 let path = path.as_ref().to_path_buf();
308 let source = open_for_seeking(&path)?;
309 Ok(IndexedFasta {
310 source,
311 index,
312 path,
313 })
314 }
315}
316
317impl<S: ReadSeek> IndexedFasta<S> {
318 pub fn from_source(source: S, index: FastaIndex) -> IndexedFasta<S> {
320 IndexedFasta {
321 source,
322 index,
323 path: PathBuf::from("<memory>"),
324 }
325 }
326
327 pub fn index(&self) -> &FastaIndex {
329 &self.index
330 }
331
332 pub fn path(&self) -> &Path {
334 &self.path
335 }
336
337 pub fn fetch(&mut self, name: &str) -> Result<Sequence> {
339 let length = self
340 .index
341 .get(name)
342 .ok_or_else(|| Error::UnknownSequence(name.to_string()))?
343 .length;
344 self.fetch_region(name, 0, length)
345 }
346
347 pub fn fetch_region(&mut self, name: &str, start: u64, end: u64) -> Result<Sequence> {
352 let record = self
353 .index
354 .get(name)
355 .ok_or_else(|| Error::UnknownSequence(name.to_string()))?
356 .clone();
357 if start > end || end > record.length {
358 return Err(Error::OutOfBounds {
359 id: name.to_string(),
360 start,
361 end,
362 length: record.length,
363 });
364 }
365 let bases = (end - start) as usize;
366 let mut seq = Vec::with_capacity(bases);
367 if bases > 0 {
368 let from = record.byte_offset(start);
369 let to = record.byte_offset(end - 1) + 1;
370 self.source.seek(SeekFrom::Start(from))?;
371 let mut raw = vec![0u8; (to - from) as usize];
372 self.source.read_exact(&mut raw).map_err(|e| {
373 if e.kind() == io::ErrorKind::UnexpectedEof {
374 Error::Index(format!(
375 "{}: index does not match the FASTA file (truncated at {name})",
376 self.path.display()
377 ))
378 } else {
379 Error::Io(e)
380 }
381 })?;
382 seq.extend(raw.into_iter().filter(|b| *b != b'\n' && *b != b'\r'));
383 }
384 if seq.len() != bases {
385 return Err(Error::Index(format!(
386 "{}: index does not match the FASTA file (got {} of {bases} bases for {name})",
387 self.path.display(),
388 seq.len()
389 )));
390 }
391 let id = if start == 0 && end == record.length {
392 name.to_string()
393 } else {
394 format!("{name}:{}-{end}", start + 1)
395 };
396 Ok(Sequence::fasta(id, seq))
397 }
398
399 pub fn fetch_locus(&mut self, locus: &str) -> Result<Sequence> {
410 let (name, range) = match locus.rsplit_once(':') {
411 None => return self.fetch(locus),
412 Some((name, range)) => (name, range),
413 };
414 let range = range.replace("..", "-").replace(',', "");
415 let (start, end) = match range.split_once('-') {
416 Some((s, e)) => (s, e),
417 None => (range.as_str(), range.as_str()),
418 };
419 let parse = |text: &str, what: &str| -> Result<u64> {
420 text.trim()
421 .parse::<u64>()
422 .map_err(|e| Error::Index(format!("bad {what} coordinate in {locus:?}: {e}")))
423 };
424 let start = parse(start, "start")?;
425 let end = parse(end, "end")?;
426 if start == 0 {
427 return Err(Error::Index(format!("{locus:?}: coordinates are 1-based")));
428 }
429 self.fetch_region(name, start - 1, end)
430 }
431}
432
433fn open_for_scanning(path: &Path) -> Result<Box<dyn Read>> {
435 let file = open_file(path)?;
436 match probe(path)? {
437 Container::Plain => Ok(Box::new(file)),
438 #[cfg(feature = "gzip")]
439 Container::Bgzf => Ok(Box::new(crate::bgzf::BgzfReader::new(file)?)),
440 #[cfg(feature = "gzip")]
441 Container::PlainGzip => Err(plain_gzip_error(path)),
442 Container::Zstd => Err(not_seekable_error(path, "Zstandard")),
443 }
444}
445
446fn open_for_seeking(path: &Path) -> Result<BoxedSource> {
448 let file = open_file(path)?;
449 match probe(path)? {
450 Container::Plain => Ok(Box::new(file)),
451 #[cfg(feature = "gzip")]
452 Container::Bgzf => {
453 let index = match crate::bgzf::GziIndex::from_path(crate::bgzf::gzi_path(path)) {
454 Ok(index) => index,
455 Err(Error::Io(e)) if e.kind() == io::ErrorKind::NotFound => {
457 crate::bgzf::GziIndex::build(open_file(path)?)?
458 }
459 Err(e) => return Err(e),
460 };
461 Ok(Box::new(
462 crate::bgzf::BgzfReader::new(file)?.with_index(index),
463 ))
464 }
465 #[cfg(feature = "gzip")]
466 Container::PlainGzip => Err(plain_gzip_error(path)),
467 Container::Zstd => Err(not_seekable_error(path, "Zstandard")),
468 }
469}
470
471enum Container {
476 Plain,
478 #[cfg(feature = "gzip")]
480 Bgzf,
481 #[cfg(feature = "gzip")]
483 PlainGzip,
484 Zstd,
486}
487
488fn probe(path: &Path) -> Result<Container> {
489 let mut file = open_file(path)?;
490 let mut head = [0u8; 128];
491 let mut filled = 0;
492 while filled < head.len() {
493 match file.read(&mut head[filled..]) {
494 Ok(0) => break,
495 Ok(n) => filled += n,
496 Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
497 Err(e) => return Err(Error::Io(e)),
498 }
499 }
500 let head = &head[..filled];
501 match crate::format::Compression::from_magic(head) {
502 crate::format::Compression::None => return Ok(Container::Plain),
503 crate::format::Compression::Zstd => return Ok(Container::Zstd),
504 _ => {}
505 }
506 #[cfg(feature = "gzip")]
507 {
508 if crate::bgzf::is_bgzf(head) {
509 Ok(Container::Bgzf)
510 } else {
511 Ok(Container::PlainGzip)
512 }
513 }
514 #[cfg(not(feature = "gzip"))]
515 Err(Error::FeatureDisabled("gzip"))
516}
517
518#[cfg(feature = "gzip")]
519#[cfg(feature = "gzip")]
520fn plain_gzip_error(path: &Path) -> Error {
521 not_seekable_error(path, "plain gzip")
522}
523
524fn not_seekable_error(path: &Path, container: &str) -> Error {
527 Error::Index(format!(
528 "{}: this is {container}, which cannot be randomly accessed. \
529 Recompress it as BGZF (`bgzip`, or this crate's own gzip output) to index it.",
530 path.display()
531 ))
532}
533
534fn open_file(path: &Path) -> Result<File> {
535 File::open(path)
536 .map_err(|e| Error::Io(io::Error::new(e.kind(), format!("{}: {e}", path.display()))))
537}
538
539fn trim_newline(line: &[u8]) -> &[u8] {
540 let mut end = line.len();
541 if end > 0 && line[end - 1] == b'\n' {
542 end -= 1;
543 }
544 if end > 0 && line[end - 1] == b'\r' {
545 end -= 1;
546 }
547 &line[..end]
548}
549
550#[cfg(test)]
551mod tests {
552 use super::*;
553
554 const FASTA: &[u8] = b">chr1 first\nACGTACGTAC\nGGGG\n>chr2\nTTTTTTTTTT\nTTTTTTTTTT\n";
555
556 #[test]
557 fn builds_index_matching_samtools_layout() {
558 let index = FastaIndex::build(FASTA).unwrap();
559 assert_eq!(index.len(), 2);
560 let chr1 = index.get("chr1").unwrap();
561 assert_eq!(chr1.name, "chr1");
562 assert_eq!(chr1.length, 14);
563 assert_eq!(chr1.offset, 12); assert_eq!(chr1.line_bases, 10);
565 assert_eq!(chr1.line_width, 11);
566
567 let chr2 = index.get("chr2").unwrap();
568 assert_eq!(chr2.length, 20);
569 assert_eq!(chr2.offset, 12 + 11 + 5 + 6);
570 assert_eq!(index.total_length(), 34);
571 }
572
573 #[test]
574 fn round_trips_fai_text() {
575 let index = FastaIndex::build(FASTA).unwrap();
576 let mut text = Vec::new();
577 index.write(&mut text).unwrap();
578 assert_eq!(text, b"chr1\t14\t12\t10\t11\nchr2\t20\t34\t10\t11\n");
579 let reparsed = FastaIndex::parse(&text[..]).unwrap();
580 assert_eq!(reparsed.records(), index.records());
581 }
582
583 #[test]
584 fn rejects_ragged_lines() {
585 let ragged = b">a\nACGT\nAC\nACGT\n";
586 assert!(matches!(
587 FastaIndex::build(&ragged[..]),
588 Err(Error::Index(_))
589 ));
590 assert!(FastaIndex::build(&b">a\nACGT\nAC\n"[..]).is_ok());
592 }
593
594 #[test]
595 fn blank_line_before_the_sequence_does_not_shift_the_offset() {
596 let index = FastaIndex::build(&b">a\n\nACGT\n>b\nAC\n"[..]).unwrap();
598 let a = index.get("a").unwrap();
599 assert_eq!((a.offset, a.length, a.line_bases), (4, 4, 4));
600
601 let index = FastaIndex::build(&b">a\n\n\n>b\nAC\n"[..]).unwrap();
603 assert_eq!(index.get("a").unwrap().length, 0);
604 assert_eq!(index.get("b").unwrap().length, 2);
605 }
606
607 #[test]
608 fn rejects_duplicate_names() {
609 assert!(matches!(
610 FastaIndex::build(&b">a\nAC\n>a\nGT\n"[..]),
611 Err(Error::Index(_))
612 ));
613 }
614
615 #[test]
616 fn names_match_what_the_reader_produces() {
617 let fasta = ">chr\u{a0}1 description\nACGT\n";
621 let index = FastaIndex::build(fasta.as_bytes()).unwrap();
622 let records = crate::read_all_from(fasta.as_bytes()).unwrap();
623 assert_eq!(records[0].id, "chr\u{a0}1");
624 assert!(index.get(&records[0].id).is_some());
625
626 assert!(matches!(
629 FastaIndex::build(&b"> a\nACGT\n"[..]),
630 Err(Error::Index(_))
631 ));
632 assert!(crate::read_all_from(&b"> a\nACGT\n"[..]).is_err());
633 }
634
635 #[test]
636 fn rejects_data_before_header() {
637 assert!(matches!(
638 FastaIndex::build(&b"ACGT\n>a\nAC\n"[..]),
639 Err(Error::Index(_))
640 ));
641 }
642
643 #[test]
644 fn fetches_regions() {
645 let dir = std::env::temp_dir().join(format!("fastx-index-{}", std::process::id()));
646 std::fs::create_dir_all(&dir).unwrap();
647 let path = dir.join("ref.fa");
648 std::fs::write(&path, FASTA).unwrap();
649
650 let index = FastaIndex::build_from_path(&path).unwrap();
651 index.write_to_path(&path).unwrap();
652 assert!(fai_path(&path).exists());
653
654 let mut fasta = IndexedFasta::open(&path).unwrap();
655 assert_eq!(fasta.fetch("chr1").unwrap().seq, b"ACGTACGTACGGGG");
656 assert_eq!(fasta.fetch("chr2").unwrap().seq, b"T".repeat(20));
657
658 assert_eq!(fasta.fetch_region("chr1", 0, 4).unwrap().seq, b"ACGT");
660 assert_eq!(fasta.fetch_region("chr1", 8, 12).unwrap().seq, b"ACGG");
662 assert_eq!(fasta.fetch_region("chr1", 10, 14).unwrap().seq, b"GGGG");
664 assert!(fasta.fetch_region("chr1", 5, 5).unwrap().seq.is_empty());
666 assert_eq!(fasta.fetch_locus("chr1:1-4").unwrap().seq, b"ACGT");
668 assert_eq!(fasta.fetch_locus("chr1:9..12").unwrap().seq, b"ACGG");
669 assert_eq!(fasta.fetch_locus("chr2").unwrap().seq.len(), 20);
670 assert_eq!(fasta.fetch_locus("chr1:1-4").unwrap().id, "chr1:1-4");
671
672 assert!(matches!(
674 fasta.fetch("nope"),
675 Err(Error::UnknownSequence(_))
676 ));
677 assert!(matches!(
678 fasta.fetch_region("chr1", 0, 99),
679 Err(Error::OutOfBounds { .. })
680 ));
681 assert!(fasta.fetch_locus("chr1:0-4").is_err());
682 assert!(fasta.fetch_locus("chr1:x-4").is_err());
683
684 std::fs::remove_dir_all(&dir).ok();
685 }
686
687 #[cfg(feature = "gzip")]
688 #[test]
689 fn indexes_and_fetches_from_bgzf() {
690 use std::io::Write;
691
692 let dir = std::env::temp_dir().join(format!("fastx-bgzf-idx-{}", std::process::id()));
693 std::fs::create_dir_all(&dir).unwrap();
694
695 let mut fasta = Vec::new();
697 for chromosome in 0..3 {
698 fasta.extend_from_slice(format!(">chr{chromosome} test\n").as_bytes());
699 for line in 0..1_500 {
700 let base = b"ACGT"[(chromosome + line) % 4];
701 fasta.extend(std::iter::repeat(base).take(60));
702 fasta.push(b'\n');
703 }
704 }
705
706 let bgzf_path = dir.join("ref.fa.gz");
707 let mut writer = crate::bgzf::BgzfWriter::create(&bgzf_path).unwrap();
708 writer.write_all(&fasta).unwrap();
709 let (_, gzi) = writer.finish_with_index().unwrap();
710 gzi.write_to_path(&bgzf_path).unwrap();
711
712 let index = FastaIndex::build_from_path(&bgzf_path).unwrap();
715 let plain_index = FastaIndex::build(&fasta[..]).unwrap();
716 assert_eq!(index.records(), plain_index.records());
717 index.write_to_path(&bgzf_path).unwrap();
718
719 let plain_path = dir.join("ref.fa");
722 std::fs::write(&plain_path, &fasta).unwrap();
723 let mut compressed = IndexedFasta::open(&bgzf_path).unwrap();
724 let mut plain = IndexedFasta::open(&plain_path).unwrap();
725
726 for name in ["chr0", "chr1", "chr2"] {
727 assert_eq!(compressed.fetch(name).unwrap(), plain.fetch(name).unwrap());
728 for (start, end) in [(0, 1), (59, 61), (1_000, 1_100), (89_000, 90_000)] {
729 assert_eq!(
730 compressed.fetch_region(name, start, end).unwrap(),
731 plain.fetch_region(name, start, end).unwrap(),
732 "{name}:{start}-{end}"
733 );
734 }
735 }
736
737 std::fs::remove_file(crate::bgzf::gzi_path(&bgzf_path)).unwrap();
739 let mut rescanned = IndexedFasta::open(&bgzf_path).unwrap();
740 assert_eq!(
741 rescanned.fetch_region("chr1", 1_000, 1_100).unwrap(),
742 plain.fetch_region("chr1", 1_000, 1_100).unwrap()
743 );
744
745 std::fs::remove_dir_all(&dir).ok();
746 }
747
748 #[cfg(feature = "gzip")]
749 #[test]
750 fn refuses_to_index_plain_gzip() {
751 use std::io::Write;
752
753 let dir = std::env::temp_dir().join(format!("fastx-plain-gz-{}", std::process::id()));
754 std::fs::create_dir_all(&dir).unwrap();
755 let path = dir.join("ref.fa.gz");
756
757 let file = std::fs::File::create(&path).unwrap();
759 let mut encoder = flate2::write::GzEncoder::new(file, flate2::Compression::fast());
760 encoder.write_all(FASTA).unwrap();
761 encoder.finish().unwrap();
762
763 let error = FastaIndex::build_from_path(&path).unwrap_err();
764 assert!(matches!(error, Error::Index(_)), "{error}");
765 assert!(error.to_string().contains("BGZF"), "{error}");
766 assert_eq!(crate::read_all(&path).unwrap().len(), 2);
768
769 std::fs::remove_dir_all(&dir).ok();
770 }
771}