Better function representation

5 files changed, 152 insertions, 107 deletions

diff --git a/src/eval/class.rs b/src/eval/class.rs
index 307bebcc..91a9daef 100644
--- a/src/eval/class.rs
+++ b/src/eval/class.rs
@@ -1,8 +1,7 @@
 use std::fmt::{self, Debug, Formatter, Write};
 
-use super::{Args, EvalContext, Node, StyleMap};
+use super::{Args, EvalContext, Func, Node, StyleMap, Value};
 use crate::diag::TypResult;
-use crate::util::EcoString;
 
 /// A class of [nodes](Node).
 ///
@@ -35,38 +34,40 @@ use crate::util::EcoString;
 /// [`set`]: Self::set
 #[derive(Clone)]
 pub struct Class {
-    name: EcoString,
-    construct: fn(&mut EvalContext, &mut Args) -> TypResult<Node>,
+    name: &'static str,
+    construct: fn(&mut EvalContext, &mut Args) -> TypResult<Value>,
     set: fn(&mut Args, &mut StyleMap) -> TypResult<()>,
 }
 
 impl Class {
     /// Create a new class.
-    pub fn new<T>(name: EcoString) -> Self
+    pub fn new<T>(name: &'static str) -> Self
     where
         T: Construct + Set + 'static,
     {
         Self {
             name,
-            construct: T::construct,
+            construct: |ctx, args| {
+                let mut styles = StyleMap::new();
+                T::set(args, &mut styles)?;
+                let node = T::construct(ctx, args)?;
+                Ok(Value::Node(node.styled_with_map(styles.scoped())))
+            },
             set: T::set,
         }
     }
 
     /// The name of the class.
-    pub fn name(&self) -> &EcoString {
-        &self.name
+    pub fn name(&self) -> &'static str {
+        self.name
     }
 
     /// Construct an instance of the class.
     ///
     /// This parses both property and data arguments (in this order) and styles
     /// the node constructed from the data with the style properties.
-    pub fn construct(&self, ctx: &mut EvalContext, args: &mut Args) -> TypResult<Node> {
-        let mut styles = StyleMap::new();
-        self.set(args, &mut styles)?;
-        let node = (self.construct)(ctx, args)?;
-        Ok(node.styled_with_map(styles.scoped()))
+    pub fn construct(&self, ctx: &mut EvalContext, args: &mut Args) -> TypResult<Value> {
+        (self.construct)(ctx, args)
     }
 
     /// Execute the class's set rule.
@@ -76,6 +77,11 @@ impl Class {
     pub fn set(&self, args: &mut Args, styles: &mut StyleMap) -> TypResult<()> {
         (self.set)(args, styles)
     }
+
+    /// Return the class constructor as a function.
+    pub fn constructor(&self) -> Func {
+        Func::native(self.name, self.construct)
+    }
 }
 
 impl Debug for Class {
diff --git a/src/eval/function.rs b/src/eval/func.rs
index 0edc1e78..ccd0932f 100644
--- a/src/eval/function.rs
+++ b/src/eval/func.rs
@@ -1,44 +1,68 @@
 use std::fmt::{self, Debug, Formatter, Write};
 use std::sync::Arc;
 
-use super::{Cast, EvalContext, Value};
+use super::{Cast, Eval, EvalContext, Scope, Value};
 use crate::diag::{At, TypResult};
+use crate::syntax::ast::Expr;
 use crate::syntax::{Span, Spanned};
 use crate::util::EcoString;
 
 /// An evaluatable function.
 #[derive(Clone)]
-pub struct Function(Arc<Inner<Func>>);
+pub struct Func(Arc<Repr>);
 
-/// The unsized structure behind the [`Arc`].
-struct Inner<T: ?Sized> {
-    name: Option<EcoString>,
-    func: T,
+/// The different kinds of function representations.
+enum Repr {
+    /// A native rust function.
+    Native(Native),
+    /// A user-defined closure.
+    Closure(Closure),
+    /// A nested function with pre-applied arguments.
+    With(Func, Args),
 }
 
-type Func = dyn Fn(&mut EvalContext, &mut Args) -> TypResult<Value>;
+impl Func {
+    /// Create a new function from a native rust function.
+    pub fn native(
+        name: &'static str,
+        func: fn(&mut EvalContext, &mut Args) -> TypResult<Value>,
+    ) -> Self {
+        Self(Arc::new(Repr::Native(Native { name, func })))
+    }
 
-impl Function {
-    /// Create a new function from a rust closure.
-    pub fn new<F>(name: Option<EcoString>, func: F) -> Self
-    where
-        F: Fn(&mut EvalContext, &mut Args) -> TypResult<Value> + 'static,
-    {
-        Self(Arc::new(Inner { name, func }))
+    /// Create a new function from a closure.
+    pub fn closure(closure: Closure) -> Self {
+        Self(Arc::new(Repr::Closure(closure)))
     }
 
     /// The name of the function.
-    pub fn name(&self) -> Option<&EcoString> {
-        self.0.name.as_ref()
+    pub fn name(&self) -> Option<&str> {
+        match self.0.as_ref() {
+            Repr::Native(native) => Some(native.name),
+            Repr::Closure(closure) => closure.name.as_deref(),
+            Repr::With(func, _) => func.name(),
+        }
     }
 
     /// Call the function in the context with the arguments.
     pub fn call(&self, ctx: &mut EvalContext, args: &mut Args) -> TypResult<Value> {
-        (&self.0.func)(ctx, args)
+        match self.0.as_ref() {
+            Repr::Native(native) => (native.func)(ctx, args),
+            Repr::Closure(closure) => closure.call(ctx, args),
+            Repr::With(wrapped, applied) => {
+                args.items.splice(.. 0, applied.items.iter().cloned());
+                wrapped.call(ctx, args)
+            }
+        }
+    }
+
+    /// Apply the given arguments to the function.
+    pub fn with(self, args: Args) -> Self {
+        Self(Arc::new(Repr::With(self, args)))
     }
 }
 
-impl Debug for Function {
+impl Debug for Func {
     fn fmt(&self, f: &mut Formatter) -> fmt::Result {
         f.write_str("<function")?;
         if let Some(name) = self.name() {
@@ -49,13 +73,65 @@ impl Debug for Function {
     }
 }
 
-impl PartialEq for Function {
+impl PartialEq for Func {
     fn eq(&self, other: &Self) -> bool {
-        // We cast to thin pointers for comparison.
-        std::ptr::eq(
-            Arc::as_ptr(&self.0) as *const (),
-            Arc::as_ptr(&other.0) as *const (),
-        )
+        Arc::ptr_eq(&self.0, &other.0)
+    }
+}
+
+/// A native rust function.
+struct Native {
+    /// The name of the function.
+    pub name: &'static str,
+    /// The function pointer.
+    pub func: fn(&mut EvalContext, &mut Args) -> TypResult<Value>,
+}
+
+/// A user-defined closure.
+pub struct Closure {
+    /// The name of the closure.
+    pub name: Option<EcoString>,
+    /// Captured values from outer scopes.
+    pub captured: Scope,
+    /// The parameter names and default values. Parameters with default value
+    /// are named parameters.
+    pub params: Vec<(EcoString, Option<Value>)>,
+    /// The name of an argument sink where remaining arguments are placed.
+    pub sink: Option<EcoString>,
+    /// The expression the closure should evaluate to.
+    pub body: Expr,
+}
+
+impl Closure {
+    /// Call the function in the context with the arguments.
+    pub fn call(&self, ctx: &mut EvalContext, args: &mut Args) -> TypResult<Value> {
+        // Don't leak the scopes from the call site. Instead, we use the
+        // scope of captured variables we collected earlier.
+        let prev_scopes = std::mem::take(&mut ctx.scopes);
+        ctx.scopes.top = self.captured.clone();
+
+        // Parse the arguments according to the parameter list.
+        for (param, default) in &self.params {
+            ctx.scopes.def_mut(param, match default {
+                None => args.expect::<Value>(param)?,
+                Some(default) => {
+                    args.named::<Value>(param)?.unwrap_or_else(|| default.clone())
+                }
+            });
+        }
+
+        // Put the remaining arguments into the sink.
+        if let Some(sink) = &self.sink {
+            ctx.scopes.def_mut(sink, args.take());
+        }
+
+        // Evaluate the body.
+        let value = self.body.eval(ctx)?;
+
+        // Restore the call site scopes.
+        ctx.scopes = prev_scopes;
+
+        Ok(value)
     }
 }
 
diff --git a/src/eval/mod.rs b/src/eval/mod.rs
index a453a357..ae680d95 100644
--- a/src/eval/mod.rs
+++ b/src/eval/mod.rs
@@ -10,7 +10,7 @@ mod value;
 mod styles;
 mod capture;
 mod class;
-mod function;
+mod func;
 mod node;
 mod ops;
 mod scope;
@@ -19,7 +19,7 @@ pub use array::*;
 pub use capture::*;
 pub use class::*;
 pub use dict::*;
-pub use function::*;
+pub use func::*;
 pub use node::*;
 pub use scope::*;
 pub use styles::*;
@@ -602,9 +602,9 @@ impl Eval for CallExpr {
 
             Value::Class(class) => {
                 let point = || Tracepoint::Call(Some(class.name().to_string()));
-                let node = class.construct(ctx, &mut args).trace(point, self.span())?;
+                let value = class.construct(ctx, &mut args).trace(point, self.span())?;
                 args.finish()?;
-                Ok(Value::Node(node))
+                Ok(value)
             }
 
             v => bail!(
@@ -669,6 +669,9 @@ impl Eval for ClosureExpr {
     type Output = Value;
 
     fn eval(&self, ctx: &mut EvalContext) -> TypResult<Self::Output> {
+        // The closure's name is defined by its let binding if there's one.
+        let name = self.name().map(Ident::take);
+
         // Collect captured variables.
         let captured = {
             let mut visitor = CapturesVisitor::new(&ctx.scopes);
@@ -676,8 +679,8 @@ impl Eval for ClosureExpr {
             visitor.finish()
         };
 
-        let mut sink = None;
         let mut params = Vec::new();
+        let mut sink = None;
 
         // Collect parameters and an optional sink parameter.
         for param in self.params() {
@@ -697,39 +700,14 @@ impl Eval for ClosureExpr {
             }
         }
 
-        // Clone the body expression so that we don't have a lifetime
-        // dependence on the AST.
-        let name = self.name().map(Ident::take);
-        let body = self.body();
-
         // Define the actual function.
-        let func = Function::new(name, move |ctx, args| {
-            // Don't leak the scopes from the call site. Instead, we use the
-            // scope of captured variables we collected earlier.
-            let prev_scopes = mem::take(&mut ctx.scopes);
-            ctx.scopes.top = captured.clone();
-
-            // Parse the arguments according to the parameter list.
-            for (param, default) in &params {
-                ctx.scopes.def_mut(param, match default {
-                    None => args.expect::<Value>(param)?,
-                    Some(default) => {
-                        args.named::<Value>(param)?.unwrap_or_else(|| default.clone())
-                    }
-                });
-            }
-
-            // Put the remaining arguments into the sink.
-            if let Some(sink) = &sink {
-                ctx.scopes.def_mut(sink, args.take());
-            }
-
-            let value = body.eval(ctx)?;
-            ctx.scopes = prev_scopes;
-            Ok(value)
-        });
-
-        Ok(Value::Func(func))
+        Ok(Value::Func(Func::closure(Closure {
+            name,
+            captured,
+            params,
+            sink,
+            body: self.body(),
+        })))
     }
 }
 
@@ -738,16 +716,9 @@ impl Eval for WithExpr {
 
     fn eval(&self, ctx: &mut EvalContext) -> TypResult<Self::Output> {
         let callee = self.callee();
-        let wrapped = callee.eval(ctx)?.cast::<Function>().at(callee.span())?;
-        let applied = self.args().eval(ctx)?;
-
-        let name = wrapped.name().cloned();
-        let func = Function::new(name, move |ctx, args| {
-            args.items.splice(.. 0, applied.items.iter().cloned());
-            wrapped.call(ctx, args)
-        });
-
-        Ok(Value::Func(func))
+        let func = callee.eval(ctx)?.cast::<Func>().at(callee.span())?;
+        let args = self.args().eval(ctx)?;
+        Ok(Value::Func(func.with(args)))
     }
 }
 
diff --git a/src/eval/scope.rs b/src/eval/scope.rs
index 34da68d4..3e79afc1 100644
--- a/src/eval/scope.rs
+++ b/src/eval/scope.rs
@@ -4,7 +4,7 @@ use std::fmt::{self, Debug, Formatter};
 use std::iter;
 use std::sync::Arc;
 
-use super::{Args, Class, Construct, EvalContext, Function, Set, Value};
+use super::{Args, Class, Construct, EvalContext, Func, Set, Value};
 use crate::diag::TypResult;
 use crate::util::EcoString;
 
@@ -98,22 +98,21 @@ impl Scope {
         self.values.insert(var.into(), slot);
     }
 
-    /// Define a constant function.
-    pub fn def_func<F>(&mut self, name: &str, f: F)
-    where
-        F: Fn(&mut EvalContext, &mut Args) -> TypResult<Value> + 'static,
-    {
-        let name = EcoString::from(name);
-        self.def_const(name.clone(), Function::new(Some(name), f));
+    /// Define a constant native function.
+    pub fn def_func(
+        &mut self,
+        name: &'static str,
+        func: fn(&mut EvalContext, &mut Args) -> TypResult<Value>,
+    ) {
+        self.def_const(name, Func::native(name, func));
     }
 
     /// Define a constant class.
-    pub fn def_class<T>(&mut self, name: &str)
+    pub fn def_class<T>(&mut self, name: &'static str)
     where
         T: Construct + Set + 'static,
     {
-        let name = EcoString::from(name);
-        self.def_const(name.clone(), Class::new::<T>(name));
+        self.def_const(name, Class::new::<T>(name));
     }
 
     /// Look up the value of a variable.
diff --git a/src/eval/value.rs b/src/eval/value.rs
index 7d65d5af..c19794df 100644
--- a/src/eval/value.rs
+++ b/src/eval/value.rs
@@ -4,7 +4,7 @@ use std::fmt::{self, Debug, Formatter};
 use std::hash::Hash;
 use std::sync::Arc;
 
-use super::{ops, Args, Array, Class, Dict, Function, Node};
+use super::{ops, Args, Array, Class, Dict, Func, Node};
 use crate::diag::StrResult;
 use crate::geom::{Angle, Color, Fractional, Length, Linear, Relative, RgbaColor};
 use crate::layout::Layout;
@@ -45,7 +45,7 @@ pub enum Value {
     /// A node value: `[*Hi* there]`.
     Node(Node),
     /// An executable function.
-    Func(Function),
+    Func(Func),
     /// Captured arguments to a function.
     Args(Args),
     /// A class of nodes.
@@ -89,7 +89,7 @@ impl Value {
             Self::Array(_) => Array::TYPE_NAME,
             Self::Dict(_) => Dict::TYPE_NAME,
             Self::Node(_) => Node::TYPE_NAME,
-            Self::Func(_) => Function::TYPE_NAME,
+            Self::Func(_) => Func::TYPE_NAME,
             Self::Args(_) => Args::TYPE_NAME,
             Self::Class(_) => Class::TYPE_NAME,
             Self::Dyn(v) => v.type_name(),
@@ -401,13 +401,7 @@ primitive! { EcoString: "string", Str }
 primitive! { Array: "array", Array }
 primitive! { Dict: "dictionary", Dict }
 primitive! { Node: "template", Node }
-primitive! { Function: "function",
-    Func,
-    Class(v) => Function::new(
-        Some(v.name().clone()),
-        move |ctx, args| v.construct(ctx, args).map(Value::Node)
-    )
-}
+primitive! { Func: "function", Func, Class(v) => v.constructor() }
 primitive! { Args: "arguments", Args }
 primitive! { Class: "class", Class }
 
@@ -554,9 +548,8 @@ mod tests {
         test(dict!["two" => false, "one" => 1], "(one: 1, two: false)");
 
         // Functions.
-        test(Function::new(None, |_, _| Ok(Value::None)), "<function>");
         test(
-            Function::new(Some("nil".into()), |_, _| Ok(Value::None)),
+            Func::native("nil", |_, _| Ok(Value::None)),
             "<function nil>",
         );