sequential_macro/
lib.rs

1//! **Do not use this**
2#![warn(clippy::pedantic)]
3extern crate proc_macro;
4
5use quote::quote;
6use std::sync::atomic::{AtomicBool, Ordering};
7
8/// Used for the attribute macros such that we only define [`__TestState`] and [`__PAIR`] once.
9static SET: AtomicBool = AtomicBool::new(false);
10/// Annotates tests which must run sequentially.
11#[proc_macro_attribute]
12pub fn sequential(
13    _attr: proc_macro::TokenStream,
14    item: proc_macro::TokenStream,
15) -> proc_macro::TokenStream {
16    inner(
17        item,
18        quote! {
19            let _ = __PAIR.1.wait_while(__PAIR.0.lock().expect("sequential-test error"), |pending|
20                match pending {
21                    __TestState::Parallel(0) => {
22                        *pending = __TestState::Sequential;
23                        false
24                    },
25                    _ => true
26                }
27            ).expect("sequential-test error");
28        },
29        quote! {
30            *__PAIR.0.lock().expect("sequential-test error") = __TestState::Parallel(0);
31        },
32    )
33}
34/// Annotates tests which may run in parallel.
35#[proc_macro_attribute]
36pub fn parallel(
37    _attr: proc_macro::TokenStream,
38    item: proc_macro::TokenStream,
39) -> proc_macro::TokenStream {
40    inner(
41        item,
42        quote! {
43            let _ = __PAIR.1.wait_while(__PAIR.0.lock().expect("sequential-test error"), |pending|
44                match pending {
45                    __TestState::Sequential => true,
46                    __TestState::Parallel(ref mut x) => {
47                        *x += 1;
48                        false
49                    }
50                }
51            ).expect("sequential-test error");
52        },
53        quote! {
54            match *__PAIR.0.lock().expect("sequential-test error") {
55                __TestState::Sequential => unreachable!("sequential-test error"),
56                __TestState::Parallel(ref mut x) => {
57                    *x -= 1;
58                }
59            }
60        },
61    )
62}
63fn inner(
64    item: proc_macro::TokenStream,
65    prefix: proc_macro2::TokenStream,
66    suffix: proc_macro2::TokenStream,
67) -> proc_macro::TokenStream {
68    // We get whether this was the first macro invocation and set first macro invocation to false.
69    let ret = SET.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst);
70    let mut iter = item.into_iter().peekable();
71    let signature = proc_macro2::TokenStream::from(
72        std::iter::from_fn(|| {
73            iter.next_if(|x| match x {
74                proc_macro::TokenTree::Group(group) => {
75                    !matches!(group.delimiter(), proc_macro::Delimiter::Brace)
76                }
77                _ => true,
78            })
79        })
80        .collect::<proc_macro::TokenStream>(),
81    );
82    let block = proc_macro2::TokenStream::from(iter.collect::<proc_macro::TokenStream>());
83    let item = quote! {
84        #signature {
85            #prefix
86            let res = std::panic::catch_unwind(|| #block );
87            #suffix
88            __PAIR.1.notify_all();
89            if let Err(err) = res {
90                std::panic::resume_unwind(err);
91            }
92        }
93    };
94    // If this was the first macro invocation define the mutex and condvar used for locking.
95    let rtn = if ret.is_ok() {
96        quote! {
97            /// The test state stored in the mutex in the pair used for ordering tests.
98            enum __TestState {
99                Sequential,
100                Parallel(u64)
101            }
102            /// The mutex and condvar pair used for ordering tests.
103            static __PAIR: (std::sync::Mutex<__TestState>, std::sync::Condvar) = (
104                std::sync::Mutex::new(__TestState::Parallel(0)), std::sync::Condvar::new()
105            );
106            #item
107        }
108    } else {
109        item
110    };
111    proc_macro::TokenStream::from(rtn)
112}