fibertools_rs/utils/bamannotations.rs
1use molecular_annotation::{AnnotationInfo, MolecularAnnotations};
2use std::cell::OnceCell;
3
4/// Extract the single per-annotation quality fibertools expects, returning
5/// 0 when the annotation has none. Debug-asserts that the type carries at
6/// most one quality — fibertools' current annotation types (`nuc`, `msp`,
7/// `fire`, `m6a`, `cpg`) are all single-quality. If you add a multi-quality
8/// type, pick the index explicitly at the call site instead of using this
9/// helper.
10#[inline]
11pub fn primary_qual(qualities: &[u8], type_name: &str) -> u8 {
12 debug_assert!(
13 qualities.len() <= 1,
14 "primary_qual: type {:?} carries {} qualities; fibertools expects \u{2264} 1",
15 type_name,
16 qualities.len(),
17 );
18 qualities.first().copied().unwrap_or(0)
19}
20
21/// View over a specific annotation type, providing the column-oriented
22/// surface fibertools historically got from `FiberAnnotations` while
23/// delegating per-annotation iteration to the library's
24/// [`AnnotationInfo`].
25///
26/// Stays valid when the type is absent: every accessor returns an empty
27/// `Vec` / zero count, so call sites avoid `Option` plumbing.
28///
29/// All accessors and the per-annotation iterator yield results in
30/// **BAM-orient ascending** order. The spec stores annotations in
31/// molecular order; for reverse-aligned reads we reverse so consumers
32/// (BED12 blocks, pileup intervals, TSV columns) get ascending output.
33#[derive(Debug)]
34pub struct AnnotationTypeView<'a> {
35 annot: &'a MolecularAnnotations,
36 type_name: &'a str,
37 /// Memoizes [`Self::bam_ordered`]. Reference-coordinate liftover for
38 /// every annotation runs in `iter_type`, so without this cache each
39 /// accessor (and each call in a hot loop) re-lifts the whole set. The
40 /// cell is per-view, so callers that reuse one view across many calls
41 /// (e.g. the FIRE per-base windowing loop) pay the liftover once.
42 ordered: OnceCell<Vec<AnnotationInfo<'a>>>,
43}
44
45impl<'a> AnnotationTypeView<'a> {
46 pub(crate) fn new(annot: &'a MolecularAnnotations, type_name: &'a str) -> Self {
47 Self {
48 annot,
49 type_name,
50 ordered: OnceCell::new(),
51 }
52 }
53
54 pub fn len(&self) -> usize {
55 self.annot
56 .get_type(self.type_name)
57 .map(|t| t.annotations.len())
58 .unwrap_or(0)
59 }
60 pub fn is_empty(&self) -> bool {
61 self.len() == 0
62 }
63
64 /// Materialize the annotation infos for this type in BAM-orient
65 /// ascending order, memoized for the life of the view. Reference-coord
66 /// liftover happens here (in `iter_type`), so caching is what keeps
67 /// repeated accessor / per-call use from re-lifting the whole set.
68 fn bam_ordered(&self) -> &[AnnotationInfo<'a>] {
69 self.ordered.get_or_init(|| {
70 let Some(it) = self.annot.iter_type(self.type_name) else {
71 return Vec::new();
72 };
73 let mut v: Vec<AnnotationInfo<'a>> = it.collect();
74 if self.annot.is_reverse_aligned() {
75 v.reverse();
76 }
77 v
78 })
79 }
80
81 /// Borrow the memoized infos directly, BAM-orient ascending. Lets hot
82 /// loops scan query coordinates without the per-call `Vec` allocation
83 /// that the column accessors incur.
84 pub fn infos(&self) -> &[AnnotationInfo<'a>] {
85 self.bam_ordered()
86 }
87
88 /// Count annotations whose BAM-orient `query_start` falls in
89 /// `[start, end)`. Allocation-free over the memoized infos.
90 pub fn count_query_in(&self, start: i64, end: i64) -> usize {
91 self.bam_ordered()
92 .iter()
93 .filter(|a| {
94 let pos = a.query_start as i64;
95 pos >= start && pos < end
96 })
97 .count()
98 }
99
100 pub fn starts(&self) -> Vec<i64> {
101 self.bam_ordered()
102 .iter()
103 .map(|a| a.query_start as i64)
104 .collect()
105 }
106 pub fn ends(&self) -> Vec<i64> {
107 self.bam_ordered()
108 .iter()
109 .map(|a| a.query_end as i64)
110 .collect()
111 }
112 pub fn option_starts(&self) -> Vec<Option<i64>> {
113 self.bam_ordered()
114 .iter()
115 .map(|a| Some(a.query_start as i64))
116 .collect()
117 }
118 pub fn option_ends(&self) -> Vec<Option<i64>> {
119 self.bam_ordered()
120 .iter()
121 .map(|a| Some(a.query_end as i64))
122 .collect()
123 }
124 pub fn option_lengths(&self) -> Vec<Option<i64>> {
125 self.bam_ordered()
126 .iter()
127 .map(|a| Some((a.query_end - a.query_start) as i64))
128 .collect()
129 }
130 pub fn lengths(&self) -> Vec<i64> {
131 self.bam_ordered()
132 .iter()
133 .map(|a| (a.query_end - a.query_start) as i64)
134 .collect()
135 }
136
137 /// Convenience over [`Self::qual_at`] for the common case where the
138 /// annotation type carries at most one quality per annotation. Returns
139 /// 0 for annotations with no qualities.
140 ///
141 /// All fibertools-rs annotation types (`nuc`, `msp+` / `msp+Q`,
142 /// `fire+P`, `m6a+Q`, `cpg+Q`) are single-quality; this method debug-
143 /// asserts that invariant. If you add a multi-quality type, call
144 /// [`Self::qual_at`] with an explicit index instead — `qual()`
145 /// silently dropping quality columns would be a footgun.
146 pub fn qual(&self) -> Vec<u8> {
147 self.qual_at(0)
148 }
149
150 /// Per-annotation quality at the given index, BAM-orient ascending.
151 /// Returns 0 when the annotation has fewer than `idx + 1` qualities.
152 ///
153 /// `qual_at(0)` is the canonical single-quality accessor and
154 /// debug-asserts that the type has at most one quality. Other indices
155 /// skip the assertion — the caller is presumed to know the type's
156 /// `QualitySpec`.
157 pub fn qual_at(&self, idx: usize) -> Vec<u8> {
158 self.bam_ordered()
159 .iter()
160 .map(|a| {
161 debug_assert!(
162 idx > 0 || a.qualities.len() <= 1,
163 "AnnotationTypeView::qual() called on multi-quality type {:?} ({} qualities); use qual_at(idx)",
164 self.type_name,
165 a.qualities.len(),
166 );
167 a.qualities.get(idx).copied().unwrap_or(0)
168 })
169 .collect()
170 }
171
172 pub fn reference_starts(&self) -> Vec<Option<i64>> {
173 self.bam_ordered()
174 .iter()
175 .map(|a| a.ref_start.map(|x| x as i64))
176 .collect()
177 }
178 pub fn reference_ends(&self) -> Vec<Option<i64>> {
179 self.bam_ordered()
180 .iter()
181 .map(|a| a.ref_end.map(|x| x as i64))
182 .collect()
183 }
184 pub fn reference_lengths(&self) -> Vec<Option<i64>> {
185 self.bam_ordered()
186 .iter()
187 .map(|a| match (a.ref_start, a.ref_end) {
188 (Some(s), Some(e)) => Some((e - s) as i64),
189 _ => None,
190 })
191 .collect()
192 }
193
194 /// Iterate per-annotation in BAM-orient ascending order, yielding the
195 /// library's [`AnnotationInfo`] view directly. Call sites read fields
196 /// like `info.query_start`, `info.ref_start`, `info.qualities`.
197 pub fn iter(&self) -> std::vec::IntoIter<AnnotationInfo<'a>> {
198 self.bam_ordered().to_vec().into_iter()
199 }
200}
201
202impl<'a> IntoIterator for &AnnotationTypeView<'a> {
203 type Item = AnnotationInfo<'a>;
204 type IntoIter = std::vec::IntoIter<AnnotationInfo<'a>>;
205 fn into_iter(self) -> Self::IntoIter {
206 self.iter()
207 }
208}