fayalite/crates/fayalite/src/build/registry.rs

// SPDX-License-Identifier: LGPL-3.0-or-later
// See Notices.txt for copyright information

use crate::{
    build::{DynJobKind, JobKind, built_in_job_kinds},
    intern::Interned,
    util::InternedStrCompareAsStr,
};
use std::{
    collections::BTreeMap,
    fmt,
    sync::{Arc, OnceLock, RwLock, RwLockWriteGuard},
};

impl DynJobKind {
    pub fn registry() -> JobKindRegistrySnapshot {
        JobKindRegistrySnapshot(JobKindRegistry::get())
    }
    #[track_caller]
    pub fn register(self) {
        JobKindRegistry::register(JobKindRegistry::lock(), self);
    }
}

#[derive(Clone, Debug)]
struct JobKindRegistry {
    job_kinds: BTreeMap<InternedStrCompareAsStr, DynJobKind>,
}

enum JobKindRegisterError {
    SameName {
        name: InternedStrCompareAsStr,
        old_job_kind: DynJobKind,
        new_job_kind: DynJobKind,
    },
}

impl fmt::Display for JobKindRegisterError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::SameName {
                name,
                old_job_kind,
                new_job_kind,
            } => write!(
                f,
                "two different `JobKind` can't share the same name:\n\
                {name:?}\n\
                old job kind:\n\
                {old_job_kind:?}\n\
                new job kind:\n\
                {new_job_kind:?}",
            ),
        }
    }
}

trait JobKindRegistryRegisterLock {
    type Locked;
    fn lock(self) -> Self::Locked;
    fn make_mut(locked: &mut Self::Locked) -> &mut JobKindRegistry;
}

impl JobKindRegistryRegisterLock for &'static RwLock<Arc<JobKindRegistry>> {
    type Locked = RwLockWriteGuard<'static, Arc<JobKindRegistry>>;
    fn lock(self) -> Self::Locked {
        self.write().expect("shouldn't be poisoned")
    }
    fn make_mut(locked: &mut Self::Locked) -> &mut JobKindRegistry {
        Arc::make_mut(locked)
    }
}

impl JobKindRegistryRegisterLock for &'_ mut JobKindRegistry {
    type Locked = Self;

    fn lock(self) -> Self::Locked {
        self
    }

    fn make_mut(locked: &mut Self::Locked) -> &mut JobKindRegistry {
        locked
    }
}

impl JobKindRegistry {
    fn lock() -> &'static RwLock<Arc<Self>> {
        static REGISTRY: OnceLock<RwLock<Arc<JobKindRegistry>>> = OnceLock::new();
        REGISTRY.get_or_init(Default::default)
    }
    fn try_register<L: JobKindRegistryRegisterLock>(
        lock: L,
        job_kind: DynJobKind,
    ) -> Result<(), JobKindRegisterError> {
        use std::collections::btree_map::Entry;
        let name = InternedStrCompareAsStr(job_kind.name());
        // run user code only outside of lock
        let mut locked = lock.lock();
        let this = L::make_mut(&mut locked);
        let result = match this.job_kinds.entry(name) {
            Entry::Occupied(entry) => Err(JobKindRegisterError::SameName {
                name,
                old_job_kind: entry.get().clone(),
                new_job_kind: job_kind,
            }),
            Entry::Vacant(entry) => {
                entry.insert(job_kind);
                Ok(())
            }
        };
        drop(locked);
        // outside of lock now, so we can test if it's the same DynJobKind
        match result {
            Err(JobKindRegisterError::SameName {
                name: _,
                old_job_kind,
                new_job_kind,
            }) if old_job_kind == new_job_kind => Ok(()),
            result => result,
        }
    }
    #[track_caller]
    fn register<L: JobKindRegistryRegisterLock>(lock: L, job_kind: DynJobKind) {
        match Self::try_register(lock, job_kind) {
            Err(e) => panic!("{e}"),
            Ok(()) => {}
        }
    }
    fn get() -> Arc<Self> {
        Self::lock().read().expect("shouldn't be poisoned").clone()
    }
}

impl Default for JobKindRegistry {
    fn default() -> Self {
        let mut retval = Self {
            job_kinds: BTreeMap::new(),
        };
        for job_kind in built_in_job_kinds() {
            Self::register(&mut retval, job_kind);
        }
        retval
    }
}

#[derive(Clone, Debug)]
pub struct JobKindRegistrySnapshot(Arc<JobKindRegistry>);

