1use proc_macro::TokenStream;
2use proc_macro2::TokenStream as Tokens;
3use quote::{format_ident, quote};
4use syn::parse::{Parse, ParseStream};
5use syn::visit::Visit;
6use syn::{Expr, Path, ReturnType, Token, Type, TypeBareFn, parse_quote};
7
8mod signature;
9
10#[proc_macro]
11pub fn __mock(input: TokenStream) -> TokenStream {
12 expand(syn::parse(input)).into()
13}
14
15#[proc_macro]
16pub fn __check_signature(input: TokenStream) -> TokenStream {
17 expand_check(syn::parse(input)).into()
18}
19
20#[proc_macro]
21pub fn __replace_local(input: TokenStream) -> TokenStream {
22 expand_replacement(syn::parse(input)).into()
23}
24
25struct Replacement {
26 mock: Input,
27 target: Expr,
28}
29
30impl Parse for Replacement {
31 fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
32 let root = input.parse()?;
33 input.parse::<Token![,]>()?;
34 let session = input.parse()?;
35 input.parse::<Token![,]>()?;
36 let source = input.parse()?;
37 input.parse::<Token![,]>()?;
38 let target = input.parse()?;
39 input.parse::<Token![,]>()?;
40 let signature = input.parse()?;
41 Ok(Self {
42 mock: Input {
43 root,
44 session,
45 source,
46 signature,
47 },
48 target,
49 })
50 }
51}
52
53fn expand_replacement(input: syn::Result<Replacement>) -> Tokens {
54 let Replacement { mock, target } = match input {
55 Ok(input) => input,
56 Err(error) => return error.into_compile_error(),
57 };
58 let Input {
59 root,
60 session,
61 source,
62 signature,
63 } = mock;
64 let abi = &signature.abi;
65 let unsafety = &signature.unsafety;
66 let types: Vec<_> = (0..signature.inputs.len())
67 .map(|index| format_ident!("__Arg{index}"))
68 .collect();
69 let names: Vec<_> = (0..signature.inputs.len())
70 .map(|index| format_ident!("__arg{index}"))
71 .collect();
72 quote! {
73 {
74 struct __Site;
75 #[allow(clippy::too_many_arguments)]
76 #abi fn __dispatch<__Marker, #(#types,)* __Return>(#(#names: #types),*) -> __Return {
77 let address = __dispatch::<__Marker, #(#types,)* __Return> as *const () as usize;
78 let target = #root::__private::route(address).expect("replacement is not active");
79 #root::__invoke!(target, #abi fn(#(#types),*) -> __Return, (#(#names),*))
80 }
81 fn __install<#(#types,)* __Return>(
82 session: &mut #root::Session,
83 source: #unsafety #abi fn(#(#types),*) -> __Return,
84 target: #unsafety #abi fn(#(#types),*) -> __Return,
85 ) -> ::std::result::Result<(), #root::Error> {
86 #root::__install_replacement!(session, source as *const (),
87 __dispatch::<__Site, #(#types,)* __Return> as *const (), target as *const ())
88 }
89 #root::__private::check(__install(#session, #source, #target))
90 }
91 }
92}
93
94fn expand_check(input: syn::Result<SignatureCheck>) -> Tokens {
95 match input {
96 Ok(check) => signature::check(
97 source_item(&check.source, &check.value),
98 &check.signature,
99 &check.target,
100 ),
101 Err(error) => error.into_compile_error(),
102 }
103}
104
105struct SignatureCheck {
106 source: Expr,
107 value: Expr,
108 target: Expr,
109 signature: TypeBareFn,
110}
111
112impl Parse for SignatureCheck {
113 fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
114 let source = input.parse()?;
115 input.parse::<Token![,]>()?;
116 let value = input.parse()?;
117 input.parse::<Token![,]>()?;
118 let target = input.parse()?;
119 input.parse::<Token![,]>()?;
120 let signature = input.parse()?;
121 Ok(Self {
122 source,
123 value,
124 target,
125 signature,
126 })
127 }
128}
129
130fn source_item<'a>(source: &'a Expr, value: &'a Expr) -> &'a Expr {
131 match source {
132 Expr::Path(_) => source,
133 Expr::Paren(paren) => source_item(&paren.expr, value),
134 Expr::Group(group) => source_item(&group.expr, value),
135 _ => value,
136 }
137}
138
139struct Input {
140 root: Path,
141 session: Expr,
142 source: Expr,
143 signature: TypeBareFn,
144}
145
146impl Parse for Input {
147 fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
148 let root = input.parse()?;
149 input.parse::<Token![,]>()?;
150 let session = input.parse()?;
151 input.parse::<Token![,]>()?;
152 let source = input.parse()?;
153 input.parse::<Token![,]>()?;
154 let signature: TypeBareFn = input.parse()?;
155 if signature.variadic.is_some() {
156 return Err(syn::Error::new_spanned(
157 signature,
158 "mock! does not support variadic functions",
159 ));
160 }
161 if input.peek(Token![,]) {
162 input.parse::<Token![,]>()?;
163 }
164 Ok(Self {
165 root,
166 session,
167 source,
168 signature,
169 })
170 }
171}
172
173#[derive(Default)]
174struct Borrowed(bool);
175
176impl<'ast> Visit<'ast> for Borrowed {
177 fn visit_type_reference(&mut self, node: &'ast syn::TypeReference) {
178 self.0 |= node.lifetime.as_ref().is_none_or(|lt| lt.ident != "static");
179 syn::visit::visit_type_reference(self, node);
180 }
181
182 fn visit_lifetime(&mut self, node: &'ast syn::Lifetime) {
183 self.0 |= node.ident != "static";
184 }
185}
186
187fn expand(input: syn::Result<Input>) -> Tokens {
188 match input {
189 Ok(input) => generate(input),
190 Err(error) => error.into_compile_error(),
191 }
192}
193
194fn generate(input: Input) -> Tokens {
195 let Input {
196 root,
197 session,
198 source,
199 signature,
200 } = input;
201 let args: Vec<_> = signature.inputs.iter().map(|arg| &arg.ty).collect();
202 let saved_source = parse_quote!(__shimforge_original);
203 let signature_check = signature::check(
204 source_item(&source, &saved_source),
205 &signature,
206 &parse_quote!(__shimforge_call),
207 );
208 let names: Vec<_> = (0..args.len())
209 .map(|index| format_ident!("__shimforge_arg_{index}"))
210 .collect();
211 let infer = args.iter().map(|_| quote!(_));
212 let output: Type = match &signature.output {
213 ReturnType::Default => parse_quote!(()),
214 ReturnType::Type(_, ty) => *ty.clone(),
215 };
216 let mut borrowed = Borrowed::default();
217 borrowed.visit_type(&output);
218 let builder_output = if borrowed.0 {
219 quote!(())
220 } else {
221 quote!(#output)
222 };
223 let constants = (!borrowed.0).then(|| {
224 quote! {
225 impl<__Output> __ShimforgeBuilder<__Output>
226 where __Output: ::std::marker::Send + 'static + ::std::convert::Into<#output> {
227 #[track_caller]
228 pub fn returns(self, value: __Output) -> #root::Expectation
229 where __Output: ::std::clone::Clone {
230 self.returning(move |#(#names),*| {
231 let _ = (#(#names),*);
232 value.clone().into()
233 })
234 }
235
236 #[track_caller]
237 pub fn return_once(self, value: __Output) -> #root::Expectation {
238 self.returning_once(move |#(#names),*| {
239 let _ = (#(#names),*);
240 value.into()
241 })
242 }
243
244 #[track_caller]
245 pub fn returns_default(self) -> #root::Expectation
246 where __Output: ::std::default::Default {
247 self.returning(|#(#names),*| {
248 let _ = (#(#names),*);
249 __Output::default().into()
250 })
251 }
252 }
253 }
254 });
255 let binder = &signature.lifetimes;
256 let unsafety = &signature.unsafety;
257 let abi = &signature.abi;
258 let mut callable = signature.clone();
259 callable.unsafety = None;
260 callable.lifetimes = None;
261 let parameters = binder.as_ref().map(|binder| &binder.lifetimes);
262 let generics = parameters.map(|parameters| quote!(<#parameters>));
263 quote! {{
264 struct __ShimforgeRule {
265 meta: ::std::sync::Arc<#root::__private::Meta>,
266 matcher: ::std::boxed::Box<dyn #binder ::std::ops::Fn(#(&#args),*) -> bool + ::std::marker::Send + ::std::marker::Sync + 'static>,
267 action: ::std::sync::Mutex<__ShimforgeAction>,
268 }
269
270 enum __ShimforgeAction {
271 Repeat(::std::boxed::Box<dyn #binder ::std::ops::FnMut(#(#args),*) -> #output + ::std::marker::Send + 'static>),
272 Once(::std::option::Option<::std::boxed::Box<dyn #binder ::std::ops::FnOnce(#(#args),*) -> #output + ::std::marker::Send + 'static>>),
273 }
274
275 impl #root::__private::Rule for __ShimforgeRule {
276 fn meta(&self) -> &::std::sync::Arc<#root::__private::Meta> { &self.meta }
277 }
278
279 ::std::thread_local! {
280 static __SHIMFORGE_LOCAL: ::std::cell::RefCell<::std::option::Option<::std::sync::Arc<#root::__private::State<__ShimforgeRule>>>> = const { ::std::cell::RefCell::new(::std::option::Option::None) };
281 }
282 static __SHIMFORGE_GLOBAL: ::std::sync::Mutex<::std::option::Option<::std::sync::Arc<#root::__private::State<__ShimforgeRule>>>> = ::std::sync::Mutex::new(::std::option::Option::None);
283 static __SHIMFORGE_GLOBAL_ACTIVE: ::std::sync::atomic::AtomicBool = ::std::sync::atomic::AtomicBool::new(false);
284
285 #[allow(clippy::too_many_arguments)]
286 #abi fn __shimforge_call #generics (#(#names: #args),*) -> #output {
287 let state = __SHIMFORGE_LOCAL.try_with(|slot| slot.try_borrow().ok().and_then(|state| state.clone())).ok().flatten()
288 .or_else(|| {
289 if __SHIMFORGE_GLOBAL_ACTIVE.load(::std::sync::atomic::Ordering::Acquire) {
290 #root::__private::lock(&__SHIMFORGE_GLOBAL).clone()
291 } else {
292 ::std::option::Option::None
293 }
294 });
295 let state = match state {
296 ::std::option::Option::Some(state) => state,
297 ::std::option::Option::None => {
298 let address = __shimforge_call as *const () as usize;
299 let target = #root::__private::route(address).expect("mock is not active");
300 return #root::__invoke!(target, #callable, (#(#names),*));
301 }
302 };
303 let _call = state.enter();
304 let rule = state.select(&|rule| (rule.matcher)(#(&#names),*));
305 let mut action = #root::__private::lock(&rule.action);
306 match &mut *action {
307 __ShimforgeAction::Repeat(action) => action(#(#names),*),
308 __ShimforgeAction::Once(action) =>
309 action.take().expect("one-use return was already used")(#(#names),*),
310 }
311 }
312
313 struct __ShimforgeMock {
314 state: ::std::sync::Arc<#root::__private::State<__ShimforgeRule>>,
315 }
316
317 impl __ShimforgeMock {
318 pub fn expect(&self) -> __ShimforgeBuilder<#builder_output> {
319 __ShimforgeBuilder {
320 state: self.state.clone(),
321 config: #root::__private::Config::default(),
322 matcher: ::std::boxed::Box::new(|#(#names),*| { let _ = (#(#names),*); true }),
323 output: ::std::marker::PhantomData,
324 }
325 }
326
327 #[track_caller]
328 pub fn verify(&self) { #root::__private::check(self.state.verify()) }
329
330 #[track_caller]
331 pub fn checkpoint(&self) { #root::__private::check(self.state.checkpoint()) }
332 }
333
334 struct __ShimforgeBuilder<__Output> {
335 state: ::std::sync::Arc<#root::__private::State<__ShimforgeRule>>,
336 config: #root::__private::Config,
337 matcher: ::std::boxed::Box<dyn #binder ::std::ops::Fn(#(&#args),*) -> bool + ::std::marker::Send + ::std::marker::Sync + 'static>,
338 output: ::std::marker::PhantomData<fn() -> __Output>,
339 }
340
341 impl<__Output> __ShimforgeBuilder<__Output> {
342 pub fn with<__Matcher>(mut self, matcher: __Matcher) -> Self
343 where __Matcher: #binder ::std::ops::Fn(#(&#args),*) -> bool + ::std::marker::Send + ::std::marker::Sync + 'static {
344 self.matcher = ::std::boxed::Box::new(matcher);
345 self
346 }
347
348 pub fn times(mut self, count: impl ::std::convert::Into<#root::CallCount>) -> Self {
349 self.config = self.config.times(count);
350 self
351 }
352
353 pub fn once(mut self) -> Self {
354 self.config = self.config.once();
355 self
356 }
357
358 pub fn in_sequence(mut self, sequence: &#root::Sequence) -> Self {
359 self.config = self.config.in_sequence(sequence);
360 self
361 }
362
363 #[track_caller]
364 pub fn returning<__Action>(self, action: __Action) -> #root::Expectation
365 where __Action: #binder ::std::ops::FnMut(#(#args),*) -> #output + ::std::marker::Send + 'static {
366 #root::__private::check(self.state.add(self.config, |meta| __ShimforgeRule {
367 meta,
368 matcher: self.matcher,
369 action: ::std::sync::Mutex::new(__ShimforgeAction::Repeat(::std::boxed::Box::new(action))),
370 }))
371 }
372
373 #[track_caller]
374 pub fn returning_once<__Action>(self, action: __Action) -> #root::Expectation
375 where __Action: #binder ::std::ops::FnOnce(#(#args),*) -> #output + ::std::marker::Send + 'static {
376 let config = #root::__private::check(self.config.for_once());
377 #root::__private::check(self.state.add(config, |meta| __ShimforgeRule {
378 meta,
379 matcher: self.matcher,
380 action: ::std::sync::Mutex::new(__ShimforgeAction::Once(::std::option::Option::Some(::std::boxed::Box::new(action)))),
381 }))
382 }
383
384 #[track_caller]
385 pub fn never(self) -> #root::Expectation {
386 self.times(0usize).panics("forbidden mock call")
387 }
388
389 #[track_caller]
390 pub fn panics(self, message: impl ::std::convert::Into<::std::string::String>) -> #root::Expectation {
391 let message = message.into();
392 self.returning(move |#(#names),*| {
393 let _ = (#(#names),*);
394 ::std::panic!("{}", message)
395 })
396 }
397 }
398
399 #constants
400
401 #root::__private::check((|| -> ::std::result::Result<__ShimforgeMock, #root::Error> {
402 let __shimforge_original = #source;
403 #signature_check
404 let __shimforge_source = __shimforge_original as #unsafety #abi fn(#(#infer),*) -> _;
405 #[allow(clippy::type_complexity)]
406 let __shimforge_target: #signature = __shimforge_call;
407 fn __shimforge_checked<T>(source: T, _: T) -> T { source }
408 let __shimforge_source = __shimforge_checked(__shimforge_source, __shimforge_target);
409 let __shimforge_session = (#session).__borrow();
410 let __shimforge_thread = __shimforge_session.__thread();
411 let __shimforge_state = #root::__private::State::new(::std::stringify!(#source));
412 let __shimforge_set = |slot: &mut ::std::option::Option<::std::sync::Arc<#root::__private::State<__ShimforgeRule>>>| {
413 if slot.is_some() {
414 return ::std::result::Result::Err(#root::Error::Expectation("mock site is already active".into()));
415 }
416 *slot = ::std::option::Option::Some(__shimforge_state.clone());
417 ::std::result::Result::Ok(())
418 };
419 if __shimforge_thread.is_some() {
420 __SHIMFORGE_LOCAL.with(|slot| __shimforge_set(&mut slot.borrow_mut()))?;
421 } else {
422 __shimforge_set(&mut #root::__private::lock(&__SHIMFORGE_GLOBAL))?;
423 __SHIMFORGE_GLOBAL_ACTIVE.store(true, ::std::sync::atomic::Ordering::Release);
424 }
425 let __shimforge_detach = ::std::boxed::Box::new(move || {
426 let state = if __shimforge_thread.is_some() {
427 __SHIMFORGE_LOCAL.with(|slot| slot.borrow_mut().take())
428 } else {
429 __SHIMFORGE_GLOBAL_ACTIVE.store(false, ::std::sync::atomic::Ordering::Release);
430 #root::__private::lock(&__SHIMFORGE_GLOBAL).take()
431 };
432 ::std::mem::drop(state);
433 });
434 #root::__install!(__shimforge_session, __shimforge_source as *const (), __shimforge_target as *const (), __shimforge_state.clone(), __shimforge_detach)?;
435 ::std::result::Result::Ok(__ShimforgeMock { state: __shimforge_state })
436 })())
437 }}
438}
439
440#[cfg(test)]
441mod tests;