Skip to main content

fgumi_lib/
reference.rs

1//! Reference genome FASTA reading with all sequences loaded into memory.
2//!
3//! This module provides thread-safe access to reference genome sequences, which is needed
4//! for tasks like NM/UQ/MD tag calculation and variant calling.
5//!
6//! Following fgbio's approach, the entire reference is loaded into memory at startup
7//! to ensure O(1) lookup performance for each read during tag regeneration.
8//!
9//! Uses FAI index for fast raw-byte reading (htsjdk-style) instead of line-by-line parsing.
10//!
11//! # Memory Usage
12//!
13//! Reference sequences are stored as raw bytes, matching htsjdk's approach.
14//! For a typical human reference (~3GB), this uses approximately 3GB of memory
15//! but provides the fastest possible load times (~9s vs ~22s for compressed storage).
16//!
17//! # Future improvement
18//!
19//! The custom FAI-based raw-byte reading (`read_sequence_raw`) could be replaced with
20//! noodles' built-in indexed reader once <https://github.com/zaeleus/noodles/pull/365>
21//! is merged and released, which adds the same optimization to noodles.
22use crate::errors::FgumiError;
23use anyhow::{Context, Result};
24use log::debug;
25use noodles::core::Position;
26use noodles::fasta::fai;
27use std::collections::HashMap;
28use std::fs::File;
29use std::io::{Read, Seek, SeekFrom};
30use std::path::{Path, PathBuf};
31use std::sync::Arc;
32
33/// Read a sequence from a FASTA file using FAI index metadata.
34/// Uses raw byte reading with mathematical newline handling (htsjdk-style).
35///
36/// Optimized: Reads entire sequence bytes in one syscall, then strips newlines in memory.
37/// Fast path: If sequence fits in a single line, skip newline stripping entirely.
38#[allow(clippy::cast_possible_truncation)]
39fn read_sequence_raw(file: &mut File, record: &fai::Record) -> Result<Vec<u8>> {
40    let line_bases = record.line_bases() as usize;
41    let line_width = record.line_width() as usize;
42    let seq_len = record.length() as usize;
43    let offset = record.offset();
44
45    // Fast path: if sequence fits in a single line, no newlines to strip
46    if seq_len <= line_bases {
47        file.seek(SeekFrom::Start(offset))?;
48        let mut sequence = vec![0u8; seq_len];
49        file.read_exact(&mut sequence)?;
50        return Ok(sequence);
51    }
52
53    // Calculate number of complete lines and remaining bases
54    let complete_lines = seq_len / line_bases;
55    let remaining_bases = seq_len % line_bases;
56
57    // Total bytes = (complete_lines * line_width) + remaining_bases
58    // But last line might not have a terminator, so we calculate conservatively
59    let total_bytes = if remaining_bases > 0 {
60        complete_lines * line_width + remaining_bases
61    } else if complete_lines > 0 {
62        // All bases fit exactly in complete lines, last line has no terminator at end of seq
63        (complete_lines - 1) * line_width + line_bases
64    } else {
65        0
66    };
67
68    // Seek and read all bytes at once
69    file.seek(SeekFrom::Start(offset))?;
70    let mut raw_bytes = vec![0u8; total_bytes];
71    file.read_exact(&mut raw_bytes)?;
72
73    // Strip newlines in memory (much faster than seeking)
74    let mut sequence = Vec::with_capacity(seq_len);
75    let terminator_len = line_width - line_bases;
76
77    let mut pos = 0;
78    while sequence.len() < seq_len && pos < raw_bytes.len() {
79        // Read up to line_bases bytes
80        let bases_to_copy = (seq_len - sequence.len()).min(line_bases).min(raw_bytes.len() - pos);
81        sequence.extend_from_slice(&raw_bytes[pos..pos + bases_to_copy]);
82        pos += bases_to_copy;
83
84        // Skip line terminator if present and we need more bases
85        if sequence.len() < seq_len && pos < raw_bytes.len() {
86            pos += terminator_len;
87        }
88    }
89
90    Ok(sequence)
91}
92
93/// Find a sibling file for a FASTA file by trying two naming conventions:
94/// 1. Replace extension (e.g., `ref.fasta` → `ref.<replace_ext>`)
95/// 2. Append extension to full path (e.g., `ref.fa` → `ref.fa.<append_ext>`)
96fn find_sibling_file(fasta_path: &Path, replace_ext: &str, append_ext: &str) -> Option<PathBuf> {
97    let replaced = fasta_path.with_extension(replace_ext);
98    if replaced.exists() {
99        return Some(replaced);
100    }
101
102    let appended = PathBuf::from(format!("{}.{append_ext}", fasta_path.display()));
103    if appended.exists() {
104        return Some(appended);
105    }
106
107    None
108}
109
110/// Find FAI index path for a FASTA file.
111///
112/// Tries multiple naming conventions:
113/// 1. Replace extension with `.fa.fai` (e.g., `ref.fasta` → `ref.fa.fai`)
114/// 2. Append `.fai` to full path (e.g., `ref.fa` → `ref.fa.fai`, `ref.fasta` → `ref.fasta.fai`)
115fn find_fai_path(fasta_path: &Path) -> Option<PathBuf> {
116    find_sibling_file(fasta_path, "fa.fai", "fai")
117}
118
119/// Find sequence dictionary path for a FASTA file.
120///
121/// Tries multiple naming conventions used by different tools:
122/// 1. Replace extension with `.dict` (fgbio/HTSJDK/Picard convention: `ref.fa` → `ref.dict`)
123/// 2. Append `.dict` to full path (GATK convention: `ref.fa` → `ref.fa.dict`)
124///
125/// # Arguments
126/// * `fasta_path` - Path to the FASTA file
127///
128/// # Returns
129/// The path to the dictionary file if found, or `None` if not found.
130///
131/// # Examples
132/// ```no_run
133/// use std::path::Path;
134/// use fgumi_lib::reference::find_dict_path;
135///
136/// // Will find either "ref.dict" or "ref.fa.dict"
137/// if let Some(dict_path) = find_dict_path(Path::new("ref.fa")) {
138///     println!("Found dictionary: {}", dict_path.display());
139/// }
140/// ```
141#[must_use]
142pub fn find_dict_path(fasta_path: &Path) -> Option<PathBuf> {
143    find_sibling_file(fasta_path, "dict", "dict")
144}
145
146/// A thread-safe reference genome reader with all sequences preloaded into memory.
147///
148/// This reader loads the entire FASTA file into memory at construction time,
149/// providing O(1) lookup performance for sequence fetches. This approach matches
150/// fgbio's `nmUqMdTagRegeneratingWriter` which reads all contigs into a Map upfront.
151///
152/// For a typical human reference (e.g., hs38DH at ~3GB), this uses approximately
153/// 3GB of memory (raw byte storage like htsjdk) and provides the fastest load times.
154#[derive(Clone)]
155pub struct ReferenceReader {
156    /// All sequences loaded into memory as raw bytes, keyed by sequence name
157    sequences: Arc<HashMap<String, Vec<u8>>>,
158}
159
160impl ReferenceReader {
161    /// Creates a new reference reader, loading all sequences into memory.
162    ///
163    /// This reads the entire FASTA file into memory at construction time.
164    /// For a typical human reference (~3GB), this takes a few seconds but
165    /// provides O(1) lookup performance for all subsequent fetches.
166    ///
167    /// # Arguments
168    /// * `path` - Path to the reference FASTA file (may be gzipped)
169    ///
170    /// # Errors
171    /// Returns an error if:
172    /// - The file does not exist
173    /// - The file cannot be read or parsed as FASTA
174    ///
175    /// # Examples
176    /// ```no_run
177    /// use fgumi_lib::reference::ReferenceReader;
178    ///
179    /// let reader = ReferenceReader::new("reference.fasta")?;
180    /// # Ok::<(), anyhow::Error>(())
181    /// ```
182    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self> {
183        let path = path.as_ref();
184
185        // Verify the file exists and is readable
186        if !path.exists() {
187            return Err(FgumiError::InvalidFileFormat {
188                file_type: "Reference FASTA".to_string(),
189                path: path.display().to_string(),
190                reason: "File does not exist".to_string(),
191            }
192            .into());
193        }
194
195        debug!("Reading reference FASTA into memory: {}", path.display());
196
197        // Try FAI-based fast reading first
198        if let Some(fai_path) = find_fai_path(path) {
199            debug!("Using FAI index for fast loading: {}", fai_path.display());
200            return Self::new_with_fai(path, &fai_path);
201        }
202
203        // Fall back to noodles for non-indexed files
204        debug!("No FAI index found, using sequential reading");
205        Self::new_sequential(path)
206    }
207
208    /// Load sequences using FAI index for fast raw-byte reading (htsjdk-style).
209    fn new_with_fai(fasta_path: &Path, fai_path: &Path) -> Result<Self> {
210        let index = fai::fs::read(fai_path)
211            .with_context(|| format!("Failed to read FAI index: {}", fai_path.display()))?;
212        let records: &[fai::Record] = index.as_ref();
213        let mut file = File::open(fasta_path)
214            .with_context(|| format!("Failed to open FASTA: {}", fasta_path.display()))?;
215
216        let mut sequences = HashMap::with_capacity(records.len());
217
218        for record in records {
219            let raw_sequence = read_sequence_raw(&mut file, record)?;
220            let name = String::from_utf8_lossy(record.name().as_ref()).into_owned();
221            sequences.insert(name, raw_sequence);
222        }
223
224        debug!("Loaded {} contigs into memory (FAI-indexed)", sequences.len());
225        Ok(Self { sequences: Arc::new(sequences) })
226    }
227
228    /// Load sequences using noodles sequential reading (fallback for non-indexed files).
229    fn new_sequential(path: &Path) -> Result<Self> {
230        use noodles::fasta;
231
232        let mut sequences = HashMap::new();
233        let mut reader = fasta::io::reader::Builder.build_from_path(path)?;
234
235        for result in reader.records() {
236            let record = result?;
237            let name = std::str::from_utf8(record.name())?.to_string();
238            let raw_sequence: Vec<u8> = record.sequence().as_ref().to_vec();
239            sequences.insert(name, raw_sequence);
240        }
241
242        debug!("Loaded {} contigs into memory (sequential)", sequences.len());
243        Ok(Self { sequences: Arc::new(sequences) })
244    }
245
246    /// Retrieves a subsequence from the reference genome.
247    ///
248    /// Since all sequences are preloaded into memory, this is an O(1) lookup
249    /// followed by a slice copy.
250    ///
251    /// # Arguments
252    /// * `chrom` - Chromosome/sequence name (e.g., "chr1", "1")
253    /// * `start` - Start position (1-based, inclusive)
254    /// * `end` - End position (1-based, inclusive)
255    ///
256    /// # Returns
257    /// The requested subsequence as a vector of bytes (preserving original case)
258    ///
259    /// # Errors
260    /// Returns an error if:
261    /// - The chromosome is not found in the reference
262    /// - The requested region exceeds the chromosome length
263    ///
264    /// # Examples
265    /// ```no_run
266    /// use fgumi_lib::reference::ReferenceReader;
267    /// use noodles::core::Position;
268    ///
269    /// let reader = ReferenceReader::new("reference.fasta")?;
270    ///
271    /// // Fetch first 100 bases of chr1
272    /// let seq = reader.fetch("chr1", Position::try_from(1)?, Position::try_from(100)?)?;
273    /// assert_eq!(seq.len(), 100);
274    /// # Ok::<(), anyhow::Error>(())
275    /// ```
276    pub fn fetch(&self, chrom: &str, start: Position, end: Position) -> Result<Vec<u8>> {
277        Ok(self.fetch_slice(chrom, start, end)?.to_vec())
278    }
279
280    /// Borrowed-slice variant of [`Self::fetch`] that returns a reference into the
281    /// in-memory sequence — same lookup semantics, no allocation.
282    ///
283    /// Prefer this over [`Self::fetch`] in hot loops (e.g., per-record reference
284    /// access in `fgumi zipper`'s `--restore-unconverted-bases` path) where
285    /// allocating a fresh `Vec<u8>` per call adds up to gigabytes of churn.
286    ///
287    /// # Errors
288    ///
289    /// Same as [`Self::fetch`]: returns an error if the chromosome is missing or the
290    /// requested region exceeds the sequence length.
291    pub fn fetch_slice(&self, chrom: &str, start: Position, end: Position) -> Result<&[u8]> {
292        let sequence = self
293            .sequences
294            .get(chrom)
295            .ok_or_else(|| FgumiError::ReferenceNotFound { ref_name: chrom.to_string() })?;
296
297        // Convert from 1-based inclusive to 0-based [start, end) indexing
298        let start_idx = usize::from(start) - 1;
299        let end_idx = usize::from(end);
300
301        if end_idx > sequence.len() || start_idx >= end_idx {
302            return Err(FgumiError::InvalidParameter {
303                parameter: "region".to_string(),
304                reason: format!(
305                    "Requested region {}:{}-{} exceeds sequence length {}",
306                    chrom,
307                    start,
308                    end,
309                    sequence.len()
310                ),
311            }
312            .into());
313        }
314
315        Ok(&sequence[start_idx..end_idx])
316    }
317
318    /// Gets a single base from the reference at the specified position.
319    ///
320    /// This is a convenience method that delegates to `fetch()` with a single-base region.
321    ///
322    /// # Arguments
323    /// * `chrom` - Chromosome/sequence name (e.g., "chr1", "1")
324    /// * `pos` - Position (1-based)
325    ///
326    /// # Returns
327    /// The base at the specified position (preserving original case from FASTA)
328    ///
329    /// # Errors
330    /// Returns an error if:
331    /// - The chromosome is not found in the reference
332    /// - The position exceeds the chromosome length
333    ///
334    /// # Examples
335    /// ```no_run
336    /// use fgumi_lib::reference::ReferenceReader;
337    /// use noodles::core::Position;
338    ///
339    /// let reader = ReferenceReader::new("reference.fasta")?;
340    ///
341    /// // Get the base at position 1000 of chr1
342    /// let base = reader.base_at("chr1", Position::try_from(1000)?)?;
343    /// assert!(matches!(base, b'A' | b'C' | b'G' | b'T' | b'N'));
344    /// # Ok::<(), anyhow::Error>(())
345    /// ```
346    pub fn base_at(&self, chrom: &str, pos: Position) -> Result<u8> {
347        let sequence = self
348            .sequences
349            .get(chrom)
350            .ok_or_else(|| FgumiError::ReferenceNotFound { ref_name: chrom.to_string() })?;
351
352        // Convert from 1-based to 0-based indexing
353        let pos_idx = usize::from(pos) - 1;
354
355        sequence.get(pos_idx).copied().ok_or_else(|| {
356            FgumiError::InvalidParameter {
357                parameter: "position".to_string(),
358                reason: format!(
359                    "Position {}:{} exceeds sequence length {}",
360                    chrom,
361                    pos,
362                    sequence.len()
363                ),
364            }
365            .into()
366        })
367    }
368}
369
370impl fgumi_sam::ReferenceProvider for ReferenceReader {
371    fn fetch(
372        &self,
373        chrom: &str,
374        start: noodles::core::Position,
375        end: noodles::core::Position,
376    ) -> anyhow::Result<Vec<u8>> {
377        self.fetch(chrom, start, end)
378    }
379
380    /// Returns a borrow into the in-memory sequence, avoiding the per-call `Vec`
381    /// allocation `fetch` performs.
382    fn fetch_borrowed(
383        &self,
384        chrom: &str,
385        start: noodles::core::Position,
386        end: noodles::core::Position,
387    ) -> anyhow::Result<std::borrow::Cow<'_, [u8]>> {
388        Ok(std::borrow::Cow::Borrowed(self.fetch_slice(chrom, start, end)?))
389    }
390}
391
392#[cfg(feature = "simplex")]
393impl fgumi_consensus::methylation::RefBaseProvider for ReferenceReader {
394    fn base_at_0based(&self, chrom: &str, pos: u64) -> Option<u8> {
395        let sequence = self.sequences.get(chrom)?;
396        sequence.get(usize::try_from(pos).ok()?).copied()
397    }
398
399    fn sequence_for(&self, chrom: &str) -> Option<&[u8]> {
400        self.sequences.get(chrom).map(Vec::as_slice)
401    }
402}
403
404#[cfg(test)]
405mod tests {
406    use super::*;
407    use crate::sam::builder::create_default_test_fasta;
408
409    /// `fetch_borrowed` must hand back a borrow, not an owned copy — that is the
410    /// entire point of the override. A `Cow::Owned` here means the default trait
411    /// implementation is being used and the per-record allocation is still happening.
412    #[test]
413    fn test_fetch_borrowed_borrows_and_matches_fetch() -> Result<()> {
414        use fgumi_sam::ReferenceProvider;
415
416        let fasta = create_default_test_fasta()?;
417        let reader = ReferenceReader::new(fasta.path())?;
418        let (start, end) = (Position::try_from(1)?, Position::try_from(4)?);
419
420        let borrowed = ReferenceProvider::fetch_borrowed(&reader, "chr1", start, end)?;
421        assert!(
422            matches!(borrowed, std::borrow::Cow::Borrowed(_)),
423            "expected a borrowed slice; an owned value means the allocation is still happening"
424        );
425
426        // Same bytes as the allocating path.
427        let owned = ReferenceProvider::fetch(&reader, "chr1", start, end)?;
428        assert_eq!(borrowed.as_ref(), owned.as_slice());
429
430        // Errors propagate identically for a missing chromosome.
431        assert!(ReferenceProvider::fetch_borrowed(&reader, "nope", start, end).is_err());
432        Ok(())
433    }
434
435    /// The production call sites (`clip.rs`, `filter.rs`) pass `&ReferenceReader`
436    /// into a `R: ReferenceProvider` generic, so they resolve through the `T: Deref`
437    /// blanket impl rather than `ReferenceReader`'s own. If that blanket stops
438    /// forwarding `fetch_borrowed`, it silently falls back to the trait default —
439    /// which returns `Cow::Owned`, reinstating the per-record allocation with no
440    /// output change and no other test failing. This pins the forwarding.
441    #[test]
442    fn test_fetch_borrowed_forwards_through_reference() -> Result<()> {
443        use fgumi_sam::ReferenceProvider;
444
445        /// Mirrors how `regenerate_alignment_tags_raw` takes its provider.
446        fn borrows_via_generic<R: ReferenceProvider>(
447            provider: R,
448            chrom: &str,
449            start: Position,
450            end: Position,
451        ) -> Result<bool> {
452            Ok(matches!(provider.fetch_borrowed(chrom, start, end)?, std::borrow::Cow::Borrowed(_)))
453        }
454
455        let fasta = create_default_test_fasta()?;
456        let reader = ReferenceReader::new(fasta.path())?;
457        let (start, end) = (Position::try_from(1)?, Position::try_from(4)?);
458
459        assert!(
460            borrows_via_generic(&reader, "chr1", start, end)?,
461            "&ReferenceReader must still borrow; an owned value means the Deref blanket \
462             impl stopped forwarding fetch_borrowed and the allocation is back"
463        );
464        Ok(())
465    }
466
467    #[test]
468    fn test_fetch_subsequence() -> Result<()> {
469        let fasta = create_default_test_fasta()?;
470        let reader = ReferenceReader::new(fasta.path())?;
471
472        // Fetch from chr1
473        let seq = reader.fetch("chr1", Position::try_from(1)?, Position::try_from(4)?)?;
474        assert_eq!(seq, b"ACGT");
475
476        // Fetch from chr2
477        let seq = reader.fetch("chr2", Position::try_from(5)?, Position::try_from(8)?)?;
478        assert_eq!(seq, b"CCCC");
479
480        Ok(())
481    }
482
483    #[test]
484    fn test_fetch_slice_returns_borrowed_bytes() -> Result<()> {
485        let fasta = create_default_test_fasta()?;
486        let reader = ReferenceReader::new(fasta.path())?;
487
488        // Borrowed alternative to `fetch` — same bytes, no allocation.
489        let slice: &[u8] =
490            reader.fetch_slice("chr1", Position::try_from(1)?, Position::try_from(4)?)?;
491        assert_eq!(slice, b"ACGT");
492
493        // Same bytes as `fetch` would return.
494        let owned = reader.fetch("chr1", Position::try_from(1)?, Position::try_from(4)?)?;
495        assert_eq!(slice, owned.as_slice());
496
497        // Bounds errors propagate just like `fetch`.
498        assert!(
499            reader
500                .fetch_slice("chr1", Position::try_from(1)?, Position::try_from(10_000)?)
501                .is_err()
502        );
503        assert!(
504            reader.fetch_slice("nope", Position::try_from(1)?, Position::try_from(2)?).is_err()
505        );
506        // Inverted interval (start > end) must also error.
507        assert!(
508            reader.fetch_slice("chr1", Position::try_from(5)?, Position::try_from(4)?).is_err()
509        );
510
511        Ok(())
512    }
513
514    #[test]
515    fn test_base_at() -> Result<()> {
516        let fasta = create_default_test_fasta()?;
517        let reader = ReferenceReader::new(fasta.path())?;
518
519        assert_eq!(reader.base_at("chr1", Position::try_from(1)?)?, b'A');
520        assert_eq!(reader.base_at("chr1", Position::try_from(2)?)?, b'C');
521        assert_eq!(reader.base_at("chr2", Position::try_from(1)?)?, b'G');
522
523        Ok(())
524    }
525
526    #[test]
527    fn test_all_sequences_loaded() -> Result<()> {
528        let fasta = create_default_test_fasta()?;
529        let reader = ReferenceReader::new(fasta.path())?;
530
531        // All sequences should be available immediately after construction
532        let seq1 = reader.fetch("chr1", Position::try_from(1)?, Position::try_from(4)?)?;
533        assert_eq!(seq1, b"ACGT");
534
535        let seq2 = reader.fetch("chr2", Position::try_from(1)?, Position::try_from(4)?)?;
536        assert_eq!(seq2, b"GGGG");
537
538        // Fetching chr1 again should still work (all in memory)
539        let seq1_again = reader.fetch("chr1", Position::try_from(1)?, Position::try_from(4)?)?;
540        assert_eq!(seq1_again, b"ACGT");
541
542        Ok(())
543    }
544
545    #[test]
546    fn test_nonexistent_sequence() {
547        let fasta = create_default_test_fasta().expect("creating test FASTA should succeed");
548        let reader =
549            ReferenceReader::new(fasta.path()).expect("creating reference reader should succeed");
550
551        let result = reader.fetch(
552            "chr999",
553            Position::try_from(1).expect("position conversion should succeed"),
554            Position::try_from(4).expect("position conversion should succeed"),
555        );
556        assert!(result.is_err());
557    }
558
559    #[test]
560    fn test_out_of_bounds() {
561        let fasta = create_default_test_fasta().expect("creating test FASTA should succeed");
562        let reader =
563            ReferenceReader::new(fasta.path()).expect("creating reference reader should succeed");
564
565        // chr1 is only 12 bases long
566        let result = reader.fetch(
567            "chr1",
568            Position::try_from(1).expect("position conversion should succeed"),
569            Position::try_from(100).expect("position conversion should succeed"),
570        );
571        assert!(result.is_err());
572    }
573
574    #[test]
575    fn test_reference_case_preserved() -> Result<()> {
576        // Test that FASTA case is preserved (important for MD tag generation)
577        // fgbio preserves case in reference sequences for proper MD tag output
578        use crate::sam::builder::create_test_fasta;
579
580        let file = create_test_fasta(&[("chr1", "AcGtNnAaCcGgTt")])?; // Mixed case sequence
581
582        let reader = ReferenceReader::new(file.path())?;
583
584        // Fetch the full sequence and verify case is preserved
585        let seq = reader.fetch("chr1", Position::try_from(1)?, Position::try_from(14)?)?;
586        assert_eq!(seq, b"AcGtNnAaCcGgTt");
587
588        // Verify individual bases preserve case
589        assert_eq!(reader.base_at("chr1", Position::try_from(1)?)?, b'A'); // uppercase A
590        assert_eq!(reader.base_at("chr1", Position::try_from(2)?)?, b'c'); // lowercase c
591        assert_eq!(reader.base_at("chr1", Position::try_from(3)?)?, b'G'); // uppercase G
592        assert_eq!(reader.base_at("chr1", Position::try_from(4)?)?, b't'); // lowercase t
593        assert_eq!(reader.base_at("chr1", Position::try_from(5)?)?, b'N'); // uppercase N
594        assert_eq!(reader.base_at("chr1", Position::try_from(6)?)?, b'n'); // lowercase n
595
596        Ok(())
597    }
598
599    #[test]
600    fn test_n_bases_at_various_positions() -> Result<()> {
601        use crate::sam::builder::create_test_fasta;
602
603        // N at start, middle, and end
604        let file = create_test_fasta(&[("chr1", "NACGTNACGTN")])?;
605        let reader = ReferenceReader::new(file.path())?;
606
607        assert_eq!(reader.base_at("chr1", Position::try_from(1)?)?, b'N');
608        assert_eq!(reader.base_at("chr1", Position::try_from(6)?)?, b'N');
609        assert_eq!(reader.base_at("chr1", Position::try_from(11)?)?, b'N');
610
611        // Fetch range including N
612        let seq = reader.fetch("chr1", Position::try_from(1)?, Position::try_from(6)?)?;
613        assert_eq!(seq, b"NACGTN");
614
615        Ok(())
616    }
617
618    #[test]
619    fn test_all_n_sequence() -> Result<()> {
620        use crate::sam::builder::create_test_fasta;
621
622        let file = create_test_fasta(&[("chrN", "NNNNNNNNNN")])?;
623        let reader = ReferenceReader::new(file.path())?;
624
625        let seq = reader.fetch("chrN", Position::try_from(1)?, Position::try_from(10)?)?;
626        assert_eq!(seq, b"NNNNNNNNNN");
627
628        for i in 1..=10 {
629            assert_eq!(reader.base_at("chrN", Position::try_from(i)?)?, b'N');
630        }
631
632        Ok(())
633    }
634
635    #[test]
636    fn test_long_sequence_boundaries() -> Result<()> {
637        use crate::sam::builder::create_test_fasta;
638
639        // Create a 100-base sequence with markers at specific positions
640        // We want to test fetching across the 32-base boundary (for bit-packed storage)
641        let mut seq = String::new();
642
643        // Positions 1-28: ACGT repeated (7 times)
644        for _ in 0..7 {
645            seq.push_str("ACGT");
646        }
647        // Positions 29-32: NNNN (straddles u64 boundary at position 32)
648        seq.push_str("NNNN");
649        // Positions 33-60: ACGT repeated (7 times)
650        for _ in 0..7 {
651            seq.push_str("ACGT");
652        }
653        // Positions 61-64: lowercase acgt (straddles second u64 boundary)
654        seq.push_str("acgt");
655        // Positions 65-100: ACGT repeated (9 times)
656        for _ in 0..9 {
657            seq.push_str("ACGT");
658        }
659
660        assert_eq!(seq.len(), 100);
661
662        let file = create_test_fasta(&[("chr1", &seq)])?;
663        let reader = ReferenceReader::new(file.path())?;
664
665        // Verify positions around first marker (NNNN at 29-32)
666        assert_eq!(reader.base_at("chr1", Position::try_from(28)?)?, b'T');
667        assert_eq!(reader.base_at("chr1", Position::try_from(29)?)?, b'N');
668        assert_eq!(reader.base_at("chr1", Position::try_from(32)?)?, b'N');
669        assert_eq!(reader.base_at("chr1", Position::try_from(33)?)?, b'A');
670
671        // Verify positions around second marker (acgt at 61-64)
672        assert_eq!(reader.base_at("chr1", Position::try_from(60)?)?, b'T');
673        assert_eq!(reader.base_at("chr1", Position::try_from(61)?)?, b'a');
674        assert_eq!(reader.base_at("chr1", Position::try_from(64)?)?, b't');
675        assert_eq!(reader.base_at("chr1", Position::try_from(65)?)?, b'A');
676
677        // Fetch across boundaries
678        let cross_first = reader.fetch("chr1", Position::try_from(27)?, Position::try_from(34)?)?;
679        assert_eq!(cross_first, b"GTNNNNAC");
680
681        let cross_second =
682            reader.fetch("chr1", Position::try_from(59)?, Position::try_from(66)?)?;
683        assert_eq!(cross_second, b"GTacgtAC");
684
685        Ok(())
686    }
687
688    #[test]
689    fn test_single_base_sequence() -> Result<()> {
690        use crate::sam::builder::create_test_fasta;
691
692        let file = create_test_fasta(&[("chr1", "A"), ("chr2", "N"), ("chr3", "g")])?;
693        let reader = ReferenceReader::new(file.path())?;
694
695        assert_eq!(reader.base_at("chr1", Position::try_from(1)?)?, b'A');
696        assert_eq!(reader.base_at("chr2", Position::try_from(1)?)?, b'N');
697        assert_eq!(reader.base_at("chr3", Position::try_from(1)?)?, b'g');
698
699        let seq = reader.fetch("chr1", Position::try_from(1)?, Position::try_from(1)?)?;
700        assert_eq!(seq, b"A");
701
702        Ok(())
703    }
704
705    #[test]
706    fn test_fetch_full_sequence() -> Result<()> {
707        use crate::sam::builder::create_test_fasta;
708
709        let original = "ACGTNacgtn";
710        let file = create_test_fasta(&[("chr1", original)])?;
711        let reader = ReferenceReader::new(file.path())?;
712
713        let seq = reader.fetch("chr1", Position::try_from(1)?, Position::try_from(10)?)?;
714        assert_eq!(seq, original.as_bytes());
715
716        Ok(())
717    }
718
719    #[test]
720    fn test_fetch_last_base() -> Result<()> {
721        use crate::sam::builder::create_test_fasta;
722
723        let file = create_test_fasta(&[("chr1", "ACGTN")])?;
724        let reader = ReferenceReader::new(file.path())?;
725
726        // Fetch last base
727        let seq = reader.fetch("chr1", Position::try_from(5)?, Position::try_from(5)?)?;
728        assert_eq!(seq, b"N");
729
730        assert_eq!(reader.base_at("chr1", Position::try_from(5)?)?, b'N');
731
732        Ok(())
733    }
734
735    #[test]
736    fn test_multiple_chromosomes_isolation() -> Result<()> {
737        use crate::sam::builder::create_test_fasta;
738
739        let file = create_test_fasta(&[
740            ("chr1", "AAAA"),
741            ("chr2", "CCCC"),
742            ("chr3", "GGGG"),
743            ("chr4", "TTTT"),
744        ])?;
745        let reader = ReferenceReader::new(file.path())?;
746
747        // Verify each chromosome has its own sequence
748        assert_eq!(reader.fetch("chr1", Position::try_from(1)?, Position::try_from(4)?)?, b"AAAA");
749        assert_eq!(reader.fetch("chr2", Position::try_from(1)?, Position::try_from(4)?)?, b"CCCC");
750        assert_eq!(reader.fetch("chr3", Position::try_from(1)?, Position::try_from(4)?)?, b"GGGG");
751        assert_eq!(reader.fetch("chr4", Position::try_from(1)?, Position::try_from(4)?)?, b"TTTT");
752
753        Ok(())
754    }
755
756    #[test]
757    fn test_mixed_case_all_bases() -> Result<()> {
758        use crate::sam::builder::create_test_fasta;
759
760        // Test all 10 possible base values: A, C, G, T, N (upper and lower)
761        let file = create_test_fasta(&[("chr1", "ACGTNacgtn")])?;
762        let reader = ReferenceReader::new(file.path())?;
763
764        let expected = [b'A', b'C', b'G', b'T', b'N', b'a', b'c', b'g', b't', b'n'];
765        for (i, &expected_base) in expected.iter().enumerate() {
766            let pos = Position::try_from(i + 1)?;
767            assert_eq!(
768                reader.base_at("chr1", pos)?,
769                expected_base,
770                "Mismatch at position {}",
771                i + 1
772            );
773        }
774
775        Ok(())
776    }
777
778    #[test]
779    fn test_runs_of_n_bases() -> Result<()> {
780        use crate::sam::builder::create_test_fasta;
781
782        // Simulate masked regions with runs of N
783        let file = create_test_fasta(&[("chr1", "ACGTNNNNNNNNACGT")])?;
784        let reader = ReferenceReader::new(file.path())?;
785
786        // Fetch run of N's
787        let n_run = reader.fetch("chr1", Position::try_from(5)?, Position::try_from(12)?)?;
788        assert_eq!(n_run, b"NNNNNNNN");
789
790        // Fetch across N boundary
791        let across = reader.fetch("chr1", Position::try_from(3)?, Position::try_from(14)?)?;
792        assert_eq!(across, b"GTNNNNNNNNAC");
793
794        Ok(())
795    }
796
797    #[test]
798    fn test_position_one_based() -> Result<()> {
799        use crate::sam::builder::create_test_fasta;
800
801        let file = create_test_fasta(&[("chr1", "ACGTN")])?;
802        let reader = ReferenceReader::new(file.path())?;
803
804        // Position 1 should be first base 'A', not second
805        assert_eq!(reader.base_at("chr1", Position::try_from(1)?)?, b'A');
806        assert_eq!(reader.base_at("chr1", Position::try_from(2)?)?, b'C');
807
808        // fetch(1, 1) should return single base 'A'
809        let seq = reader.fetch("chr1", Position::try_from(1)?, Position::try_from(1)?)?;
810        assert_eq!(seq, b"A");
811
812        // fetch(1, 2) should return "AC"
813        let seq = reader.fetch("chr1", Position::try_from(1)?, Position::try_from(2)?)?;
814        assert_eq!(seq, b"AC");
815
816        Ok(())
817    }
818
819    #[test]
820    fn test_find_dict_path_replacing_convention() -> Result<()> {
821        // Test the fgbio/HTSJDK/Picard convention: ref.fa -> ref.dict
822        let temp_dir = tempfile::tempdir()?;
823        let fasta_path = temp_dir.path().join("ref.fa");
824        let dict_path = temp_dir.path().join("ref.dict");
825
826        // Create empty files
827        std::fs::write(&fasta_path, "")?;
828        std::fs::write(&dict_path, "")?;
829
830        let found = find_dict_path(&fasta_path);
831        assert!(found.is_some());
832        assert_eq!(found.expect("should find dictionary path"), dict_path);
833
834        Ok(())
835    }
836
837    #[test]
838    fn test_find_dict_path_appending_convention() -> Result<()> {
839        // Test the GATK convention: ref.fa -> ref.fa.dict
840        let temp_dir = tempfile::tempdir()?;
841        let fasta_path = temp_dir.path().join("ref.fa");
842        let dict_path = temp_dir.path().join("ref.fa.dict");
843
844        // Create empty files
845        std::fs::write(&fasta_path, "")?;
846        std::fs::write(&dict_path, "")?;
847
848        let found = find_dict_path(&fasta_path);
849        assert!(found.is_some());
850        assert_eq!(found.expect("should find dictionary path"), dict_path);
851
852        Ok(())
853    }
854
855    #[test]
856    fn test_find_dict_path_prefers_replacing_convention() -> Result<()> {
857        // When both exist, prefer the fgbio/HTSJDK convention (ref.dict)
858        let temp_dir = tempfile::tempdir()?;
859        let fasta_path = temp_dir.path().join("ref.fa");
860        let dict_replacing = temp_dir.path().join("ref.dict");
861        let dict_appending = temp_dir.path().join("ref.fa.dict");
862
863        // Create all files
864        std::fs::write(&fasta_path, "")?;
865        std::fs::write(&dict_replacing, "")?;
866        std::fs::write(&dict_appending, "")?;
867
868        let found = find_dict_path(&fasta_path);
869        assert!(found.is_some());
870        // Should find ref.dict first (fgbio/HTSJDK convention)
871        assert_eq!(found.expect("should find dictionary path"), dict_replacing);
872
873        Ok(())
874    }
875
876    #[test]
877    fn test_find_dict_path_not_found() {
878        // When no dict file exists, return None
879        let temp_dir = tempfile::tempdir().expect("creating temp file/dir should succeed");
880        let fasta_path = temp_dir.path().join("ref.fa");
881
882        // Create only the FASTA file, no dict
883        std::fs::write(&fasta_path, "").expect("writing file should succeed");
884
885        let found = find_dict_path(&fasta_path);
886        assert!(found.is_none());
887    }
888
889    #[test]
890    fn test_find_dict_path_fasta_extension() -> Result<()> {
891        // Test with .fasta extension: ref.fasta -> ref.dict
892        let temp_dir = tempfile::tempdir()?;
893        let fasta_path = temp_dir.path().join("ref.fasta");
894        let dict_path = temp_dir.path().join("ref.dict");
895
896        std::fs::write(&fasta_path, "")?;
897        std::fs::write(&dict_path, "")?;
898
899        let found = find_dict_path(&fasta_path);
900        assert!(found.is_some());
901        assert_eq!(found.expect("should find dictionary path"), dict_path);
902
903        Ok(())
904    }
905
906    #[test]
907    fn test_find_dict_path_fasta_appending_convention() -> Result<()> {
908        // Test with .fasta extension: ref.fasta -> ref.fasta.dict
909        let temp_dir = tempfile::tempdir()?;
910        let fasta_path = temp_dir.path().join("ref.fasta");
911        let dict_path = temp_dir.path().join("ref.fasta.dict");
912
913        std::fs::write(&fasta_path, "")?;
914        std::fs::write(&dict_path, "")?;
915
916        let found = find_dict_path(&fasta_path);
917        assert!(found.is_some());
918        assert_eq!(found.expect("should find dictionary path"), dict_path);
919
920        Ok(())
921    }
922}