henad_compute/entry/
set.rs

1//! The set of models a host offers.
2//!
3//! A model id is a lowercase ASCII letter, followed by lowercase ASCII letters, digits and underscores. Ids in that
4//! form stay safe as TOML values, CLI arguments and CSV cells.
5
6use std::fmt;
7use std::iter::FusedIterator;
8
9use henad_core::provenance::{BuildInfo, ModelSource};
10
11use super::ModelEntry;
12use crate::gpu::{GpuContext, GpuNeeds};
13
14/// Models a host offers, each id once, in the order the host lists them.
15#[derive(Debug, Clone)]
16pub struct ModelSet {
17    /// Build recorded on every entry inserted without one.
18    build: BuildInfo,
19    entries: Vec<ModelEntry>,
20}
21
22impl ModelSet {
23    /// Returns an empty set that records `build` as the build of every entry inserted without one.
24    pub fn new(build: BuildInfo) -> Self {
25        Self {
26            build,
27            entries: Vec::new(),
28        }
29    }
30
31    /// Adds `entry`. An entry whose source records no build yet receives the set's build, and any other entry keeps
32    /// the build it records.
33    ///
34    /// Note that an entry registered from a type in a third-party crate also receives the set's build. Its source's
35    /// type path still refers to the real crate. A model library therefore exports a set built from its own
36    /// [`build_info!`](henad_core::build_info), and its entries keep its name and version.
37    ///
38    /// # Errors
39    ///
40    /// Returns [`ModelSetError::DuplicateId`] for an id the set holds, and [`ModelSetError::InvalidId`] for an id
41    /// outside the grammar. The set is unchanged on an error.
42    pub fn insert(&mut self, entry: ModelEntry) -> Result<&mut Self, ModelSetError> {
43        check_id(entry.id())?;
44        self.check_free(&entry)?;
45        let entry = if entry.source().build().is_some() {
46            entry
47        } else {
48            let source = entry.source().clone().__with_build(self.build);
49            entry.with_source(source)
50        };
51        self.entries.push(entry);
52        Ok(self)
53    }
54
55    /// Adds every entry of `other`, each keeping its own source.
56    ///
57    /// # Errors
58    ///
59    /// Returns [`ModelSetError::DuplicateId`] for an id both sets hold. Checks every id before it adds any, and
60    /// leaves the set unchanged on an error.
61    pub fn extend(&mut self, other: Self) -> Result<&mut Self, ModelSetError> {
62        for entry in &other.entries {
63            self.check_free(entry)?;
64        }
65        self.entries.extend(other.entries);
66        Ok(self)
67    }
68
69    /// Returns entry `id`, whether or not a host can run it.
70    pub fn get(&self, id: &str) -> Option<&ModelEntry> {
71        self.entries.iter().find(|entry| entry.id() == id)
72    }
73
74    /// Returns entry `id`, or the reason a host with `gpu` cannot run it.
75    ///
76    /// # Errors
77    ///
78    /// Returns [`ModelLookupError::NotInSet`] for an id missing from the set, and
79    /// [`ModelLookupError::NeedsGpu`] for a GPU model when `gpu` is `None`.
80    pub fn lookup(&self, id: &str, gpu: Option<&GpuContext>) -> Result<&ModelEntry, ModelLookupError> {
81        if let Some(entry) = self.runnable(gpu).find(|entry| entry.id() == id) {
82            return Ok(entry);
83        }
84        let id = id.to_owned();
85        Err(if self.get(&id).is_some() {
86            ModelLookupError::NeedsGpu { id }
87        } else {
88            ModelLookupError::NotInSet { id }
89        })
90    }
91
92    /// Returns the entries a host with `gpu` can run, in the set's order. A GPU model runs only with a device.
93    pub fn runnable<'a>(&'a self, gpu: Option<&GpuContext>) -> impl Iterator<Item = &'a ModelEntry> + use<'a> {
94        let device = gpu.is_some();
95        self.entries
96            .iter()
97            .filter(move |entry| device || entry.gpu_needs().is_none())
98    }
99
100    /// Returns an iterator over the entries, in the set's order.
101    pub fn iter(&self) -> ModelSetIter<'_> {
102        ModelSetIter(self.entries.iter())
103    }
104
105    /// Number of entries.
106    pub fn len(&self) -> usize {
107        self.entries.len()
108    }
109
110    /// Returns whether the set holds no entry.
111    pub fn is_empty(&self) -> bool {
112        self.entries.is_empty()
113    }
114
115    /// Returns the merged needs of every GPU entry. A host requests a device that meets them before any model builds.
116    pub fn gpu_needs(&self) -> GpuNeeds {
117        self.entries
118            .iter()
119            .filter_map(ModelEntry::gpu_needs)
120            .fold(GpuNeeds::default(), GpuNeeds::merge)
121    }
122
123    /// Returns an error when the set already holds an entry with the same id as `entry`.
124    fn check_free(&self, entry: &ModelEntry) -> Result<(), ModelSetError> {
125        match self.get(entry.id()) {
126            Some(existing) => Err(ModelSetError::DuplicateId {
127                id: entry.id().to_owned(),
128                existing: Box::new(existing.source().clone()),
129                added: Box::new(entry.source().clone()),
130            }),
131            None => Ok(()),
132        }
133    }
134}
135
136impl<'a> IntoIterator for &'a ModelSet {
137    type Item = &'a ModelEntry;
138    type IntoIter = ModelSetIter<'a>;
139
140    fn into_iter(self) -> Self::IntoIter {
141        self.iter()
142    }
143}
144
145/// Iterator over the entries of a [`ModelSet`], in the set's order.
146#[derive(Debug, Clone)]
147pub struct ModelSetIter<'a>(std::slice::Iter<'a, ModelEntry>);
148
149impl<'a> Iterator for ModelSetIter<'a> {
150    type Item = &'a ModelEntry;
151
152    fn next(&mut self) -> Option<Self::Item> {
153        self.0.next()
154    }
155
156    fn size_hint(&self) -> (usize, Option<usize>) {
157        self.0.size_hint()
158    }
159}
160
161impl ExactSizeIterator for ModelSetIter<'_> {}
162
163impl FusedIterator for ModelSetIter<'_> {}
164
165/// Returns an error when `id` is outside the model id grammar.
166fn check_id(id: &str) -> Result<(), ModelSetError> {
167    let invalid = |reason| {
168        Err(ModelSetError::InvalidId {
169            id: id.to_owned(),
170            reason,
171        })
172    };
173    let mut chars = id.chars();
174    match chars.next() {
175        None => invalid("is empty"),
176        Some(first) if !first.is_ascii_lowercase() => invalid("does not start with a lowercase ASCII letter"),
177        Some(_)
178            if !chars
179                .all(|character| character.is_ascii_lowercase() || character.is_ascii_digit() || character == '_') =>
180        {
181            invalid("holds a character other than a lowercase ASCII letter, a digit or an underscore")
182        }
183        Some(_) => Ok(()),
184    }
185}
186
187/// Reason a model set rejects an entry.
188#[derive(Debug)]
189#[non_exhaustive]
190pub enum ModelSetError {
191    /// The set already holds a model with this id.
192    DuplicateId {
193        /// Id the two entries share.
194        id: String,
195        /// Source of the entry the set holds.
196        existing: Box<ModelSource>,
197        /// Source of the rejected entry.
198        added: Box<ModelSource>,
199    },
200    /// The id is outside the model id grammar.
201    InvalidId {
202        /// Id the entry declares.
203        id: String,
204        /// Clause saying how `id` breaks the grammar, such as "is empty".
205        reason: &'static str,
206    },
207}
208
209impl fmt::Display for ModelSetError {
210    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
211        match self {
212            Self::DuplicateId { id, existing, added } => write!(
213                f,
214                "the set already holds model '{id}' ({}), and cannot add it again ({})",
215                existing.type_path(),
216                added.type_path()
217            ),
218            Self::InvalidId { id, reason } => write!(
219                f,
220                "model id '{id}' {reason} (an id is a lowercase ASCII letter, then lowercase letters, digits and underscores)"
221            ),
222        }
223    }
224}
225
226impl std::error::Error for ModelSetError {}
227
228/// Reason a host cannot run a model that it requested from a set.
229#[derive(Debug)]
230#[non_exhaustive]
231pub enum ModelLookupError {
232    /// The set holds no model with this id.
233    NotInSet {
234        /// Id the host requested.
235        id: String,
236    },
237    /// A GPU model, requested with no compute device.
238    NeedsGpu {
239        /// Id of the GPU model.
240        id: String,
241    },
242}
243
244impl fmt::Display for ModelLookupError {
245    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
246        match self {
247            Self::NotInSet { id } => write!(f, "this build does not include model '{id}'"),
248            Self::NeedsGpu { id } => write!(
249                f,
250                "model '{id}' needs a GPU with compute support, and this device has none"
251            ),
252        }
253    }
254}
255
256impl std::error::Error for ModelLookupError {}
257
258#[cfg(test)]
259mod tests {
260    use henad_core::authoring::model::grid_model::GridModel;
261    use henad_core::grid::Grid2D;
262    use henad_core::params::{ParamDescriptor, ParamValue};
263    use henad_core::provenance::BuildInfo;
264    use henad_core::topology::NeighborhoodKind;
265    use henad_core::view::{StatDescriptor, StatValue};
266
267    use super::{ModelLookupError, ModelSet, ModelSetError};
268    use crate::entry::{ModelEntry, register_grid_model};
269
270    /// A grid model that never changes, registered under the id `$id`.
271    macro_rules! still_model {
272        ($ty:ident, $id:literal) => {
273            struct $ty;
274
275            impl GridModel for $ty {
276                const NAME: &'static str = $id;
277                const ID: &'static str = $id;
278                const DESCRIPTION: &'static str = $id;
279                const PALETTE: &'static [[u8; 4]] = &[[0, 0, 0, 255]];
280                const NEIGHBORHOOD: NeighborhoodKind = NeighborhoodKind::Moore;
281                const STATS: &'static [StatDescriptor] = &[];
282                type Params = ();
283
284                fn param_descriptors() -> Vec<ParamDescriptor> {
285                    Vec::new()
286                }
287
288                fn from_params(_params: &[ParamValue]) -> Self::Params {}
289
290                fn init(_grid: &mut Grid2D<u8>, _params: &[ParamValue], _rng: &mut u64) {}
291
292                fn step_cell(cell: u8, _neighbors: &[u8], _params: &Self::Params, _rng: &mut u64) -> u8 {
293                    cell
294                }
295
296                fn stats(_grid: &Grid2D<u8>) -> Vec<StatValue> {
297                    Vec::new()
298                }
299            }
300        };
301    }
302
303    still_model!(Still, "still");
304    still_model!(StillAgain, "still");
305    still_model!(Calm, "calm_2");
306    still_model!(Capital, "Still");
307    still_model!(Hyphen, "still-life");
308    still_model!(Digit, "2still");
309    still_model!(Empty, "");
310
311    fn build(package: &'static str) -> BuildInfo {
312        BuildInfo::__from_env(package, "1.0.0", None, None, None, None, false)
313    }
314
315    fn ids(models: &ModelSet) -> Vec<&str> {
316        models.iter().map(ModelEntry::id).collect()
317    }
318
319    #[test]
320    fn a_model_set_refuses_a_duplicate_id() {
321        let mut models = ModelSet::new(build("host"));
322        models.insert(register_grid_model::<Still>()).expect("a fresh id");
323        let Err(error) = models.insert(register_grid_model::<StillAgain>()) else {
324            panic!("a second 'still' joined the set");
325        };
326        let ModelSetError::DuplicateId { id, existing, added } = &error else {
327            panic!("expected a duplicate id, got {error:?}");
328        };
329        assert_eq!(id, "still");
330        assert!(existing.type_path().ends_with("::Still"), "{error}");
331        assert!(added.type_path().ends_with("::StillAgain"), "{error}");
332        assert_eq!(ids(&models), ["still"]);
333
334        let mut other = ModelSet::new(build("library"));
335        other.insert(register_grid_model::<Calm>()).expect("a fresh id");
336        other.insert(register_grid_model::<StillAgain>()).expect("a fresh id");
337        assert!(matches!(models.extend(other), Err(ModelSetError::DuplicateId { .. })));
338        assert_eq!(ids(&models), ["still"], "a refused extend adds nothing");
339    }
340
341    #[test]
342    fn a_set_iterates_in_insertion_order_either_way() {
343        let mut models = ModelSet::new(build("host"));
344        models.insert(register_grid_model::<Still>()).expect("a fresh id");
345        models.insert(register_grid_model::<Calm>()).expect("a fresh id");
346        let mut through_iter = models.iter();
347        assert_eq!(through_iter.len(), 2);
348        assert_eq!(through_iter.next().map(ModelEntry::id), Some("still"));
349        assert_eq!(through_iter.len(), 1);
350        let through_loop: Vec<&str> = (&models).into_iter().map(ModelEntry::id).collect();
351        assert_eq!(through_loop, ["still", "calm_2"]);
352    }
353
354    #[test]
355    fn a_model_set_refuses_an_id_outside_the_grammar() {
356        let mut models = ModelSet::new(build("host"));
357        for entry in [
358            register_grid_model::<Capital>(),
359            register_grid_model::<Hyphen>(),
360            register_grid_model::<Digit>(),
361            register_grid_model::<Empty>(),
362        ] {
363            let id = entry.id().to_owned();
364            let Err(error) = models.insert(entry) else {
365                panic!("'{id}' joined the set");
366            };
367            assert!(matches!(error, ModelSetError::InvalidId { .. }), "{error:?}");
368            assert!(error.to_string().contains(&format!("'{id}'")), "{error}");
369        }
370        assert!(models.is_empty());
371        models
372            .insert(register_grid_model::<Calm>())
373            .expect("letters, digits and an underscore");
374    }
375
376    #[test]
377    fn a_set_keeps_the_source_of_an_entry_from_another_set() {
378        let mut library = ModelSet::new(build("library"));
379        library.insert(register_grid_model::<Still>()).expect("a fresh id");
380        let mut host = ModelSet::new(build("host"));
381        host.insert(register_grid_model::<Calm>()).expect("a fresh id");
382        host.insert(library.get("still").expect("still is in the library").clone())
383            .expect("a fresh id");
384        assert_eq!(host.get("still").map(|entry| entry.source().package()), Some("library"));
385        assert_eq!(host.get("calm_2").map(|entry| entry.source().package()), Some("host"));
386
387        let mut combined = ModelSet::new(build("other"));
388        combined.extend(host).expect("no shared id");
389        assert_eq!(ids(&combined), ["calm_2", "still"]);
390        assert_eq!(
391            combined.get("still").map(|entry| entry.source().package()),
392            Some("library")
393        );
394        assert_eq!(
395            combined.get("calm_2").map(|entry| entry.source().package()),
396            Some("host")
397        );
398    }
399
400    #[test]
401    fn a_cpu_model_needs_no_device() {
402        let mut models = ModelSet::new(build("host"));
403        models.insert(register_grid_model::<Still>()).expect("a fresh id");
404        assert_eq!(models.gpu_needs(), crate::gpu::GpuNeeds::default());
405        assert!(models.lookup("still", None).is_ok());
406        assert!(matches!(
407            models.lookup("calm_2", None),
408            Err(ModelLookupError::NotInSet { .. })
409        ));
410    }
411}
henad_compute/entry/set.rs

henad_compute/entry/
set.rs
