Repository: florimondmanca/fountain-lang

Issue:
Functions
Add function capability, based on https://www.craftinginterpreters.com/functions.html

* Call expressions: `somefunc()`, `somefunc(a, b, c)`
* Function definitions: `fn f(x, y) return x + y end`
* `return` support


Discussion:

Patch:
diff --git a/DESIGN.md b/DESIGN.md
index 9a96d5e..c6e5a9a 100644
--- a/DESIGN.md
+++ b/DESIGN.md
@@ -26,6 +26,8 @@ simple_stmt:
   | assert_stmt
   | assign_stmt
   | expr_stmt
+  | fn_stmt
+  | "return" expression?
   | "break"
   | "continue"
 block:
@@ -42,6 +44,10 @@ assign_stmt:
   | IDENTIFIER "=" expression
 expr_stmt:
   | expression
+fn_stmt:
+  | "fn" IDENTIFIER "(" parameters? ")" stmt* "end"
+parameters:
+  | IDENTIFIER ("," IDENTIFIER)*
 
 # Expressions
 expression:
@@ -62,7 +68,12 @@ factor:
   | unary ( ( "*" | "/" ) unary )*
 unary:
   | ( "-" | "not" ) unary
-  | primary;
+  | call
+call:
+  | expression "(" arguments? ")"
+  | primary
+arguments:
+  | expression ("," expression)*
 primary:
   | IDENTIFIER
   | TRUE
