From 46b96d03b1541db7c1c94cc689f54c9c56c65b44 Mon Sep 17 00:00:00 2001
From: ConnorSkees <39542938+ConnorSkees@users.noreply.github.com>
Date: Fri, 28 Feb 2020 01:02:11 -0500
Subject: [PATCH] Refactor how @return is parsed

---
 src/atrule.rs      | 23 ++++++++++++++++++-----
 tests/functions.rs | 15 ++++++++++-----
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/src/atrule.rs b/src/atrule.rs
index e99566c..cd5b279 100644
--- a/src/atrule.rs
+++ b/src/atrule.rs
@@ -73,11 +73,24 @@ impl AtRule {
                 let (name, func) = Function::decl_from_tokens(toks, scope)?;
                 AtRule::Function(name, Box::new(func))
             }
-            AtRuleKind::Return => AtRule::Return(
-                // todo: return may not end in semicolon
-                toks.take_while(|t| t.kind != TokenKind::Symbol(Symbol::SemiColon))
-                    .collect(),
-            ),
+            AtRuleKind::Return => {
+                let mut t = Vec::new();
+                let mut n = 0;
+                while let Some(tok) = toks.peek() {
+                    match tok.kind {
+                        TokenKind::Symbol(Symbol::OpenCurlyBrace) => n += 1,
+                        TokenKind::Symbol(Symbol::CloseCurlyBrace) => n -= 1,
+                        TokenKind::Interpolation => n += 1,
+                        TokenKind::Symbol(Symbol::SemiColon) => break,
+                        _ => {}
+                    }
+                    if n < 0 {
+                        break;
+                    }
+                    t.push(toks.next().unwrap());
+                }
+                AtRule::Return(t)
+            },
             AtRuleKind::Use => todo!("@use not yet implemented"),
             AtRuleKind::Annotation => todo!("@annotation not yet implemented"),
             AtRuleKind::AtRoot => todo!("@at-root not yet implemented"),
diff --git a/tests/functions.rs b/tests/functions.rs
index 8b966ed..d4aac96 100644
--- a/tests/functions.rs
+++ b/tests/functions.rs
@@ -38,8 +38,13 @@ test!(
     "@function a($a) {\n  @return $a;\n}\n\nb {\ncolor: a(red,);\n}\n",
     "b {\n  color: red;\n}\n"
 );
-// test!(
-//     return_no_semicolon,
-//     "@function a() {\n  @return 1\n}\n\nb {\ncolor: a();\n}\n",
-//     "b {\n  color: 1;\n}\n"
-// );
+test!(
+    return_no_semicolon,
+    "@function a() {\n  @return 1\n}\n\nb {\ncolor: a();\n}\n",
+    "b {\n  color: 1;\n}\n"
+);
+test!(
+    two_returns,
+    "@function a() {\n  @return 1; @return 2;\n}\n\nb {\ncolor: a();\n}\n",
+    "b {\n  color: 1;\n}\n"
+);