Fix function sinks (typst#638)

Author: Marmare314

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

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),
+                        _ => {}
                     }
                 }
 

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(),
         };
 

src/syntax/ast.rs

@@ -1570,15 +1570,15 @@ pub enum Param {
     /// A named parameter with a default value: `draw: false`.
     Named(Named),
     /// An argument sink: `..args`.
-    Sink(Ident),
+    Sink(Option<Ident>),
 }
 
 impl AstNode for Param {
     fn from_untyped(node: &SyntaxNode) -> Option<Self> {
         match node.kind() {
             SyntaxKind::Ident => node.cast().map(Self::Pos),
             SyntaxKind::Named => node.cast().map(Self::Named),
-            SyntaxKind::Spread => node.cast_first_match().map(Self::Sink),
+            SyntaxKind::Spread => Some(Self::Sink(node.cast_first_match())),
             _ => Option::None,
         }
     }
@@ -1587,7 +1587,7 @@ impl AstNode for Param {
         match self {
             Self::Pos(v) => v.as_untyped(),
             Self::Named(v) => v.as_untyped(),
-            Self::Sink(v) => v.as_untyped(),
+            Self::Sink(_) => self.as_untyped(),
         }
     }
 }

src/syntax/parser.rs

@@ -1069,6 +1069,7 @@ fn validate_dict(p: &mut Parser, m: Marker) {
 }
 
 fn validate_params(p: &mut Parser, m: Marker) {
+    let mut used_spread = false;
     let mut used = HashSet::new();
     for child in p.post_process(m) {
         match child.kind() {
@@ -1086,12 +1087,24 @@ fn validate_params(p: &mut Parser, m: Marker) {
             }
             SyntaxKind::Spread => {
                 let Some(within) = child.children_mut().last_mut() else { continue };
-                if within.kind() != SyntaxKind::Ident {
+                if used_spread {
+                    child.convert_to_error("only one argument sink is allowed");
+                    continue;
+                }
+                used_spread = true;
+                if within.kind() == SyntaxKind::Dots {
+                    continue;
+                } else if within.kind() != SyntaxKind::Ident {
                     within.convert_to_error(eco_format!(
                         "expected identifier, found {}",
                         within.kind().name(),
                     ));
                     child.make_erroneous();
+                    continue;
+                }
+                if !used.insert(within.text().clone()) {
+                    within.convert_to_error("duplicate parameter");
+                    child.make_erroneous();
                 }
             }
             SyntaxKind::LeftParen | SyntaxKind::RightParen | SyntaxKind::Comma => {}

tests/typ/compiler/closure.typ

@@ -155,6 +155,10 @@
 // Error: 35-36 duplicate parameter
 #let f(a, b, a: none, b: none, c, b) = none
 
+---
+// Error: 13-14 duplicate parameter
+#let f(a, ..a) = none
+
 ---
 // Error: 7-17 expected identifier, named pair or argument sink, found keyed pair
 #((a, "named": b) => none)

tests/typ/compiler/spread.typ

@@ -27,7 +27,7 @@
 #{
   let save(..args) = {
     test(type(args), "arguments")
-    test(repr(args), "(1, 2, three: true)")
+    test(repr(args), "(three: true, 1, 2)")
   }
 
   save(1, 2, three: true)
@@ -55,6 +55,11 @@
 #f(..if false {})
 #f(..for x in () [])
 
+---
+// unnamed spread
+#let f(.., a) = a
+#test(f(1, 2, 3), 3)
+
 ---
 // Error: 13-19 cannot spread string
 #calc.min(.."nope")
@@ -64,7 +69,7 @@
 #let f(..true) = none
 
 ---
-// Error: 15-16 only one argument sink is allowed
+// Error: 13-16 only one argument sink is allowed
 #let f(..a, ..b) = none
 
 ---
@@ -91,3 +96,28 @@
 ---
 // Error: 5-11 cannot spread array into dictionary
 #(..(1, 2), a: 1)
+
+---
+// Spread at beginning.
+#{
+  let f(..a, b) = (a, b)
+  test(repr(f(1)), "((), 1)")
+  test(repr(f(1, 2, 3)), "((1, 2), 3)")
+  test(repr(f(1, 2, 3, 4, 5)), "((1, 2, 3, 4), 5)")
+}
+
+---
+// Spread in the middle.
+#{
+  let f(a, ..b, c) = (a, b, c)
+  test(repr(f(1, 2)), "(1, (), 2)")
+  test(repr(f(1, 2, 3, 4, 5)), "(1, (2, 3, 4), 5)")
+}
+
+---
+#{
+  let f(..a, b, c, d) = none
+
+  // Error: 4-10 missing argument: d
+  f(1, 2)
+}