use std::iter::FusedIterator;
use molframe_core::index::{AtomIndex, ChainIndex, ModelIndex, ResidueIndex};
use molframe_core::structure::{
AtomRef, ChainRef, ModelRef, ResidueRef, Structure as CoreStructure,
};
macro_rules! hierarchy_iter {
($(#[$meta:meta])* $name:ident, $item:ty, $index:ty, $handle:ident) => {
$(#[$meta])*
#[derive(Clone, Debug)]
pub struct $name<'a> {
structure: &'a CoreStructure,
position: usize,
end: usize,
}
impl<'a> Iterator for $name<'a> {
type Item = $item;
fn next(&mut self) -> Option<Self::Item> {
if self.position >= self.end {
return None;
}
let ordinal = u32::try_from(self.position).ok()?;
self.position += 1;
self.structure.$handle(<$index>::new(ordinal))
}
fn size_hint(&self) -> (usize, Option<usize>) {
let remaining = self.end - self.position;
(remaining, Some(remaining))
}
}
impl ExactSizeIterator for $name<'_> {
fn len(&self) -> usize {
self.end - self.position
}
}
impl FusedIterator for $name<'_> {}
};
}
macro_rules! hierarchy_collection {
(
$(#[$meta:meta])* $name:ident, $iter:ident, $item:ty, $index:ty, $handle:ident
) => {
$(#[$meta])*
#[derive(Clone, Copy, Debug)]
pub struct $name<'a> {
structure: &'a CoreStructure,
}
impl<'a> $name<'a> {
pub(super) fn new(structure: &'a CoreStructure) -> Self {
Self { structure }
}
#[must_use]
pub fn len(&self) -> usize {
self.iteration_end()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.iteration_end() == 0
}
#[must_use]
pub fn get(&self, position: usize) -> Option<$item> {
let ordinal = u32::try_from(position).ok()?;
self.structure.$handle(<$index>::new(ordinal))
}
#[must_use]
pub fn iter(&self) -> $iter<'a> {
self.structure_iter()
}
fn structure_iter(&self) -> $iter<'a> {
$iter {
structure: self.structure,
position: 0,
end: self.iteration_end(),
}
}
}
impl<'a> IntoIterator for $name<'a> {
type Item = $item;
type IntoIter = $iter<'a>;
fn into_iter(self) -> Self::IntoIter {
self.structure_iter()
}
}
impl<'a> IntoIterator for &$name<'a> {
type Item = $item;
type IntoIter = $iter<'a>;
fn into_iter(self) -> Self::IntoIter {
self.structure_iter()
}
}
};
}
hierarchy_iter!(
ChainsIter,
ChainRef<'a>,
ChainIndex,
chain
);
hierarchy_collection!(
Chains,
ChainsIter,
ChainRef<'a>,
ChainIndex,
chain
);
impl Chains<'_> {
fn iteration_end(self) -> usize {
self.structure.chain_count()
}
}
hierarchy_iter!(
ResiduesIter,
ResidueRef<'a>,
ResidueIndex,
residue
);
hierarchy_collection!(
Residues,
ResiduesIter,
ResidueRef<'a>,
ResidueIndex,
residue
);
impl Residues<'_> {
fn iteration_end(self) -> usize {
self.structure.residue_count()
}
}
hierarchy_iter!(
AtomsIter,
AtomRef<'a>,
AtomIndex,
atom
);
hierarchy_collection!(
Atoms,
AtomsIter,
AtomRef<'a>,
AtomIndex,
atom
);
impl Atoms<'_> {
fn iteration_end(self) -> usize {
self.structure.atom_count() as usize
}
}
hierarchy_iter!(
ModelsIter,
ModelRef<'a>,
ModelIndex,
model
);
hierarchy_collection!(
Models,
ModelsIter,
ModelRef<'a>,
ModelIndex,
model
);
impl Models<'_> {
fn iteration_end(self) -> usize {
self.structure.model_count()
}
}