kcode-kweb-db 0.1.0

use crate::{
    Error, NodeId, ObjectId, Result, TransactionId, TransactionPackage, wire::ParsedTransaction,
};
use std::collections::{BTreeMap, BTreeSet};

#[derive(Clone)]
pub(crate) struct Record {
    pub parsed: ParsedTransaction,
    pub package: TransactionPackage,
}

#[derive(Clone, Default)]
pub(crate) struct Graph {
    records: BTreeMap<TransactionId, Record>,
}

impl Graph {
    pub fn insert(&mut self, record: Record) -> Result<bool> {
        let id = record.parsed.id;
        if let Some(existing) = self.records.get(&id) {
            if existing.package == record.package {
                return Ok(false);
            }
            return Err(Error::corrupt(
                "one transaction ID names two different retained packages",
            ));
        }
        self.records.insert(id, record);
        Ok(true)
    }

    pub fn get(&self, id: TransactionId) -> Option<&Record> {
        self.records.get(&id)
    }

    pub fn records(&self) -> impl Iterator<Item = (&TransactionId, &Record)> {
        self.records.iter()
    }

    pub fn heads(&self) -> Vec<TransactionId> {
        let parents = self
            .records
            .values()
            .flat_map(|record| record.parsed.unsigned.heads.iter().copied())
            .collect::<BTreeSet<_>>();
        self.records
            .keys()
            .copied()
            .filter(|id| !parents.contains(id))
            .collect()
    }

    pub fn active(&self) -> BTreeSet<TransactionId> {
        let mut active = BTreeSet::new();
        loop {
            let before = active.len();
            for (id, record) in &self.records {
                if record
                    .parsed
                    .unsigned
                    .heads
                    .iter()
                    .all(|parent| active.contains(parent))
                {
                    active.insert(*id);
                }
            }
            if active.len() == before {
                break;
            }
        }
        active
    }

    pub fn generations(&self, active: &BTreeSet<TransactionId>) -> BTreeMap<TransactionId, u64> {
        let mut generations = BTreeMap::new();
        loop {
            let before = generations.len();
            for id in active {
                let record = &self.records[id];
                if record
                    .parsed
                    .unsigned
                    .heads
                    .iter()
                    .all(|parent| generations.contains_key(parent))
                {
                    let generation = record
                        .parsed
                        .unsigned
                        .heads
                        .iter()
                        .filter_map(|parent| generations.get(parent))
                        .copied()
                        .max()
                        .map_or(0, |value| value + 1);
                    generations.insert(*id, generation);
                }
            }
            if generations.len() == before {
                break;
            }
        }
        generations
    }

    pub fn is_ancestor(&self, ancestor: TransactionId, descendant: TransactionId) -> bool {
        if ancestor == descendant {
            return false;
        }
        let mut pending = vec![descendant];
        let mut visited = BTreeSet::new();
        while let Some(id) = pending.pop() {
            if !visited.insert(id) {
                continue;
            }
            let Some(record) = self.records.get(&id) else {
                continue;
            };
            for parent in &record.parsed.unsigned.heads {
                if *parent == ancestor {
                    return true;
                }
                pending.push(*parent);
            }
        }
        false
    }

    pub fn known_node_ids(&self) -> BTreeSet<NodeId> {
        self.records
            .values()
            .flat_map(|record| {
                record
                    .parsed
                    .unsigned
                    .creates
                    .iter()
                    .chain(record.parsed.unsigned.updates.iter())
                    .map(|operation| operation.id)
            })
            .collect()
    }

    pub fn known_object_ids(&self) -> BTreeSet<ObjectId> {
        self.records
            .values()
            .flat_map(|record| {
                record
                    .parsed
                    .unsigned
                    .objects
                    .iter()
                    .map(|object| object.id)
            })
            .collect()
    }
}