1 | //! Basic operations useful for building a testsuite |
2 | |
3 | use color_eyre::eyre::Result; |
4 | use crossbeam_channel::unbounded; |
5 | use crossbeam_channel::Receiver; |
6 | use crossbeam_channel::Sender; |
7 | use regex::bytes::RegexSet; |
8 | use std::num::NonZeroUsize; |
9 | use std::path::Component; |
10 | use std::path::Path; |
11 | use std::path::Prefix; |
12 | use std::sync::OnceLock; |
13 | use std::thread; |
14 | |
15 | /// Remove the common prefix of this path and the `root_dir`. |
16 | pub(crate) fn strip_path_prefix<'a>( |
17 | path: &'a Path, |
18 | prefix: &Path, |
19 | ) -> impl Iterator<Item = Component<'a>> { |
20 | let mut components: Components<'_> = path.components(); |
21 | for c: Component<'_> in prefix.components() { |
22 | // Windows has some funky paths. This is probably wrong, but works well in practice. |
23 | let deverbatimize: impl Fn(Component<'_>) -> … = |c: Component<'_>| match c { |
24 | Component::Prefix(prefix: PrefixComponent<'_>) => Err(match prefix.kind() { |
25 | Prefix::VerbatimUNC(a: &OsStr, b: &OsStr) => Prefix::UNC(a, b), |
26 | Prefix::VerbatimDisk(d: u8) => Prefix::Disk(d), |
27 | other: Prefix<'_> => other, |
28 | }), |
29 | c: Component<'_> => Ok(c), |
30 | }; |
31 | let c2: Option> = components.next(); |
32 | if Some(deverbatimize(c)) == c2.map(deverbatimize) { |
33 | continue; |
34 | } |
35 | return c2.into_iter().chain(components); |
36 | } |
37 | None.into_iter().chain(components) |
38 | } |
39 | |
40 | impl CrateType { |
41 | /// Heuristic: |
42 | /// * [`CrateType::ProcMacro`] if the file contains a [proc macro attribute] |
43 | /// * [`CrateType::Test`] if the file contains `#[test]` |
44 | /// * [`CrateType::Bin`] if the file contains `fn main()` or `#[start]` |
45 | /// * otherwise [`CrateType::Lib`] |
46 | /// |
47 | /// [proc macro attribute]: https://doc.rust-lang.org/reference/procedural-macros.html |
48 | pub fn from_file_contents(file_contents: &[u8]) -> CrateType { |
49 | static RE: OnceLock<RegexSet> = OnceLock::new(); |
50 | let re = RE.get_or_init(|| { |
51 | RegexSet::new([ |
52 | r"#\[proc_macro(_derive|_attribute)?[\](]" , |
53 | r"#\[test\]" , |
54 | r"fn main()|#\[start\]" , |
55 | ]) |
56 | .unwrap() |
57 | }); |
58 | |
59 | match re.matches(file_contents).iter().next() { |
60 | Some(0) => CrateType::ProcMacro, |
61 | Some(1) => CrateType::Test, |
62 | Some(2) => CrateType::Bin, |
63 | _ => CrateType::Lib, |
64 | } |
65 | } |
66 | } |
67 | |
68 | /// The kind of crate we're building here. Corresponds to `--crate-type` flags of rustc |
69 | pub enum CrateType { |
70 | /// A proc macro |
71 | ProcMacro, |
72 | /// A file containing unit tests |
73 | Test, |
74 | /// A binary file containing a main function or start function |
75 | Bin, |
76 | /// A library crate |
77 | Lib, |
78 | } |
79 | |
80 | /// A generic multithreaded runner that has a thread for producing work, |
81 | /// a thread for collecting work, and `num_threads` threads for doing the work. |
82 | pub fn run_and_collect<const N: usize, SUBMISSION: Send, RESULT: Send>( |
83 | num_threads: NonZeroUsize, |
84 | submitter: impl FnOnce([Sender<SUBMISSION>; N]) + Send, |
85 | runner: impl Sync + Fn(&[Receiver<SUBMISSION>; N], Sender<RESULT>) -> Result<()>, |
86 | collector: impl FnOnce(Receiver<RESULT>) + Send, |
87 | ) -> Result<()> { |
88 | // A channel for files to process |
89 | let (submit, receive): (Vec<_>, Vec<_>) = std::iter::repeat_with(unbounded).take(N).unzip(); |
90 | let receive = receive[..].try_into().unwrap(); |
91 | let mut submit = submit.into_iter(); |
92 | let submit = std::array::from_fn(|_| submit.next().unwrap()); |
93 | |
94 | thread::scope(|s| { |
95 | // Create a thread that is in charge of walking the directory and submitting jobs. |
96 | // It closes the channel when it is done. |
97 | s.spawn(|| submitter(submit)); |
98 | |
99 | // A channel for the messages emitted by the individual test threads. |
100 | // Used to produce live updates while running the tests. |
101 | let (finished_files_sender, finished_files_recv) = unbounded(); |
102 | |
103 | s.spawn(|| collector(finished_files_recv)); |
104 | |
105 | let mut threads = vec![]; |
106 | |
107 | // Create N worker threads that receive files to test. |
108 | for _ in 0..num_threads.get() { |
109 | let finished_files_sender = finished_files_sender.clone(); |
110 | threads.push(s.spawn(|| runner(receive, finished_files_sender))); |
111 | } |
112 | |
113 | for thread in threads { |
114 | thread.join().unwrap()?; |
115 | } |
116 | Ok(()) |
117 | }) |
118 | } |
119 | |