variadic functions

3 files changed, 26 insertions, 8 deletions

diff --git a/interpreter.py b/interpreter.py
index 1911c05..f5a8562 100644
--- a/interpreter.py
+++ b/interpreter.py
@@ -47,7 +47,7 @@ class Function:
             for arg in self.args:
                 out.append(f"{arg}")
         if self.many is not None:
-            out.append(f"*{self.many}")
+            out.append(f"{self.many}")
         return " ".join(out) + ")"
 
     def arity_check(self, symbol, params):
@@ -102,26 +102,36 @@ class UserFunction(Function):
     
     def __init__(self, name, params, body):
         # TODO this doesn't do type checking, or optional, or lazy
-        newparams, args = self.process_params(name, params)
-        super().__init__(name, newparams, body, args)
+        newparams, args, many = self.process_params(name, params)
+        super().__init__(name, newparams, body, args, many)
 
     def process_params(self, name, params):
         newparams = []
         args = []
+        many = None
         prev_type = False
         first = True
         for param in params:
             if isinstance(param, Symbol):
-                newparams.append(param)
-                args.append(Arg(param.name, TypeEnum.ANY))
+                if many is not None:
+                    raise NebPanic("& must be last argument")
+                if param.name == "&":
+                    many = Arg(param.name, TypeEnum.ANY)
+                else:
+                    newparams.append(param)
+                    args.append(Arg(param.name, TypeEnum.ANY))
                 prev_type = False
             elif isinstance(param, Type) and not prev_type and not first:
-                args[-1].type_ = TypeEnum.__getattr__(param.name[1:].upper())
+                typ = TypeEnum.__getattr__(param.name[1:].upper())
+                if many is None:
+                    args[-1].type_ = typ
+                else:
+                    many.type_ = typ
                 prev_type = True
             else:
                 raise NebPanic("invalid :func signature", param)
             first = False
-        return newparams, args
+        return newparams, args, many
 
     def call(self, expr, env, ns):
         self.arity_check(expr.args[0], expr.args[1:])
@@ -129,6 +139,11 @@ class UserFunction(Function):
         this_env = Environment(env)
         for idx, param in enumerate(self.params):
             this_env.register(param.name, evaluated_args[idx])
+
+        # if we got "many", wrap the rest in a list
+        if self.many:
+            this_env.register(self.many.name, List(evaluated_args[len(self.params):], True))
+
         return interpret(self.body, env=this_env, ns=ns)
 
 class Environment:
diff --git a/lexer.py b/lexer.py
index ea4fefb..bce3f51 100644
--- a/lexer.py
+++ b/lexer.py
@@ -34,7 +34,8 @@ keywords = {
     "for-count": TokenType.FOR_COUNT,
     "|": TokenType.PIPE,
     "def": TokenType.DEF,
-    "lambda": TokenType.LAMBDA }
+    "lambda": TokenType.LAMBDA,
+    "&": TokenType.MANY }
 
 
 WHITESPACE = [" ", "\n", "\t"]
diff --git a/structs.py b/structs.py
index adfd0e4..02693be 100644
--- a/structs.py
+++ b/structs.py
@@ -62,6 +62,8 @@ class TokenType(Enum):
     LITERAL_TYPE = auto()
     BOOL_TYPE = auto()
 
+    MANY = auto()
+
 @dataclass
 class Token:
     type_: TokenType