diff --git a/docs/README.md b/docs/README.md
index 81d1b08..0237636 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -124,7 +124,7 @@ end
 ### Functions
 
 ```lua
-fun f(x, y)
+fn f(x, y)
   return x + y
 end
 
diff --git a/examples/statements.ft b/examples/statements.ft
index a0215a3..f656ecf 100644
--- a/examples/statements.ft
+++ b/examples/statements.ft
@@ -35,3 +35,36 @@ for do
   end
   i = i + 1
 end
+
+print clock()
+
+fn fibo(n)
+  a = 0
+  b = 1
+  temp = 0
+  i = 0
+
+  for do
+    if i >= n do
+      break
+    end
+    temp = a
+    a = b
+    b = temp + b
+    i = i + 1
+  end
+
+  return a
+end
+
+print fibo(1)
+
+fn adder(y)
+  fn add(x)
+    return x + y
+  end
+  return add
+end
+
+add2 = adder(2)
+print add2(1)
diff --git a/src/fountain/_ast/__init__.py b/src/fountain/_ast/__init__.py
index 4233beb..ccf8f79 100644
--- a/src/fountain/_ast/__init__.py
+++ b/src/fountain/_ast/__init__.py
@@ -4,16 +4,19 @@ from .nodes import (
     Binary,
     Block,
     Break,
+    Call,
     Conjunction,
     Continue,
     Disjunction,
     Expr,
     Expression,
     For,
+    Function,
     Group,
     If,
     Literal,
     Print,
+    Return,
     Stmt,
     Unary,
     Variable,
@@ -28,12 +31,14 @@ __all__ = [
     "Binary",
     "Block",
     "Break",
+    "Call",
     "Conjunction",
     "Continue",
     "Disjunction",
     "Expr",
     "Expression",
     "For",
+    "Function",
     "Group",
     "If",
     "Literal",
@@ -41,6 +46,7 @@ __all__ = [
     "Variable",
     "parse",
     "Print",
+    "Return",
     "Stmt",
     "Token",
     "TokenType",
diff --git a/src/fountain/_ast/nodes.py b/src/fountain/_ast/nodes.py
index 03d918b..9143069 100644
--- a/src/fountain/_ast/nodes.py
+++ b/src/fountain/_ast/nodes.py
@@ -47,6 +47,13 @@ class Variable(Expr):
     name: Token
 
 
+@dataclass
+class Call(Expr):
+    callee: Expr
+    arguments: list[Expr]
+    closing: Token
+
+
 @dataclass
 class Stmt:
     pass
@@ -100,3 +107,16 @@ class Assert(Stmt):
 @dataclass
 class Block(Stmt):
     statements: list[Stmt]
+
+
+@dataclass
+class Function(Stmt):
+    name: Token
+    parameters: list[Token]
+    body: list[Stmt]
+
+
+@dataclass
+class Return(Stmt):
+    op: Token
+    expr: Optional[Expr]
diff --git a/src/fountain/_ast/parse.py b/src/fountain/_ast/parse.py
index 6870454..42852d9 100644
--- a/src/fountain/_ast/parse.py
+++ b/src/fountain/_ast/parse.py
@@ -5,16 +5,19 @@ from .nodes import (
     Binary,
     Block,
     Break,
+    Call,
     Conjunction,
     Continue,
     Disjunction,
     Expr,
     Expression,
     For,
+    Function,
     Group,
     If,
     Literal,
     Print,
+    Return,
     Stmt,
     Unary,
     Variable,
@@ -25,6 +28,7 @@ from .tokens import Token, TokenType
 def parse(tokens: list[Token]) -> list[Stmt]:
     current = 0
     loop = False
+    inside_function = False
 
     # Helpers.
 
@@ -54,10 +58,10 @@ def parse(tokens: list[Token]) -> list[Stmt]:
             return False
         return peek().type == t
 
-    def consume(t: TokenType, error_message: str) -> None:
+    def consume(t: TokenType, error_message: str) -> Token:
         if check(t):
             movenext()
-            return
+            return previous()
 
         raise ParseError(peek(), error_message)
 
@@ -91,6 +95,12 @@ def parse(tokens: list[Token]) -> list[Stmt]:
                 raise ParseError(previous(), "continue outside loop")
             return Continue(previous())
 
+        if match(TokenType.FN):
+            return function_statement()
+
+        if match(TokenType.RETURN):
+            return return_statement()
+
         if match(TokenType.ASSERT):
             return assert_statement()
 
@@ -154,6 +164,41 @@ def parse(tokens: list[Token]) -> list[Stmt]:
 
         return For(body)
 
+    def function_statement() -> Stmt:
+        nonlocal inside_function
+        prev = inside_function
+        inside_function = True
+
+        name = consume(TokenType.IDENTIFIER, "expected function name")
+
+        consume(TokenType.LEFT_PARENS, "expected '(' after function name")
+
+        parameters = []
+        if not check(TokenType.RIGHT_PARENS):
+            while True:
+                parameters.append(
+                    consume(TokenType.IDENTIFIER, "expected parameter name")
+                )
+                if not match(TokenType.COMMA):
+                    break
+
+        consume(TokenType.RIGHT_PARENS, "expected ')' after parameters")
+
+        body = []
+        while not check(TokenType.END) and not done():
+            body.append(statement())
+        consume(TokenType.END, "expected 'end' to close function")
+
+        inside_function = prev
+
+        return Function(name, parameters, body)
+
+    def return_statement() -> Stmt:
+        if not inside_function:
+            raise ParseError(previous(), "return outside function")
+        expr = None if check(TokenType.END) else expression()
+        return Return(previous(), expr)
+
     def assert_statement() -> Stmt:
         op = previous()
         test = expression()
@@ -230,7 +275,32 @@ def parse(tokens: list[Token]) -> list[Stmt]:
             op = previous()
             right = unary()
             return Unary(op, right)
-        return primary()
+        return call()
+
+    def call() -> Expr:
+        expr = primary()
+
+        # NOTE: there may be multiple calls, e.g. f(a, b)(c), where each
+        # call becomes the callee of the next one.
+
+        def finish_call(callee: Expr) -> Expr:
+            arguments = []
+
+            if not check(TokenType.RIGHT_PARENS):
+                arguments.append(expression())
+                while match(TokenType.COMMA):
+                    if len(arguments) >= 255:
+                        raise ParseError(previous(), "more than 255 arguments")
+                    arguments.append(expression())
+
+            closing = consume(TokenType.RIGHT_PARENS, "expected ')' after arguments")
+
+            return Call(callee, arguments, closing)
+
+        while match(TokenType.LEFT_PARENS):
+            expr = finish_call(callee=expr)
+
+        return expr
 
     def primary() -> Expr:
         if match(TokenType.FALSE):
@@ -267,7 +337,7 @@ def parse(tokens: list[Token]) -> list[Stmt]:
                 return
 
             if peek().type in (
-                TokenType.FUN,
+                TokenType.FN,
                 TokenType.FOR,
                 TokenType.IF,
                 TokenType.RETURN,
diff --git a/src/fountain/_ast/tokens.py b/src/fountain/_ast/tokens.py
index 73e0253..2bda9a5 100644
--- a/src/fountain/_ast/tokens.py
+++ b/src/fountain/_ast/tokens.py
@@ -47,7 +47,7 @@ class TokenType(Enum):
     CONTINUE = auto()  # `continue`
     PASS = auto()  # `pass`
     IN = auto()  # `in`
-    FUN = auto()  # `fun`
+    FN = auto()  # `fn`
     RETURN = auto()  # `return`
     PRINT = auto()  # `print`  # TODO: make it a function later
     ASSERT = auto()  # `assert`  # TODO: make it a function later
@@ -78,7 +78,7 @@ KEYWORDS = {
     "continue": TokenType.CONTINUE,
     "pass": TokenType.PASS,
     "in": TokenType.IN,
-    "fun": TokenType.FUN,
+    "fn": TokenType.FN,
     "print": TokenType.PRINT,
     "assert": TokenType.ASSERT,
     "return": TokenType.RETURN,
diff --git a/src/fountain/_builtins.py b/src/fountain/_builtins.py
new file mode 100644
index 0000000..3d5ac04
--- /dev/null
+++ b/src/fountain/_builtins.py
@@ -0,0 +1,23 @@
+import time
+from typing import TYPE_CHECKING, Any
+
+from ._functions import FunctionType
+
+if TYPE_CHECKING:
+    from ._interpreter import Interpreter
+
+
+class Clock(FunctionType):
+    def arity(self) -> int:
+        return 0
+
+    def call(self, interpreter: "Interpreter", *arguments: Any) -> Any:
+        return time.time()
+
+    def str(self) -> str:
+        return "<built-in function clock>"
+
+
+BUILTINS = [
+    ("clock", Clock()),
+]
diff --git a/src/fountain/_exceptions.py b/src/fountain/_exceptions.py
index 1096121..7f49d60 100644
--- a/src/fountain/_exceptions.py
+++ b/src/fountain/_exceptions.py
@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
     from ._ast import Token
@@ -37,3 +37,9 @@ class BreakExc(Exception):
 
 class ContinueExc(Exception):
     pass
+
+
+class ReturnExc(Exception):
+    def __init__(self, value: Any) -> None:
+        super().__init__()
+        self.value = value
diff --git a/src/fountain/_functions.py b/src/fountain/_functions.py
new file mode 100644
index 0000000..04eb4f5
--- /dev/null
+++ b/src/fountain/_functions.py
@@ -0,0 +1,48 @@
+from typing import TYPE_CHECKING, Any
+
+from ._ast import Function
+from ._exceptions import ReturnExc
+from ._scope import Scope
+
+if TYPE_CHECKING:
+    from ._interpreter import Interpreter
+
+
+class FunctionType:
+    def arity(self) -> int:
+        raise NotImplementedError  # pragma: no cover
+
+    def call(self, interpreter: "Interpreter", *arguments: Any) -> Any:
+        raise NotImplementedError  # pragma: no cover
+
+    def __str__(self) -> str:
+        raise NotImplementedError  # pragma: no cover
+
+
+class UserFunction(FunctionType):
+    def __init__(self, stmt: Function, closure: Scope) -> None:
+        self._stmt = stmt
+        self._closure = closure
+
+    @property
+    def name(self) -> str:
+        return self._stmt.name.lexeme
+
+    def arity(self) -> int:
+        return len(self._stmt.parameters)
+
+    def call(self, interpreter: "Interpreter", *arguments: Any) -> Any:
+        scope = Scope(self._closure)
+
+        for i, param in enumerate(self._stmt.parameters):
+            scope.assign(param.lexeme, arguments[i])
+
+        try:
+            interpreter.execute_scoped(self._stmt.body, scope)
+        except ReturnExc as exc:
+            return exc.value
+        else:
+            return None
+
+    def __str__(self) -> str:
+        return f"<function {self.name} at {hex(id(self))}>"
diff --git a/src/fountain/_interpreter.py b/src/fountain/_interpreter.py
index a4d61fc..566e332 100644
--- a/src/fountain/_interpreter.py
+++ b/src/fountain/_interpreter.py
@@ -6,30 +6,39 @@ from ._ast import (
     Binary,
     Block,
     Break,
+    Call,
     Conjunction,
     Continue,
     Disjunction,
     Expression,
     For,
+    Function,
     Group,
     If,
     Literal,
     NodeVisitor,
     Print,
+    Return,
     Stmt,
     Token,
     TokenType,
     Unary,
     Variable,
 )
-from ._exceptions import BreakExc, ContinueExc, EvalError
+from ._builtins import BUILTINS
+from ._exceptions import BreakExc, ContinueExc, EvalError, ReturnExc
+from ._functions import FunctionType, UserFunction
+from ._scope import Scope
 
 __all__ = ["Interpreter"]
 
 
 class Interpreter(NodeVisitor[Any]):
     def __init__(self) -> None:
-        self._scope = Scope()
+        scope = Scope()
+        for name, value in BUILTINS:
+            scope.assign(name, value)
+        self._scope = scope
 
     def interpret(self, statements: list[Stmt]) -> None:
         try:
@@ -77,6 +86,10 @@ class Interpreter(NodeVisitor[Any]):
     def execute_Continue(self, stmt: Continue) -> None:
         raise ContinueExc()
 
+    def execute_Return(self, stmt: Return) -> None:
+        value = self.evaluate(stmt.expr) if stmt.expr is not None else None
+        raise ReturnExc(value)
+
     def execute_Assert(self, stmt: Assert) -> None:
         test = self.evaluate(stmt.test)
 
@@ -92,15 +105,22 @@ class Interpreter(NodeVisitor[Any]):
         raise EvalError(stmt.op, message)
 
     def execute_Block(self, stmt: Block) -> None:
+        scope = Scope(self._scope)
+        self.execute_scoped(stmt.statements, scope)
+
+    def execute_scoped(self, statements: list[Stmt], scope: Scope) -> None:
         previous_scope = self._scope
-        scope = Scope(parent=previous_scope)
         try:
             self._scope = scope
-            for statement in stmt.statements:
+            for statement in statements:
                 self.execute(statement)
         finally:
             self._scope = previous_scope
 
+    def execute_Function(self, stmt: Function) -> Any:
+        func = UserFunction(stmt, closure=self._scope)
+        self._scope.assign(func.name, func)
+
     def evaluate_Literal(self, expr: Literal) -> Any:
         return expr.value
 
@@ -176,22 +196,20 @@ class Interpreter(NodeVisitor[Any]):
     def evaluate_Variable(self, expr: Variable) -> Any:
         return self._scope.get(expr.name)
 
+    def evaluate_Call(self, expr: Call) -> Any:
+        callee = self.evaluate(expr.callee)
+        if not isinstance(callee, FunctionType):
+            raise EvalError(expr.closing, "can only call functions")
 
-class Scope:
-    def __init__(self, parent: "Scope" = None) -> None:
-        self._parent = parent
-        self._values: dict[str, Any] = {}
+        numargs, arity = len(expr.arguments), callee.arity()
+        if numargs != arity:
+            plural = "" if arity == 1 else "s"
+            message = f"expected {arity} argument{plural}, got {numargs}"
+            raise EvalError(expr.closing, message)
 
-    def assign(self, name: str, value: Any) -> None:
-        self._values[name] = value
+        arguments = [self.evaluate(arg) for arg in expr.arguments]
 
-    def get(self, name: Token) -> Any:
-        try:
-            return self._values[name.lexeme]
-        except KeyError:
-            if self._parent is not None:
-                return self._parent.get(name)
-            raise EvalError(name, f"name {name.lexeme!r} is not defined") from None
+        return callee.call(self, *arguments)
 
 
 def check_number_operand(op: Token, value: Any) -> None:
diff --git a/src/fountain/_scope.py b/src/fountain/_scope.py
new file mode 100644
index 0000000..c13b545
--- /dev/null
+++ b/src/fountain/_scope.py
@@ -0,0 +1,21 @@
+from typing import Any
+
+from ._ast import Token
+from ._exceptions import EvalError
+
+
+class Scope:
+    def __init__(self, parent: "Scope" = None) -> None:
+        self._parent = parent
+        self._values: dict[str, Any] = {}
+
+    def assign(self, name: str, value: Any) -> None:
+        self._values[name] = value
+
+    def get(self, name: Token) -> Any:
+        try:
+            return self._values[name.lexeme]
+        except KeyError:
+            if self._parent is not None:
+                return self._parent.get(name)
+            raise EvalError(name, f"name {name.lexeme!r} is not defined") from None
diff --git a/tools/astgen.py b/tools/astgen.py
index 2f08467..37bf294 100644
--- a/tools/astgen.py
+++ b/tools/astgen.py
@@ -55,6 +55,7 @@ if __name__ == "__main__":
         "Disjunction   = expressions: list[Expr]",
         "Conjunction   = expressions: list[Expr]",
         "Variable      = name: Token",
+        "Call          = callee: Expr, arguments: list[Expr], closing: Token",
     ]
     stmt_types = [
         "Expression = expression: Expr",
@@ -66,6 +67,8 @@ if __name__ == "__main__":
         "Continue   = op: Token",
         "Assert     = op: Token, test: Expr, message: Optional[Expr]",
         "Block      = statements: list[Stmt]",
+        "Function   = name: Token, parameters: list[Token], body: list[Stmt]",
+        "Return     = op: Token, expr: Optional[Expr]",
     ]
 
     ast = make_ast(


Test patch:
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 9d1ec2d..b8a3bef 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -115,6 +115,37 @@ def test_cli_repl(monkeypatch: Any, capsys: Any) -> None:
             """,
             "0\n1\n2\n4\n5\n",
         ),
+        (
+            "fn f() print 'OK' end; f()",
+            # TODO: drop 'nil' after:
+            # https://github.com/florimondmanca/fountain-lang/issues/1
+            "OK\nnil\n",
+        ),
+        (
+            """
+            fn add(x, y)
+                return x + y
+            end
+
+            print add(1, 2)
+            """,
+            "3\n",
+        ),
+        (
+            """
+            fn adder(y)
+                fn add(x)
+                    return x + y
+                end
+                return add
+            end
+
+            add2 = adder(2)
+            print add2(1)
+            print adder(3)(2)
+            """,
+            "3\n5\n",
+        ),
         ("assert true", ""),
         ("-- Comment", ""),
         ("", ""),
@@ -205,6 +236,26 @@ def test_cli_eval(source: str, result: str, capsys: Any) -> None:
             "[line 4] error: at end: expected 'end' to close 'for'\n",
             65,
         ),
+        (
+            "fn f end",
+            "[line 1] error: at 'end': expected '(' after function name\n",
+            65,
+        ),
+        (
+            "fn f( end",
+            "[line 1] error: at 'end': expected parameter name\n",
+            65,
+        ),
+        (
+            "fn f(x end",
+            "[line 1] error: at 'end': expected ')' after parameters\n",
+            65,
+        ),
+        (
+            "return 0",
+            "[line 1] error: at 'return': return outside function\n",
+            65,
+        ),
         # Runtime errors
         (
             "1/0",
@@ -236,6 +287,21 @@ def test_cli_eval(source: str, result: str, capsys: Any) -> None:
             "[line 1] error: at 'assert': false!\n",
             70,
         ),
+        (
+            "fn f(x, y) return x end; f(1)",
+            "[line 1] error: at ')': expected 2 arguments, got 1\n",
+            70,
+        ),
+        (
+            "fn f(x) return x end; f(1, 2)",
+            "[line 1] error: at ')': expected 1 argument, got 2\n",
+            70,
+        ),
+        (
+            "1()",
+            "[line 1] error: at ')': can only call functions\n",
+            70,
+        ),
     ],
 )
 def test_cli_eval_error(source: str, err: str, exit_code: int, capsys: Any) -> None: