From 3341faeb91bd06be899958d664f79ba11e2adb99 Mon Sep 17 00:00:00 2001
From: Bradlee Speice <bradlee@speice.io>
Date: Sat, 17 Nov 2018 10:32:58 -0500
Subject: [PATCH] Handle `return` showing up in closures

---
 qadapt-macro/src/lib.rs | 13 +++++++++++--
 tests/macros.rs         | 17 +++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

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<TokenTree> = 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()
+}