Support functions of multiple arities in compiler (#80)

* Support functions of multiple arities in compiler

* Refactoring function AST code slightly

* Test new function arities

* Remove leading space

* Remove unused function def

Author: chrisrink10

basilisp/compiler.py

@@ -40,6 +40,7 @@
 _DO_PREFIX = 'lisp_do'
 _FN_PREFIX = 'lisp_fn'
 _IF_PREFIX = 'lisp_if'
+_MULTI_ARITY_ARG_NAME = 'multi_arity_args'
 _THROW_PREFIX = 'lisp_throw'
 _TRY_PREFIX = 'lisp_try'
 _NS_VAR = '__NS'
@@ -479,9 +480,11 @@ def _do_ast(ctx: CompilerContext, form: llist.List) -> ASTStream:
         func=ast.Name(id=do_fn_name, ctx=ast.Load()), args=[], keywords=[]))
 
 
+FunctionDefDetails = Tuple[List[ast.arg], ASTStream, Optional[ast.arg]]
+
+
 def _fn_args_body(ctx: CompilerContext, arg_vec: vec.Vector,
-                  body_exprs: llist.List
-                  ) -> Tuple[List[ast.arg], ASTStream, Optional[ast.arg]]:
+                  body_exprs: llist.List) -> FunctionDefDetails:
     """Generate the Python AST Nodes for a Lisp function argument vector
     and body expressions. Return a tuple of arg nodes and body AST nodes."""
     st = ctx.symbol_table
@@ -522,22 +525,189 @@ def _fn_args_body(ctx: CompilerContext, arg_vec: vec.Vector,
     return args, cast(ASTStream, body), vargs
 
 
-def _fn_ast(ctx: CompilerContext, form: llist.List) -> ASTStream:
-    """Generate a Python AST Node for an anonymous function."""
-    assert form.first == _FN
-    has_name = isinstance(form[1], sym.Symbol)
-    name = genname("__" + (munge(form[1].name) if has_name else _FN_PREFIX))
+FunctionArityDetails = Tuple[int, bool, llist.List]
+
+
+def _fn_arities(form: llist.List) -> Iterable[FunctionArityDetails]:
+    """Return the arities of a function definition and some additional details about
+    the argument vector. Verify that all arities are compatible. In particular, this
+    function will throw a CompilerException if any of the following are true:
+     - two function definitions have the same number of arguments
+     - two function definitions have a rest parameter
+     - any function definition has the same number of arguments as a definition
+       with a rest parameter
+
+    Given a function such as this:
+
+        (fn a
+          ([] :a)
+          ([a] a))
+
+    Returns a generator yielding: '(([] :a) ([a] a))
+
+    Single arity functions yield the rest:
+
+        (fn a [] :a) ;=> '(([] :a))"""
+    if not all(map(lambda f: isinstance(f, llist.List) and isinstance(f.first, vec.Vector), form)):
+        assert isinstance(form.first, vec.Vector)
+        yield len(form.first), False, form
+        return
 
-    arg_idx = 1 + int(has_name)
-    body_idx = 2 + int(has_name)
+    arg_counts: Dict[int, llist.List] = {}
+    has_vargs = False
+    vargs_len = None
+    for arity in form:
+        # Verify each arity is unique
+        arg_count = len(arity.first)
+        if arg_count in arg_counts:
+            raise CompilerException("Each arity in multi-arity fn must be unique",
+                                    [arity, arg_counts[arg_count]])
+
+        # Verify that only one arity contains a rest-param
+        is_rest = False
+        for arg in arity.first:
+            if arg == _AMPERSAND:
+                if has_vargs:
+                    raise CompilerException("Only one arity in multi-arity fn may have rest param")
+                is_rest = True
+                has_vargs = True
+                arg_count -= 1
+                vargs_len = arg_count
+
+        # Verify that arities do not exceed rest-param arity
+        if vargs_len is not None and any([c >= vargs_len for c in arg_counts.keys()]):
+            raise CompilerException("No arity in multi-arity fn may exceed the rest param arity")
+
+        # Put this in last so it does not conflict with the above checks
+        arg_counts[arg_count] = arity
+
+        yield arg_count, is_rest, arity
+
+
+def _compose_ifs(if_stmts: List[Dict[str, ast.AST]], orelse: List[ast.AST] = None) -> ast.If:
+    """Compose a series of If statements into nested elifs, with
+    an optional terminating else."""
+    first = if_stmts[0]
+    try:
+        rest = if_stmts[1:]
+        return ast.If(test=first["test"],
+                      body=[first["body"]],
+                      orelse=[_compose_ifs(rest, orelse=orelse)])
+    except IndexError:
+        return ast.If(test=first["test"],
+                      body=[first["body"]],
+                      orelse=Maybe(orelse).or_else_get([]))
 
-    assert isinstance(form[arg_idx], vec.Vector)
 
+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):
-        args, body, vargs = _fn_args_body(ctx, form[arg_idx], form[body_idx:])
+        args, body, vargs = _fn_args_body(ctx, fndef.first, fndef.rest)
 
         yield _dependency(_expressionize(body, name, args=args, vargs=vargs))
         yield _node(ast.Name(id=name, ctx=ast.Load()))
