From 078ab91be01d866ae1f698655987304d59299c51 Mon Sep 17 00:00:00 2001 From: Jacob Lifshay Date: Sun, 28 Sep 2025 23:05:24 -0700 Subject: [PATCH] WIP refactoring to have JobKind be internal jobs --- crates/fayalite/src/build.rs | 529 ++++++++++---------------- crates/fayalite/src/build/registry.rs | 241 ++++++++++++ 2 files changed, 442 insertions(+), 328 deletions(-) create mode 100644 crates/fayalite/src/build/registry.rs diff --git a/crates/fayalite/src/build.rs b/crates/fayalite/src/build.rs index 9005f785..0e03287c 100644 --- a/crates/fayalite/src/build.rs +++ b/crates/fayalite/src/build.rs @@ -7,13 +7,16 @@ use crate::{ module::Module, util::{HashMap, HashSet, job_server::AcquiredJob}, }; -use hashbrown::hash_map::Entry; use petgraph::{ algo::{DfsSpace, kosaraju_scc, toposort}, graph::DiGraph, visit::{GraphBase, Visitable}, }; -use serde::{Deserialize, Deserializer, Serialize, Serializer, de::Error, ser::SerializeSeq}; +use serde::{ + Deserialize, Deserializer, Serialize, Serializer, + de::{DeserializeOwned, Error}, + ser::SerializeSeq, +}; use std::{ any::{Any, TypeId}, borrow::Cow, @@ -25,13 +28,14 @@ use std::{ marker::PhantomData, panic, rc::Rc, - sync::{Arc, OnceLock, RwLock, RwLockWriteGuard, mpsc}, + sync::{Arc, OnceLock, mpsc}, thread::{self, ScopedJoinHandle}, }; use tempfile::TempDir; pub mod external; pub mod firrtl; +pub mod registry; macro_rules! write_str { ($s:expr, $($rest:tt)*) => { @@ -93,43 +97,85 @@ impl Ord for JobItemName { } } -pub trait JobKind: 'static + Send + Sync + Hash + Eq + fmt::Debug { - type Job: 'static + Send + Sync + Hash + Eq + fmt::Debug; - fn inputs_and_direct_dependencies<'a>( - &'a self, - job: &'a Self::Job, - ) -> Cow<'a, BTreeMap>>; - fn outputs(&self, job: &Self::Job) -> Interned<[JobItemName]>; - /// gets the part of the command line that is common for all members of this job kind -- usually the executable name/path and any global options and/or subcommands - fn command_line_prefix(&self) -> Interned<[Interned]>; - fn to_command_line(&self, job: &Self::Job) -> Interned<[Interned]>; - /// return the subcommand if this is an internal JobKind - fn subcommand(&self) -> Option; - /// Parse from [`ArgMatches`], this should only be called with the results of parsing [`subcommand()`]. - /// If [`subcommand()`] returned [`None`], you should not call this function since it will panic. - /// - /// [`ArgMatches`]: clap::ArgMatches - /// [`subcommand()`]: JobKind::subcommand - fn from_arg_matches(&self, matches: &mut clap::ArgMatches) -> clap::error::Result; - fn debug_name(&self, job: &Self::Job) -> String { - let name = self - .command_line_prefix() - .last() - .copied() - .or_else(|| self.to_command_line(job).first().copied()) - .unwrap_or_default(); - let name = match name.rsplit_once(['/', '\\']) { - Some((_, name)) if name.trim() != "" => name, - _ => &*name, - }; - format!("job:{name}") - } - fn parse_command_line( - &self, - command_line: Interned<[Interned]>, - ) -> clap::error::Result; +pub trait JobArgs: + clap::Args + 'static + Send + Sync + Hash + Eq + fmt::Debug + Clone + Serialize + DeserializeOwned +{ + fn to_args>(&self, args: &mut Args); +} + +pub trait JobDependencies: 'static + Send + Sync + Hash + Eq + fmt::Debug + Copy { + type Args: 'static + Send + Sync + Hash + Eq + fmt::Debug + Clone; + type Jobs: 'static + Send + Sync + Hash + Eq + fmt::Debug; + fn kinds_dyn(self) -> Vec; + fn jobs_dyn(self, jobs: Self::Jobs) -> Vec; + #[track_caller] + fn from_dyn_args(args: Vec) -> Self::Args; +} + +macro_rules! impl_job_dependencies { + (@impl $(($a:ident, $b:ident: $T:ident),)*) => { + impl<$($T: JobKind),*> JobDependencies for ($($T,)*) { + type Args = ($($T::Args,)*); + type Jobs = ($($T::Job,)*); + + fn kinds_dyn(self) -> Vec { + let ($($a,)*) = self; + vec![$(DynJobKind::new($a),)*] + } + + fn jobs_dyn(self, jobs: Self::Jobs) -> Vec { + let ($($a,)*) = self; + let ($($b,)*) = jobs; + vec![$(DynJob::new($a, $b),)*] + } + + #[track_caller] + fn from_dyn_args(args: Vec) -> Self::Args { + let Ok([$($a,)*]) = args.try_into() else { + panic!("wrong number of dependencies"); + }; + $(let Some($a) = $a.downcast_ref::<$T>().cloned() else { + panic!("wrong type of dependency, expected {} got:\n{:?}", std::any::type_name::<$T>(), $a); + };)* + ($($a,)*) + } + } + }; + ($($first:tt, $($rest:tt,)*)?) => { + impl_job_dependencies!(@impl $($first, $($rest,)*)?); + $(impl_job_dependencies!($($rest,)*);)? + }; +} + +impl_job_dependencies! { + (a0, b0: T0), + (a1, b1: T1), + (a2, b2: T2), + (a3, b3: T3), + (a4, b4: T4), + (a5, b5: T5), + (a6, b6: T6), + (a7, b7: T7), + (a8, b8: T8), + (a9, b9: T9), + (va0, b10: T10), + (va1, b11: T11), +} + +pub trait JobKind: 'static + Send + Sync + Hash + Eq + fmt::Debug + Sized + Copy { + type Args: JobArgs; + type Job: AsRef + 'static + Send + Sync + Hash + Eq + fmt::Debug; + type Dependencies: JobDependencies; + fn dependencies(self) -> Self::Dependencies; + fn args_to_jobs( + self, + args: Self::Args, + dependencies_args: ::Args, + ) -> eyre::Result<(Self::Job, ::Jobs)>; + fn outputs(self, job: &Self::Job) -> Interned<[JobItemName]>; + fn name(self) -> Interned; fn run( - &self, + self, job: &Self::Job, inputs: &[JobItem], acquired_job: &mut AcquiredJob, @@ -141,16 +187,15 @@ trait DynJobKindTrait: 'static + Send + Sync + fmt::Debug { fn as_arc_any(self: Arc) -> Arc; fn eq_dyn(&self, other: &dyn DynJobKindTrait) -> bool; fn hash_dyn(&self, state: &mut dyn Hasher); - fn command_line_prefix_dyn(&self) -> Interned<[Interned]>; - fn subcommand_dyn(&self) -> Option; + fn dependencies_kinds_dyn(&self) -> Vec; + fn args_group_id_dyn(&self) -> Option; + fn augment_args_dyn(&self, cmd: clap::Command) -> clap::Command; + fn augment_args_for_update_dyn(&self, cmd: clap::Command) -> clap::Command; fn from_arg_matches_dyn( self: Arc, matches: &mut clap::ArgMatches, - ) -> clap::error::Result; - fn parse_command_line_dyn( - self: Arc, - command_line: Interned<[Interned]>, - ) -> clap::error::Result; + ) -> clap::error::Result; + fn name_dyn(&self) -> Interned; } impl DynJobKindTrait for T { @@ -165,36 +210,42 @@ impl DynJobKindTrait for T { fn eq_dyn(&self, other: &dyn DynJobKindTrait) -> bool { other .as_any() - .downcast_ref::() + .downcast_ref::() .is_some_and(|other| self == other) } fn hash_dyn(&self, mut state: &mut dyn Hasher) { - self.hash(&mut state) + self.hash(&mut state); } - fn command_line_prefix_dyn(&self) -> Interned<[Interned]> { - self.command_line_prefix() + fn dependencies_kinds_dyn(&self) -> Vec { + self.dependencies().kinds_dyn() } - fn subcommand_dyn(&self) -> Option { - self.subcommand() + fn args_group_id_dyn(&self) -> Option { + ::group_id() + } + + fn augment_args_dyn(&self, cmd: clap::Command) -> clap::Command { + ::augment_args(cmd) + } + + fn augment_args_for_update_dyn(&self, cmd: clap::Command) -> clap::Command { + ::augment_args_for_update(cmd) } fn from_arg_matches_dyn( - self: Arc, + self: Arc, matches: &mut clap::ArgMatches, - ) -> clap::error::Result { - let job = self.from_arg_matches(matches)?; - Ok(DynJob::from_arc(self, job)) + ) -> clap::error::Result { + Ok(DynJobArgs::from_arc( + self, + ::from_arg_matches_mut(matches)?, + )) } - fn parse_command_line_dyn( - self: Arc, - command_line: Interned<[Interned]>, - ) -> clap::error::Result { - let job = self.parse_command_line(command_line)?; - Ok(DynJob::from_arc(self, job)) + fn name_dyn(&self) -> Interned { + self.name() } } @@ -203,31 +254,19 @@ pub struct DynJobKind(Arc); impl DynJobKind { pub fn from_arc(job_kind: Arc) -> Self { - if TypeId::of::() == TypeId::of::() { - Self::clone( - &Arc::downcast::(job_kind.as_arc_any()) - .ok() - .expect("already checked type"), - ) - } else { - Self(job_kind) - } + Self(job_kind) } pub fn new(job_kind: T) -> Self { - if let Some(job_kind) = DynJobKindTrait::as_any(&job_kind).downcast_ref::() { - job_kind.clone() - } else { - Self(Arc::new(job_kind)) - } + Self(Arc::new(job_kind)) } pub fn type_id(&self) -> TypeId { DynJobKindTrait::as_any(&*self.0).type_id() } - pub fn downcast_ref(&self) -> Option<&T> { - DynJobKindTrait::as_any(&*self.0).downcast_ref() + pub fn downcast(&self) -> Option { + DynJobKindTrait::as_any(&*self.0).downcast_ref().copied() } pub fn downcast_arc(self) -> Result, Self> { - if self.downcast_ref::().is_some() { + if self.downcast::().is_some() { Ok(Arc::downcast::(self.0.as_arc_any()) .ok() .expect("already checked type")) @@ -235,12 +274,26 @@ impl DynJobKind { Err(self) } } - pub fn registry() -> JobKindRegistrySnapshot { - JobKindRegistrySnapshot(JobKindRegistry::get()) + pub fn dependencies_kinds(&self) -> Vec { + DynJobKindTrait::dependencies_kinds_dyn(&*self.0) } - #[track_caller] - pub fn register(self) { - JobKindRegistry::register(JobKindRegistry::lock(), self); + pub fn args_group_id(&self) -> Option { + DynJobKindTrait::args_group_id_dyn(&*self.0) + } + pub fn augment_args(&self, cmd: clap::Command) -> clap::Command { + DynJobKindTrait::augment_args_dyn(&*self.0, cmd) + } + pub fn augment_args_for_update(&self, cmd: clap::Command) -> clap::Command { + DynJobKindTrait::augment_args_for_update_dyn(&*self.0, cmd) + } + pub fn from_arg_matches( + self, + matches: &mut clap::ArgMatches, + ) -> clap::error::Result { + DynJobKindTrait::from_arg_matches_dyn(self.0, matches) + } + pub fn name(&self) -> Interned { + DynJobKindTrait::name_dyn(&*self.0) } } @@ -270,7 +323,7 @@ impl Serialize for DynJobKind { where S: Serializer, { - self.command_line_prefix().serialize(serializer) + self.name().serialize(serializer) } } @@ -279,226 +332,90 @@ impl<'de> Deserialize<'de> for DynJobKind { where D: Deserializer<'de>, { - let command_line_prefix: Cow<'_, [Interned]> = Cow::deserialize(deserializer)?; - match Self::registry().get_by_command_line_prefix(&command_line_prefix) { + let name = Cow::::deserialize(deserializer)?; + match Self::registry().get_by_name(&name) { Some(retval) => Ok(retval.clone()), None => Err(D::Error::custom(format_args!( - "unknown job kind: command line prefix not found in registry: {command_line_prefix:?}" + "unknown job kind: name not found in registry: {name:?}" ))), } } } -#[derive(Clone, Debug)] -struct JobKindRegistry { - command_line_prefix_to_job_kind_map: HashMap<&'static [Interned], DynJobKind>, - job_kinds: Vec, - subcommand_names: BTreeMap<&'static str, DynJobKind>, +trait DynJobArgsTrait: 'static + Send + Sync + fmt::Debug { + fn as_any(&self) -> &dyn Any; + fn as_arc_any(self: Arc) -> Arc; + fn eq_dyn(&self, other: &dyn DynJobArgsTrait) -> bool; + fn hash_dyn(&self, state: &mut dyn Hasher); + fn kind(&self) -> DynJobKind; } -enum JobKindRegisterError { - SameCommandLinePrefix { - command_line_prefix: &'static [Interned], - old_job_kind: DynJobKind, - new_job_kind: DynJobKind, - }, - CommandLinePrefixDoesNotMatchSubcommandName { - command_line_prefix: &'static [Interned], - expected: [Interned; 2], - job_kind: DynJobKind, - }, -} - -impl fmt::Display for JobKindRegisterError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::SameCommandLinePrefix { - command_line_prefix, - old_job_kind, - new_job_kind, - } => write!( - f, - "two different `DynJobKind` can't share the same `command_line_prefix` of:\n\ - {command_line_prefix:?}\n\ - old job kind:\n\ - {old_job_kind:?}\n\ - new job kind:\n\ - {new_job_kind:?}", - ), - Self::CommandLinePrefixDoesNotMatchSubcommandName { - command_line_prefix, - expected, - job_kind, - } => write!( - f, - "`JobKind::subcommand()` returned `Some` but the `command_line_prefix` is not as expected\n\ - (it should be `[program_name_for_internal_jobs(), subcommand_name]`):\n\ - command_line_prefix:\n\ - {command_line_prefix:?}\n\ - expected:\n\ - {expected:?}\n\ - job kind:\n\ - {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> { - type Locked = RwLockWriteGuard<'static, Arc>; - 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 { +impl DynJobArgsTrait for inner::DynJobArgs { + fn as_any(&self) -> &dyn Any { self } - fn make_mut(locked: &mut Self::Locked) -> &mut JobKindRegistry { - locked + fn as_arc_any(self: Arc) -> Arc { + self + } + + fn eq_dyn(&self, other: &dyn DynJobArgsTrait) -> bool { + other + .as_any() + .downcast_ref::() + .is_some_and(|other| self == other) + } + + fn hash_dyn(&self, mut state: &mut dyn Hasher) { + self.hash(&mut state); + } + + fn kind(&self) -> DynJobKind { + DynJobKind::from_arc(self.kind.clone()) } } -impl JobKindRegistry { - fn lock() -> &'static RwLock> { - static REGISTRY: OnceLock>> = OnceLock::new(); - REGISTRY.get_or_init(Default::default) +#[derive(Clone)] +pub struct DynJobArgs(Arc); + +impl DynJobArgs { + pub fn from_arc(kind: Arc, args: K::Args) -> Self { + Self(Arc::new(inner::DynJobArgs { kind, args })) } - fn try_register( - lock: L, - job_kind: DynJobKind, - ) -> Result<(), JobKindRegisterError> { - let command_line_prefix = Interned::into_inner(job_kind.command_line_prefix()); - let subcommand_name = job_kind - .subcommand() - .map(|subcommand| subcommand.get_name().intern()); - if let Some(subcommand_name) = subcommand_name { - let expected = [program_name_for_internal_jobs(), subcommand_name]; - if command_line_prefix != &expected { - return Err( - JobKindRegisterError::CommandLinePrefixDoesNotMatchSubcommandName { - command_line_prefix, - expected, - job_kind, - }, - ); - } - } - // run user code only outside of lock - let mut locked = lock.lock(); - let this = L::make_mut(&mut locked); - let result = match this - .command_line_prefix_to_job_kind_map - .entry(command_line_prefix) - { - Entry::Occupied(entry) => Err(JobKindRegisterError::SameCommandLinePrefix { - command_line_prefix, - old_job_kind: entry.get().clone(), - new_job_kind: job_kind, - }), - Entry::Vacant(entry) => { - this.job_kinds.push(job_kind.clone()); - if let Some(subcommand_name) = subcommand_name { - this.subcommand_names - .insert(Interned::into_inner(subcommand_name), job_kind.clone()); - } - 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::SameCommandLinePrefix { - command_line_prefix: _, - old_job_kind, - new_job_kind, - }) if old_job_kind == new_job_kind => Ok(()), - result => result, - } - } - #[track_caller] - fn register(lock: L, job_kind: DynJobKind) { - match Self::try_register(lock, job_kind) { - Err(e) => panic!("{e}"), - Ok(()) => {} - } - } - fn get() -> Arc { - Self::lock().read().expect("shouldn't be poisoned").clone() + pub fn new(kind: K, args: K::Args) -> Self { + Self::from_arc(Arc::new(kind), args) } } -impl Default for JobKindRegistry { - fn default() -> Self { - let mut retval = Self { - command_line_prefix_to_job_kind_map: HashMap::default(), - job_kinds: Vec::new(), - subcommand_names: BTreeMap::new(), - }; - for job_kind in [] { - Self::register(&mut retval, job_kind); - } - retval - } -} - -#[derive(Clone, Debug)] -pub struct JobKindRegistrySnapshot(Arc); - -impl JobKindRegistrySnapshot { - pub fn get() -> Self { - JobKindRegistrySnapshot(JobKindRegistry::get()) - } - pub fn get_by_command_line_prefix<'a>( - &'a self, - command_line_prefix: &[Interned], - ) -> Option<&'a DynJobKind> { - self.0 - .command_line_prefix_to_job_kind_map - .get(command_line_prefix) - } - pub fn job_kinds(&self) -> &[DynJobKind] { - &self.0.job_kinds - } -} - -#[derive(Clone, Debug, Hash, PartialEq, Eq)] -pub struct AnyInternalJob(pub DynJob); - -impl clap::Subcommand for AnyInternalJob { +impl clap::Subcommand for DynJobArgs { fn augment_subcommands(mut cmd: clap::Command) -> clap::Command { - for job_kind in JobKindRegistrySnapshot::get().0.subcommand_names.values() { - let Some(subcommand) = job_kind.subcommand() else { - // shouldn't happen, ignore it - continue; - }; - cmd = cmd.subcommand(subcommand); + let snapshot = JobKindRegistrySnapshot::get(); + for (&name, job_kind) in &snapshot.0.job_kinds { + let mut subcommand = clap::Command::new(name); + for dependency in job_kind.dependencies_kinds() { + subcommand = dependency.augment_args(subcommand); + } + cmd = cmd.subcommand(job_kind.augment_args(subcommand)); } cmd } - fn augment_subcommands_for_update(cmd: clap::Command) -> clap::Command { - Self::augment_subcommands(cmd) + fn augment_subcommands_for_update(mut cmd: clap::Command) -> clap::Command { + let snapshot = JobKindRegistrySnapshot::get(); + for (&name, job_kind) in &snapshot.0.job_kinds { + let mut subcommand = clap::Command::new(name); + for dependency in job_kind.dependencies_kinds() { + subcommand = dependency.augment_args_for_update(subcommand); + } + cmd = cmd.subcommand(job_kind.augment_args_for_update(subcommand)); + } + cmd } fn has_subcommand(name: &str) -> bool { JobKindRegistrySnapshot::get() .0 - .subcommand_names + .job_kinds .contains_key(name) } } @@ -528,16 +445,13 @@ impl clap::FromArgMatches for AnyInternalJob { } fn update_from_arg_matches(&mut self, matches: &clap::ArgMatches) -> Result<(), clap::Error> { - *self = Self::from_arg_matches(matches)?; - Ok(()) + Self::update_from_arg_matches_mut(self, &mut matches.clone()) } fn update_from_arg_matches_mut( &mut self, matches: &mut clap::ArgMatches, ) -> Result<(), clap::Error> { - *self = Self::from_arg_matches_mut(matches)?; - Ok(()) } } @@ -558,6 +472,12 @@ trait DynJobTrait: 'static + Send + Sync + fmt::Debug { mod inner { use super::*; + #[derive(Debug, PartialEq, Eq, Hash)] + pub(crate) struct DynJobArgs { + pub(crate) kind: Arc, + pub(crate) args: T::Args, + } + #[derive(Debug, PartialEq, Eq, Hash)] pub(crate) struct DynJob { pub(crate) kind: Arc, @@ -575,7 +495,7 @@ impl DynJobTrait for inner::DynJob { fn eq_dyn(&self, other: &dyn DynJobTrait) -> bool { other .as_any() - .downcast_ref::>() + .downcast_ref::() .is_some_and(|other| self == other) } @@ -727,53 +647,6 @@ impl<'de> Deserialize<'de> for DynJob { } } -impl JobKind for DynJobKind { - type Job = DynJob; - - fn inputs_and_direct_dependencies<'a>( - &'a self, - job: &'a Self::Job, - ) -> Cow<'a, BTreeMap>> { - Cow::Borrowed(job.inputs_and_direct_dependencies()) - } - - fn outputs(&self, job: &Self::Job) -> Interned<[JobItemName]> { - job.outputs() - } - - fn command_line_prefix(&self) -> Interned<[Interned]> { - self.0.command_line_prefix_dyn() - } - - fn to_command_line(&self, job: &Self::Job) -> Interned<[Interned]> { - job.to_command_line() - } - - fn subcommand(&self) -> Option { - self.0.subcommand_dyn() - } - - fn from_arg_matches(&self, matches: &mut clap::ArgMatches) -> clap::error::Result { - self.0.clone().from_arg_matches_dyn(matches) - } - - fn parse_command_line( - &self, - command_line: Interned<[Interned]>, - ) -> clap::error::Result { - self.0.clone().parse_command_line_dyn(command_line) - } - - fn run( - &self, - job: &Self::Job, - inputs: &[JobItem], - acquired_job: &mut AcquiredJob, - ) -> eyre::Result> { - job.run(inputs, acquired_job) - } -} - #[derive(Clone, Debug)] enum JobGraphNode { Job(DynJob), diff --git a/crates/fayalite/src/build/registry.rs b/crates/fayalite/src/build/registry.rs new file mode 100644 index 00000000..b2ae9397 --- /dev/null +++ b/crates/fayalite/src/build/registry.rs @@ -0,0 +1,241 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// See Notices.txt for copyright information + +use crate::{ + build::{DynJobKind, JobKind}, + intern::Interned, +}; +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<&'static str, DynJobKind>, +} + +enum JobKindRegisterError { + SameName { + name: &'static str, + 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> { + type Locked = RwLockWriteGuard<'static, Arc>; + 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> { + static REGISTRY: OnceLock>> = OnceLock::new(); + REGISTRY.get_or_init(Default::default) + } + fn try_register( + lock: L, + job_kind: DynJobKind, + ) -> Result<(), JobKindRegisterError> { + use std::collections::btree_map::Entry; + let name = Interned::into_inner(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(lock: L, job_kind: DynJobKind) { + match Self::try_register(lock, job_kind) { + Err(e) => panic!("{e}"), + Ok(()) => {} + } + } + fn get() -> Arc { + 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 [] { + Self::register(&mut retval, job_kind); + } + retval + } +} + +#[derive(Clone, Debug)] +pub struct JobKindRegistrySnapshot(Arc); + +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(&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, &'static str, DynJobKind>, +); + +impl<'a> Iterator for JobKindRegistryIter<'a> { + type Item = &'a DynJobKind; + + fn next(&mut self) -> Option { + self.0.next() + } + + fn size_hint(&self) -> (usize, Option) { + self.0.size_hint() + } + + fn count(self) -> usize + where + Self: Sized, + { + self.0.count() + } + + fn last(self) -> Option { + self.0.last() + } + + fn nth(&mut self, n: usize) -> Option { + self.0.nth(n) + } + + fn fold(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.0.next_back() + } + + fn nth_back(&mut self, n: usize) -> Option { + self.0.nth_back(n) + } + + fn rfold(self, init: B, f: F) -> B + where + F: FnMut(B, Self::Item) -> B, + { + self.0.rfold(init, f) + } +} + +#[track_caller] +pub fn register_job_kind(kind: K) { + DynJobKind::new(kind).register(); +}