1use proc_macro2::{Ident, TokenStream};
2use quote::{format_ident, quote};
3use core::panic;
4use std::collections::HashMap;
5use std::path::{Path, PathBuf};
6use syn::parse::{Parse, ParseStream};
7use syn::{Expr, Result, Token};
8
9struct MacroInput {
10 prefix: Ident,
11 tag: String,
12 operation: Option<Ident>,
13 processed: Expr,
14}
15
16impl Parse for MacroInput {
17 fn parse(input: ParseStream) -> Result<Self> {
18 let prefix = input.parse::<Ident>()?;
19 input.parse::<Token![,]>()?;
20 let tag = input.parse::<Ident>()?.to_string();
21 input.parse::<Token![,]>()?;
22 let operation = input.parse::<Ident>().ok();
23 if operation.is_some() {
24 let _ = input.parse::<Token![,]>();
25 }
26 let processed = input.parse::<Expr>()?;
27 Ok(MacroInput {
28 prefix,
29 tag,
30 operation,
31 processed,
32 })
33 }
34}
35
36#[proc_macro]
37pub fn generate_code_tests(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
38 let macro_input = syn::parse_macro_input!(input as MacroInput);
39
40 let tests = walkdir::WalkDir::new("example-code")
41 .into_iter()
42 .map(|entry| entry.expect("entry to exist"))
43 .filter(|entry| entry.file_type().is_file())
44 .filter(|entry| !entry.file_name().to_string_lossy().starts_with('_'))
45 .map(|entry| entry.path().to_path_buf())
46 .map(|path| (module_hierarchy(&path), generate_test(path, ¯o_input)))
47 .collect::<Vec<(Vec<String>, TokenStream)>>();
48
49 let mut modules = HashMap::<String, Vec<TokenStream>>::new();
50 for (hierarchy, test) in tests {
51 let module_name = hierarchy.join("_");
53
54 if let Some(tests) = modules.get_mut(&module_name) {
55 tests.push(test);
56 continue;
57 } else {
58 modules.insert(module_name, vec![test]);
59 }
60 }
61
62 modules
63 .into_iter()
64 .map(|(module_name, tests)| {
65 if module_name.is_empty() {
66 quote! {
67 #(#tests)*
68 }
69 } else {
70 let module_name = format_ident!("{}", module_name);
71 quote! {
72 mod #module_name {
73 use super::*;
74 #(#tests)*
75 }
76 }
77 }
78 })
79 .collect::<TokenStream>()
80 .into()
81}
82
83fn module_hierarchy(path: &Path) -> Vec<String> {
84 let relative = &path.strip_prefix("example-code").unwrap();
85
86 let mut hierarchy = relative
87 .components()
88 .map(|c| c.as_os_str().to_str().unwrap().to_string())
89 .collect::<Vec<_>>();
90 hierarchy.pop();
91 hierarchy
92}
93
94fn generate_test(path: PathBuf, macro_input: &MacroInput) -> TokenStream {
95 let MacroInput {
96 prefix,
97 tag,
98 operation,
99 processed,
100 } = macro_input;
101 let name = path.file_stem().unwrap().to_str().unwrap();
102
103 let Ok(test_file) = std::fs::read_to_string(&path) else {
104 panic!("Test file not found!")
105 };
106
107 let Ok(expected_struct) = expected_result(&test_file, tag).map_err(|e| match e {
108 TestError::TagNotClosed => panic!("Tag not closed in test file!"),
109 TestError::InvalidRustCode => panic!("Invalid Rust code in test file '{}'!", path.to_string_lossy()),
110 TestError::TagNotFound => e,
111 }) else {
112 return quote! {};
113 };
114
115 let code = test_file;
116
117 let test_name = format_ident!("{prefix}_{name}");
118 if let Some(operation) = operation {
119 quote! {
120 #[test]
121 fn #test_name() {
122 let code = #code;
123 let expected_struct = #expected_struct.#operation();
124 let actual_struct = #processed.#operation();
125 assert_eq!(expected_struct, actual_struct);
126 }
127 }
128 } else {
129 quote! {
130 #[test]
131 fn #test_name() {
132 let code = #code;
133 let expected_struct = #expected_struct;
134 let actual_struct = #processed;
135 assert_eq!(expected_struct, actual_struct);
136 }
137 }
138 }
139}
140
141enum TestError {
142 TagNotFound,
143 TagNotClosed,
144 InvalidRustCode,
145}
146
147fn expected_result(test_file: &str, tag: &str) -> std::result::Result<TokenStream, TestError> {
148 let prefix = "//@";
149 let mut lines = vec![];
150 let mut iter = test_file.lines();
151
152 loop {
153 match iter.next() {
154 Some(line)
155 if line
156 .strip_prefix(prefix)
157 .is_some_and(|s| s.trim().starts_with(tag)) =>
158 {
159 break
160 }
161 Some(_) => continue,
162 None => return Err(TestError::TagNotFound),
163 }
164 }
165
166 loop {
167 match iter.next() {
168 Some(line) if line.starts_with("//@end") => break,
169 Some(line) => lines.push(line.strip_prefix("//").expect(&format!("line to start with '//': {line} in \n{test_file}"))),
170 None => return Err(TestError::TagNotClosed),
171 }
172 }
173
174 let code = lines.join("\n");
175 code.parse::<TokenStream>()
176 .map_err(|_| TestError::InvalidRustCode)
177}