molframe-query 0.3.0

Typed and textual structural selection plans for molframe
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114

//! Ordered selection traversal shared by column predicates.

use super::AtomContext;
use molframe_core::selection::AtomSelection;
use molframe_core::structure::Structure;

pub(super) struct SelectionCursor<I> {
    iter: I,
    next: Option<u32>,
}

impl<I> SelectionCursor<I>
where
    I: Iterator<Item = u32>,
{
    pub(super) fn new(mut iter: I) -> Self {
        let next = iter.next();
        Self { iter, next }
    }

    pub(super) fn is_exhausted(&self) -> bool {
        self.next.is_none()
    }

    pub(super) fn skip_before(&mut self, minimum: u32) {
        while self.next.is_some_and(|atom| atom < minimum) {
            self.next = self.iter.next();
        }
    }

    pub(super) fn matches(&mut self, atom: u32) -> bool {
        self.skip_before(atom);
        if self.next != Some(atom) {
            return false;
        }
        self.next = self.iter.next();
        true
    }

    pub(super) fn next_before(&mut self, end: u32) -> Option<u32> {
        let atom = self.next?;
        if atom >= end {
            return None;
        }
        self.next = self.iter.next();
        Some(atom)
    }
}

pub(crate) fn scan(
    structure: &Structure,
    universe: &AtomSelection,
    mut accepts: impl FnMut(AtomContext<'_>) -> bool,
) -> AtomSelection {
    if universe.is_empty() {
        return AtomSelection::Empty;
    }
    // The universe's own count is the ceiling on what a predicate can accept,
    // so reserving it once replaces a growth sequence that reallocated as the
    // matches accumulated. A predicate accepting few atoms wastes the
    // reservation only until the buffer is dropped, while one accepting most of
    // the molecule — the common case for a broad selection — would otherwise
    // rebuild the buffer repeatedly as it grows.
    //
    // A universe larger than the address space reserves nothing and grows as it
    // did before: the reservation is an optimisation, and the selection it is
    // sized for cannot exist on a platform that cannot index it.
    let mut selected = match usize::try_from(universe.len()) {
        Ok(ceiling) => Vec::with_capacity(ceiling),
        Err(_) => Vec::new(),
    };
    visit(structure, universe, |context| {
        if accepts(context) {
            selected.push(context.atom.index().get());
        }
    });
    selection_from_sorted(selected)
}

pub(super) fn visit<'a>(
    structure: &'a Structure,
    universe: &AtomSelection,
    mut visitor: impl FnMut(AtomContext<'a>),
) {
    let mut universe = SelectionCursor::new(universe.iter());
    if universe.is_exhausted() {
        return;
    }
    for chain in structure.data().chains() {
        for residue in chain.residues() {
            for atom in residue.atoms() {
                let index = atom.index().get();
                if universe.matches(index) {
                    visitor(AtomContext {
                        atom,
                        residue,
                        chain,
                    });
                    if universe.is_exhausted() {
                        return;
                    }
                }
            }
        }
    }
}

pub(super) fn selection_from_sorted(selected: Vec<u32>) -> AtomSelection {
    if selected.is_empty() {
        AtomSelection::Empty
    } else {
        AtomSelection::from_sorted(selected)
    }
}