pfx_mod/
set.rs

1use wasmtime::component::{Instance, TypedFunc};
2use wasmtime::{Store, Trap, WasmBacktrace};
3
4use crate::ctx::ModCtx;
5use crate::depth::Site;
6use crate::error::{Cause, Frame, ModError};
7use crate::fingerprint::Fingerprint;
8use crate::host::{Host, Package};
9use crate::limits::Fuel;
10use crate::manifest::Manifest;
11
12pub type Bind<T, B> = fn(&mut Store<ModCtx<T>>, &Instance) -> wasmtime::Result<B>;
13
14#[derive(Clone, Copy, Debug, PartialEq, Eq)]
15pub struct Called<R> {
16    pub value: R,
17    pub fuel: u64,
18}
19
20pub const SAVE: &str = "save";
21pub const RESTORE: &str = "restore";
22
23#[derive(Clone, Debug, PartialEq, Eq)]
24pub enum Swapped {
25    Kept { bytes: usize },
26    Fresh(Fresh),
27}
28
29#[derive(Clone, Debug, PartialEq, Eq)]
30pub enum Fresh {
31    Stopped,
32    NoSave,
33    NoRestore,
34    Mismatch { export: &'static str },
35    Save(Box<ModError>),
36}
37
38impl std::fmt::Display for Swapped {
39    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40        match self {
41            Self::Kept { bytes } => {
42                write!(f, "kept its state ({bytes} bytes through save and restore)")
43            }
44            Self::Fresh(Fresh::Stopped) => write!(
45                f,
46                "restarted clean: the old module was stopped, so there was no state to save"
47            ),
48            Self::Fresh(Fresh::NoSave) => {
49                write!(f, "restarted clean: the old module exports no `save`")
50            }
51            Self::Fresh(Fresh::NoRestore) => {
52                write!(f, "restarted clean: the new module exports no `restore`")
53            }
54            Self::Fresh(Fresh::Mismatch { export }) => write!(
55                f,
56                "restarted clean: `{export}` has the wrong type; save is func() -> list<u8> and restore is func(state: list<u8>)"
57            ),
58            Self::Fresh(Fresh::Save(error)) => write!(f, "restarted clean: save failed, {error}"),
59        }
60    }
61}
62
63struct Running<T: 'static, B> {
64    store: Store<ModCtx<T>>,
65    instance: Instance,
66    bindings: B,
67}
68
69pub struct Mod<T: 'static, B> {
70    package: Package<T>,
71    bind: Bind<T, B>,
72    fuel: Fuel,
73    running: Option<Running<T, B>>,
74    stopped: Option<(ModError, ModCtx<T>)>,
75}
76
77impl<T: 'static, B> Mod<T, B> {
78    pub fn start(host: &Host<T>, package: Package<T>, bind: Bind<T, B>, game: T) -> Self {
79        let ctx = ModCtx::new(game, host.limits());
80        let mut this = Self {
81            package,
82            bind,
83            fuel: host.limits().fuel,
84            running: None,
85            stopped: None,
86        };
87        this.launch(host, ctx);
88        this
89    }
90
91    fn launch(&mut self, host: &Host<T>, ctx: ModCtx<T>) {
92        let mut store = Store::new(host.engine(), ctx);
93        store.limiter(|cx| &mut cx.limiter);
94        let budget = self.fuel.budget();
95        if let Err(e) = store.set_fuel(budget) {
96            let error = self.broke(&e);
97            self.stopped = Some((error, store.into_data()));
98            return;
99        }
100        let made = self
101            .package
102            .pre
103            .instantiate(&mut store)
104            .and_then(|instance| Ok((instance, (self.bind)(&mut store, &instance)?)));
105        match made {
106            Ok((instance, bindings)) => {
107                self.running = Some(Running {
108                    store,
109                    instance,
110                    bindings,
111                });
112                self.stopped = None;
113            }
114            Err(e) => {
115                let error = classify(&self.package, e, &store, budget);
116                self.stopped = Some((error, store.into_data()));
117            }
118        }
119    }
120
121    fn error(&self, cause: Cause, fuel: u64, at: Option<Frame>) -> ModError {
122        ModError {
123            id: self.package.manifest.id.clone(),
124            cause,
125            fuel,
126            at,
127        }
128    }
129
130    fn broke(&self, e: &wasmtime::Error) -> ModError {
131        self.error(Cause::Error(format!("{e:#}")), 0, None)
132    }
133
134    pub fn call<R>(
135        &mut self,
136        f: impl FnOnce(&B, &mut Store<ModCtx<T>>) -> wasmtime::Result<R>,
137    ) -> Result<Called<R>, ModError> {
138        let Some(running) = self.running.as_mut() else {
139            return Err(self.error(Cause::Disabled, 0, None));
140        };
141        let store = &mut running.store;
142        if let Fuel::PerCall(budget) = self.fuel
143            && let Err(e) = store.set_fuel(budget)
144        {
145            let error = self.broke(&e);
146            return Err(self.stop(error));
147        }
148        let before = store.get_fuel().unwrap_or(0);
149        store.data_mut().limiter.hit = None;
150        match f(&running.bindings, store) {
151            Ok(value) => Ok(Called {
152                value,
153                fuel: before.saturating_sub(store.get_fuel().unwrap_or(0)),
154            }),
155            Err(e) => {
156                let error = classify(&self.package, e, &running.store, before);
157                Err(self.stop(error))
158            }
159        }
160    }
161
162    fn stop(&mut self, error: ModError) -> ModError {
163        if let Some(running) = self.running.take() {
164            self.stopped = Some((error.clone(), running.store.into_data()));
165        }
166        error
167    }
168
169    pub fn tick(&mut self) {
170        let fuel = self.fuel;
171        if let Some(running) = self.running.as_mut() {
172            running.store.data_mut().log.tick();
173            if let Fuel::PerTick(budget) = fuel
174                && let Err(e) = running.store.set_fuel(budget)
175            {
176                let error = self.broke(&e);
177                self.stop(error);
178            }
179        }
180    }
181
182    pub fn swap(&mut self, host: &Host<T>, package: Package<T>) -> Result<(), ModError> {
183        if package.manifest.id != self.package.manifest.id {
184            return Err(self.error(
185                Cause::Error(format!(
186                    "a swap must keep the id, but the new module has id `{}`",
187                    package.manifest.id
188                )),
189                0,
190                None,
191            ));
192        }
193        let ctx = match (self.running.take(), self.stopped.take()) {
194            (Some(running), _) => running.store.into_data(),
195            (None, Some((_, ctx))) => ctx,
196            (None, None) => return Err(self.error(Cause::Disabled, 0, None)),
197        };
198        self.package = package;
199        self.launch(host, ctx);
200        match &self.stopped {
201            Some((error, _)) => Err(error.clone()),
202            None => Ok(()),
203        }
204    }
205
206    pub fn reload(&mut self, host: &Host<T>, package: Package<T>) -> Result<Swapped, ModError> {
207        if package.manifest.id != self.package.manifest.id {
208            return self
209                .swap(host, package)
210                .map(|()| Swapped::Fresh(Fresh::Stopped));
211        }
212        let saved = self.save();
213        self.swap(host, package)?;
214        let state = match saved {
215            Ok(state) => state,
216            Err(fresh) => return Ok(Swapped::Fresh(fresh)),
217        };
218        let budget = self.fuel.budget();
219        let Some(running) = self.running.as_mut() else {
220            return Ok(Swapped::Fresh(Fresh::Stopped));
221        };
222        let restore = match export::<T, (Vec<u8>,), ()>(running, RESTORE) {
223            Ok(Some(restore)) => restore,
224            Ok(None) => return Ok(Swapped::Fresh(Fresh::NoRestore)),
225            Err(fresh) => return Ok(Swapped::Fresh(fresh)),
226        };
227        let bytes = state.len();
228        let store = &mut running.store;
229        let result = store
230            .set_fuel(budget)
231            .and_then(|()| restore.call(&mut *store, (state,)));
232        let refill = store.set_fuel(budget);
233        match result.and(refill) {
234            Ok(()) => Ok(Swapped::Kept { bytes }),
235            Err(e) => {
236                let error = classify(&self.package, e, &running.store, budget);
237                Err(self.stop(error))
238            }
239        }
240    }
241
242    fn save(&mut self) -> Result<Vec<u8>, Fresh> {
243        let budget = self.fuel.budget();
244        let Some(running) = self.running.as_mut() else {
245            return Err(Fresh::Stopped);
246        };
247        let save = match export::<T, (), (Vec<u8>,)>(running, SAVE)? {
248            Some(save) => save,
249            None => return Err(Fresh::NoSave),
250        };
251        let store = &mut running.store;
252        match store
253            .set_fuel(budget)
254            .and_then(|()| save.call(&mut *store, ()))
255        {
256            Ok((state,)) => Ok(state),
257            Err(e) => {
258                let error = classify(&self.package, e, &running.store, budget);
259                Err(Fresh::Save(Box::new(self.stop(error))))
260            }
261        }
262    }
263
264    pub fn id(&self) -> &str {
265        &self.package.manifest.id
266    }
267
268    pub fn manifest(&self) -> &Manifest {
269        &self.package.manifest
270    }
271
272    pub fn fingerprint(&self) -> Fingerprint {
273        self.package.fingerprint
274    }
275
276    pub fn running(&self) -> bool {
277        self.running.is_some()
278    }
279
280    pub fn stopped(&self) -> Option<&ModError> {
281        self.stopped.as_ref().map(|(error, _)| error)
282    }
283
284    pub fn fuel_left(&self) -> Option<u64> {
285        self.running
286            .as_ref()
287            .and_then(|running| running.store.get_fuel().ok())
288    }
289
290    fn ctx(&self) -> &ModCtx<T> {
291        match (&self.running, &self.stopped) {
292            (Some(running), _) => running.store.data(),
293            (None, Some((_, ctx))) => ctx,
294            (None, None) => unreachable!("a mod is always running or stopped"),
295        }
296    }
297
298    fn ctx_mut(&mut self) -> &mut ModCtx<T> {
299        match (&mut self.running, &mut self.stopped) {
300            (Some(running), _) => running.store.data_mut(),
301            (None, Some((_, ctx))) => ctx,
302            (None, None) => unreachable!("a mod is always running or stopped"),
303        }
304    }
305
306    pub fn game(&self) -> &T {
307        &self.ctx().game
308    }
309
310    pub fn game_mut(&mut self) -> &mut T {
311        &mut self.ctx_mut().game
312    }
313
314    pub fn take_log(&mut self) -> Vec<String> {
315        self.ctx_mut().log.take()
316    }
317
318    pub fn dropped_lines(&self) -> u64 {
319        self.ctx().log.dropped()
320    }
321}
322
323fn export<T: 'static, P, R>(
324    running: &mut Running<T, impl Sized>,
325    name: &'static str,
326) -> Result<Option<TypedFunc<P, R>>, Fresh>
327where
328    P: wasmtime::component::ComponentNamedList + wasmtime::component::Lower,
329    R: wasmtime::component::ComponentNamedList + wasmtime::component::Lift,
330{
331    let Some(func) = running.instance.get_func(&mut running.store, name) else {
332        return Ok(None);
333    };
334    func.typed::<P, R>(&running.store)
335        .map(Some)
336        .map_err(|_| Fresh::Mismatch { export: name })
337}
338
339fn classify<T>(
340    package: &Package<T>,
341    e: wasmtime::Error,
342    store: &Store<ModCtx<T>>,
343    before: u64,
344) -> ModError {
345    let fuel = before.saturating_sub(store.get_fuel().unwrap_or(0));
346    let frames: Vec<Frame> = e
347        .downcast_ref::<WasmBacktrace>()
348        .map(|trace| {
349            trace
350                .frames()
351                .iter()
352                .map(|frame| Frame {
353                    func: frame.func_index(),
354                    offset: frame.module_offset(),
355                })
356                .collect()
357        })
358        .unwrap_or_default();
359    let deep = matches!(e.downcast_ref::<Trap>(), Some(Trap::UnreachableCodeReached))
360        && frames.first().is_some_and(|frame| {
361            frame.offset.is_some_and(|offset| {
362                package.sites.contains(&Site {
363                    func: frame.func,
364                    offset,
365                })
366            })
367        });
368    let at = if deep {
369        frames.get(1).or(frames.first()).copied()
370    } else {
371        frames.first().copied()
372    };
373    let cause = match e.downcast_ref::<Trap>() {
374        _ if deep => Cause::Depth,
375        Some(Trap::OutOfFuel) => Cause::Fuel,
376        Some(Trap::StackOverflow) => Cause::Stack,
377        Some(trap) => Cause::Trap(*trap),
378        None => match store.data().limiter.hit {
379            Some(limit) => Cause::Limit(limit),
380            None => Cause::Error(format!("{e:#}")),
381        },
382    };
383    ModError {
384        id: package.manifest.id.clone(),
385        cause,
386        fuel,
387        at,
388    }
389}
390
391pub struct ModSet<T: 'static, B> {
392    mods: Vec<Mod<T, B>>,
393}
394
395impl<T: 'static, B> ModSet<T, B> {
396    pub fn start(
397        host: &Host<T>,
398        packages: Vec<Package<T>>,
399        bind: Bind<T, B>,
400        mut game: impl FnMut(&Manifest) -> T,
401    ) -> (Self, Vec<ModError>) {
402        let mut packages = packages;
403        packages.sort_by(|a, b| a.manifest.id.cmp(&b.manifest.id));
404        packages.dedup_by(|a, b| a.manifest.id == b.manifest.id);
405        let mods: Vec<Mod<T, B>> = packages
406            .into_iter()
407            .map(|package| {
408                let data = game(&package.manifest);
409                Mod::start(host, package, bind, data)
410            })
411            .collect();
412        let failed = mods.iter().filter_map(|m| m.stopped().cloned()).collect();
413        (Self { mods }, failed)
414    }
415
416    pub fn is_empty(&self) -> bool {
417        self.mods.is_empty()
418    }
419
420    pub fn len(&self) -> usize {
421        self.mods.len()
422    }
423
424    pub fn mods(&self) -> &[Mod<T, B>] {
425        &self.mods
426    }
427
428    pub fn mods_mut(&mut self) -> &mut [Mod<T, B>] {
429        &mut self.mods
430    }
431
432    pub fn get_mut(&mut self, id: &str) -> Option<&mut Mod<T, B>> {
433        self.mods.iter_mut().find(|m| m.id() == id)
434    }
435
436    pub fn tick(&mut self) {
437        for m in &mut self.mods {
438            m.tick();
439        }
440    }
441
442    pub fn fingerprint(&self) -> Fingerprint {
443        Fingerprint::of_set(self.mods.iter().map(|m| (m.id(), m.fingerprint())))
444    }
445}
pfx_mod/set.rs

pfx_mod/
set.rs
