Capture variable by slot instead of value 🎣

7 files changed, 64 insertions, 61 deletions

diff --git a/src/eval/capture.rs b/src/eval/capture.rs
index b7052c70..054b64ab 100644
--- a/src/eval/capture.rs
+++ b/src/eval/capture.rs
@@ -1,3 +1,5 @@
+use std::rc::Rc;
+
 use super::*;
 use crate::syntax::visit::*;
 
@@ -25,7 +27,7 @@ impl<'a> Visitor<'a> for CapturesVisitor<'a> {
     }
 
     fn visit_def(&mut self, id: &mut Ident) {
-        self.internal.define(id.as_str(), Value::None);
+        self.internal.def_mut(id.as_str(), Value::None);
     }
 
     fn visit_expr(&mut self, expr: &'a mut Expr) {
@@ -34,7 +36,7 @@ impl<'a> Visitor<'a> for CapturesVisitor<'a> {
             // 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::CapturedValue(value.clone());
+                    *expr = Expr::Captured(Rc::clone(&value));
                 }
             }
         } else {
diff --git a/src/eval/mod.rs b/src/eval/mod.rs
index af3303af..996c908c 100644
--- a/src/eval/mod.rs
+++ b/src/eval/mod.rs
@@ -160,7 +160,7 @@ impl Eval for Spanned<&Expr> {
         match self.v {
             Expr::None => Value::None,
             Expr::Ident(v) => match ctx.scopes.get(v) {
-                Some(value) => value.clone(),
+                Some(slot) => slot.borrow().clone(),
                 None => {
                     ctx.diag(error!(self.span, "unknown variable"));
                     Value::Error
@@ -185,7 +185,7 @@ impl Eval for Spanned<&Expr> {
             Expr::Let(v) => v.with_span(self.span).eval(ctx),
             Expr::If(v) => v.with_span(self.span).eval(ctx),
             Expr::For(v) => v.with_span(self.span).eval(ctx),
-            Expr::CapturedValue(v) => v.clone(),
+            Expr::Captured(v) => v.borrow().clone(),
         }
     }
 }
@@ -341,30 +341,29 @@ impl Spanned<&ExprBinary> {
         let rhs = self.v.rhs.eval(ctx);
         let span = self.v.lhs.span;
 
-        match &self.v.lhs.v {
-            Expr::Ident(id) => {
-                if let Some(slot) = ctx.scopes.get_mut(id) {
-                    *slot = op(std::mem::take(slot), rhs);
-                    return Value::None;
-                } else if ctx.scopes.is_const(id) {
-                    ctx.diag(error!(span, "cannot assign to a constant"));
-                } else {
+        let slot = match &self.v.lhs.v {
+            Expr::Ident(id) => match ctx.scopes.get(id) {
+                Some(slot) => slot,
+                None => {
                     ctx.diag(error!(span, "unknown variable"));
+                    return Value::Error;
                 }
-            }
+            },
 
-            Expr::CapturedValue(_) => {
-                ctx.diag(error!(
-                    span,
-                    "cannot assign to captured expression in a template",
-                ));
-            }
+            Expr::Captured(slot) => slot,
 
             _ => {
                 ctx.diag(error!(span, "cannot assign to this expression"));
+                return Value::Error;
             }
+        };
+
+        if let Ok(mut slot) = slot.try_borrow_mut() {
+            *slot = op(std::mem::take(&mut slot), rhs);
+            return Value::None;
         }
 
+        ctx.diag(error!(span, "cannot assign to a constant"));
         Value::Error
     }
 }
@@ -377,7 +376,7 @@ impl Eval for Spanned<&ExprLet> {
             Some(expr) => expr.eval(ctx),
             None => Value::None,
         };
-        ctx.scopes.define(self.v.pat.v.as_str(), value);
+        ctx.scopes.def_mut(self.v.pat.v.as_str(), value);
         Value::None
     }
 }
@@ -418,7 +417,7 @@ impl Eval for Spanned<&ExprFor> {
             (for ($($binding:ident => $value:ident),*) in $iter:expr) => {
                 #[allow(unused_parens)]
                 for ($($value),*) in $iter {
-                    $(ctx.scopes.define($binding.as_str(), $value);)*
+                    $(ctx.scopes.def_mut($binding.as_str(), $value);)*
 
                     if let Value::Template(new) = self.v.body.eval(ctx) {
                         output.extend(new);
diff --git a/src/eval/scope.rs b/src/eval/scope.rs
index 9c966a24..70338041 100644
--- a/src/eval/scope.rs
+++ b/src/eval/scope.rs
@@ -1,9 +1,14 @@
+use std::cell::RefCell;
 use std::collections::HashMap;
 use std::fmt::{self, Debug, Formatter};
 use std::iter;
+use std::rc::Rc;
 
 use super::Value;
 
+/// A slot where a variable is stored.
+pub type Slot = Rc<RefCell<Value>>;
+
 /// A stack of scopes.
 #[derive(Debug, Default, Clone, PartialEq)]
 pub struct Scopes<'a> {
@@ -34,38 +39,29 @@ impl<'a> Scopes<'a> {
         self.top = self.scopes.pop().expect("no pushed scope");
     }
 
-    /// Define a variable in the active scope.
-    pub fn define(&mut self, var: impl Into<String>, value: impl Into<Value>) {
-        self.top.define(var, value);
+    /// Define a constant variable in the active scope.
+    pub fn def_const(&mut self, var: impl Into<String>, value: impl Into<Value>) {
+        self.top.def_const(var, value);
     }
 
-    /// Look up the value of a variable.
-    pub fn get(&self, var: &str) -> Option<&Value> {
+    /// Define a mutable variable in the active scope.
+    pub fn def_mut(&mut self, var: impl Into<String>, value: impl Into<Value>) {
+        self.top.def_mut(var, value);
+    }
+
+    /// Look up the slot of a variable.
+    pub fn get(&self, var: &str) -> Option<&Slot> {
         iter::once(&self.top)
             .chain(self.scopes.iter().rev())
             .chain(self.base.into_iter())
             .find_map(|scope| scope.get(var))
     }
-
-    /// Get a mutable reference to a variable.
-    pub fn get_mut(&mut self, var: &str) -> Option<&mut Value> {
-        iter::once(&mut self.top)
-            .chain(self.scopes.iter_mut().rev())
-            .find_map(|scope| scope.get_mut(var))
-    }
-
-    /// Return whether the variable is constant (not writable).
-    ///
-    /// Defaults to `false` if the variable does not exist.
-    pub fn is_const(&self, var: &str) -> bool {
-        self.base.map_or(false, |base| base.get(var).is_some())
-    }
 }
 
-/// A map from variable names to values.
+/// A map from variable names to variable slots.
 #[derive(Default, Clone, PartialEq)]
 pub struct Scope {
-    values: HashMap<String, Value>,
+    values: HashMap<String, Slot>,
 }
 
 impl Scope {
@@ -74,19 +70,25 @@ impl Scope {
         Self::default()
     }
 
-    /// Define a new variable.
-    pub fn define(&mut self, var: impl Into<String>, value: impl Into<Value>) {
-        self.values.insert(var.into(), value.into());
+    /// Define a constant variable.
+    pub fn def_const(&mut self, var: impl Into<String>, value: impl Into<Value>) {
+        let cell = RefCell::new(value.into());
+
+        // Make it impossible to write to this value again.
+        // FIXME: Use Ref::leak once stable.
+        std::mem::forget(cell.borrow());
+
+        self.values.insert(var.into(), Rc::new(cell));
     }
 
-    /// Look up the value of a variable.
-    pub fn get(&self, var: &str) -> Option<&Value> {
-        self.values.get(var)
+    /// Define a mutable variable.
+    pub fn def_mut(&mut self, var: impl Into<String>, value: impl Into<Value>) {
+        self.values.insert(var.into(), Rc::new(RefCell::new(value.into())));
     }
 
-    /// Get a mutable reference to a variable.
-    pub fn get_mut(&mut self, var: &str) -> Option<&mut Value> {
-        self.values.get_mut(var)
+    /// Look up the value of a variable.
+    pub fn get(&self, var: &str) -> Option<&Slot> {
+        self.values.get(var)
     }
 }
 
diff --git a/src/library/mod.rs b/src/library/mod.rs
index ed83270d..48da093b 100644
--- a/src/library/mod.rs
+++ b/src/library/mod.rs
@@ -23,10 +23,10 @@ pub fn new() -> Scope {
     let mut std = Scope::new();
     macro_rules! set {
         (func: $name:expr, $func:expr) => {
-            std.define($name, ValueFunc::new($name, $func))
+            std.def_const($name, ValueFunc::new($name, $func))
         };
         (any: $var:expr, $any:expr) => {
-            std.define($var, ValueAny::new($any))
+            std.def_const($var, ValueAny::new($any))
         };
     }
 
diff --git a/src/syntax/expr.rs b/src/syntax/expr.rs
index 7a055cc7..ebe82199 100644
--- a/src/syntax/expr.rs
+++ b/src/syntax/expr.rs
@@ -1,6 +1,6 @@
 use super::*;
 use crate::color::RgbaColor;
-use crate::eval::Value;
+use crate::eval::Slot;
 use crate::geom::{AngularUnit, LengthUnit};
 
 /// An expression.
@@ -51,11 +51,11 @@ pub enum Expr {
     If(ExprIf),
     /// A for expression: `#for x #in y { z }`.
     For(ExprFor),
-    /// A captured value.
+    /// A captured variable slot.
     ///
     /// This node is never created by parsing. It only results from an in-place
-    /// transformation of an identifier to a captured value.
-    CapturedValue(Value),
+    /// transformation of an identifier to a captured variable.
+    Captured(Slot),
 }
 
 impl Pretty for Expr {
@@ -88,7 +88,7 @@ impl Pretty for Expr {
             Self::Let(v) => v.pretty(p),
             Self::If(v) => v.pretty(p),
             Self::For(v) => v.pretty(p),
-            Self::CapturedValue(v) => v.pretty(p),
+            Self::Captured(v) => v.borrow().pretty(p),
         }
     }
 }
diff --git a/src/syntax/visit.rs b/src/syntax/visit.rs
index e9e5dad7..6f0e7ef3 100644
--- a/src/syntax/visit.rs
+++ b/src/syntax/visit.rs
@@ -94,7 +94,7 @@ pub fn walk_expr<'a, V: Visitor<'a>>(v: &mut V, expr: &'a mut Expr) {
         Expr::Let(e) => v.visit_let(e),
         Expr::If(e) => v.visit_if(e),
         Expr::For(e) => v.visit_for(e),
-        Expr::CapturedValue(_) => {}
+        Expr::Captured(_) => {}
     }
 }
 
diff --git a/tests/typeset.rs b/tests/typeset.rs
index cd58bdab..4b9d26a4 100644
--- a/tests/typeset.rs
+++ b/tests/typeset.rs
@@ -317,8 +317,8 @@ fn register_helpers(scope: &mut Scope, panics: Rc<RefCell<Vec<Panic>>>) {
         }
     };
 
-    scope.define("f", ValueFunc::new("f", f));
-    scope.define("test", ValueFunc::new("test", test));
+    scope.def_const("f", ValueFunc::new("f", f));
+    scope.def_const("test", ValueFunc::new("test", test));
 }
 
 fn print_diag(diag: &Spanned<Diag>, map: &LineMap, lines: u32) {