refactor @ for to be used in @ function

src/atrule/for_rule.rs

@@ -1,3 +1,5 @@
+use std::iter::Iterator;
+
 use codemap::{Span, Spanned};
 
 use peekmore::{PeekMore, PeekMoreIterator};
@@ -15,7 +17,44 @@ use crate::utils::{
     devour_whitespace, eat_ident, read_until_closing_curly_brace, read_until_open_curly_brace,
 };
 use crate::value::{Number, Value};
-use crate::Token;
+use crate::{Stmt, Token};
+
+#[derive(Debug, Clone)]
+pub(crate) struct For {
+    pub var: Spanned<String>,
+    // TODO: optimization: this could be a generic or &dyn Iterator maybe?
+    pub iter: Vec<usize>,
+    pub body: Vec<Token>,
+}
+
+impl For {
+    pub fn ruleset_eval(
+        self,
+        scope: &mut Scope,
+        super_selector: &Selector,
+    ) -> SassResult<Vec<Spanned<Stmt>>> {
+        let mut stmts = Vec::new();
+        for i in self.iter {
+            scope.insert_var(
+                &self.var.node,
+                Spanned {
+                    node: Value::Dimension(Number::from(i), Unit::None),
+                    span: self.var.span,
+                },
+            )?;
+            stmts.extend(eat_stmts(
+                &mut self.body.clone().into_iter().peekmore(),
+                scope,
+                super_selector,
+            )?);
+        }
+        Ok(stmts)
+    }
+
+    pub fn iter(&self) -> std::slice::Iter<'_, usize> {
+        self.iter.iter()
+    }
+}
 
 pub(crate) fn parse_for<I: Iterator<Item = Token>>(
     toks: &mut PeekMoreIterator<I>,
@@ -23,7 +62,6 @@ pub(crate) fn parse_for<I: Iterator<Item = Token>>(
     super_selector: &Selector,
     span: Span,
 ) -> SassResult<AtRule> {
-    let mut stmts = Vec::new();
     devour_whitespace(toks);
     let var = match toks.next().ok_or(("expected \"$\".", span))?.kind {
         '$' => eat_ident(toks, scope, super_selector)?,
@@ -137,29 +175,11 @@ pub(crate) fn parse_for<I: Iterator<Item = Token>>(
 
     devour_whitespace(toks);
 
-    let (mut x, mut y);
+    let iter = if from < to {
-    let iter: &mut dyn std::iter::Iterator<Item = usize> = if from < to {
+        (from..(to + through)).collect()
-        x = from..(to + through);
-        &mut x
     } else {
-        y = ((to - through)..(from + 1)).skip(1).rev();
+        ((to - through)..(from + 1)).skip(1).rev().collect()
-        &mut y
     };
 
-    for i in iter {
+    Ok(AtRule::For(For { iter, body, var }))
-        scope.insert_var(
-            &var,
-            Spanned {
-                node: Value::Dimension(Number::from(i), Unit::None),
-                span: var.span,
-            },
-        )?;
-        stmts.extend(eat_stmts(
-            &mut body.clone().into_iter().peekmore(),
-            scope,
-            super_selector,
-        )?);
-    }
-
-    Ok(AtRule::For(stmts))
 }

src/atrule/function.rs

@@ -9,8 +9,9 @@ use crate::atrule::AtRule;
 use crate::error::SassResult;
 use crate::scope::Scope;
 use crate::selector::Selector;
+use crate::unit::Unit;
 use crate::utils::{devour_whitespace, eat_ident, read_until_closing_curly_brace};
-use crate::value::Value;
+use crate::value::{Number, Value};
 use crate::{Stmt, Token};
 
 #[derive(Debug, Clone)]
@@ -118,32 +119,50 @@ impl Function {
     ) -> SassResult<Value> {
         self.args(args, scope, super_selector)?;
         let stmts = self.eval_body(super_selector)?;
-        self.call(super_selector, stmts)
+        self.call(super_selector, stmts)?
+            .ok_or(("Function finished without @return.", self.pos).into())
     }
 
-    pub fn call(&self, super_selector: &Selector, stmts: Vec<Spanned<Stmt>>) -> SassResult<Value> {
+    pub fn call(
+        &mut self,
+        super_selector: &Selector,
+        stmts: Vec<Spanned<Stmt>>,
+    ) -> SassResult<Option<Value>> {
         for stmt in stmts {
             match stmt.node {
                 Stmt::AtRule(AtRule::Return(toks)) => {
-                    return Ok(Value::from_tokens(
+                    return Ok(Some(
-                        &mut toks.into_iter().peekmore(),
+                        Value::from_vec(toks, &self.scope, super_selector)?.node,
-                        &self.scope,
+                    ));
-                        super_selector,
-                    )?
-                    .node)
                 }
-                Stmt::AtRule(AtRule::For(..)) => todo!("@for in function"),
+                Stmt::AtRule(AtRule::For(f)) => {
-                Stmt::AtRule(AtRule::If(i)) => {
+                    for i in f.iter().cloned() {
-                    if let Ok(v) = self.call(
+                        self.scope.insert_var(
+                            &f.var.node,
+                            Spanned {
+                                node: Value::Dimension(Number::from(i), Unit::None),
+                                span: f.var.span,
+                            },
+                        )?;
+                    }
+                    let for_stmts = eat_stmts(
+                        &mut f.body.clone().into_iter().peekmore(),
+                        &mut self.scope,
                         super_selector,
-                        i.eval(&mut self.scope.clone(), super_selector)?,
+                    )?;
-                    ) {
+                    if let Some(v) = self.call(super_selector, for_stmts)? {
-                        return Ok(v);
+                        return Ok(Some(v));
+                    }
+                }
+                Stmt::AtRule(AtRule::If(i)) => {
+                    let if_stmts = i.eval(&mut self.scope, super_selector)?;
+                    if let Some(v) = self.call(super_selector, if_stmts)? {
+                        return Ok(Some(v));
                     }
                 }
                 _ => return Err(("This at-rule is not allowed here.", stmt.span).into()),
             }
         }
-        Err(("Function finished without @return.", self.pos).into())
+        Ok(None)
     }
 }

src/atrule/mixin.rs

@@ -123,9 +123,10 @@ impl Mixin {
             let span = expr.span;
             match expr.node {
                 Expr::AtRule(a) => match a {
-                    AtRule::Include(s) | AtRule::While(s) | AtRule::Each(s) | AtRule::For(s) => {
+                    AtRule::For(f) => {
-                        stmts.extend(s)
+                        stmts.extend(f.ruleset_eval(&mut self.scope, super_selector)?)
                     }
+                    AtRule::Include(s) | AtRule::While(s) | AtRule::Each(s) => stmts.extend(s),
                     AtRule::If(i) => stmts.extend(i.eval(&mut self.scope.clone(), super_selector)?),
                     AtRule::Content => stmts.extend(self.content.clone()),
                     AtRule::Return(..) => {
@@ -193,7 +194,7 @@ pub(crate) fn eat_include<I: Iterator<Item = Token>>(
         match tok.kind {
             ';' => CallArgs::new(name.span),
             '(' => {
-                let tmp = eat_call_args(toks, scope, super_selector)?;
+                let tmp = eat_call_args(toks)?;
                 devour_whitespace_or_comment(toks)?;
                 if let Some(tok) = toks.next() {
                     match tok.kind {

src/atrule/mod.rs

@@ -13,6 +13,7 @@ use crate::utils::{
 use crate::value::Value;
 use crate::{RuleSet, Stmt, Token};
 
+use for_rule::For;
 pub(crate) use function::Function;
 pub(crate) use if_rule::If;
 pub(crate) use kind::AtRuleKind;
@@ -38,7 +39,7 @@ pub(crate) enum AtRule {
     Charset,
     Content,
     Unknown(UnknownAtRule),
-    For(Vec<Spanned<Stmt>>),
+    For(For),
     Each(Vec<Spanned<Stmt>>),
     While(Vec<Spanned<Stmt>>),
     Include(Vec<Spanned<Stmt>>),

src/lib.rs

@@ -378,10 +378,10 @@ impl<'a> StyleSheetParser<'a> {
                                         ("This at-rule is not allowed here.", rule.span).into()
                                     )
                                 }
+                                AtRule::For(f) => rules.extend(f.ruleset_eval(&mut Scope::new(), &Selector::new())?),
                                 AtRule::Include(s)
                                 | AtRule::While(s)
-                                | AtRule::Each(s)
+                                | AtRule::Each(s) => rules.extend(s),
-                                | AtRule::For(s) => rules.extend(s),
                                 AtRule::Content => return Err(
                                     ("@content is only allowed within mixin declarations.", rule.span
                                 ).into()),
@@ -425,9 +425,8 @@ impl<'a> StyleSheetParser<'a> {
                     span,
                 }),
                 Expr::AtRule(a) => match a {
-                    AtRule::Include(s) | AtRule::While(s) | AtRule::Each(s) | AtRule::For(s) => {
+                    AtRule::For(f) => stmts.extend(f.ruleset_eval(scope, super_selector)?),
-                        stmts.extend(s)
+                    AtRule::Include(s) | AtRule::While(s) | AtRule::Each(s) => stmts.extend(s),
-                    }
                     AtRule::If(i) => stmts.extend(i.eval(scope, super_selector)?),
                     AtRule::Content => {
                         return Err((

tests/for.rs

@@ -43,3 +43,13 @@ test!(
     "a {\n  $a: red;\n  @for $i from 1 to 3 {\n    $a: blue;\n  }\n  color: $a;\n}\n",
     "a {\n  color: blue;\n}\n"
 );
+test!(
+    simple_return_in_function,
+    "@function foo() {\n  @for $i from 1 to 2 {\n    @return $i;\n  }\n}\na {\n  color: foo();\n}\n",
+    "a {\n  color: 1;\n}\n"
+);
+test!(
+    return_gated_by_if_in_function,
+    "@function foo() {\n  @for $i from 1 through 2 {\n    @if $i==2 {\n      @return $i;\n    }\n  }\n}\na {\n  color: foo();\n}\n",
+    "a {\n  color: 2;\n}\n"
+);