go_defer/
lib.rs

1// This file is part of Substrate.
2
3// Copyright (C) Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: Apache-2.0
5
6// Licensed under the Apache License, Version 2.0 (the "License");
7// you may not use this file except in compliance with the License.
8// You may obtain a copy of the License at
9//
10// 	http://www.apache.org/licenses/LICENSE-2.0
11//
12// Unless required by applicable law or agreed to in writing, software
13// distributed under the License is distributed on an "AS IS" BASIS,
14// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15// See the License for the specific language governing permissions and
16// limitations under the License.
17
18//! Contains the [`crate::defer!`] macro for *deferring* the execution
19//! of code until the current scope is dropped.
20//! This helps with *always* executing cleanup code.
21
22/// Executes the wrapped closure on drop.
23///
24/// Should be used together with the [`crate::defer!`] macro.
25#[must_use]
26pub struct DeferGuard<F: FnOnce()>(pub Option<F>);
27
28impl<F: FnOnce()> Drop for DeferGuard<F> {
29	fn drop(&mut self) {
30		self.0.take().map(|f| f());
31	}
32}
33
34/// Executes the given code when the current scope is dropped.
35///
36/// Multiple calls to [`crate::defer!`] will execute the passed codes in reverse order.
37/// This also applies to panic stack unwinding.
38///
39/// # Example
40///
41/// ```rust
42/// use sp_core::defer;
43///
44/// let message = std::cell::RefCell::new("".to_string());
45/// {
46/// 	defer!(
47/// 		message.borrow_mut().push_str("world!");
48/// 	);
49/// 	defer!(
50/// 		message.borrow_mut().push_str("Hello ");
51/// 	);
52/// }
53/// assert_eq!(*message.borrow(), "Hello world!");
54/// ```
55#[macro_export]
56macro_rules! defer(
57	( $( $code:tt )* ) => {
58		let _guard = $crate::DeferGuard(Some(|| { $( $code )* }));
59	};
60);
61
62#[cfg(test)]
63mod test {
64	#[test]
65	fn defer_guard_works() {
66		let mut called = false;
67		{
68			defer!(
69				called = true;
70			);
71		}
72		assert!(called, "DeferGuard should have executed the closure");
73	}
74
75	#[test]
76	/// `defer` executes the code in reverse order of being called.
77	fn defer_guard_order_works() {
78		let called = std::cell::RefCell::new(1);
79
80		defer!(
81			assert_eq!(*called.borrow(), 3);
82		);
83		defer!(
84			assert_eq!(*called.borrow(), 2);
85			*called.borrow_mut() = 3;
86		);
87		defer!({
88			assert_eq!(*called.borrow(), 1);
89			*called.borrow_mut() = 2;
90		});
91	}
92
93	#[test]
94	#[allow(unused_braces)]
95	#[allow(clippy::unnecessary_operation)]
96	fn defer_guard_syntax_works() {
97		let called = std::cell::RefCell::new(0);
98		{
99			defer!(*called.borrow_mut() += 1);
100			defer!(*called.borrow_mut() += 1;); // With ;
101			defer!({ *called.borrow_mut() += 1 });
102			defer!({ *called.borrow_mut() += 1 };); // With ;
103		}
104		assert_eq!(*called.borrow(), 4);
105	}
106
107	#[test]
108	/// `defer` executes the code even in case of a panic.
109	fn defer_guard_panic_unwind_works() {
110		use std::panic::{catch_unwind, AssertUnwindSafe};
111		let mut called = false;
112
113		let should_panic = catch_unwind(AssertUnwindSafe(|| {
114			defer!(called = true);
115			panic!();
116		}));
117
118		assert!(should_panic.is_err(), "DeferGuard should have panicked");
119		assert!(called, "DeferGuard should have executed the closure");
120	}
121
122	#[test]
123	/// `defer` executes the code even in case another `defer` panics.
124	fn defer_guard_defer_panics_unwind_works() {
125		use std::panic::{catch_unwind, AssertUnwindSafe};
126		let counter = std::cell::RefCell::new(0);
127
128		let should_panic = catch_unwind(AssertUnwindSafe(|| {
129			defer!(*counter.borrow_mut() += 1);
130			defer!(
131				*counter.borrow_mut() += 1;
132				panic!();
133			);
134			defer!(*counter.borrow_mut() += 1);
135		}));
136
137		assert!(should_panic.is_err(), "DeferGuard should have panicked");
138		assert_eq!(*counter.borrow(), 3, "DeferGuard should have executed the closure");
139	}
140}