Skip to main content

henad_compute/
simulation.rs

1//! Checked setups for one build of a model, and the simulations they build.
2//!
3//! A [`RunSetup`] holds a model's parameter values by id, a seed and scheduled actions, each checked against the
4//! model's descriptors when set. [`RunSetup::build`] returns a [`Simulation`], one built model that its caller steps,
5//! and a [`StatSample`] is one sample of the model's stats, read by label.
6//!
7//! Actions fire under [`henad_core::action::Fire::AfterStep`]. Tick 0's actions fire inside the build, and each
8//! later tick's actions fire after the step that reaches it. Every call returns a model's panic, a device error or a
9//! lost device as a [`Fault`], and a simulation that returned one rejects every later call that runs model code.
10//! [`Simulation`] lists the calls that enter the rayon pool.
11
12use std::fmt;
13use std::io::{self, Write};
14use std::ops::ControlFlow;
15
16#[cfg(not(target_arch = "wasm32"))]
17use henad_core::action::Fire;
18use henad_core::action::{RefusedActions, Schedule, Scheduled};
19use henad_core::explore::replay::Replay;
20use henad_core::explore::value::{ValueError, check_value, parse_value};
21use henad_core::export::state as state_export;
22use henad_core::model::SimState;
23use henad_core::params::ParamValue;
24use henad_core::send_sync::WasmNotSend;
25use henad_core::view::{EdgeView, GridView, PointView, StatEntry, StatValue};
26
27use crate::entry::{ModelEntry, ModelState};
28use crate::fault::{BUILDING, Fault, STEPPING, catching};
29use crate::gpu::GpuContext;
30#[cfg(not(target_arch = "wasm32"))]
31use crate::gpu::{GpuSimState, fault::catching_on, stepping};
32
33/// Values, seed and schedule for one build of an entry, each checked against its descriptors when set.
34///
35/// [`ModelEntry::setup`] returns one at the declared defaults, the default seed and no actions.
36#[derive(Debug, Clone)]
37pub struct RunSetup {
38    entry: ModelEntry,
39    /// One value per parameter in descriptor order.
40    values: Vec<ParamValue>,
41    seed: Option<u64>,
42    schedule: Schedule,
43}
44
45impl RunSetup {
46    /// Returns a setup of `entry` at its declared defaults, the default seed and no actions.
47    pub(crate) fn new(entry: ModelEntry) -> Self {
48        let values = entry
49            .param_descriptors()
50            .iter()
51            .map(|descriptor| descriptor.kind.default_value())
52            .collect();
53        Self {
54            entry,
55            values,
56            seed: None,
57            schedule: Schedule::default(),
58        }
59    }
60
61    /// Returns the setup with parameter `id` set to `value`.
62    ///
63    /// `value` is an `f32`, a `u32`, a `bool`, or a [`ParamValue`] itself, as `ParamValue::Choice(index)` is for a
64    /// choice parameter. A `u32` is never accepted as a choice index.
65    ///
66    /// # Errors
67    ///
68    /// Returns [`SetupError::Param`] for an unknown id, a value of another kind, or a value out of bounds.
69    pub fn set(mut self, id: &str, value: impl Into<ParamValue>) -> Result<Self, SetupError> {
70        let value = value.into();
71        let index = param_index(&self.entry, id)?;
72        check_value(&self.entry.param_descriptors()[index].kind, &value).map_err(|error| param_error(id, error))?;
73        self.values[index] = value;
74        Ok(self)
75    }
76
77    /// Returns the setup with parameter `id` read from `text`, as `--set` reads it. A choice parameter accepts an
78    /// option name.
79    ///
80    /// # Errors
81    ///
82    /// Returns [`SetupError::Param`] for an unknown id, text that cannot be parsed as the parameter's kind, or a value
83    /// out of bounds.
84    pub fn set_text(mut self, id: &str, text: &str) -> Result<Self, SetupError> {
85        let index = param_index(&self.entry, id)?;
86        let value =
87            parse_value(&self.entry.param_descriptors()[index].kind, text).map_err(|error| param_error(id, error))?;
88        self.values[index] = value;
89        Ok(self)
90    }
91
92    /// Returns the setup with `seed` instead of the model's default seed.
93    pub fn with_seed(mut self, seed: u64) -> Self {
94        self.seed = Some(seed);
95        self
96    }
97
98    /// Returns the setup with action `action_id` scheduled at `tick`. Two actions due at one tick fire in the order
99    /// they were scheduled.
100    ///
101    /// # Errors
102    ///
103    /// Returns [`SetupError::UnknownAction`] for an id the model does not declare.
104    pub fn act_at(mut self, action_id: &str, tick: u64) -> Result<Self, SetupError> {
105        let index = action_index(&self.entry, action_id)?;
106        let mut entries = self.schedule.entries().to_vec();
107        entries.push(Scheduled {
108            index,
109            id: action_id.to_owned(),
110            tick,
111        });
112        self.schedule = Schedule::from_entries(entries);
113        Ok(self)
114    }
115
116    /// Returns the setup a recorded run was built from.
117    ///
118    /// # Errors
119    ///
120    /// Returns [`SetupError::WrongModel`] for a run of another model, and the errors of [`Self::from_parts`] for
121    /// values or actions the entry rejects.
122    pub fn from_replay(entry: &ModelEntry, replay: &Replay) -> Result<Self, SetupError> {
123        if replay.model != entry.id() {
124            return Err(SetupError::WrongModel {
125                expected: entry.id().to_owned(),
126                found: replay.model.clone(),
127            });
128        }
129        Self::from_parts(entry, &replay.params, Some(replay.seed), replay.schedule.clone())
130    }
131
132    /// Returns a setup holding `values`, `seed` and `schedule`, each checked as [`Self::set`] and [`Self::act_at`]
133    /// check them.
134    ///
135    /// `values` holds one value per parameter, in the order of [`ModelEntry::param_descriptors`], as
136    /// [`Self::values`] returns them.
137    ///
138    /// # Errors
139    ///
140    /// Returns a [`SetupError`] for a value count other than the entry's parameter count, a value of another kind or
141    /// out of bounds, or a schedule entry for an action that the entry does not declare.
142    pub fn from_parts(
143        entry: &ModelEntry,
144        values: &[ParamValue],
145        seed: Option<u64>,
146        schedule: Schedule,
147    ) -> Result<Self, SetupError> {
148        let descriptors = entry.param_descriptors();
149        if values.len() != descriptors.len() {
150            return Err(SetupError::ParamCount {
151                expected: descriptors.len(),
152                found: values.len(),
153            });
154        }
155        for (descriptor, value) in descriptors.iter().zip(values) {
156            check_value(&descriptor.kind, value).map_err(|error| param_error(descriptor.id, error))?;
157        }
158        let actions = entry.action_descriptors();
159        for scheduled in schedule.entries() {
160            if actions
161                .get(scheduled.index)
162                .is_none_or(|action| action.id != scheduled.id)
163            {
164                return Err(SetupError::UnknownAction {
165                    id: scheduled.id.clone(),
166                });
167            }
168        }
169        Ok(Self {
170            entry: entry.clone(),
171            values: values.to_vec(),
172            seed,
173            schedule,
174        })
175    }
176
177    /// Entry the setup builds.
178    pub fn entry(&self) -> &ModelEntry {
179        &self.entry
180    }
181
182    /// One value per parameter, in the order of [`ModelEntry::param_descriptors`].
183    pub fn values(&self) -> &[ParamValue] {
184        &self.values
185    }
186
187    /// Seed of the build. `None` is the model's default seed, and no `Some` value reproduces it.
188    pub fn seed(&self) -> Option<u64> {
189        self.seed
190    }
191
192    /// Actions scheduled for the build, by tick.
193    pub fn schedule(&self) -> &Schedule {
194        &self.schedule
195    }
196
197    /// Builds the model, on `gpu` for a GPU entry, then fires the schedule's tick-0 entries.
198    ///
199    /// The returned simulation already holds tick 0's actions, and records tick 0 as fired. On native targets a CPU
200    /// model's `init` runs outside the rayon pool, and the build enters the pool only to fire tick 0's actions.
201    ///
202    /// # Errors
203    ///
204    /// Returns a [`Fault`] when the build or a tick-0 action panics, the device rejects the build or the action, a GPU
205    /// entry has no device, or a GPU entry is built on wasm32.
206    pub fn build(&self, gpu: Option<&GpuContext>) -> Result<Simulation, Fault> {
207        #[cfg(target_arch = "wasm32")]
208        if self.entry.gpu_needs().is_some() {
209            return Err(gpu_in_browser(self.entry.id()));
210        }
211        let schedule = &self.schedule;
212        let engine = match self.entry.build(&self.values, self.seed, gpu)? {
213            ModelState::Cpu(mut state) => {
214                if !schedule.is_empty() {
215                    cpu_call(BUILDING, || refusal(&schedule.run_due(&mut *state), BUILDING))?;
216                }
217                Engine::Cpu(state)
218            }
219            #[cfg(not(target_arch = "wasm32"))]
220            ModelState::Gpu(mut state) => {
221                let ctx = gpu
222                    .ok_or_else(|| Fault::refused(BUILDING, "a GPU model was built with no device"))?
223                    .clone();
224                if !schedule.is_empty() {
225                    gpu_call(&ctx, BUILDING, || {
226                        refusal(&stepping::run_due(&mut *state, &ctx, schedule), BUILDING)?;
227                        stepping::wait(&ctx)
228                    })?;
229                }
230                Engine::Gpu { state, ctx }
231            }
232            #[cfg(target_arch = "wasm32")]
233            ModelState::Gpu(_) => return Err(gpu_in_browser(self.entry.id())),
234        };
235        Ok(Simulation {
236            setup: self.clone(),
237            engine,
238            earlier_fault: None,
239        })
240    }
241}
242
243/// Returns the position of parameter `id` in the descriptors of `entry`.
244fn param_index(entry: &ModelEntry, id: &str) -> Result<usize, SetupError> {
245    entry.param_index(id).ok_or_else(|| {
246        SetupError::Param(ValueError::UnknownParam {
247            id: id.to_owned(),
248            known: entry
249                .param_descriptors()
250                .iter()
251                .map(|descriptor| descriptor.id)
252                .collect(),
253        })
254    })
255}
256
257/// Returns the position of action `id` in the descriptors of `entry`.
258fn action_index(entry: &ModelEntry, id: &str) -> Result<usize, SetupError> {
259    entry
260        .action_index(id)
261        .ok_or_else(|| SetupError::UnknownAction { id: id.to_owned() })
262}
263
264/// Returns the error for a value that parameter `id` rejects.
265fn param_error(id: &str, error: ValueError) -> SetupError {
266    SetupError::Param(ValueError::Param {
267        id: id.to_owned(),
268        source: Box::new(error),
269    })
270}
271
272/// Returns the fault for a scheduled action the state rejected.
273///
274/// Every engine rejects only an index past its model's actions, and a setup checks each id before it stores an entry.
275/// A rejection is therefore an engine contract violation.
276fn refusal(refused: &RefusedActions<'_>, during: &'static str) -> Result<(), Fault> {
277    match refused.first() {
278        None => Ok(()),
279        Some(action) => Err(Fault::refused(
280            during,
281            format!(
282                "the model refused its own action '{}' at tick {}",
283                action.id, action.tick
284            ),
285        )),
286    }
287}
288
289/// Returns the fault for a GPU entry built in a browser.
290#[cfg(target_arch = "wasm32")]
291fn gpu_in_browser(id: &str) -> Fault {
292    Fault::refused(
293        BUILDING,
294        format!("model '{id}' runs on the GPU, and a browser steps a GPU model through GpuSimThread alone"),
295    )
296}
297
298/// Runs `task` inside the current rayon pool and catches a panic out of it.
299fn cpu_call<T: WasmNotSend>(
300    during: &'static str,
301    task: impl FnOnce() -> Result<T, Fault> + WasmNotSend,
302) -> Result<T, Fault> {
303    in_pool(|| catching(during, task)?)
304}
305
306/// Runs `task` inside the current rayon pool.
307///
308/// Called from outside the pool, the body runs on a worker. Each parallel pass then starts from a worker instead of
309/// being injected from outside and parking the caller.
310///
311/// Note that the body is one job on the pool. A worker waiting in another host's join can run it nested, and that
312/// join returns only once the body ends.
313#[cfg(not(target_arch = "wasm32"))]
314fn in_pool<T: Send>(task: impl FnOnce() -> T + Send) -> T {
315    rayon::scope(|_| task())
316}
317
318/// Runs `task` on the calling thread. The frame driver pumps from outside the pool too.
319#[cfg(target_arch = "wasm32")]
320fn in_pool<T>(task: impl FnOnce() -> T) -> T {
321    task()
322}
323
324/// Reason a setup, a live edit or a live action was rejected.
325#[derive(Debug)]
326#[non_exhaustive]
327pub enum SetupError {
328    /// An unknown parameter id, or a value its descriptor rejects.
329    Param(ValueError),
330    /// [`RunSetup::from_parts`] received `found` values where the entry declares `expected` parameters.
331    ParamCount {
332        /// Number of parameters the entry declares.
333        expected: usize,
334        /// Number of values passed.
335        found: usize,
336    },
337    /// An action id the model does not declare.
338    UnknownAction {
339        /// Action id as given, by a call or a schedule entry.
340        id: String,
341    },
342    /// A live edit of a parameter that applies only on a rebuild.
343    ReloadOnly {
344        /// Id of the parameter edited.
345        id: String,
346    },
347    /// A recorded run of model `found`, passed to an entry of model `expected`.
348    WrongModel {
349        /// Id of the entry's model.
350        expected: String,
351        /// Id of the model the run records.
352        found: String,
353    },
354    /// Model code panicked, or the device rejected an action pass, during a live edit or action.
355    Fault(Fault),
356}
357
358impl fmt::Display for SetupError {
359    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
360        match self {
361            Self::Param(_) => f.write_str("cannot set the parameter"),
362            Self::ParamCount { expected, found } => {
363                write!(f, "expected {expected} parameter values, found {found}")
364            }
365            Self::UnknownAction { id } => write!(f, "model has no action '{id}'"),
366            Self::ReloadOnly { id } => write!(f, "parameter '{id}' applies only when the model is built"),
367            Self::WrongModel { expected, found } => {
368                write!(f, "the run is of model '{found}', not '{expected}'")
369            }
370            Self::Fault(_) => f.write_str("the model faulted"),
371        }
372    }
373}
374
375impl std::error::Error for SetupError {
376    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
377        match self {
378            Self::Param(error) => Some(error),
379            Self::Fault(fault) => Some(fault),
380            Self::ParamCount { .. }
381            | Self::UnknownAction { .. }
382            | Self::ReloadOnly { .. }
383            | Self::WrongModel { .. } => None,
384        }
385    }
386}
387
388impl From<Fault> for SetupError {
389    fn from(fault: Fault) -> Self {
390        Self::Fault(fault)
391    }
392}
393
394/// State a [`Simulation`] steps, with the device a GPU state lives on.
395enum Engine {
396    Cpu(Box<dyn SimState>),
397    #[cfg(not(target_arch = "wasm32"))]
398    Gpu {
399        state: Box<dyn GpuSimState>,
400        ctx: GpuContext,
401    },
402}
403
404impl Engine {
405    fn state(&self) -> &dyn SimState {
406        match self {
407            Self::Cpu(state) => &**state,
408            #[cfg(not(target_arch = "wasm32"))]
409            Self::Gpu { state, .. } => &**state,
410        }
411    }
412}
413
414/// One built model and its schedule, stepped by its caller.
415///
416/// On native targets a CPU simulation runs the model code of each stepping, sampling, view, layout and action call
417/// inside one `rayon::scope` on the current pool, and `pool.install(|| simulation.run_for(n))` picks the pool.
418/// `run_for(0)` and a `run_to` at or behind the current tick step nothing and enter no pool. [`Self::write_state`]
419/// enters it for its view preparation alone, and [`Self::set_param`] runs no parallel pass and stays on the calling
420/// thread. A GPU simulation runs each call on the calling thread, inside wgpu error scopes that keep a device error on
421/// the call that raised it.
422///
423/// Note that a call made from outside the pool runs as one job on it, and a worker waiting in another host's join
424/// can run that job nested. The other host's join then waits for the whole call. A host that steps independently
425/// steps inside `install` on its own pool.
426///
427/// An error no scope catches is stored in the context's [`FaultSink`](crate::fault::FaultSink), and whichever holder of
428/// the context waits next reports it. A lost device fails the next call that waits for the device with
429/// [`FaultKind::DeviceLost`]. Simulations that step at once on several threads each take their own device, for
430/// example from one `henad::gpu::acquire_headless` call per thread. A second [`GpuContext::new`] on one device takes
431/// over its error handling from the first context.
432///
433/// A call that returns a [`Fault`] can stop part way through a step or an action, with some of its edits applied.
434/// Every later stepping, sampling, live edit, action, view and layout call then returns a [`FaultKind::Refused`]
435/// fault that carries the first fault's message, or [`FaultKind::DeviceLost`] once the device is lost. Rebuild from
436/// [`Self::setup`] to go on.
437///
438/// [`FaultKind::DeviceLost`]: crate::fault::FaultKind::DeviceLost
439/// [`FaultKind::Refused`]: crate::fault::FaultKind::Refused
440pub struct Simulation {
441    setup: RunSetup,
442    engine: Engine,
443    /// First fault a call returned, as its message.
444    earlier_fault: Option<String>,
445}
446
447// A host steps a simulation on its own thread.
448#[cfg(not(target_arch = "wasm32"))]
449const _: fn() = || {
450    fn send<T: Send>() {}
451    send::<Simulation>();
452};
453
454/// Prints the setup and the tick, and leaves out the state.
455impl fmt::Debug for Simulation {
456    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
457        f.debug_struct("Simulation")
458            .field("setup", &self.setup)
459            .field("tick", &self.tick())
460            .finish_non_exhaustive()
461    }
462}
463
464impl Simulation {
465    /// Setup the simulation was built from. Live [`Self::set_param`] edits and immediate [`Self::act`] calls leave it
466    /// unchanged.
467    pub fn setup(&self) -> &RunSetup {
468        &self.setup
469    }
470
471    /// Number of ticks stepped so far.
472    pub fn tick(&self) -> u64 {
473        self.engine.state().tick()
474    }
475
476    /// Size of the population: its cells, its agents or its live nodes.
477    pub fn population(&self) -> u64 {
478        self.engine.state().population()
479    }
480
481    /// Approximate memory the state owns: host heap for a CPU state, device buffers and textures for a GPU state.
482    pub fn heap_bytes(&self) -> usize {
483        self.engine.state().heap_bytes()
484    }
485
486    /// Number of jobs one step splits into. `None` for a backend with no such split.
487    pub fn parallel_jobs(&self) -> Option<usize> {
488        self.engine.state().parallel_jobs()
489    }
490
491    /// Steps one tick. Same as `run_for(1)`.
492    ///
493    /// Note that each call enters the rayon pool once on the CPU and waits for the device once on the GPU. A loop of
494    /// `step()` pays that per tick, where [`Self::run_for`] pays it once. A loop of `step()` from outside the pool
495    /// belongs inside `rayon::scope` or `pool.install`.
496    ///
497    /// # Errors
498    ///
499    /// Returns a [`Fault`] when the model panics or the device reports an error. Once an earlier call returned a
500    /// fault, the fault is [`FaultKind::Refused`], or [`FaultKind::DeviceLost`] once the device is lost.
501    ///
502    /// [`FaultKind::DeviceLost`]: crate::fault::FaultKind::DeviceLost
503    /// [`FaultKind::Refused`]: crate::fault::FaultKind::Refused
504    pub fn step(&mut self) -> Result<(), Fault> {
505        self.run_for(1)
506    }
507
508    /// Steps `ticks` ticks.
509    ///
510    /// Note that without the template's profile block, a debug build runs the kernels at opt-level 0 in the crate that
511    /// registers the model, or in henad-models for an example entry, and its timings mean little. `--release` is the
512    /// measured configuration.
513    ///
514    /// # Errors
515    ///
516    /// Returns a [`Fault`] when the model panics or the device reports an error. Once an earlier call returned a
517    /// fault, the fault is [`FaultKind::Refused`], or [`FaultKind::DeviceLost`] once the device is lost.
518    ///
519    /// [`FaultKind::DeviceLost`]: crate::fault::FaultKind::DeviceLost
520    /// [`FaultKind::Refused`]: crate::fault::FaultKind::Refused
521    pub fn run_for(&mut self, ticks: u64) -> Result<(), Fault> {
522        let end = self.tick().saturating_add(ticks);
523        self.advance(end)
524    }
525
526    /// Steps up to `tick` and never past it. A tick at or behind the current one steps nothing.
527    ///
528    /// # Errors
529    ///
530    /// Returns a [`Fault`] when the model panics or the device reports an error. Once an earlier call returned a
531    /// fault, the fault is [`FaultKind::Refused`], or [`FaultKind::DeviceLost`] once the device is lost.
532    ///
533    /// [`FaultKind::DeviceLost`]: crate::fault::FaultKind::DeviceLost
534    /// [`FaultKind::Refused`]: crate::fault::FaultKind::Refused
535    pub fn run_to(&mut self, tick: u64) -> Result<(), Fault> {
536        self.advance(tick)
537    }
538
539    /// Runs to `end_tick`, calling `on_sample` at the current tick, at each later multiple of `interval`, and at
540    /// `end_tick`.
541    ///
542    /// Returns the `Break` of the first sample that asks to stop, or `Continue` once `end_tick` is reached. Each
543    /// sample is taken as [`Self::stats`] takes it. An `end_tick` at or behind the current tick steps nothing and
544    /// takes the one sample at the current tick. Each call samples the tick it starts on, and two calls back to back
545    /// both sample the tick they share.
546    ///
547    /// On native targets a CPU model enters the pool once for the whole call, and `on_sample` runs inside it on a
548    /// pool worker. It needs `Send`, and it holds that worker while it runs. A GPU model calls it on the calling
549    /// thread. Either way `on_sample` runs outside the fault scopes, and a panic in it unwinds to the caller as an
550    /// ordinary panic.
551    ///
552    /// Note that a CPU model's whole call can run nested under another host's join on the same pool, as the type's
553    /// docs describe. An `on_sample` that waits on another user of the pool, for example through a bounded channel, can
554    /// then deadlock.
555    ///
556    /// # Errors
557    ///
558    /// Returns a [`Fault`] when the model panics or the device reports an error. Once an earlier call returned a
559    /// fault, the fault is [`FaultKind::Refused`], or [`FaultKind::DeviceLost`] once the device is lost.
560    ///
561    /// [`FaultKind::DeviceLost`]: crate::fault::FaultKind::DeviceLost
562    /// [`FaultKind::Refused`]: crate::fault::FaultKind::Refused
563    ///
564    /// # Panics
565    ///
566    /// Panics when `interval` is 0.
567    pub fn run_sampled<B: WasmNotSend>(
568        &mut self,
569        end_tick: u64,
570        interval: u64,
571        mut on_sample: impl FnMut(&StatSample) -> ControlFlow<B> + WasmNotSend,
572    ) -> Result<ControlFlow<B>, Fault> {
573        assert!(interval > 0, "a sampling interval is at least one tick");
574        self.check_earlier_fault()?;
575        let Self { setup, engine, .. } = self;
576        let schedule = &setup.schedule;
577        let sampled = match engine {
578            // One pool entry for the whole call. A pool entry per sample costs more than a small model's tick.
579            Engine::Cpu(state) => in_pool(|| {
580                let first = catching(STEPPING, || sample_cpu(&mut **state))?;
581                if let ControlFlow::Break(stop) = on_sample(&first) {
582                    return Ok(ControlFlow::Break(stop));
583                }
584                while state.tick() < end_tick {
585                    let next = next_sample_tick(state.tick(), interval, end_tick);
586                    let sample = catching(STEPPING, || {
587                        step_cpu(&mut **state, schedule, next)?;
588                        Ok(sample_cpu(&mut **state))
589                    })??;
590                    if let ControlFlow::Break(stop) = on_sample(&sample) {
591                        return Ok(ControlFlow::Break(stop));
592                    }
593                }
594                Ok(ControlFlow::Continue(()))
595            }),
596            #[cfg(not(target_arch = "wasm32"))]
597            Engine::Gpu { state, ctx } => run_sampled_gpu(&mut **state, ctx, schedule, end_tick, interval, on_sample),
598        };
599        self.record(sampled)
600    }
601
602    /// Returns the stats of the current tick, sampled as a sweep track samples them: `prepare_view` first on the CPU, a
603    /// blocking stats-only readback on the GPU.
604    ///
605    /// # Errors
606    ///
607    /// Returns a [`Fault`] when the model panics or the device reports an error. Once an earlier call returned a
608    /// fault, the fault is [`FaultKind::Refused`], or [`FaultKind::DeviceLost`] once the device is lost.
609    ///
610    /// [`FaultKind::DeviceLost`]: crate::fault::FaultKind::DeviceLost
611    /// [`FaultKind::Refused`]: crate::fault::FaultKind::Refused
612    pub fn stats(&mut self) -> Result<StatSample, Fault> {
613        self.check_earlier_fault()?;
614        let sample = match &mut self.engine {
615            Engine::Cpu(state) => cpu_call(STEPPING, || Ok(sample_cpu(&mut **state))),
616            #[cfg(not(target_arch = "wasm32"))]
617            Engine::Gpu { state, ctx } => sample_gpu(&mut **state, ctx),
618        };
619        self.record(sample)
620    }
621
622    /// Sets parameter `id` on the running model.
623    ///
624    /// # Errors
625    ///
626    /// Returns [`SetupError::Param`] for an unknown id or a value the descriptor rejects, [`SetupError::ReloadOnly`]
627    /// for a parameter that takes effect only when the model is built, and [`SetupError::Fault`] when the model panics
628    /// or the device reports an error. Once an earlier call returned a fault, a value the descriptor accepts gets
629    /// [`FaultKind::Refused`], or [`FaultKind::DeviceLost`] once the device is lost.
630    ///
631    /// [`FaultKind::DeviceLost`]: crate::fault::FaultKind::DeviceLost
632    /// [`FaultKind::Refused`]: crate::fault::FaultKind::Refused
633    pub fn set_param(&mut self, id: &str, value: impl Into<ParamValue>) -> Result<(), SetupError> {
634        let value = value.into();
635        let index = param_index(&self.setup.entry, id)?;
636        let descriptor = &self.setup.entry.param_descriptors()[index];
637        check_value(&descriptor.kind, &value).map_err(|error| param_error(id, error))?;
638        self.check_earlier_fault()?;
639        // The state rejects a reload-only index itself, from the descriptors it was built with.
640        let accepted = match &mut self.engine {
641            Engine::Cpu(state) => catching(STEPPING, || state.set_param(index, &value)),
642            #[cfg(not(target_arch = "wasm32"))]
643            Engine::Gpu { state, ctx } => gpu_call(ctx, STEPPING, || Ok(state.set_param(index, &value))),
644        };
645        if self.record(accepted)? {
646            Ok(())
647        } else {
648            Err(SetupError::ReloadOnly { id: id.to_owned() })
649        }
650    }
651
652    /// Runs action `action_id` now, between ticks.
653    ///
654    /// # Errors
655    ///
656    /// Returns [`SetupError::UnknownAction`] for an id the model does not declare, and [`SetupError::Fault`] when the
657    /// model panics or the device rejects the action's pass. Once an earlier call returned a fault, a declared action
658    /// gets [`FaultKind::Refused`], or [`FaultKind::DeviceLost`] once the device is lost.
659    ///
660    /// [`FaultKind::DeviceLost`]: crate::fault::FaultKind::DeviceLost
661    /// [`FaultKind::Refused`]: crate::fault::FaultKind::Refused
662    pub fn act(&mut self, action_id: &str) -> Result<(), SetupError> {
663        let index = action_index(&self.setup.entry, action_id)?;
664        self.check_earlier_fault()?;
665        let refused = || Fault::refused(STEPPING, format!("the model refused its own action '{action_id}'"));
666        let acted = match &mut self.engine {
667            Engine::Cpu(state) => cpu_call(STEPPING, || if state.act(index) { Ok(()) } else { Err(refused()) }),
668            #[cfg(not(target_arch = "wasm32"))]
669            Engine::Gpu { state, ctx } => gpu_call(ctx, STEPPING, || {
670                let mut encoder = ctx.device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
671                    label: Some("henad_simulation_action"),
672                });
673                if !state.encode_action(&mut encoder, index) {
674                    return Err(refused());
675                }
676                ctx.queue.submit(Some(encoder.finish()));
677                stepping::wait(ctx)
678            }),
679        };
680        Ok(self.record(acted)?)
681    }
682
683    /// Prepares the views as a publish would, then borrows all three together. A GPU model's views stay on the
684    /// device, and each view is `None`.
685    ///
686    /// # Errors
687    ///
688    /// Returns a [`Fault`] when the model panics while preparing its views. Once an earlier call returned a fault, the
689    /// fault is [`FaultKind::Refused`], or [`FaultKind::DeviceLost`] once the device is lost.
690    ///
691    /// [`FaultKind::DeviceLost`]: crate::fault::FaultKind::DeviceLost
692    /// [`FaultKind::Refused`]: crate::fault::FaultKind::Refused
693    pub fn views(&mut self) -> Result<SimulationViews<'_>, Fault> {
694        self.check_earlier_fault()?;
695        let prepared = match &mut self.engine {
696            Engine::Cpu(state) => cpu_call(STEPPING, || {
697                state.prepare_view();
698                Ok(())
699            }),
700            #[cfg(not(target_arch = "wasm32"))]
701            Engine::Gpu { .. } => Ok(()),
702        };
703        self.record(prepared)?;
704        Ok(match &self.engine {
705            Engine::Cpu(state) => SimulationViews {
706                grid: state.grid_view(),
707                points: state.point_view(),
708                edges: state.edge_view(),
709            },
710            #[cfg(not(target_arch = "wasm32"))]
711            Engine::Gpu { .. } => SimulationViews {
712                grid: None,
713                points: None,
714                edges: None,
715            },
716        })
717    }
718
719    /// Relaxes a network model's layout for `budget_ms` milliseconds. A model without a layout keeps its points where
720    /// they are.
721    ///
722    /// # Errors
723    ///
724    /// Returns a [`Fault`] when the model panics. Once an earlier call returned a fault, the fault is
725    /// [`FaultKind::Refused`], or [`FaultKind::DeviceLost`] once the device is lost.
726    ///
727    /// [`FaultKind::DeviceLost`]: crate::fault::FaultKind::DeviceLost
728    /// [`FaultKind::Refused`]: crate::fault::FaultKind::Refused
729    pub fn relax_layout(&mut self, budget_ms: f32) -> Result<(), Fault> {
730        self.check_earlier_fault()?;
731        let relaxed = match &mut self.engine {
732            Engine::Cpu(state) => cpu_call(STEPPING, || {
733                if state.set_layout(true, budget_ms) {
734                    state.relax_layout();
735                    state.set_layout(false, budget_ms);
736                }
737                Ok(())
738            }),
739            #[cfg(not(target_arch = "wasm32"))]
740            Engine::Gpu { .. } => Ok(()),
741        };
742        self.record(relaxed)
743    }
744
745    /// Writes the grid, points or edges as `--export` does.
746    ///
747    /// Every section the model has is written, so a model with a field under its agents writes both the grid and the
748    /// points. The views are prepared first, as [`Self::views`] prepares them.
749    ///
750    /// # Errors
751    ///
752    /// Returns [`ExportError::Fault`] when the model panics while preparing its views. Once an earlier call returned a
753    /// fault, it returns [`ExportError::Fault`] for every model, holding [`FaultKind::Refused`], or
754    /// [`FaultKind::DeviceLost`] once the device is lost. Otherwise it returns [`ExportError::GpuState`] for a GPU
755    /// model, whose views stay on the device, and [`ExportError::Io`] when `writer` fails.
756    ///
757    /// [`FaultKind::DeviceLost`]: crate::fault::FaultKind::DeviceLost
758    /// [`FaultKind::Refused`]: crate::fault::FaultKind::Refused
759    pub fn write_state(&mut self, writer: &mut dyn Write) -> Result<(), ExportError> {
760        self.check_earlier_fault()?;
761        match &mut self.engine {
762            Engine::Cpu(state) => {
763                // Only the prepare is scoped. The writer is not `Send`.
764                let prepared = cpu_call(STEPPING, || {
765                    state.prepare_view();
766                    Ok(())
767                });
768                self.record(prepared)?;
769                write_views(self.engine.state(), writer)?;
770                Ok(())
771            }
772            #[cfg(not(target_arch = "wasm32"))]
773            Engine::Gpu { .. } => Err(ExportError::GpuState),
774        }
775    }
776
777    /// Steps to `end` under [`Fire::AfterStep`], firing each tick's actions after the step that reaches it.
778    fn advance(&mut self, end: u64) -> Result<(), Fault> {
779        self.check_earlier_fault()?;
780        let Self { setup, engine, .. } = self;
781        let schedule = &setup.schedule;
782        let advanced = match engine {
783            Engine::Cpu(state) => {
784                if state.tick() >= end {
785                    return Ok(());
786                }
787                cpu_call(STEPPING, || step_cpu(&mut **state, schedule, end))
788            }
789            #[cfg(not(target_arch = "wasm32"))]
790            Engine::Gpu { state, ctx } => advance_gpu(&mut **state, ctx, schedule, end),
791        };
792        self.record(advanced)
793    }
794
795    /// Returns the fault every call returns once an earlier call faulted.
796    fn check_earlier_fault(&self) -> Result<(), Fault> {
797        let Some(earlier) = &self.earlier_fault else {
798            return Ok(());
799        };
800        #[cfg(not(target_arch = "wasm32"))]
801        if let Engine::Gpu { ctx, .. } = &self.engine
802            && ctx.is_lost()
803        {
804            return Err(Fault::device_lost(STEPPING));
805        }
806        Err(Fault::refused(
807            STEPPING,
808            format!("an earlier call faulted and the simulation has to be rebuilt ({earlier})"),
809        ))
810    }
811
812    /// Returns `result`, and records its fault for [`Self::check_earlier_fault`] to report on every later call.
813    fn record<T>(&mut self, result: Result<T, Fault>) -> Result<T, Fault> {
814        if let Err(fault) = &result
815            && self.earlier_fault.is_none()
816        {
817            self.earlier_fault = Some(fault.to_string());
818        }
819        result
820    }
821}
822
823/// Runs a GPU state to `end_tick` as [`Simulation::run_sampled`] does.
824#[cfg(not(target_arch = "wasm32"))]
825fn run_sampled_gpu<B>(
826    state: &mut dyn GpuSimState,
827    ctx: &GpuContext,
828    schedule: &Schedule,
829    end_tick: u64,
830    interval: u64,
831    mut on_sample: impl FnMut(&StatSample) -> ControlFlow<B>,
832) -> Result<ControlFlow<B>, Fault> {
833    let first = sample_gpu(state, ctx)?;
834    if let ControlFlow::Break(stop) = on_sample(&first) {
835        return Ok(ControlFlow::Break(stop));
836    }
837    while state.tick() < end_tick {
838        let next = next_sample_tick(state.tick(), interval, end_tick);
839        advance_gpu(state, ctx, schedule, next)?;
840        let sample = sample_gpu(state, ctx)?;
841        if let ControlFlow::Break(stop) = on_sample(&sample) {
842            return Ok(ControlFlow::Break(stop));
843        }
844    }
845    Ok(ControlFlow::Continue(()))
846}
847
848/// Runs `f` inside the fault scopes. Once the device is lost, the loss is returned instead of the errors it caused.
849#[cfg(not(target_arch = "wasm32"))]
850fn gpu_call<T>(ctx: &GpuContext, during: &'static str, f: impl FnOnce() -> Result<T, Fault>) -> Result<T, Fault> {
851    match catching_on(ctx, during, f) {
852        Ok(Ok(value)) => Ok(value),
853        Ok(Err(fault)) | Err(fault) => Err(if ctx.is_lost() {
854            Fault::device_lost(during)
855        } else {
856            fault
857        }),
858    }
859}
860
861/// Steps a GPU state to `end` under [`Fire::AfterStep`], at most [`crate::gpu::MAX_STEPS_PER_SUBMISSION`] steps per
862/// command buffer, and waits for the device.
863#[cfg(not(target_arch = "wasm32"))]
864fn advance_gpu(state: &mut dyn GpuSimState, ctx: &GpuContext, schedule: &Schedule, end: u64) -> Result<(), Fault> {
865    let count = end.saturating_sub(state.tick());
866    if count == 0 {
867        return Ok(());
868    }
869    gpu_call(ctx, STEPPING, || {
870        let refused = stepping::run_steps_acting(state, ctx, count, schedule, Fire::AfterStep)?;
871        refusal(&refused, STEPPING)
872    })
873}
874
875/// Returns the tick of the sample after `tick`: the next multiple of `interval`, or `end_tick` when that comes first.
876fn next_sample_tick(tick: u64, interval: u64, end_tick: u64) -> u64 {
877    (tick / interval)
878        .saturating_add(1)
879        .saturating_mul(interval)
880        .min(end_tick)
881}
882
883/// Steps `state` to `end`, testing the schedule once per stretch between two due ticks.
884fn step_cpu(state: &mut dyn SimState, schedule: &Schedule, end: u64) -> Result<(), Fault> {
885    if schedule.is_empty() {
886        for _ in state.tick()..end {
887            state.step();
888        }
889        return Ok(());
890    }
891    while state.tick() < end {
892        let stop = schedule.next_due_after(state.tick()).map_or(end, |due| due.min(end));
893        for _ in state.tick()..stop {
894            state.step();
895        }
896        refusal(&schedule.run_due(state), STEPPING)?;
897    }
898    Ok(())
899}
900
901/// Returns the stats of a CPU state's current tick, prepared as a publish prepares them.
902fn sample_cpu(state: &mut dyn SimState) -> StatSample {
903    state.prepare_view();
904    StatSample {
905        tick: state.tick(),
906        entries: state.stats(),
907    }
908}
909
910/// Returns the stats of a GPU state's current tick, blocking on a stats-only readback.
911#[cfg(not(target_arch = "wasm32"))]
912fn sample_gpu(state: &mut dyn GpuSimState, ctx: &GpuContext) -> Result<StatSample, Fault> {
913    let entries = gpu_call(ctx, STEPPING, || stepping::sample_stats(state, ctx))?;
914    // Reports a fault the sample raised outside the scopes. No later call might wait.
915    stepping::wait(ctx)?;
916    Ok(StatSample {
917        tick: state.tick(),
918        entries,
919    })
920}
921
922/// Writes every view of a prepared CPU state to `writer`.
923fn write_views(state: &dyn SimState, mut writer: &mut dyn Write) -> io::Result<()> {
924    if let Some(grid) = state.grid_view() {
925        state_export::write_grid(&mut writer, grid.width, grid.height, grid.cells)?;
926    }
927    let points = state.point_view();
928    if let Some(points) = &points {
929        state_export::write_points(&mut writer, points.pos_x, points.pos_y, points.color)?;
930    }
931    if let (Some(points), Some(edges)) = (&points, state.edge_view()) {
932        let rows = state_export::point_rows(points.pos_x, points.pos_y);
933        state_export::write_edges(&mut writer, edges.src, edges.dst, edges.color.unwrap_or(&[]), &rows)?;
934    }
935    Ok(())
936}
937
938/// Views of one model, prepared together and borrowed from its state.
939#[derive(Debug)]
940#[non_exhaustive]
941pub struct SimulationViews<'a> {
942    /// Grid of a CPU grid model, or the field layer of a CPU agent model that has one.
943    pub grid: Option<GridView<'a>>,
944    /// Agents or nodes of a CPU agent or network model.
945    pub points: Option<PointView<'a>>,
946    /// Edges of a CPU network model.
947    pub edges: Option<EdgeView<'a>>,
948}
949
950/// Reason [`Simulation::write_state`] wrote nothing, or stopped part way.
951#[derive(Debug)]
952#[non_exhaustive]
953pub enum ExportError {
954    /// The model is a GPU model.
955    GpuState,
956    /// The writer failed.
957    Io(io::Error),
958    /// The model faulted while preparing its views, or an earlier call faulted.
959    Fault(Fault),
960}
961
962impl fmt::Display for ExportError {
963    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
964        match self {
965            Self::GpuState => f.write_str("a GPU model's state stays on the device and cannot be exported"),
966            Self::Io(_) => f.write_str("cannot write the state"),
967            Self::Fault(_) => f.write_str("the model faulted while preparing its views"),
968        }
969    }
970}
971
972impl std::error::Error for ExportError {
973    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
974        match self {
975            Self::GpuState => None,
976            Self::Io(error) => Some(error),
977            Self::Fault(fault) => Some(fault),
978        }
979    }
980}
981
982impl From<io::Error> for ExportError {
983    fn from(error: io::Error) -> Self {
984        Self::Io(error)
985    }
986}
987
988impl From<Fault> for ExportError {
989    fn from(fault: Fault) -> Self {
990        Self::Fault(fault)
991    }
992}
993
994/// One sample of a model's stats.
995#[derive(Debug, Clone)]
996pub struct StatSample {
997    tick: u64,
998    entries: Vec<StatEntry>,
999}
1000
1001impl StatSample {
1002    /// Tick the sample was taken at.
1003    pub fn tick(&self) -> u64 {
1004        self.tick
1005    }
1006
1007    /// Every stat the model reports, in the order it reports them.
1008    pub fn entries(&self) -> &[StatEntry] {
1009        &self.entries
1010    }
1011
1012    /// Returns the value of stat `label`, `None` for a label the model does not report.
1013    pub fn get(&self, label: &str) -> Option<&StatValue> {
1014        self.entries
1015            .iter()
1016            .find(|entry| entry.label == label)
1017            .map(|entry| &entry.value)
1018    }
1019
1020    /// Returns [`StatValue::scalar`] of stat `label`: a scalar's value, a vector's magnitude, a histogram's total
1021    /// count. `None` for a label the model does not report.
1022    pub fn scalar(&self, label: &str) -> Option<f64> {
1023        self.get(label).map(StatValue::scalar)
1024    }
1025}