Recur correctly handles variadic arg lists (#116)

* Recur is only fixed to the current arity arglist and correctly handles variadic args

* Add recur test

* Test trampoline arguments

* Fix lint error

* More trampoline test cases

Author: chrisrink10

basilisp/compiler.py

@@ -130,10 +130,11 @@ def new_frame(self, name):
 
 
 class RecurPoint:
-    __slots__ = ('name', 'has_recur')
+    __slots__ = ('name', 'args', 'has_recur')
 
-    def __init__(self, name: str) -> None:
+    def __init__(self, name: str, args: vec.Vector) -> None:
         self.name = name
+        self.args = args
         self.has_recur = False
 
 
@@ -159,8 +160,8 @@ def recur_point(self):
         return self._recur_points[-1]
 
     @contextlib.contextmanager
-    def new_recur_point(self, name: str):
-        self._recur_points.append(RecurPoint(name))
+    def new_recur_point(self, name: str, args: vec.Vector):
+        self._recur_points.append(RecurPoint(name, args))
         yield
         self._recur_points.pop()
 
@@ -715,7 +716,7 @@ def _compose_ifs(if_stmts: List[Dict[str, ast.AST]], orelse: List[ast.AST] = Non
 
 def _single_arity_fn_ast(ctx: CompilerContext, name: str, fndef: llist.List) -> ASTStream:
     """Generate Python AST nodes for a single-arity function."""
-    with ctx.new_symbol_table(name):
+    with ctx.new_symbol_table(name), ctx.new_recur_point(name, fndef.first):
         args, body, vargs = _fn_args_body(ctx, fndef.first, fndef.rest)
 
         yield _dependency(_expressionize(body, name, args=args, vargs=vargs))
@@ -771,13 +772,26 @@ def __f_68(*multi_arity_args):
     has_rest = False
 
     for arg_count, is_rest, arity in arities:
-        has_rest = any([has_rest, is_rest])
-        arity_name = f"{name}__arity{'_rest' if is_rest else arg_count}"
-
-        with ctx.new_symbol_table(arity_name):
-            args, body, vargs = _fn_args_body(ctx, arity.first, arity.rest)
-
-            yield _dependency(_expressionize(body, arity_name, args=args, vargs=vargs))
+        with ctx.new_recur_point(name, arity.first):
+            has_rest = any([has_rest, is_rest])
+            arity_name = f"{name}__arity{'_rest' if is_rest else arg_count}"
+
+            with ctx.new_symbol_table(arity_name):
+                # Generate the arity function
+                args, body, vargs = _fn_args_body(ctx, arity.first, arity.rest)
+                yield _dependency(_expressionize(body, arity_name, args=args, vargs=vargs))
+
+            # If a recur point was established, we generate a trampoline version of the
+            # generated function to allow repeated recursive calls without blowing up the
+            # stack size.
+            if ctx.recur_point.has_recur:
+                yield _dependency(ast.Assign(targets=[ast.Name(id=arity_name, ctx=ast.Store())],
+                                             value=ast.Call(func=_TRAMPOLINE_FN_NAME,
+                                                            args=[
+                                                                ast.Name(id=arity_name, ctx=ast.Load())],
+                                                            keywords=[])))
+
+            # Generate an if-statement branch for the arity-dispatch function
             compare_op = ast.GtE() if is_rest else ast.Eq()
             if_stmts.append(
                 {"test": ast.Compare(left=ast.Call(func=_load_attr('len'),
@@ -808,15 +822,7 @@ def __f_68(*multi_arity_args):
         decorator_list=[],
         returns=None))
 
-    # If a recur point was established, we generate a trampoline version of the
-    # generated function to allow repeated recursive calls without blowing up the
-    # stack size.
-    if ctx.recur_point.has_recur:
-        yield _node(ast.Call(func=_TRAMPOLINE_FN_NAME,
-                             args=[ast.Name(id=ctx.recur_point.name, ctx=ast.Load())],
-                             keywords=[]))
-    else:
-        yield _node(ast.Name(id=name, ctx=ast.Load()))
+    yield _node(ast.Name(id=name, ctx=ast.Load()))
 
 
 def _fn_ast(ctx: CompilerContext, form: llist.List) -> ASTStream:
@@ -825,18 +831,17 @@ def _fn_ast(ctx: CompilerContext, form: llist.List) -> ASTStream:
     has_name = isinstance(form[1], sym.Symbol)
     name = genname("__" + (munge(form[1].name) if has_name else _FN_PREFIX))
 
-    with ctx.new_recur_point(name):
-        rest_idx = 1 + int(has_name)
-        arities = list(_fn_arities(ctx, form[rest_idx:]))
-        if len(arities) == 0:
-            raise CompilerException("Function def must have argument vector")
-        elif len(arities) == 1:
-            _, _, fndef = arities[0]
-            yield from _single_arity_fn_ast(ctx, name, fndef)
-            return
-        else:
-            yield from _multi_arity_fn_ast(ctx, name, arities)
-            return
+    rest_idx = 1 + int(has_name)
+    arities = list(_fn_arities(ctx, form[rest_idx:]))
+    if len(arities) == 0:
+        raise CompilerException("Function def must have argument vector")
+    elif len(arities) == 1:
+        _, _, fndef = arities[0]
+        yield from _single_arity_fn_ast(ctx, name, fndef)
+        return
+    else:
+        yield from _multi_arity_fn_ast(ctx, name, arities)
+        return
 
 
 def _if_ast(ctx: CompilerContext, form: llist.List) -> ASTStream:
@@ -1080,8 +1085,10 @@ def _recur_ast(ctx: CompilerContext, form: llist.List) -> ASTStream:
         expr_deps, exprs = _collection_literal_ast(ctx, form.rest)
         yield from expr_deps
 
+        has_vargs = any([s == _AMPERSAND for s in ctx.recur_point.args])
         yield _node(ast.Call(func=_TRAMPOLINE_ARGS_FN_NAME,
-                             args=list(_unwrap_nodes(exprs)),
+                             args=list(itertools.chain([ast.NameConstant(has_vargs)],
+                                                       _unwrap_nodes(exprs))),
                              keywords=[]))
     except IndexError:
         raise CompilerException("Attempting to recur without recur point") from None

basilisp/core/__init__.lpy

@@ -66,6 +66,18 @@
   (fn string? [o]
     (instance? builtins/str o)))
 
+(def
+  ^{:doc "Return true if obj is a symbol."}
+  symbol?
+  (fn symbol? [o]
+    (instance? basilisp.lang.symbol/Symbol o)))
+
+(def
+  ^{:doc "Return true if obj is a keyword."}
+  keyword?
+  (fn keyword? [o]
+    (instance? basilisp.lang.keyword/Keyword o)))
+
 (def
   ^{:doc "Return true if o is a list."}
   list?
@@ -134,25 +146,35 @@
 
 (def
   ^{:macro true
-    :doc   "Define a new function."}
+    :doc   "Define a new function with an optional docstring."}
   defn
   (fn defn [&form name & body]
-    (let [body (concat body)
-          doc  (if (string? (first body))
-                 (first body)
-                 nil)
-          fname (if doc
-                  (with-meta name {:doc doc})
-                  name)
-          body (if doc
-                 (rest body)
-                 body)
-          args (if (vector? (first body))
-                 (first body)
-                 nil)  ;; Should throw here!
-          body (rest body)]
+    (if (symbol? name)
+      nil  ;; Do nothing!
+      (throw (ex-info "First argument to defn must be a symbol"
+                      {:found name :type (builtins/type name)})))
+    (let [body   (concat body)
+          doc    (if (string? (first body))
+                   (first body)
+                   nil)
+          fname  (if doc
+                   (with-meta name {:doc doc})
+                   name)
+          body   (if doc
+                   (rest body)
+                   body)
+          multi? (list? (first body))
+          body   (if multi?
+                   body
+                   (cons
+                    (if (vector? (first body))
+                      (first body)
+                      (throw
+                       (ex-info "Expected an argument vector"
+                                {:found (first body)})))
+                    (rest body)))]
       `(def ~fname
-         (fn* ~fname ~args
+         (fn* ~fname
               ~@body)))))
 
 (defn nth
@@ -185,3 +207,12 @@
   (if (seq (rest s))
     (recur (rest s))
     (first s)))
+
+(defn +
+  "Sum the arguments together. If no arguments given, returns 0."
+  ([] 0)
+  ([x] x)
+  ([x & args]
+   (if (seq (rest args))
+     (recur (operator/add x (first args)) (rest args))
+     (operator/add x (first args)))))

basilisp/lang/runtime.py

@@ -2,7 +2,7 @@
 import itertools
 import threading
 import types
-from typing import Optional, List, Dict
+from typing import Optional, Dict, Tuple
 
 from functional import seq
 from pyrsistent import pmap, PMap, PSet, pset
@@ -506,15 +506,29 @@ def _collect_args(args) -> lseq.Seq:
 
 
 class _TrampolineArgs:
-    __slots__ = ('_args', '_kwargs')
+    __slots__ = ('_has_varargs', '_args', '_kwargs')
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, has_varargs: bool, *args, **kwargs) -> None:
+        self._has_varargs = has_varargs
         self._args = args
         self._kwargs = kwargs
 
     @property
-    def args(self) -> List:
-        return self._args
+    def args(self) -> Tuple:
+        """Return the arguments for a trampolined function. If the function
+        that is being trampolined has varargs, unroll the final argument if
+        it is a sequence."""
+        if not self._has_varargs:
+            return self._args
+
+        try:
+            final = self._args[-1]
+            if isinstance(final, lseq.Seq):
+                inits = self._args[:-1]
+                return tuple(itertools.chain(inits, final))
+            return self._args
+        except IndexError:
+            return ()
 
     @property
     def kwargs(self) -> Dict:

tests/compiler_test.py

@@ -414,6 +414,45 @@ def test_recur(ns_var: Var):
     assert "3ba" == lcompile("(rev-str \"a\" :b 3)")
 
 
+def test_recur_arity(ns_var: Var):
+    # Single arity function
+    code = """
+    (def ++
+      (fn ++ [x & args]
+        (if (seq (rest args))
+          (recur (operator/add x (first args)) (rest args))
+          (operator/add x (first args)))))
+    """
+
+    lcompile(code)
+
+    assert 3 == lcompile("(++ 1 2)")
+    assert 6 == lcompile("(++ 1 2 3)")
+    assert 10 == lcompile("(++ 1 2 3 4)")
+    assert 15 == lcompile("(++ 1 2 3 4 5)")
+
+    # Multi-arity function
+    code = """
+    (def +++
+      (fn +++ 
+        ([] 0)
+        ([x] x)
+        ([x & args]
+          (if (seq (rest args))
+            (recur (operator/add x (first args)) (rest args))
+            (operator/add x (first args))))))
+    """
+
+    lcompile(code)
+
+    assert 0 == lcompile("(+++)")
+    assert 1 == lcompile("(+++ 1)")
+    assert 3 == lcompile("(+++ 1 2)")
+    assert 6 == lcompile("(+++ 1 2 3)")
+    assert 10 == lcompile("(+++ 1 2 3 4)")
+    assert 15 == lcompile("(+++ 1 2 3 4 5)")
+
+
 def test_disallow_recur_in_special_forms(ns_var: Var):
     with pytest.raises(compiler.CompilerException):
         lcompile("(fn [a] (def b (recur \"a\")))")

tests/runtime_test.py

@@ -201,3 +201,23 @@ def test_conj():
 
     with pytest.raises(TypeError):
         runtime.conj("b", 1, "a")
+
+
+def test_trampoline_args():
+    args = runtime._TrampolineArgs(True)
+    assert () == args.args
+
+    args = runtime._TrampolineArgs(False, llist.l(2, 3, 4))
+    assert (llist.l(2, 3, 4),) == args.args
+
+    args = runtime._TrampolineArgs(True, llist.l(2, 3, 4))
+    assert (2, 3, 4) == args.args
+
+    args = runtime._TrampolineArgs(False, 1, 2, 3, llist.l(4, 5, 6))
+    assert (1, 2, 3, llist.l(4, 5, 6)) == args.args
+
+    args = runtime._TrampolineArgs(True, 1, 2, 3, llist.l(4, 5, 6))
+    assert (1, 2, 3, 4, 5, 6) == args.args
+
+    args = runtime._TrampolineArgs(True, 1, llist.l(2, 3, 4), 5, 6)
+    assert (1, llist.l(2, 3, 4), 5, 6) == args.args