+        return
+
+
+def _multi_arity_fn_ast(ctx: CompilerContext, name: str, arities: List[FunctionArityDetails]) -> ASTStream:
+    """Generate Python AST nodes for multi-arity Basilisp function definitions.
+
+    For example, a multi-arity function like this:
+
+        (def f
+          (fn f
+            ([] (print "No args"))
+            ([arg]
+              (print arg))
+            ([arg & rest]
+              (print (concat [arg] rest)))))
+
+    Would yield a function definition in Python code like this:
+
+        def __f_68__arity0():
+            return print_('No args')
+
+
+        def __f_68__arity1(arg_69):
+            return print_(arg_69)
+
+
+        def __f_68__arity_rest(arg_70, *rest_71):
+            rest_72 = runtime._collect_args(rest_71)
+            return print_(concat(vec.vector([arg_70], meta=None), rest_72))
+
+
+        def __f_68(*multi_arity_args):
+            if len(multi_arity_args) == 0:
+                return __f_68__arity0(*multi_arity_args)
+            elif len(multi_arity_args) == 1:
+                return __f_68__arity1(*multi_arity_args)
+            elif len(multi_arity_args) >= 2:
+                return __f_68__arity2(*multi_arity_args)
+
+
+        f = __f_68"""
+    if_stmts: List[Dict[str, ast.AST]] = []
+    multi_arity_args_arg = _load_attr(_MULTI_ARITY_ARG_NAME)
+    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))
+            compare_op = ast.GtE() if is_rest else ast.Eq()
+            if_stmts.append(
+                {"test": ast.Compare(left=ast.Call(func=_load_attr('len'),
+                                                   args=[multi_arity_args_arg],
+                                                   keywords=[]),
+                                     ops=[compare_op],
+                                     comparators=[ast.Num(arg_count)]),
+                 "body": ast.Return(value=ast.Call(func=_load_attr(arity_name),
+                                                   args=[ast.Starred(value=multi_arity_args_arg, ctx=ast.Load())],
+                                                   keywords=[]))})
+
+    assert len(if_stmts) == len(arities)
+
+    yield _dependency(ast.FunctionDef(
+        name=name,
+        args=ast.arguments(
+            args=[],
+            kwarg=None,
+            vararg=ast.arg(arg=_MULTI_ARITY_ARG_NAME, annotation=None),
+            kwonlyargs=[],
+            defaults=[],
+            kw_defaults=[]),
+        body=[_compose_ifs(if_stmts),
+              ast.Raise(exc=ast.Call(func=_load_attr('basilisp.lang.runtime.RuntimeException'),
+                                     args=[ast.Str(f"Wrong number of args passed to function: {name}")],
+                                     keywords=[]),
+                        cause=None)],
+        decorator_list=[],
+        returns=None))
+    yield _node(ast.Name(id=name, ctx=ast.Load()))
+
+
+def _fn_ast(ctx: CompilerContext, form: llist.List) -> ASTStream:
+    """Generate a Python AST Nodes for function definitions."""
+    assert form.first == _FN
+    has_name = isinstance(form[1], sym.Symbol)
+    name = genname("__" + (munge(form[1].name) if has_name else _FN_PREFIX))
+
+    rest_idx = 1 + int(has_name)
+    arities = list(_fn_arities(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:

tests/compiler_test.py

@@ -160,7 +160,7 @@ def test_do(ns_var: Var):
     assert lcompile("last-name") == "Vader"
 
 
-def test_fn(ns_var: Var):
+def test_single_arity_fn(ns_var: Var):
     code = """
     (def string-upper (fn* string-upper [s] (.upper s)))
     """
@@ -171,6 +171,16 @@ def test_fn(ns_var: Var):
     assert callable(fvar.value)
     assert fvar.value("lower") == "LOWER"
 
+    code = """
+    (def string-upper (fn* string-upper ([s] (.upper s))))
+    """
+    ns_name = ns_var.value.name
+    fvar = lcompile(code)
+    assert fvar == Var.find_in_ns(
+        sym.symbol(ns_name), sym.symbol('string-upper'))
+    assert callable(fvar.value)
+    assert fvar.value("lower") == "LOWER"
+
     code = """
     (def string-lower #(.lower %))
     """
@@ -182,6 +192,63 @@ def test_fn(ns_var: Var):
     assert fvar.value("UPPER") == "upper"
 
 
+def test_multi_arity_fn(ns_var: Var):
+    with pytest.raises(compiler.CompilerException):
+        lcompile('(fn f)')
+
+    with pytest.raises(compiler.CompilerException):
+        lcompile("""
+            (def f
+              (fn* f
+                ([] :no-args)
+                ([] :also-no-args)))
+            """)
+
+    with pytest.raises(compiler.CompilerException):
+        lcompile("""
+            (def f
+              (fn* f
+                ([& args] (concat [:no-starter] args))
+                ([s & args] (concat [s] args))))
+            """)
+
+    with pytest.raises(compiler.CompilerException):
+        lcompile("""
+            (def f
+              (fn* f
+                ([s] (concat [s] :one-arg))
+                ([& args] (concat [:rest-params] args))))
+            """)
+
+    code = """
+    (def multi-fn
+      (fn* multi-fn
+        ([] :no-args)
+        ([s] s)
+        ([s & args] (concat [s] args))))
+    """
+    ns_name = ns_var.value.name
+    fvar = lcompile(code)
+    assert fvar == Var.find_in_ns(
+        sym.symbol(ns_name), sym.symbol('multi-fn'))
+    assert callable(fvar.value)
+    assert fvar.value() == kw.keyword('no-args')
+    assert fvar.value('STRING') == 'STRING'
+    assert fvar.value(kw.keyword('first-arg'), 'second-arg', 3) == llist.l(kw.keyword('first-arg'), 'second-arg', 3)
+
+    with pytest.raises(runtime.RuntimeException):
+        code = """
+            (def angry-multi-fn
+              (fn* angry-multi-fn
+                ([] :send-me-an-arg!)
+                ([i] i)
+                ([i j] (concat [i] [j]))))
+            """
+        ns_name = ns_var.value.name
+        fvar = lcompile(code)
+        fvar.value(1, 2, 3)
+
+
 def test_fn_call(ns_var: Var):
     code = """
     (def string-upper (fn* [s] (.upper s)))