Dry-clean visitor with a macro 🏜

2 files changed, 45 insertions, 16 deletions

diff --git a/src/eval/capture.rs b/src/eval/capture.rs
index 054b64ab..42b617a7 100644
--- a/src/eval/capture.rs
+++ b/src/eval/capture.rs
@@ -15,32 +15,60 @@ impl<'a> CapturesVisitor<'a> {
     pub fn new(external: &'a Scopes) -> Self {
         Self { external, internal: Scopes::default() }
     }
+
+    /// Define an internal variable.
+    fn define(&mut self, ident: &Ident) {
+        self.internal.def_mut(ident.as_str(), Value::None);
+    }
 }
 
-impl<'a> Visitor<'a> for CapturesVisitor<'a> {
-    fn visit_scope_pre(&mut self) {
-        self.internal.push();
+impl<'ast> Visit<'ast> for CapturesVisitor<'_> {
+    fn visit_expr(&mut self, item: &'ast mut Expr) {
+        match item {
+            Expr::Ident(ident) => {
+                // Find out whether the identifier is not locally defined, but
+                // captured, and if so, replace it with its value.
+                if self.internal.get(ident).is_none() {
+                    if let Some(value) = self.external.get(ident) {
+                        *item = Expr::Captured(Rc::clone(&value));
+                    }
+                }
+            }
+            expr => visit_expr(self, expr),
+        }
     }
 
-    fn visit_scope_post(&mut self) {
+    fn visit_block(&mut self, item: &'ast mut ExprBlock) {
+        // Blocks create a scope except if directly in a template.
+        if item.scopes {
+            self.internal.push();
+        }
+        visit_block(self, item);
+        if item.scopes {
+            self.internal.pop();
+        }
+    }
+
+    fn visit_template(&mut self, item: &'ast mut ExprTemplate) {
+        // Templates always create a scope.
+        self.internal.push();
+        visit_template(self, item);
         self.internal.pop();
     }
 
-    fn visit_def(&mut self, id: &mut Ident) {
-        self.internal.def_mut(id.as_str(), Value::None);
+    fn visit_let(&mut self, item: &'ast mut ExprLet) {
+        self.define(&item.pat.v);
+        visit_let(self, item);
     }
 
-    fn visit_expr(&mut self, expr: &'a mut Expr) {
-        if let Expr::Ident(ident) = expr {
-            // Find out whether the identifier is not locally defined, but
-            // captured, and if so, replace it with its value.
-            if self.internal.get(ident).is_none() {
-                if let Some(value) = self.external.get(ident) {
-                    *expr = Expr::Captured(Rc::clone(&value));
-                }
+    fn visit_for(&mut self, item: &'ast mut ExprFor) {
+        match &mut item.pat.v {
+            ForPattern::Value(value) => self.define(value),
+            ForPattern::KeyValue(key, value) => {
+                self.define(key);
+                self.define(value);
             }
-        } else {
-            walk_expr(self, expr);
         }
+        visit_for(self, item);
     }
 }
diff --git a/src/eval/mod.rs b/src/eval/mod.rs
index 4f8961cc..cfd1c808 100644
--- a/src/eval/mod.rs
+++ b/src/eval/mod.rs
@@ -23,6 +23,7 @@ use crate::diag::Pass;
 use crate::env::Env;
 use crate::geom::{Angle, Length, Relative, Spec};
 use crate::layout::{self, Expansion, NodeSpacing, NodeStack};
+use crate::syntax::visit::Visit;
 use crate::syntax::*;
 
 /// Evaluate a syntax tree into a layout tree.