impl JobKindRegistrySnapshot {
    pub fn get() -> Self {
        JobKindRegistrySnapshot(JobKindRegistry::get())
    }
    pub fn get_by_name<'a>(&'a self, name: &str) -> Option<&'a DynJobKind> {
        self.0.job_kinds.get(name)
    }
    pub fn iter_with_names(&self) -> JobKindRegistryIterWithNames<'_> {
        JobKindRegistryIterWithNames(self.0.job_kinds.iter())
    }
    pub fn iter(&self) -> JobKindRegistryIter<'_> {
        JobKindRegistryIter(self.0.job_kinds.values())
    }
}

impl<'a> IntoIterator for &'a JobKindRegistrySnapshot {
    type Item = &'a DynJobKind;
    type IntoIter = JobKindRegistryIter<'a>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

impl<'a> IntoIterator for &'a mut JobKindRegistrySnapshot {
    type Item = &'a DynJobKind;
    type IntoIter = JobKindRegistryIter<'a>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

#[derive(Clone, Debug)]
pub struct JobKindRegistryIter<'a>(
    std::collections::btree_map::Values<'a, InternedStrCompareAsStr, DynJobKind>,
);

impl<'a> Iterator for JobKindRegistryIter<'a> {
    type Item = &'a DynJobKind;

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next()
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.0.size_hint()
    }

    fn count(self) -> usize
    where
        Self: Sized,
    {
        self.0.count()
    }

    fn last(self) -> Option<Self::Item> {
        self.0.last()
    }

    fn nth(&mut self, n: usize) -> Option<Self::Item> {
        self.0.nth(n)
    }

    fn fold<B, F>(self, init: B, f: F) -> B
    where
        F: FnMut(B, Self::Item) -> B,
    {
        self.0.fold(init, f)
    }
}

impl<'a> std::iter::FusedIterator for JobKindRegistryIter<'a> {}

impl<'a> ExactSizeIterator for JobKindRegistryIter<'a> {}

impl<'a> DoubleEndedIterator for JobKindRegistryIter<'a> {
    fn next_back(&mut self) -> Option<Self::Item> {
        self.0.next_back()
    }

    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
        self.0.nth_back(n)
    }

    fn rfold<B, F>(self, init: B, f: F) -> B
    where
        F: FnMut(B, Self::Item) -> B,
    {
        self.0.rfold(init, f)
    }
}

#[derive(Clone, Debug)]
pub struct JobKindRegistryIterWithNames<'a>(
    std::collections::btree_map::Iter<'a, InternedStrCompareAsStr, DynJobKind>,
);

impl<'a> Iterator for JobKindRegistryIterWithNames<'a> {
    type Item = (Interned<str>, &'a DynJobKind);

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next().map(|(name, job_kind)| (name.0, job_kind))
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.0.size_hint()
    }

    fn count(self) -> usize
    where
        Self: Sized,
    {
        self.0.count()
    }

    fn last(self) -> Option<Self::Item> {
        self.0.last().map(|(name, job_kind)| (name.0, job_kind))
    }

    fn nth(&mut self, n: usize) -> Option<Self::Item> {
        self.0.nth(n).map(|(name, job_kind)| (name.0, job_kind))
    }

    fn fold<B, F>(self, init: B, f: F) -> B
    where
        F: FnMut(B, Self::Item) -> B,
    {
        self.0
            .map(|(name, job_kind)| (name.0, job_kind))
            .fold(init, f)
    }
}

impl<'a> std::iter::FusedIterator for JobKindRegistryIterWithNames<'a> {}

impl<'a> ExactSizeIterator for JobKindRegistryIterWithNames<'a> {}

impl<'a> DoubleEndedIterator for JobKindRegistryIterWithNames<'a> {
    fn next_back(&mut self) -> Option<Self::Item> {
        self.0
            .next_back()
            .map(|(name, job_kind)| (name.0, job_kind))
    }

    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
        self.0
            .nth_back(n)
            .map(|(name, job_kind)| (name.0, job_kind))
    }

    fn rfold<B, F>(self, init: B, f: F) -> B
    where
        F: FnMut(B, Self::Item) -> B,
    {
        self.0
            .map(|(name, job_kind)| (name.0, job_kind))
            .rfold(init, f)
    }
}

#[track_caller]
pub fn register_job_kind<K: JobKind>(kind: K) {
    DynJobKind::new(kind).register();
}