molframe-core 0.5.0

Core types, storage and analysis contracts for molframe
Documentation
//! Allocation-bounded position iteration and sparse merging.

use super::representation::AtomSelection;
use crate::column::Ones;
use std::iter::Peekable;
use std::ops::Range;
use std::slice;

enum MergeCursor<'a> {
    Empty,
    Range(Range<u32>),
    Ranges {
        runs: &'a [Range<u32>],
        run: usize,
        next: u32,
    },
    Sparse(slice::Iter<'a, u32>),
    Dense(Ones<'a>),
}

impl<'a> MergeCursor<'a> {
    fn new(selection: &'a AtomSelection) -> Self {
        match selection {
            AtomSelection::Empty => Self::Empty,
            AtomSelection::All(count) => Self::Range(0..*count),
            AtomSelection::Range(run) => Self::Range(run.start..run.end),
            AtomSelection::Ranges(runs) => Self::Ranges {
                runs,
                run: 0,
                next: runs.first().map_or(0, |run| run.start),
            },
            AtomSelection::Sparse(positions) => Self::Sparse(positions.iter()),
            AtomSelection::Dense(mask) => Self::Dense(mask.ones()),
        }
    }
}

impl Iterator for MergeCursor<'_> {
    type Item = u32;

    fn next(&mut self) -> Option<Self::Item> {
        match self {
            Self::Empty => None,
            Self::Range(range) => range.next(),
            Self::Sparse(positions) => positions.next().copied(),
            Self::Dense(positions) => positions.next(),
            Self::Ranges { runs, run, next } => loop {
                let current = runs.get(*run)?;
                *next = (*next).max(current.start);
                if *next < current.end {
                    let position = *next;
                    *next += 1;
                    return Some(position);
                }
                *run += 1;
                *next = runs.get(*run)?.start;
            },
        }
    }
}

pub(super) fn copy_selection(selection: &AtomSelection) -> AtomSelection {
    match selection {
        AtomSelection::Empty => AtomSelection::Empty,
        AtomSelection::All(count) => AtomSelection::All(*count),
        AtomSelection::Range(run) => AtomSelection::Range(run.start..run.end),
        AtomSelection::Ranges(runs) => {
            AtomSelection::Ranges(runs.iter().map(|run| run.start..run.end).collect())
        }
        AtomSelection::Sparse(positions) => AtomSelection::Sparse(positions.clone()),
        // The mask is reference-counted, so copying one is a count bump
        // rather than a per-bit rebuild of the same words.
        AtomSelection::Dense(mask) => AtomSelection::Dense(mask.clone()),
    }
}

pub(super) fn visit_positions(selection: &AtomSelection, mut visit: impl FnMut(u32)) {
    match selection {
        AtomSelection::Empty => {}
        AtomSelection::All(count) => (0..*count).for_each(visit),
        AtomSelection::Range(run) => (run.start..run.end).for_each(visit),
        AtomSelection::Ranges(runs) => {
            for run in runs {
                for position in run.start..run.end {
                    visit(position);
                }
            }
        }
        AtomSelection::Sparse(positions) => positions.iter().copied().for_each(visit),
        AtomSelection::Dense(mask) => mask.ones().for_each(visit),
    }
}

pub(super) fn merge_positions(left: &AtomSelection, right: &AtomSelection) -> AtomSelection {
    let mut merged = Vec::new();
    merge_cursors(
        MergeCursor::new(left).peekable(),
        MergeCursor::new(right).peekable(),
        &mut merged,
    );
    AtomSelection::from_sorted(merged)
}

fn merge_cursors(
    mut left: Peekable<MergeCursor<'_>>,
    mut right: Peekable<MergeCursor<'_>>,
    merged: &mut Vec<u32>,
) {
    loop {
        let next = match (left.peek().copied(), right.peek().copied()) {
            (None, None) => break,
            (Some(left_position), Some(right_position)) if left_position == right_position => {
                right.next();
                left.next()
            }
            (Some(left_position), Some(right_position)) if left_position > right_position => {
                right.next()
            }
            (Some(_), _) => left.next(),
            (None, Some(_)) => right.next(),
        };
        if let Some(position) = next {
            merged.push(position);
        }
    }
}