From 0105eb7382801b56781308ea94b3aeffa6fd867f Mon Sep 17 00:00:00 2001
From: Marmare314 <49279081+Marmare314@users.noreply.github.com>
Date: Thu, 13 Apr 2023 16:07:58 +0200
Subject: Fix function sinks (#638)

---
 src/eval/args.rs | 11 ++++++++++
 src/eval/func.rs | 61 +++++++++++++++++++++++++++++++++++++++-----------------
 src/eval/mod.rs  | 16 ++++-----------
 3 files changed, 58 insertions(+), 30 deletions(-)

(limited to 'src/eval')

diff --git a/src/eval/args.rs b/src/eval/args.rs
index 129da463..66ca4a85 100644
--- a/src/eval/args.rs
+++ b/src/eval/args.rs
@@ -61,6 +61,17 @@ impl Args {
         Ok(None)
     }
 
+    /// Consume n positional arguments if possible.
+    pub fn consume(&mut self, n: usize) -> SourceResult<EcoVec<Arg>> {
+        if n > self.items.len() {
+            bail!(self.span, "not enough arguments");
+        }
+        let vec = self.items.to_vec();
+        let (left, right) = vec.split_at(n);
+        self.items = right.into();
+        return Ok(left.into());
+    }
+
     /// Consume and cast the first positional argument.
     ///
     /// Returns a `missing argument: {what}` error if no positional argument is
diff --git a/src/eval/func.rs b/src/eval/func.rs
index d64ca732..582eabd4 100644
--- a/src/eval/func.rs
+++ b/src/eval/func.rs
@@ -260,15 +260,22 @@ pub(super) struct Closure {
     pub name: Option<Ident>,
     /// Captured values from outer scopes.
     pub captured: Scope,
-    /// The parameter names and default values. Parameters with default value
-    /// are named parameters.
-    pub params: Vec<(Ident, Option<Value>)>,
-    /// The name of an argument sink where remaining arguments are placed.
-    pub sink: Option<Ident>,
+    /// The list of parameters.
+    pub params: Vec<Param>,
     /// The expression the closure should evaluate to.
     pub body: Expr,
 }
 
+#[derive(Hash)]
+pub enum Param {
+    /// A positional parameter: `x`.
+    Pos(Ident),
+    /// A named parameter with a default value: `draw: false`.
+    Named(Ident, Value),
+    /// An argument sink: `..args`.
+    Sink(Option<Ident>),
+}
+
 impl Closure {
     /// Call the function in the context with the arguments.
     #[allow(clippy::too_many_arguments)]
@@ -304,21 +311,38 @@ impl Closure {
         }
 
         // Parse the arguments according to the parameter list.
-        for (param, default) in &closure.params {
-            vm.define(
-                param.clone(),
-                match default {
-                    Some(default) => {
-                        args.named::<Value>(param)?.unwrap_or_else(|| default.clone())
+        let num_pos_params =
+            closure.params.iter().filter(|p| matches!(p, Param::Pos(_))).count();
+        let num_pos_args = args.to_pos().len() as usize;
+        let sink_size = num_pos_args.checked_sub(num_pos_params);
+
+        let mut sink = None;
+        let mut sink_pos_values = None;
+        for p in &closure.params {
+            match p {
+                Param::Pos(ident) => {
+                    vm.define(ident.clone(), args.expect::<Value>(ident)?);
+                }
+                Param::Sink(ident) => {
+                    sink = ident.clone();
+                    if let Some(sink_size) = sink_size {
+                        sink_pos_values = Some(args.consume(sink_size)?);
                     }
-                    None => args.expect::<Value>(param)?,
-                },
-            );
+                }
+                Param::Named(ident, default) => {
+                    let value =
+                        args.named::<Value>(ident)?.unwrap_or_else(|| default.clone());
+                    vm.define(ident.clone(), value);
+                }
+            }
         }
 
-        // Put the remaining arguments into the sink.
-        if let Some(sink) = &closure.sink {
-            vm.define(sink.clone(), args.take());
+        if let Some(sink) = sink {
+            let mut remaining_args = args.take();
+            if let Some(sink_pos_values) = sink_pos_values {
+                remaining_args.items.extend(sink_pos_values);
+            }
+            vm.define(sink, remaining_args);
         }
 
         // Ensure all arguments have been used.
@@ -407,7 +431,8 @@ impl<'a> CapturesVisitor<'a> {
                     match param {
                         ast::Param::Pos(ident) => self.bind(ident),
                         ast::Param::Named(named) => self.bind(named.name()),
-                        ast::Param::Sink(ident) => self.bind(ident),
+                        ast::Param::Sink(Some(ident)) => self.bind(ident),
+                        _ => {}
                     }
                 }
 
diff --git a/src/eval/mod.rs b/src/eval/mod.rs
index 5a450481..1b3c6ea3 100644
--- a/src/eval/mod.rs
+++ b/src/eval/mod.rs
@@ -1149,24 +1149,17 @@ impl Eval for ast::Closure {
             visitor.finish()
         };
 
-        let mut params = Vec::new();
-        let mut sink = None;
-
         // Collect parameters and an optional sink parameter.
+        let mut params = Vec::new();
         for param in self.params().children() {
             match param {
                 ast::Param::Pos(name) => {
-                    params.push((name, None));
+                    params.push(Param::Pos(name));
                 }
                 ast::Param::Named(named) => {
-                    params.push((named.name(), Some(named.expr().eval(vm)?)));
-                }
-                ast::Param::Sink(name) => {
-                    if sink.is_some() {
-                        bail!(name.span(), "only one argument sink is allowed");
-                    }
-                    sink = Some(name);
+                    params.push(Param::Named(named.name(), named.expr().eval(vm)?));
                 }
+                ast::Param::Sink(name) => params.push(Param::Sink(name)),
             }
         }
 
@@ -1176,7 +1169,6 @@ impl Eval for ast::Closure {
             name,
             captured,
             params,
-            sink,
             body: self.body(),
         };
 
-- 
cgit v1.2.3