diff --git a/qadapt-macro/src/lib.rs b/qadapt-macro/src/lib.rs index 04d2b9a..c7418d1 100644 --- a/qadapt-macro/src/lib.rs +++ b/qadapt-macro/src/lib.rs @@ -120,17 +120,26 @@ fn protected_body(fn_body: Group) -> TokenTree { /// to `return` with an exit guard. fn escape_return(ts: TokenStream) -> TokenStream { let mut protected: Vec = Vec::new(); + let mut in_closure: bool = false; let mut tt_iter = ts.into_iter(); while let Some(tt) = tt_iter.next() { let mut tokens = match tt { - TokenTree::Group(ref g) if g.delimiter() == Delimiter::Brace => { + TokenTree::Group(ref g) if g.delimiter() == Delimiter::Brace && !in_closure => { vec![group!(Delimiter::Brace, escape_return(g.stream()))] } - TokenTree::Ident(ref i) if i.to_string() == "return" => vec![ + TokenTree::Ident(ref i) if i.to_string() == "return" && !in_closure => vec![ group!(Delimiter::Brace, release_guard("exit_protected")), tt.clone(), ], + TokenTree::Punct(ref p) if p.as_char() == '|' => { + in_closure = true; + vec![tt.clone()] + } + TokenTree::Punct(ref p) if p.as_char() == ';' => { + in_closure = false; + vec![tt.clone()] + } t => vec![t], }; diff --git a/tests/macros.rs b/tests/macros.rs index f537142..6aae496 100644 --- a/tests/macros.rs +++ b/tests/macros.rs @@ -132,3 +132,20 @@ fn macro_branching_return() { assert_eq!(7, branching_return(false, false, true)); assert_eq!(8, branching_return(false, false, false)); } + +fn run_closure(x: impl Fn(bool, bool) -> bool) -> bool { + x(true, false) +} + +#[allocate_panic] +fn example_closure() { + let c = run_closure(|a: bool, b| return a && b); + assert!(!c); + let x = || return true; + assert!(x()); +} + +#[test] +fn macro_closure() { + example_closure() +}