bio_forge/model/
chain.rs

1//! Ordered residue collections representing biomolecular chains.
2//!
3//! Chains own a sequence of `Residue` records, enforce uniqueness on residue identifiers,
4//! and expose iterator helpers that cascade access down to atoms. Higher-level operations
5//! such as topology repair or solvation treat chains as the primary unit when traversing a
6//! structure.
7
8use super::residue::Residue;
9use std::fmt;
10
11/// Polymer chain containing an ordered list of residues.
12///
13/// The chain preserves the order residues were added, mirrors the chain identifier seen in
14/// structure files (e.g., `"A"`), and provides iteration helpers for both residues and
15/// atoms while keeping the internal storage private.
16#[derive(Debug, Clone, PartialEq)]
17pub struct Chain {
18    /// Chain identifier matching the source structure (usually a single character).
19    pub id: String,
20    /// Internal storage preserving insertion order for residues.
21    residues: Vec<Residue>,
22}
23
24impl Chain {
25    /// Creates an empty chain with the provided identifier.
26    ///
27    /// Use this constructor when building structures procedurally or when importing from
28    /// file formats that enumerate chains.
29    ///
30    /// # Arguments
31    ///
32    /// * `id` - Label matching the original structure's chain identifier.
33    ///
34    /// # Returns
35    ///
36    /// A new `Chain` with no residues.
37    pub fn new(id: &str) -> Self {
38        Self {
39            id: id.to_string(),
40            residues: Vec::new(),
41        }
42    }
43
44    /// Adds a residue to the chain while preventing duplicate identifiers.
45    ///
46    /// Residues are appended in insertion order. Duplicate `(id, insertion_code)` pairs are
47    /// rejected during debug builds to guard against malformed inputs.
48    ///
49    /// # Arguments
50    ///
51    /// * `residue` - The residue to append.
52    pub fn add_residue(&mut self, residue: Residue) {
53        // Avoids ambiguous lookups by ensuring each (id, insertion_code) pair is unique.
54        debug_assert!(
55            self.residue(residue.id, residue.insertion_code).is_none(),
56            "Attempted to add a duplicate residue ID '{}' (ic: {:?}) to chain '{}'",
57            residue.id,
58            residue.insertion_code,
59            self.id
60        );
61        self.residues.push(residue);
62    }
63
64    /// Looks up a residue by identifier and optional insertion code.
65    ///
66    /// # Arguments
67    ///
68    /// * `id` - Residue sequence number.
69    /// * `insertion_code` - Optional insertion code used in PDB/mmCIF records.
70    ///
71    /// # Returns
72    ///
73    /// `Some(&Residue)` if a matching residue exists; otherwise `None`.
74    pub fn residue(&self, id: i32, insertion_code: Option<char>) -> Option<&Residue> {
75        self.residues
76            .iter()
77            .find(|r| r.id == id && r.insertion_code == insertion_code)
78    }
79
80    /// Fetches a mutable reference to a residue by identifier.
81    ///
82    /// # Arguments
83    ///
84    /// * `id` - Residue sequence number.
85    /// * `insertion_code` - Optional insertion code qualifier.
86    ///
87    /// # Returns
88    ///
89    /// `Some(&mut Residue)` when present; otherwise `None`.
90    pub fn residue_mut(&mut self, id: i32, insertion_code: Option<char>) -> Option<&mut Residue> {
91        self.residues
92            .iter_mut()
93            .find(|r| r.id == id && r.insertion_code == insertion_code)
94    }
95
96    /// Returns an immutable slice containing all residues in order.
97    ///
98    /// Useful for bulk analysis or when interfacing with APIs that operate on slices.
99    ///
100    /// # Returns
101    ///
102    /// Slice view of the underlying residue list.
103    pub fn residues(&self) -> &[Residue] {
104        &self.residues
105    }
106
107    /// Reports the number of residues stored in the chain.
108    ///
109    /// # Returns
110    ///
111    /// Count of residues currently tracked.
112    pub fn residue_count(&self) -> usize {
113        self.residues.len()
114    }
115
116    /// Indicates whether the chain contains no residues.
117    ///
118    /// # Returns
119    ///
120    /// `true` when the chain is empty.
121    pub fn is_empty(&self) -> bool {
122        self.residues.is_empty()
123    }
124
125    /// Provides an iterator over immutable residue references.
126    ///
127    /// This mirrors `residues()` but avoids exposing slice internals and composes nicely with
128    /// iterator adaptors.
129    ///
130    /// # Returns
131    ///
132    /// A standard slice iterator over `Residue` references.
133    pub fn iter_residues(&self) -> std::slice::Iter<'_, Residue> {
134        self.residues.iter()
135    }
136
137    /// Provides an iterator over mutable residue references.
138    ///
139    /// # Returns
140    ///
141    /// A mutable slice iterator that allows in-place modifications.
142    pub fn iter_residues_mut(&mut self) -> std::slice::IterMut<'_, Residue> {
143        self.residues.iter_mut()
144    }
145
146    /// Iterates over all atoms contained in the chain.
147    ///
148    /// Residues are traversed in order and their atom iterators flattened, yielding atoms in
149    /// the same relative ordering seen in the original structure.
150    ///
151    /// # Returns
152    ///
153    /// Iterator that yields immutable `Atom` references.
154    pub fn iter_atoms(&self) -> impl Iterator<Item = &super::atom::Atom> {
155        self.residues.iter().flat_map(|r| r.iter_atoms())
156    }
157
158    /// Iterates over all atoms with mutable access.
159    ///
160    /// # Returns
161    ///
162    /// Iterator producing mutable `Atom` references for bulk editing operations.
163    pub fn iter_atoms_mut(&mut self) -> impl Iterator<Item = &mut super::atom::Atom> {
164        self.residues.iter_mut().flat_map(|r| r.iter_atoms_mut())
165    }
166
167    /// Retains only residues that satisfy the provided predicate.
168    ///
169    /// # Arguments
170    ///
171    /// * `f` - Predicate invoked for each residue; keep the residue when it returns `true`.
172    pub fn retain_residues<F>(&mut self, mut f: F)
173    where
174        F: FnMut(&Residue) -> bool,
175    {
176        self.residues.retain(|residue| f(residue));
177    }
178
179    /// Removes a residue by identifier and returns ownership if found.
180    ///
181    /// # Arguments
182    ///
183    /// * `id` - Residue number to remove.
184    /// * `insertion_code` - Optional insertion qualifier.
185    ///
186    /// # Returns
187    ///
188    /// `Some(Residue)` containing the removed residue; otherwise `None`.
189    pub fn remove_residue(&mut self, id: i32, insertion_code: Option<char>) -> Option<Residue> {
190        if let Some(index) = self
191            .residues
192            .iter()
193            .position(|r| r.id == id && r.insertion_code == insertion_code)
194        {
195            Some(self.residues.remove(index))
196        } else {
197            None
198        }
199    }
200}
201
202impl fmt::Display for Chain {
203    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
204        write!(
205            f,
206            "Chain {{ id: \"{}\", residues: {} }}",
207            self.id,
208            self.residue_count()
209        )
210    }
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216    use crate::model::atom::Atom;
217    use crate::model::types::{Element, Point, ResidueCategory, StandardResidue};
218
219    fn sample_residue(id: i32, name: &str) -> Residue {
220        Residue::new(
221            id,
222            None,
223            name,
224            Some(StandardResidue::ALA),
225            ResidueCategory::Standard,
226        )
227    }
228
229    #[test]
230    fn chain_new_creates_correct_chain() {
231        let chain = Chain::new("A");
232
233        assert_eq!(chain.id, "A");
234        assert!(chain.residues.is_empty());
235    }
236
237    #[test]
238    fn chain_add_residue_adds_residue_correctly() {
239        let mut chain = Chain::new("A");
240        let residue = Residue::new(
241            1,
242            None,
243            "ALA",
244            Some(StandardResidue::ALA),
245            ResidueCategory::Standard,
246        );
247
248        chain.add_residue(residue);
249
250        assert_eq!(chain.residue_count(), 1);
251        assert!(chain.residue(1, None).is_some());
252        assert_eq!(chain.residue(1, None).unwrap().name, "ALA");
253    }
254
255    #[test]
256    fn chain_residue_returns_correct_residue() {
257        let mut chain = Chain::new("A");
258        let residue = Residue::new(
259            1,
260            None,
261            "ALA",
262            Some(StandardResidue::ALA),
263            ResidueCategory::Standard,
264        );
265        chain.add_residue(residue);
266
267        let retrieved = chain.residue(1, None);
268
269        assert!(retrieved.is_some());
270        assert_eq!(retrieved.unwrap().id, 1);
271        assert_eq!(retrieved.unwrap().name, "ALA");
272    }
273
274    #[test]
275    fn chain_residue_returns_none_for_nonexistent_residue() {
276        let chain = Chain::new("A");
277
278        let retrieved = chain.residue(999, None);
279
280        assert!(retrieved.is_none());
281    }
282
283    #[test]
284    fn chain_residue_mut_returns_correct_mutable_residue() {
285        let mut chain = Chain::new("A");
286        let residue = Residue::new(
287            1,
288            None,
289            "ALA",
290            Some(StandardResidue::ALA),
291            ResidueCategory::Standard,
292        );
293        chain.add_residue(residue);
294
295        let retrieved = chain.residue_mut(1, None);
296
297        assert!(retrieved.is_some());
298        assert_eq!(retrieved.unwrap().id, 1);
299    }
300
301    #[test]
302    fn chain_residue_mut_returns_none_for_nonexistent_residue() {
303        let mut chain = Chain::new("A");
304
305        let retrieved = chain.residue_mut(999, None);
306
307        assert!(retrieved.is_none());
308    }
309
310    #[test]
311    fn chain_residues_returns_correct_slice() {
312        let mut chain = Chain::new("A");
313        let residue1 = Residue::new(
314            1,
315            None,
316            "ALA",
317            Some(StandardResidue::ALA),
318            ResidueCategory::Standard,
319        );
320        let residue2 = Residue::new(
321            2,
322            None,
323            "GLY",
324            Some(StandardResidue::GLY),
325            ResidueCategory::Standard,
326        );
327        chain.add_residue(residue1);
328        chain.add_residue(residue2);
329
330        let residues = chain.residues();
331
332        assert_eq!(residues.len(), 2);
333        assert_eq!(residues[0].id, 1);
334        assert_eq!(residues[1].id, 2);
335    }
336
337    #[test]
338    fn chain_residue_count_returns_correct_count() {
339        let mut chain = Chain::new("A");
340
341        assert_eq!(chain.residue_count(), 0);
342
343        let residue = Residue::new(
344            1,
345            None,
346            "ALA",
347            Some(StandardResidue::ALA),
348            ResidueCategory::Standard,
349        );
350        chain.add_residue(residue);
351
352        assert_eq!(chain.residue_count(), 1);
353    }
354
355    #[test]
356    fn chain_is_empty_returns_true_for_empty_chain() {
357        let chain = Chain::new("A");
358
359        assert!(chain.is_empty());
360    }
361
362    #[test]
363    fn chain_is_empty_returns_false_for_non_empty_chain() {
364        let mut chain = Chain::new("A");
365        let residue = Residue::new(
366            1,
367            None,
368            "ALA",
369            Some(StandardResidue::ALA),
370            ResidueCategory::Standard,
371        );
372        chain.add_residue(residue);
373
374        assert!(!chain.is_empty());
375    }
376
377    #[test]
378    fn chain_iter_residues_iterates_correctly() {
379        let mut chain = Chain::new("A");
380        let residue1 = Residue::new(
381            1,
382            None,
383            "ALA",
384            Some(StandardResidue::ALA),
385            ResidueCategory::Standard,
386        );
387        let residue2 = Residue::new(
388            2,
389            None,
390            "GLY",
391            Some(StandardResidue::GLY),
392            ResidueCategory::Standard,
393        );
394        chain.add_residue(residue1);
395        chain.add_residue(residue2);
396
397        let mut ids = Vec::new();
398        for residue in chain.iter_residues() {
399            ids.push(residue.id);
400        }
401
402        assert_eq!(ids, vec![1, 2]);
403    }
404
405    #[test]
406    fn chain_iter_residues_mut_iterates_correctly() {
407        let mut chain = Chain::new("A");
408        let residue = Residue::new(
409            1,
410            None,
411            "ALA",
412            Some(StandardResidue::ALA),
413            ResidueCategory::Standard,
414        );
415        chain.add_residue(residue);
416
417        for residue in chain.iter_residues_mut() {
418            residue.position = crate::model::types::ResiduePosition::Internal;
419        }
420
421        assert_eq!(
422            chain.residue(1, None).unwrap().position,
423            crate::model::types::ResiduePosition::Internal
424        );
425    }
426
427    #[test]
428    fn chain_iter_atoms_iterates_over_all_atoms() {
429        let mut chain = Chain::new("A");
430        let mut residue1 = Residue::new(
431            1,
432            None,
433            "ALA",
434            Some(StandardResidue::ALA),
435            ResidueCategory::Standard,
436        );
437        let mut residue2 = Residue::new(
438            2,
439            None,
440            "GLY",
441            Some(StandardResidue::GLY),
442            ResidueCategory::Standard,
443        );
444
445        let atom1 = Atom::new("CA1", Element::C, Point::new(0.0, 0.0, 0.0));
446        let atom2 = Atom::new("CB1", Element::C, Point::new(1.0, 0.0, 0.0));
447        let atom3 = Atom::new("CA2", Element::C, Point::new(2.0, 0.0, 0.0));
448
449        residue1.add_atom(atom1);
450        residue1.add_atom(atom2);
451        residue2.add_atom(atom3);
452
453        chain.add_residue(residue1);
454        chain.add_residue(residue2);
455
456        let mut atom_names = Vec::new();
457        for atom in chain.iter_atoms() {
458            atom_names.push(atom.name.clone());
459        }
460
461        assert_eq!(atom_names, vec!["CA1", "CB1", "CA2"]);
462    }
463
464    #[test]
465    fn chain_iter_atoms_mut_iterates_over_all_atoms_mutably() {
466        let mut chain = Chain::new("A");
467        let mut residue = Residue::new(
468            1,
469            None,
470            "ALA",
471            Some(StandardResidue::ALA),
472            ResidueCategory::Standard,
473        );
474        let atom = Atom::new("CA", Element::C, Point::new(0.0, 0.0, 0.0));
475        residue.add_atom(atom);
476        chain.add_residue(residue);
477
478        for atom in chain.iter_atoms_mut() {
479            atom.translate_by(&nalgebra::Vector3::new(1.0, 0.0, 0.0));
480        }
481
482        let translated_atom = chain.residue(1, None).unwrap().atom("CA").unwrap();
483        assert!((translated_atom.pos.x - 1.0).abs() < 1e-10);
484    }
485
486    #[test]
487    fn chain_iter_atoms_returns_empty_iterator_for_empty_chain() {
488        let chain = Chain::new("A");
489
490        let count = chain.iter_atoms().count();
491
492        assert_eq!(count, 0);
493    }
494
495    #[test]
496    fn chain_iter_atoms_mut_returns_empty_iterator_for_empty_chain() {
497        let mut chain = Chain::new("A");
498
499        let count = chain.iter_atoms_mut().count();
500
501        assert_eq!(count, 0);
502    }
503
504    #[test]
505    fn chain_display_formats_correctly() {
506        let mut chain = Chain::new("A");
507        let residue = Residue::new(
508            1,
509            None,
510            "ALA",
511            Some(StandardResidue::ALA),
512            ResidueCategory::Standard,
513        );
514        chain.add_residue(residue);
515
516        let display = format!("{}", chain);
517        let expected = "Chain { id: \"A\", residues: 1 }";
518
519        assert_eq!(display, expected);
520    }
521
522    #[test]
523    fn chain_display_formats_empty_chain_correctly() {
524        let chain = Chain::new("B");
525
526        let display = format!("{}", chain);
527        let expected = "Chain { id: \"B\", residues: 0 }";
528
529        assert_eq!(display, expected);
530    }
531
532    #[test]
533    fn chain_clone_creates_identical_copy() {
534        let mut chain = Chain::new("A");
535        let residue = Residue::new(
536            1,
537            None,
538            "ALA",
539            Some(StandardResidue::ALA),
540            ResidueCategory::Standard,
541        );
542        chain.add_residue(residue);
543
544        let cloned = chain.clone();
545
546        assert_eq!(chain, cloned);
547        assert_eq!(chain.id, cloned.id);
548        assert_eq!(chain.residues, cloned.residues);
549    }
550
551    #[test]
552    fn chain_partial_eq_compares_correctly() {
553        let mut chain1 = Chain::new("A");
554        let mut chain2 = Chain::new("A");
555        let residue = Residue::new(
556            1,
557            None,
558            "ALA",
559            Some(StandardResidue::ALA),
560            ResidueCategory::Standard,
561        );
562        chain1.add_residue(residue.clone());
563        chain2.add_residue(residue);
564
565        let chain3 = Chain::new("B");
566
567        assert_eq!(chain1, chain2);
568        assert_ne!(chain1, chain3);
569    }
570
571    #[test]
572    fn chain_with_multiple_residues_and_atoms() {
573        let mut chain = Chain::new("A");
574
575        for i in 1..=3 {
576            let mut residue = Residue::new(
577                i,
578                None,
579                &format!("RES{}", i),
580                None,
581                ResidueCategory::Standard,
582            );
583            let atom = Atom::new(
584                &format!("ATOM{}", i),
585                Element::C,
586                Point::new(i as f64, 0.0, 0.0),
587            );
588            residue.add_atom(atom);
589            chain.add_residue(residue);
590        }
591
592        assert_eq!(chain.residue_count(), 3);
593        assert_eq!(chain.iter_atoms().count(), 3);
594
595        let residue = chain.residue(2, None).unwrap();
596        assert_eq!(residue.name, "RES2");
597        assert_eq!(residue.atom("ATOM2").unwrap().name, "ATOM2");
598    }
599
600    #[test]
601    fn chain_handles_insertion_codes_correctly() {
602        let mut chain = Chain::new("A");
603        let residue1 = Residue::new(
604            1,
605            None,
606            "ALA",
607            Some(StandardResidue::ALA),
608            ResidueCategory::Standard,
609        );
610        let residue2 = Residue::new(
611            1,
612            Some('A'),
613            "ALA",
614            Some(StandardResidue::ALA),
615            ResidueCategory::Standard,
616        );
617
618        chain.add_residue(residue1);
619        chain.add_residue(residue2);
620
621        assert_eq!(chain.residue_count(), 2);
622        assert!(chain.residue(1, None).is_some());
623        assert!(chain.residue(1, Some('A')).is_some());
624        assert_eq!(
625            chain.residue(1, Some('A')).unwrap().insertion_code,
626            Some('A')
627        );
628    }
629
630    #[test]
631    fn chain_retain_residues_filters_using_predicate() {
632        let mut chain = Chain::new("A");
633        chain.add_residue(sample_residue(1, "ALA"));
634        chain.add_residue(sample_residue(2, "GLY"));
635        chain.add_residue(sample_residue(3, "SER"));
636
637        chain.retain_residues(|residue| residue.id % 2 == 1);
638
639        let ids: Vec<i32> = chain.iter_residues().map(|r| r.id).collect();
640        assert_eq!(ids, vec![1, 3]);
641    }
642
643    #[test]
644    fn chain_remove_residue_returns_removed_value() {
645        let mut chain = Chain::new("A");
646        chain.add_residue(sample_residue(5, "ALA"));
647        chain.add_residue(sample_residue(6, "GLY"));
648
649        let removed = chain.remove_residue(5, None);
650
651        assert!(removed.is_some());
652        assert_eq!(removed.unwrap().id, 5);
653        assert!(chain.residue(5, None).is_none());
654        assert_eq!(chain.residue_count(), 1);
655    }
656
657    #[test]
658    fn chain_remove_residue_returns_none_for_missing_entry() {
659        let mut chain = Chain::new("A");
660        chain.add_residue(sample_residue(5, "ALA"));
661
662        assert!(chain.remove_residue(42, None).is_none());
663        assert_eq!(chain.residue_count(), 1);
664    }
665}