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
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
//! This crate allows cleanly exiting a program using a custom exit code,
//! without the drawbacks of [`process::exit`]. Destructors will be called as
//! usual, and the stack will be completely unwound.
//!
//! It is always required to attach [`#[main]`][attribute] to the main
//! function. Then, [`with_code`] can be called from almost anywhere in the
//! program. Restrictions are noted in the documentation for that function.
//!
//! # Implementation
//!
//! Internally, this crate uses panicking to unwind the stack. Thus, if
//! panicking were set to "abort" instead of the default "unwind", setting the
//! exit status would not work correctly. This crate will cause a compile error
//! in that case, to avoid silent incorrect behavior. Further information can
//! be found in the [Rustc Development Guide][panic_runtime].
//!
//! Additionally, the program will not exit if [`with_code`] is called from a
//! spawned thread, unless panics are propagated from that thread. However,
//! propagating panics is usually recommended regardless.
//!
//! # Examples
//!
//! ```
//! use std::env;
//!
//! fn read_args() {
//!     if env::args_os().nth(1).is_some() {
//!         eprintln!("too many arguments");
//!         quit::with_code(1);
//!     }
//! }
//!
//! #[quit::main]
//! fn main() {
//!     read_args();
//! }
//! ```
//!
//! [attribute]: main
//! [panic_runtime]: https://rustc-dev-guide.rust-lang.org/panic-implementation.html#step-2-the-panic-runtime

#![forbid(unsafe_code)]
#![warn(unused_results)]

#[cfg(not(panic = "unwind"))]
compile_error!(
    r#"Quit requires unwinding panics:
https://docs.rs/quit/latest/quit/#implementation"#
);

use std::panic;
use std::panic::UnwindSafe;
use std::process;
use std::process::Termination;
use std::sync::atomic;
use std::sync::atomic::AtomicBool;

#[cfg(all(feature = "__unstable_tests", not(quit_docs_rs)))]
#[doc(hidden)]
pub mod tests_common;

/// Modifies the main function to exit with the code passed to [`with_code`].
///
/// This attribute should always be attached to the main function. Otherwise,
/// the exit code of the program may be incorrect.
///
/// # Examples
///
/// ```
/// #[quit::main]
/// fn main() {}
/// ```
pub use quit_macros::main;

struct ExitCode(process::ExitCode);

enum ResultInner<T>
where
    T: Termination,
{
    Result(T),
    ExitCode(process::ExitCode),
}

#[doc(hidden)]
pub struct __Result<T>(ResultInner<T>)
where
    T: Termination;

impl<T> Termination for __Result<T>
where
    T: Termination,
{
    #[inline]
    fn report(self) -> process::ExitCode {
        use ResultInner as Inner;

        match self.0 {
            Inner::Result(result) => result.report(),
            Inner::ExitCode(exit_code) => exit_code,
        }
    }
}

#[doc(hidden)]
#[inline]
pub fn __run<F, R>(main_fn: F) -> __Result<R>
where
    F: FnOnce() -> R + UnwindSafe,
    R: Termination,
{
    panic::catch_unwind(main_fn)
        .map(|x| __Result(ResultInner::Result(x)))
        .unwrap_or_else(|payload| {
            if let Some(&ExitCode(exit_code)) = payload.downcast_ref() {
                __Result(ResultInner::ExitCode(exit_code))
            } else {
                panic::resume_unwind(payload);
            }
        })
}

#[doc(hidden)]
pub static __ATTACHED: AtomicBool = AtomicBool::new(false);

#[doc(hidden)]
#[macro_export]
macro_rules! __main {
    (
        ( $($signature_token:tt)+ )
        ( $($args_token:tt)* ) -> $return_type:ty { $($body_token:tt)* }
    ) => {
        $($signature_token)+ (
            $($args_token)*
        ) -> $crate::__Result<$return_type> {
            $crate::__ATTACHED.store(
                true,
                ::std::sync::atomic::Ordering::Release,
            );
            $crate::__run(|| { $($body_token)* })
        }
    };
    ( $signature:tt $args:tt $body:tt ) => {
        $crate::__main!($signature $args -> () $body);
    }
}

/// Cleanly exits the program with an exit code.
///
/// Calling this function from within an FFI boundary invokes undefined
/// behavior. Because panics are used internally to unwind the stack, the exit
/// code cannot be passed safely. [`process::exit`] should be used instead in
/// that case.
///
/// This function will not behave as expected unless [`#[main]`][attribute] is
/// attached to the main function. Other implementation notes are mentioned in
/// [the module-level documentation][implementation].
///
/// # Examples
///
/// ```
/// fn exit() -> ! {
///     quit::with_code(1);
/// }
/// ```
///
/// [attribute]: main
/// [implementation]: self#implementation
#[inline]
pub fn with_code<T>(exit_code: T) -> !
where
    T: Into<process::ExitCode>,
{
    assert!(
        __ATTACHED.load(atomic::Ordering::Acquire),
        "`#[quit::main]` has not been attached to `main`",
    );
    panic::resume_unwind(Box::new(ExitCode(exit_code.into())));
}