Fix return statement

smnp/ast/node/function.py

@@ -1,10 +1,8 @@
-from smnp.ast.node.block import Block
+from smnp.ast.node.block import BlockParser
 from smnp.ast.node.identifier import IdentifierLiteralParser
 from smnp.ast.node.iterable import abstractIterableParser
 from smnp.ast.node.model import Node
 from smnp.ast.node.none import NoneNode
-from smnp.ast.node.ret import ReturnParser
-from smnp.ast.node.statement import StatementParser
 from smnp.ast.node.type import TypeParser, Type
 from smnp.ast.parser import Parser
 from smnp.token.type import TokenType
@@ -124,23 +122,8 @@ def FunctionDefinitionParser(input):
         Parser.terminal(TokenType.FUNCTION),
         Parser.doAssert(IdentifierLiteralParser, "function/method name"),
         Parser.doAssert(ArgumentsDeclarationParser, "function/method arguments"),
-        Parser.doAssert(MethodBodyParser, "function/method body"),
+        Parser.doAssert(BlockParser, "function/method body"),
         createNode=lambda _, name, args, body: FunctionDefinition.withValues(name, args, body),
         name="function definition"
     )(input)
 
-
-def MethodBodyParser(input):
-    bodyItem = Parser.oneOf(
-        ReturnParser,
-        StatementParser,
-        name="function body item"
-    )
-
-    return Parser.loop(
-        Parser.terminal(TokenType.OPEN_CURLY),
-        Parser.doAssert(bodyItem, f"statement or '{TokenType.CLOSE_CURLY.key}'"),
-        Parser.terminal(TokenType.CLOSE_CURLY),
-        createNode=lambda open, statements, close: Block.withChildren(statements, open.pos),
-        name="function body"
-    )(input)

smnp/ast/node/statement.py

@@ -11,9 +11,11 @@ def StatementParser(input):
     from smnp.ast.node.condition import IfElseStatementParser
     from smnp.ast.node.expression import ExpressionParser
 
+    from smnp.ast.node.ret import ReturnParser
     return Parser.oneOf(
         IfElseStatementParser,
         ExpressionParser,
         BlockParser,
+        ReturnParser,
         name="statement"
     )(input)