molframe-cif 0.6.2

//! First-pass metadata projection and coordinate-layout classification.

use super::keep_lowering_category;
use crate::document::{CifValueRef, Document};
use crate::lexer::Quoting;
use crate::parser::{ColumnId, DocumentSink, ValueSink};
use molframe_core::io::CategoryFilter;
use molframe_core::span::ByteSpan;
use std::collections::HashSet;

pub(super) struct Projection {
    pub(super) metadata: Document,
    pub(super) layout: AtomLayout,
}

#[derive(Default)]
pub(super) struct AtomLayout {
    pub(super) has_atom_site: bool,
    atom_loops: usize,
    atom_rows: usize,
    models: Vec<i64>,
    supported: bool,
}

impl AtomLayout {
    pub(super) fn supports_direct(&self) -> bool {
        self.supported && self.atom_loops == 1 && self.atom_rows > 0
    }

    pub(super) fn model_count(&self) -> usize {
        self.models.len()
    }
}

pub(super) struct ProjectionSink {
    metadata: DocumentSink,
    layout: AtomLayout,
    blocks: usize,
    in_atom_loop: bool,
    only_atom_site: bool,
    values_in_row: usize,
    row_model: Option<i64>,
    keep_category: fn(&str) -> bool,
    categories: CategoryFilter,
    model_column: Option<ColumnId>,
}

impl ProjectionSink {
    pub(super) fn new(keep_category: fn(&str) -> bool, categories: CategoryFilter) -> Self {
        Self {
            metadata: DocumentSink::with_filter(|_| true),
            layout: AtomLayout {
                supported: true,
                ..AtomLayout::default()
            },
            blocks: 0,
            in_atom_loop: false,
            only_atom_site: false,
            values_in_row: 0,
            row_model: None,
            keep_category,
            categories,
            model_column: None,
        }
    }

    fn in_first_block(&self) -> bool {
        self.blocks == 1
    }

    pub(super) const fn in_atom_loop(&self) -> bool {
        self.in_atom_loop
    }

    /// Whether whole rows may be fed through [`Self::bulk_row`] instead of
    /// value by value: a loop of `atom_site` items only, in the first block,
    /// none of whose values the metadata document would keep.
    pub(super) fn can_bulk(&self) -> bool {
        self.in_atom_loop
            && self.only_atom_site
            && self.in_first_block()
            && self.values_in_row == 0
            && !(keep_lowering_category("atom_site") && self.categories.keeps_category("atom_site"))
            && !(self.keep_category)("atom_site")
    }

    /// Counts a finished row whose values were read in bulk.
    pub(super) fn bulk_row(&mut self, model: Option<i64>) {
        self.record_row(model_or_one(model));
    }

    fn record_row(&mut self, model: i64) {
        self.layout.atom_rows += 1;
        if self.layout.models.last().copied() != Some(model) {
            self.layout.models.push(model);
        }
    }

    pub(super) fn metadata(&self) -> &Document {
        self.metadata.document()
    }
}

impl<'input> ValueSink<'input> for ProjectionSink {
    type Output = Projection;

    fn block(&mut self, name: &str) {
        self.blocks += 1;
        if self.in_first_block() {
            self.metadata.block(name);
        }
    }

    fn begin_loop(&mut self, tags: &[(Box<str>, Box<str>)]) {
        self.in_atom_loop = false;
        self.values_in_row = 0;
        self.row_model = None;
        self.model_column = None;
        if !self.in_first_block() || !tags.iter().any(|(category, _)| &**category == "atom_site") {
            return;
        }

        self.layout.has_atom_site = true;
        self.layout.atom_loops += 1;
        self.in_atom_loop = true;
        self.only_atom_site = tags.iter().all(|(category, _)| &**category == "atom_site");
        let mut items = HashSet::with_capacity(tags.len());
        for (category, item) in tags {
            if &**category == "atom_site" && !items.insert(&**item) {
                self.layout.supported = false;
            }
        }
        self.model_column = tags
            .iter()
            .position(|(category, item)| {
                category.as_ref() == "atom_site" && item.as_ref() == "pdbx_PDB_model_num"
            })
            .and_then(ColumnId::from_position);
    }

    fn value(
        &mut self,
        column: ColumnId,
        category: &str,
        item: &str,
        text: &'input str,
        quoting: Quoting,
        span: ByteSpan,
    ) {
        if !self.in_first_block() {
            return;
        }
        if (keep_lowering_category(category) && self.categories.keeps_category(category))
            || (self.keep_category)(category)
        {
            self.metadata.value(category, item, text, quoting, span);
        }
        if category != "atom_site" {
            if self.in_atom_loop {
                self.values_in_row += 1;
            }
            return;
        }

        self.layout.has_atom_site = true;
        if !self.in_atom_loop {
            self.layout.supported = false;
            return;
        }
        self.values_in_row += 1;
        if self.model_column == Some(column) {
            self.row_model = CifValueRef::parse(text, quoting).as_integer();
        }
    }

    fn end_row(&mut self) {
        if !self.in_atom_loop || !self.in_first_block() {
            return;
        }
        self.record_row(model_or_one(self.row_model));
        self.values_in_row = 0;
        self.row_model = None;
    }

    fn end_loop(&mut self) {
        if self.in_atom_loop && self.values_in_row != 0 {
            self.layout.supported = false;
        }
        self.in_atom_loop = false;
        self.values_in_row = 0;
        self.row_model = None;
    }

    fn finish(self) -> Self::Output {
        Projection {
            metadata: self.metadata.finish(),
            layout: self.layout,
        }
    }
}

pub(super) fn model_or_one(model: Option<i64>) -> i64 {
    let Some(model) = model else {
        return 1;
    };
    model
}