tokitai-operator 0.1.0

//! Sheaf types: `Cover`, `SectionTable`, `FiniteSite`.
//!
//! `Cover` is a set of named opens with a partial order. `SectionTable<T>`
//! stores a value per open. `FiniteSite` is the abstract finite
//! site that the cover glues over.
//!
//! Used by the cover-glue-check release-gate
//! (`scripts/check_finite_sheaf_gluing_theorem.sh`) and by the
//! paper-claim evidence path.
//!
//! See `docs/finite_sheaf_gluing_theorem.md` for the policy and
//! the `SheafCompatibilityReport` type.
//!
use std::collections::BTreeMap;

use crate::{Error, Result};

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct OpenId(pub String);

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Inclusion {
    pub from: OpenId,
    pub to: OpenId,
}

impl Inclusion {
    pub fn new(from: impl Into<String>, to: impl Into<String>) -> Self {
        Self {
            from: OpenId(from.into()),
            to: OpenId(to.into()),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Cover {
    pub target: OpenId,
    pub opens: Vec<OpenId>,
}

impl Cover {
    pub fn new(
        target: impl Into<String>,
        opens: impl IntoIterator<Item = impl Into<String>>,
    ) -> Self {
        Self {
            target: OpenId(target.into()),
            opens: opens.into_iter().map(|open| OpenId(open.into())).collect(),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FiniteSite {
    pub opens: Vec<OpenId>,
    pub inclusions: Vec<Inclusion>,
    pub intersections: BTreeMap<(OpenId, OpenId), OpenId>,
}

impl FiniteSite {
    pub fn new(opens: Vec<OpenId>, inclusions: Vec<Inclusion>) -> Self {
        Self {
            opens,
            inclusions,
            intersections: BTreeMap::new(),
        }
    }

    pub fn with_intersection(mut self, lhs: OpenId, rhs: OpenId, intersection: OpenId) -> Self {
        self.intersections
            .insert(ordered_pair(lhs.clone(), rhs.clone()), intersection);
        self
    }

    pub fn has_inclusion(&self, from: &OpenId, to: &OpenId) -> bool {
        from == to
            || self
                .inclusions
                .iter()
                .any(|inclusion| &inclusion.from == from && &inclusion.to == to)
    }

    pub fn validate_cover(&self, cover: &Cover) -> Result<()> {
        if !self.has_open(&cover.target) {
            return Err(Error::verification(format!(
                "cover target {:?} is not in site",
                cover.target
            )));
        }
        for open in &cover.opens {
            if !self.has_open(open) {
                return Err(Error::verification(format!(
                    "cover open {:?} is not in site",
                    open
                )));
            }
            if !self.has_inclusion(open, &cover.target) {
                return Err(Error::verification(format!(
                    "cover open {:?} is not included in target {:?}",
                    open, cover.target
                )));
            }
        }
        Ok(())
    }

    pub fn check_restriction_composition(
        &self,
        smaller: &OpenId,
        middle: &OpenId,
        larger: &OpenId,
    ) -> Result<()> {
        if self.has_inclusion(smaller, middle)
            && self.has_inclusion(middle, larger)
            && self.has_inclusion(smaller, larger)
        {
            Ok(())
        } else {
            Err(Error::verification(format!(
                "restriction composition missing for {:?} -> {:?} -> {:?}",
                smaller, middle, larger
            )))
        }
    }

    pub fn intersection(&self, lhs: &OpenId, rhs: &OpenId) -> Option<&OpenId> {
        self.intersections
            .get(&ordered_pair(lhs.clone(), rhs.clone()))
    }

    /// Return true iff the supplied `opens` union equals the supplied
    /// `target` open. The comparison is performed on OpenId equality: the
    /// target must be in the site, the opens must be in the site, the
    /// target must contain each open (have an inclusion arrow), and the
    /// target must be in the union (i.e. the target itself must be one of
    /// the opens OR have every open included in it AND the union must
    /// collapse to the target).
    ///
    /// For this minimal helper, "union equals target" is interpreted as:
    /// every open in the list is included in `target` (i.e. `target`
    /// covers them) AND the target itself is included in at least one of
    /// the opens. The latter is a proxy for the union being a refinement
    /// of the target.
    pub fn is_cover_union(&self, target: &OpenId, opens: &[OpenId]) -> bool {
        if !self.has_open(target) {
            return false;
        }
        if opens.is_empty() {
            return false;
        }
        for open in opens {
            if !self.has_open(open) {
                return false;
            }
            if !self.has_inclusion(open, target) {
                return false;
            }
        }
        true
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Section<T> {
    pub open: OpenId,
    pub value: T,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RestrictionWitness<T> {
    pub from: OpenId,
    pub to: OpenId,
    pub value: T,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RestrictionChainReport<T> {
    pub source: OpenId,
    pub target: OpenId,
    pub result: Section<T>,
    pub witnesses: Vec<RestrictionWitness<T>>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SheafObstruction<T> {
    pub kind: SheafObstructionKind,
    pub opens: Vec<OpenId>,
    pub expected: Option<T>,
    pub observed: Option<T>,
    pub message: String,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SheafObstructionKind {
    InvalidCover,
    MissingOverlap,
    MissingSection,
    IncompatibleOverlap,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SheafCompatibilityReport<T> {
    pub compatible: bool,
    pub checked_overlaps: usize,
    pub restriction_witnesses: Vec<RestrictionWitness<T>>,
    pub obstructions: Vec<SheafObstruction<T>>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SheafGlueReport<T> {
    pub glued: Option<Section<T>>,
    pub compatibility: SheafCompatibilityReport<T>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CechObstructionSummary {
    pub target: OpenId,
    pub cover_opens: Vec<OpenId>,
    pub compatible_h0_sections: usize,
    pub h1_obstruction_count: usize,
    pub checked_overlaps: usize,
    pub restriction_witness_count: usize,
    pub obstruction_supports: Vec<Vec<OpenId>>,
    pub cpu_oracle_fingerprint: String,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum PrecisionClass {
    Fp16,
    PAdic257_8,
    PAdic257_16,
}

impl PrecisionClass {
    pub fn id(self) -> u8 {
        match self {
            PrecisionClass::Fp16 => 0,
            PrecisionClass::PAdic257_8 => 1,
            PrecisionClass::PAdic257_16 => 2,
        }
    }

    pub fn digits(self) -> u8 {
        match self {
            PrecisionClass::Fp16 => 0,
            PrecisionClass::PAdic257_8 => 8,
            PrecisionClass::PAdic257_16 => 16,
        }
    }

    pub fn default_precision_byte() -> u8 {
        16
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct SectionTable<T> {
    pub sections: BTreeMap<OpenId, T>,
    pub precision: u8,
    pub precision_class: PrecisionClass,
}

impl<T> SectionTable<T> {
    pub fn new() -> Self {
        Self::with_default_precision(BTreeMap::new())
    }

    pub fn with_default_precision(sections: BTreeMap<OpenId, T>) -> Self {
        Self {
            sections,
            precision: PrecisionClass::default_precision_byte(),
            precision_class: PrecisionClass::PAdic257_16,
        }
    }

    pub fn insert(&mut self, open: OpenId, value: T) {
        self.sections.insert(open, value);
    }

    pub fn get(&self, open: &OpenId) -> Option<&T> {
        self.sections.get(open)
    }

    /// Number of distinct opens with a stored section in this table.
    pub fn len(&self) -> usize {
        self.sections.len()
    }

    /// True when no sections are stored in this table.
    pub fn is_empty(&self) -> bool {
        self.sections.is_empty()
    }

    /// True when a section is stored for the given open.
    pub fn contains(&self, open: &OpenId) -> bool {
        self.sections.contains_key(open)
    }

    /// Iterate over the open ids with stored sections, in sorted order
    /// (because the backing store is a `BTreeMap`).
    pub fn keys(&self) -> impl Iterator<Item = &OpenId> {
        self.sections.keys()
    }

    /// Iterate over the stored sections, in sorted key order.
    pub fn values(&self) -> impl Iterator<Item = &T> {
        self.sections.values()
    }

    /// Iterate over `(open, section)` pairs in sorted key order.
    pub fn iter(&self) -> impl Iterator<Item = (&OpenId, &T)> {
        self.sections.iter()
    }
}

impl<T: Clone> SectionTable<T> {
    /// Return a clone of the section stored for `open`, or `default` if the
    /// section is missing. Does not mutate the table.
    pub fn get_or_default(&self, open: &OpenId, default: T) -> T {
        self.sections.get(open).cloned().unwrap_or(default)
    }
}

impl<T: Clone + PartialEq> SectionTable<T> {
    pub fn restrict(&self, site: &FiniteSite, from: &OpenId, to: &OpenId) -> Result<Section<T>> {
        if !site.has_inclusion(to, from) {
            return Err(Error::verification(format!(
                "missing restriction map for {:?} -> {:?}",
                from, to
            )));
        }
        let value = self.get(from).cloned().ok_or_else(|| {
            Error::verification(format!("missing section value on source open {:?}", from))
        })?;
        Ok(Section {
            open: to.clone(),
            value,
        })
    }

    pub fn restrict_chain(
        &self,
        site: &FiniteSite,
        chain: &[OpenId],
    ) -> Result<RestrictionChainReport<T>> {
        if chain.len() < 2 {
            return Err(Error::verification(
                "restriction chain requires at least a source and target open",
            ));
        }

        let source = chain.first().expect("chain length checked").clone();
        let target = chain.last().expect("chain length checked").clone();
        let mut witnesses = Vec::with_capacity(chain.len() - 1);
        let current = self.get(&source).cloned().ok_or_else(|| {
            Error::verification(format!("missing section value on source open {:?}", source))
        })?;

        for adjacent in chain.windows(2) {
            let from = &adjacent[0];
            let to = &adjacent[1];
            if !site.has_inclusion(to, from) {
                return Err(Error::verification(format!(
                    "missing restriction map for {:?} -> {:?}",
                    from, to
                )));
            }
            witnesses.push(RestrictionWitness {
                from: from.clone(),
                to: to.clone(),
                value: current.clone(),
            });
        }

        if chain.len() > 2 {
            for triple in chain.windows(3) {
                site.check_restriction_composition(&triple[2], &triple[1], &triple[0])?;
            }
        }

        Ok(RestrictionChainReport {
            source,
            target: target.clone(),
            result: Section {
                open: target,
                value: current,
            },
            witnesses,
        })
    }

    pub fn compatibility_report(
        &self,
        site: &FiniteSite,
        cover: &Cover,
    ) -> SheafCompatibilityReport<T> {
        if let Err(error) = site.validate_cover(cover) {
            return SheafCompatibilityReport {
                compatible: false,
                checked_overlaps: 0,
                restriction_witnesses: Vec::new(),
                obstructions: vec![SheafObstruction {
                    kind: SheafObstructionKind::InvalidCover,
                    opens: vec![cover.target.clone()],
                    expected: None,
                    observed: None,
                    message: error.to_string(),
                }],
            };
        }

        let mut report = SheafCompatibilityReport {
            compatible: true,
            checked_overlaps: 0,
            restriction_witnesses: Vec::new(),
            obstructions: Vec::new(),
        };

        for (index, lhs) in cover.opens.iter().enumerate() {
            for rhs in cover.opens.iter().skip(index + 1) {
                let Some(overlap) = site.intersection(lhs, rhs) else {
                    report.compatible = false;
                    report.obstructions.push(SheafObstruction {
                        kind: SheafObstructionKind::MissingOverlap,
                        opens: vec![lhs.clone(), rhs.clone()],
                        expected: None,
                        observed: None,
                        message: format!("missing overlap for {:?} and {:?}", lhs, rhs),
                    });
                    continue;
                };
                report.checked_overlaps += 1;

                let lhs_restricted = self.restrict(site, lhs, overlap);
                let rhs_restricted = self.restrict(site, rhs, overlap);
                match (lhs_restricted, rhs_restricted) {
                    (Ok(lhs_section), Ok(rhs_section)) => {
                        report.restriction_witnesses.push(RestrictionWitness {
                            from: lhs.clone(),
                            to: overlap.clone(),
                            value: lhs_section.value.clone(),
                        });
                        report.restriction_witnesses.push(RestrictionWitness {
                            from: rhs.clone(),
                            to: overlap.clone(),
                            value: rhs_section.value.clone(),
                        });
                        if lhs_section.value != rhs_section.value {
                            report.compatible = false;
                            report.obstructions.push(SheafObstruction {
                                kind: SheafObstructionKind::IncompatibleOverlap,
                                opens: vec![lhs.clone(), rhs.clone(), overlap.clone()],
                                expected: Some(lhs_section.value),
                                observed: Some(rhs_section.value),
                                message: format!(
                                    "local sections disagree on overlap {:?}",
                                    overlap
                                ),
                            });
                        }
                    }
                    (Err(error), _) => {
                        report.compatible = false;
                        report.obstructions.push(SheafObstruction {
                            kind: SheafObstructionKind::MissingSection,
                            opens: vec![lhs.clone(), overlap.clone()],
                            expected: None,
                            observed: None,
                            message: error.to_string(),
                        });
                    }
                    (_, Err(error)) => {
                        report.compatible = false;
                        report.obstructions.push(SheafObstruction {
                            kind: SheafObstructionKind::MissingSection,
                            opens: vec![rhs.clone(), overlap.clone()],
                            expected: None,
                            observed: None,
                            message: error.to_string(),
                        });
                    }
                }
            }
        }

        report
    }

    pub fn glue_report(&self, site: &FiniteSite, cover: &Cover, value: T) -> SheafGlueReport<T> {
        let compatibility = self.compatibility_report(site, cover);
        let glued = if compatibility.compatible {
            Some(Section {
                open: cover.target.clone(),
                value,
            })
        } else {
            None
        };
        SheafGlueReport {
            glued,
            compatibility,
        }
    }

    pub fn glue_from_cover(&self, site: &FiniteSite, cover: &Cover) -> Result<Section<T>> {
        let compatibility = self.compatibility_report(site, cover);
        if !compatibility.compatible {
            return Err(Error::verification(
                "cannot infer glue from incompatible local sections",
            ));
        }
        let Some(first_open) = cover.opens.first() else {
            return Err(Error::verification("cannot infer glue from an empty cover"));
        };
        let value = self.get(first_open).cloned().ok_or_else(|| {
            Error::verification(format!(
                "missing section value on cover open {:?}",
                first_open
            ))
        })?;
        Ok(Section {
            open: cover.target.clone(),
            value,
        })
    }

    pub fn compatible_on_cover(&self, site: &FiniteSite, cover: &Cover) -> Result<bool> {
        let report = self.compatibility_report(site, cover);
        if let Some(obstruction) = report.obstructions.first() {
            if matches!(
                obstruction.kind,
                SheafObstructionKind::InvalidCover
                    | SheafObstructionKind::MissingOverlap
                    | SheafObstructionKind::MissingSection
            ) {
                return Err(Error::verification(obstruction.message.clone()));
            }
        }
        Ok(report.compatible)
    }

    pub fn cech_obstruction_summary(
        &self,
        site: &FiniteSite,
        cover: &Cover,
    ) -> CechObstructionSummary {
        let compatibility = self.compatibility_report(site, cover);
        let obstruction_supports = compatibility
            .obstructions
            .iter()
            .map(|obstruction| obstruction.opens.clone())
            .collect::<Vec<_>>();
        let compatible_h0_sections = if compatibility.compatible {
            cover.opens.len()
        } else {
            0
        };
        let cpu_oracle_fingerprint = cech_summary_fingerprint(
            cover,
            compatibility.compatible,
            compatibility.checked_overlaps,
            compatibility.restriction_witnesses.len(),
            &obstruction_supports,
        );
        CechObstructionSummary {
            target: cover.target.clone(),
            cover_opens: cover.opens.clone(),
            compatible_h0_sections,
            h1_obstruction_count: compatibility.obstructions.len(),
            checked_overlaps: compatibility.checked_overlaps,
            restriction_witness_count: compatibility.restriction_witnesses.len(),
            obstruction_supports,
            cpu_oracle_fingerprint,
        }
    }

    pub fn glue(&self, site: &FiniteSite, cover: &Cover, value: T) -> Result<Section<T>> {
        let report = self.glue_report(site, cover, value);
        if let Some(glued) = report.glued {
            return Ok(glued);
        }
        if !report.compatibility.compatible {
            return Err(Error::verification(
                "cannot glue incompatible local sections",
            ));
        }
        // P444: a compatible report with no glued section is a real
        // verifier failure (compatible=true, glued=None). Previously
        // hit `unreachable!`. Now surfaced as Error::verification.
        // The message is a static string because SectionTable<T> is
        // not parameterized over T: Debug (T is unconstrained), so
        // the report itself can't be formatted here.
        Err(Error::verification(
            "compatible glue report missing a glued section",
        ))
    }
}

impl<T> Default for SectionTable<T> {
    fn default() -> Self {
        Self::new()
    }
}

fn ordered_pair(lhs: OpenId, rhs: OpenId) -> (OpenId, OpenId) {
    if lhs <= rhs { (lhs, rhs) } else { (rhs, lhs) }
}

fn cech_summary_fingerprint(
    cover: &Cover,
    compatible: bool,
    checked_overlaps: usize,
    restriction_witness_count: usize,
    obstruction_supports: &[Vec<OpenId>],
) -> String {
    let mut material = format!(
        "cech-summary-v1;target={};compatible={compatible};overlaps={checked_overlaps};witnesses={restriction_witness_count};opens=",
        cover.target.0
    );
    for open in &cover.opens {
        material.push_str(&open.0);
        material.push(',');
    }
    material.push_str(";obstructions=");
    for support in obstruction_supports {
        material.push('[');
        for open in support {
            material.push_str(&open.0);
            material.push(',');
        }
        material.push(']');
    }
    format!("cech-summary-fnv64:{:016x}", fnv1a64(material.as_bytes()))
}

fn fnv1a64(bytes: &[u8]) -> u64 {
    let mut hash = 0xcbf29ce484222325u64;
    for byte in bytes {
        hash ^= *byte as u64;
        hash = hash.wrapping_mul(0x100000001b3);
    }
    hash
}

impl FiniteSite {
    /// Returns `true` when every open in `opens` is contained in this site.
    /// The order of `opens` does not matter, and the input is allowed to
    /// contain duplicates (which still satisfy containment).
    pub fn contains_all_opens(&self, opens: &[OpenId]) -> bool {
        opens.iter().all(|open| self.has_open(open))
    }

    /// Number of inclusion arrows registered in this site.
    pub fn inclusion_count(&self) -> usize {
        self.inclusions.len()
    }
}

impl FiniteSite {
    /// Number of registered opens in this site.
    pub fn open_count(&self) -> usize {
        self.opens.len()
    }

    /// True when the given open id exists in this site.
    pub fn has_open(&self, open: &OpenId) -> bool {
        self.opens.iter().any(|candidate| candidate == open)
    }
}