mingling/src/setups/exit_code.rs


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

use std::marker::PhantomData;

use mingling_core::{
    ProgramCollect,
    hook::{ProgramControlUnit, ProgramControls, ProgramHook},
    setup::ProgramSetup,
    this,
};

use crate::res::ResExitCode;

/// Provides the ability to control the program's exit code, which is returned when the program ends.
///
/// - Use `mingling::update_exit_code` to update the exit code.
/// - Use `mingling::current_exit_code` to query the current exit code.
pub struct ExitCodeSetup<C> {
    _collect: PhantomData<C>,
}

impl<C> Default for ExitCodeSetup<C>
where
    C: ProgramCollect<Enum = C> + 'static,
{
    fn default() -> Self {
        Self {
            _collect: PhantomData,
        }
    }
}

impl<C> ProgramSetup<C> for ExitCodeSetup<C>
where
    C: ProgramCollect<Enum = C> + 'static,
{
    fn setup(self, program: &mut crate::Program<C>) {
        // Insert resource
        program.with_resource(ResExitCode { exit_code: 0 });

        // Insert hook to override exit code before program ends
        program.with_hook(ProgramHook::empty().on_finish(|_| {
            let this = this::<C>().res_or_default::<ResExitCode>();
            let ec = this.exit_code;

            // Only override when ResExitCode has been modified
            if ec != 0 {
                ProgramControlUnit::OverrideExitCode(this.exit_code).into()
            } else {
                ProgramControls::Empty
            }
        }));
    }
}