Fine-grained capturing

3 files changed, 185 insertions, 41 deletions

diff --git a/src/eval/capture.rs b/src/eval/capture.rs
index 786da36e..4e24bc90 100644
--- a/src/eval/capture.rs
+++ b/src/eval/capture.rs
@@ -1,27 +1,118 @@
 use std::rc::Rc;
 
-use super::{Scope, Scopes};
-use crate::syntax::{NodeKind, RedRef};
+use super::{Scope, Scopes, Value};
+use crate::syntax::ast::{ClosureParam, Expr, Ident, Imports, TypedNode};
+use crate::syntax::RedRef;
 
 /// A visitor that captures variable slots.
 pub struct CapturesVisitor<'a> {
     external: &'a Scopes<'a>,
+    internal: Scopes<'a>,
     captures: Scope,
 }
 
 impl<'a> CapturesVisitor<'a> {
     /// Create a new visitor for the given external scopes.
     pub fn new(external: &'a Scopes) -> Self {
-        Self { external, captures: Scope::new() }
+        Self {
+            external,
+            internal: Scopes::new(None),
+            captures: Scope::new(),
+        }
+    }
+
+    /// Return the scope of captured variables.
+    pub fn finish(self) -> Scope {
+        self.captures
+    }
+
+    /// Bind a new internal variable.
+    pub fn bind(&mut self, ident: Ident) {
+        self.internal.def_mut(ident.take(), Value::None);
+    }
+
+    /// Capture a variable if it isn't internal.
+    pub fn capture(&mut self, ident: Ident) {
+        if self.internal.get(&ident).is_none() {
+            if let Some(slot) = self.external.get(&ident) {
+                self.captures.def_slot(ident.take(), Rc::clone(slot));
+            }
+        }
     }
 
+    /// Visit any node and collect all captured variables.
     pub fn visit(&mut self, node: RedRef) {
-        match node.kind() {
-            NodeKind::Ident(ident) => {
-                if let Some(slot) = self.external.get(ident.as_str()) {
-                    self.captures.def_slot(ident.as_str(), Rc::clone(slot));
+        match node.cast() {
+            // Every identifier is a potential variable that we need to capture.
+            // Identifiers that shouldn't count as captures because they
+            // actually bind a new name are handled further below (individually
+            // through the expressions that contain them).
+            Some(Expr::Ident(ident)) => self.capture(ident),
+
+            // A closure contains parameter bindings, which are bound before the
+            // body is evaluated. Take must be taken so that the default values
+            // of named parameters cannot access previous parameter bindings.
+            Some(Expr::Closure(expr)) => {
+                for param in expr.params() {
+                    if let ClosureParam::Named(named) = param {
+                        self.visit(named.expr().as_red());
+                    }
+                }
+
+                for param in expr.params() {
+                    match param {
+                        ClosureParam::Pos(ident) => self.bind(ident),
+                        ClosureParam::Named(named) => self.bind(named.name()),
+                        ClosureParam::Sink(ident) => self.bind(ident),
+                    }
+                }
+
+                self.visit(expr.body().as_red());
+            }
+
+            // A let expression contains a binding, but that binding is only
+            // active after the body is evaluated.
+            Some(Expr::Let(expr)) => {
+                if let Some(init) = expr.init() {
+                    self.visit(init.as_red());
+                }
+                self.bind(expr.binding());
+            }
+
+            // A for loop contains one or two bindings in its pattern. These are
+            // active after the iterable is evaluated but before the body is
+            // evaluated.
+            Some(Expr::For(expr)) => {
+                self.visit(expr.iter().as_red());
+                let pattern = expr.pattern();
+                if let Some(key) = pattern.key() {
+                    self.bind(key);
+                }
+                self.bind(pattern.value());
+                self.visit(expr.body().as_red());
+            }
+
+            // An import contains items, but these are active only after the
+            // path is evaluated.
+            Some(Expr::Import(expr)) => {
+                self.visit(expr.path().as_red());
+                if let Imports::Items(items) = expr.imports() {
+                    for item in items {
+                        self.bind(item);
+                    }
+                }
+            }
+
+            // Blocks and templates create a scope.
+            Some(Expr::Block(_) | Expr::Template(_)) => {
+                self.internal.enter();
+                for child in node.children() {
+                    self.visit(child);
                 }
+                self.internal.exit();
             }
+
+            // Everything else is traversed from left to right.
             _ => {
                 for child in node.children() {
                     self.visit(child);
@@ -29,9 +120,57 @@ impl<'a> CapturesVisitor<'a> {
             }
         }
     }
+}
 
-    /// Return the scope of captured variables.
-    pub fn finish(self) -> Scope {
-        self.captures
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::parse::parse;
+    use crate::source::SourceId;
+    use crate::syntax::RedNode;
+
+    #[track_caller]
+    fn test(src: &str, result: &[&str]) {
+        let green = parse(src);
+        let red = RedNode::from_root(green, SourceId::from_raw(0));
+
+        let mut scopes = Scopes::new(None);
+        scopes.def_const("x", 0);
+        scopes.def_const("y", 0);
+        scopes.def_const("z", 0);
+
+        let mut visitor = CapturesVisitor::new(&scopes);
+        visitor.visit(red.as_ref());
+
+        let captures = visitor.finish();
+        let mut names: Vec<_> = captures.iter().map(|(k, _)| k).collect();
+        names.sort();
+
+        assert_eq!(names, result);
+    }
+
+    #[test]
+    fn test_captures() {
+        // Let binding and function definition.
+        test("#let x = x", &["x"]);
+        test("#let x; {x + y}", &["y"]);
+        test("#let f(x, y) = x + y", &[]);
+
+        // Closure with different kinds of params.
+        test("{(x, y) => x + z}", &["z"]);
+        test("{(x: y, z) => x + z}", &["y"]);
+        test("{(..x) => x + y}", &["y"]);
+        test("{(x, y: x + z) => x + y}", &["x", "z"]);
+
+        // For loop.
+        test("#for x in y { x + z }", &["y", "z"]);
+        test("#for x, y in y { x + y }", &["y"]);
+
+        // Import.
+        test("#import x, y from z", &["z"]);
+        test("#import x, y, z from x + y", &["x", "y"]);
+
+        // Scoping.
+        test("{ let x = 1; { let y = 2; y }; x + y }", &["y"]);
     }
 }
diff --git a/src/eval/mod.rs b/src/eval/mod.rs
index 7c984691..540b58b9 100644
--- a/src/eval/mod.rs
+++ b/src/eval/mod.rs
@@ -174,8 +174,8 @@ impl Eval for Expr {
 
     fn eval(&self, ctx: &mut EvalContext) -> TypResult<Self::Output> {
         match self {
-            Self::Ident(v) => v.eval(ctx),
             Self::Lit(v) => v.eval(ctx),
+            Self::Ident(v) => v.eval(ctx),
             Self::Array(v) => v.eval(ctx).map(Value::Array),
             Self::Dict(v) => v.eval(ctx).map(Value::Dict),
             Self::Template(v) => v.eval(ctx).map(Value::Template),
@@ -200,17 +200,17 @@ impl Eval for Lit {
     type Output = Value;
 
     fn eval(&self, _: &mut EvalContext) -> TypResult<Self::Output> {
-        Ok(match *self {
-            Self::None(_) => Value::None,
-            Self::Auto(_) => Value::Auto,
-            Self::Bool(_, v) => Value::Bool(v),
-            Self::Int(_, v) => Value::Int(v),
-            Self::Float(_, v) => Value::Float(v),
-            Self::Length(_, v, unit) => Value::Length(Length::with_unit(v, unit)),
-            Self::Angle(_, v, unit) => Value::Angle(Angle::with_unit(v, unit)),
-            Self::Percent(_, v) => Value::Relative(Relative::new(v / 100.0)),
-            Self::Fractional(_, v) => Value::Fractional(Fractional::new(v)),
-            Self::Str(_, ref v) => Value::Str(v.into()),
+        Ok(match self.kind() {
+            LitKind::None => Value::None,
+            LitKind::Auto => Value::Auto,
+            LitKind::Bool(v) => Value::Bool(v),
+            LitKind::Int(v) => Value::Int(v),
+            LitKind::Float(v) => Value::Float(v),
+            LitKind::Length(v, unit) => Value::Length(Length::with_unit(v, unit)),
+            LitKind::Angle(v, unit) => Value::Angle(Angle::with_unit(v, unit)),
+            LitKind::Percent(v) => Value::Relative(Relative::new(v / 100.0)),
+            LitKind::Fractional(v) => Value::Fractional(Fractional::new(v)),
+            LitKind::Str(ref v) => Value::Str(v.into()),
         })
     }
 }
@@ -219,9 +219,9 @@ impl Eval for Ident {
     type Output = Value;
 
     fn eval(&self, ctx: &mut EvalContext) -> TypResult<Self::Output> {
-        match ctx.scopes.get(&self.string) {
+        match ctx.scopes.get(self) {
             Some(slot) => Ok(slot.borrow().clone()),
-            None => bail!(self.span, "unknown variable"),
+            None => bail!(self.span(), "unknown variable"),
         }
     }
 }
@@ -239,7 +239,7 @@ impl Eval for DictExpr {
 
     fn eval(&self, ctx: &mut EvalContext) -> TypResult<Self::Output> {
         self.items()
-            .map(|x| Ok((x.name().string.into(), x.expr().eval(ctx)?)))
+            .map(|x| Ok((x.name().take().into(), x.expr().eval(ctx)?)))
             .collect()
     }
 }
@@ -401,7 +401,7 @@ impl Eval for CallArgs {
                 CallArg::Named(x) => {
                     items.push(Arg {
                         span,
-                        name: Some(x.name().string.into()),
+                        name: Some(x.name().take().into()),
                         value: Spanned::new(x.expr().eval(ctx)?, x.expr().span()),
                     });
                 }
@@ -457,23 +457,23 @@ impl Eval for ClosureExpr {
         for param in self.params() {
             match param {
                 ClosureParam::Pos(name) => {
-                    params.push((name.string, None));
+                    params.push((name.take(), None));
                 }
                 ClosureParam::Named(named) => {
-                    params.push((named.name().string, Some(named.expr().eval(ctx)?)));
+                    params.push((named.name().take(), Some(named.expr().eval(ctx)?)));
                 }
                 ClosureParam::Sink(name) => {
                     if sink.is_some() {
-                        bail!(name.span, "only one argument sink is allowed");
+                        bail!(name.span(), "only one argument sink is allowed");
                     }
-                    sink = Some(name.string);
+                    sink = Some(name.take());
                 }
             }
         }
 
         // Clone the body expression so that we don't have a lifetime
         // dependence on the AST.
-        let name = self.name().map(|name| name.string);
+        let name = self.name().map(Ident::take);
         let body = self.body();
 
         // Define the actual function.
@@ -533,7 +533,7 @@ impl Eval for LetExpr {
             Some(expr) => expr.eval(ctx)?,
             None => Value::None,
         };
-        ctx.scopes.def_mut(self.binding().string, value);
+        ctx.scopes.def_mut(self.binding().take(), value);
         Ok(Value::None)
     }
 }
@@ -589,7 +589,7 @@ impl Eval for ForExpr {
 
                 #[allow(unused_parens)]
                 for ($($value),*) in $iter {
-                    $(ctx.scopes.def_mut(&$binding.string, $value);)*
+                    $(ctx.scopes.def_mut(&$binding, $value);)*
 
                     let value = self.body().eval(ctx)?;
                     output = ops::join(output, value)
@@ -603,7 +603,10 @@ impl Eval for ForExpr {
 
         let iter = self.iter().eval(ctx)?;
         let pattern = self.pattern();
-        match (pattern.key(), pattern.value(), iter) {
+        let key = pattern.key().map(Ident::take);
+        let value = pattern.value().take();
+
+        match (key, value, iter) {
             (None, v, Value::Str(string)) => iter!(for (v => value) in string.iter()),
             (None, v, Value::Array(array)) => {
                 iter!(for (v => value) in array.into_iter())
@@ -644,10 +647,10 @@ impl Eval for ImportExpr {
             }
             Imports::Items(idents) => {
                 for ident in idents {
-                    if let Some(slot) = module.scope.get(&ident.string) {
-                        ctx.scopes.def_mut(ident.string, slot.borrow().clone());
+                    if let Some(slot) = module.scope.get(&ident) {
+                        ctx.scopes.def_mut(ident.take(), slot.borrow().clone());
                     } else {
-                        bail!(ident.span, "unresolved import");
+                        bail!(ident.span(), "unresolved import");
                     }
                 }
             }
@@ -691,12 +694,12 @@ impl Access for Expr {
 
 impl Access for Ident {
     fn access<'a>(&self, ctx: &'a mut EvalContext) -> TypResult<RefMut<'a, Value>> {
-        match ctx.scopes.get(&self.string) {
+        match ctx.scopes.get(self) {
             Some(slot) => match slot.try_borrow_mut() {
                 Ok(guard) => Ok(guard),
-                Err(_) => bail!(self.span, "cannot mutate a constant"),
+                Err(_) => bail!(self.span(), "cannot mutate a constant"),
             },
-            None => bail!(self.span, "unknown variable"),
+            None => bail!(self.span(), "unknown variable"),
         }
     }
 }
diff --git a/src/eval/scope.rs b/src/eval/scope.rs
index eb057ae3..2290affd 100644
--- a/src/eval/scope.rs
+++ b/src/eval/scope.rs
@@ -120,6 +120,8 @@ impl Scope {
 
 impl Debug for Scope {
     fn fmt(&self, f: &mut Formatter) -> fmt::Result {
-        self.values.fmt(f)
+        f.debug_map()
+            .entries(self.values.iter().map(|(k, v)| (k, v.borrow())))
+            .finish()
     }
 }