LibJS: Add support for arrow functions

5 files changed, 144 insertions, 1 deletions

diff --git a/Libraries/LibJS/Lexer.cpp b/Libraries/LibJS/Lexer.cpp
index 04ea59ef5a..f4ffaa3fba 100644
--- a/Libraries/LibJS/Lexer.cpp
+++ b/Libraries/LibJS/Lexer.cpp
@@ -86,6 +86,7 @@ Lexer::Lexer(StringView source)
     }
 
     if (s_two_char_tokens.is_empty()) {
+        s_two_char_tokens.set("=>", TokenType::Arrow);
         s_two_char_tokens.set("+=", TokenType::PlusEquals);
         s_two_char_tokens.set("-=", TokenType::MinusEquals);
         s_two_char_tokens.set("*=", TokenType::AsteriskEquals);
diff --git a/Libraries/LibJS/Parser.cpp b/Libraries/LibJS/Parser.cpp
index 94a69bdb37..59db60e5ee 100644
--- a/Libraries/LibJS/Parser.cpp
+++ b/Libraries/LibJS/Parser.cpp
@@ -223,6 +223,78 @@ NonnullRefPtr<Statement> Parser::parse_statement()
     return statement;
 }
 
+RefPtr<FunctionExpression> Parser::try_parse_arrow_function_expression(bool expect_parens)
+{
+    save_state();
+
+    Vector<FlyString> parameters;
+    bool parse_failed = false;
+    while (true) {
+        if (match(TokenType::Comma)) {
+            consume(TokenType::Comma);
+        } else if (match(TokenType::Identifier)) {
+            auto token = consume(TokenType::Identifier);
+            parameters.append(token.value());
+        } else if (match(TokenType::ParenClose)) {
+            if (expect_parens) {
+                consume(TokenType::ParenClose);
+                if (match(TokenType::Arrow)) {
+                    consume(TokenType::Arrow);
+                } else {
+                    parse_failed = true;
+                }
+                break;
+            }
+            parse_failed = true;
+            break;
+        } else if (match(TokenType::Arrow)) {
+            if (!expect_parens) {
+                consume(TokenType::Arrow);
+                break;
+            }
+            parse_failed = true;
+            break;
+        } else {
+            parse_failed = true;
+            break;
+        }
+    }
+
+    if (parse_failed) {
+        load_state();
+        return nullptr;
+    }
+
+    auto function_body_result = [this]() -> RefPtr<BlockStatement> {
+        if (match(TokenType::CurlyOpen)) {
+            // Parse a function body with statements
+            return parse_block_statement();
+        }
+        if (match_expression()) {
+            // Parse a function body which returns a single expression
+
+            // FIXME: We synthesize a block with a return statement
+            // for arrow function bodies which are a single expression.
+            // Esprima generates a single "ArrowFunctionExpression"
+            // with a "body" property.
+            auto return_expression = parse_expression(0);
+            auto return_block = create_ast_node<BlockStatement>();
+            return_block->append<ReturnStatement>(move(return_expression));
+            return return_block;
+        }
+        // Invalid arrow function body
+        return nullptr;
+    }();
+
+    if (!function_body_result.is_null()) {
+        auto body = function_body_result.release_nonnull();
+        return create_ast_node<FunctionExpression>("", move(body), move(parameters));
+    }
+
+    load_state();
+    return nullptr;
+}
+
 NonnullRefPtr<Expression> Parser::parse_primary_expression()
 {
     if (match_unary_prefixed_expression())
@@ -231,12 +303,23 @@ NonnullRefPtr<Expression> Parser::parse_primary_expression()
     switch (m_parser_state.m_current_token.type()) {
     case TokenType::ParenOpen: {
         consume(TokenType::ParenOpen);
+        if (match(TokenType::ParenClose) || match(TokenType::Identifier)) {
+            auto arrow_function_result = try_parse_arrow_function_expression(true);
+            if (!arrow_function_result.is_null()) {
+                return arrow_function_result.release_nonnull();
+            }
+        }
         auto expression = parse_expression(0);
         consume(TokenType::ParenClose);
         return expression;
     }
-    case TokenType::Identifier:
+    case TokenType::Identifier: {
+        auto arrow_function_result = try_parse_arrow_function_expression(false);
+        if (!arrow_function_result.is_null()) {
+            return arrow_function_result.release_nonnull();
+        }
         return create_ast_node<Identifier>(consume().value());
+    }
     case TokenType::NumericLiteral:
         return create_ast_node<NumericLiteral>(consume().double_value());
     case TokenType::BoolLiteral:
diff --git a/Libraries/LibJS/Parser.h b/Libraries/LibJS/Parser.h
index c6fa879b59..c929672fbe 100644
--- a/Libraries/LibJS/Parser.h
+++ b/Libraries/LibJS/Parser.h
@@ -67,6 +67,7 @@ public:
     NonnullRefPtr<Expression> parse_secondary_expression(NonnullRefPtr<Expression>, int min_precedence, Associativity associate = Associativity::Right);
     NonnullRefPtr<CallExpression> parse_call_expression(NonnullRefPtr<Expression>);
     NonnullRefPtr<NewExpression> parse_new_expression();
+    RefPtr<FunctionExpression> try_parse_arrow_function_expression(bool expect_parens);
 
     bool has_errors() const { return m_parser_state.m_has_errors; }
 
diff --git a/Libraries/LibJS/Tests/arrow-functions.js b/Libraries/LibJS/Tests/arrow-functions.js
new file mode 100644
index 0000000000..3d2569d099
--- /dev/null
+++ b/Libraries/LibJS/Tests/arrow-functions.js
@@ -0,0 +1,57 @@
+function assert(x) { if (!x) throw 1; }
+try {
+  let getNumber = () => 42;
+  assert(getNumber() === 42);
+
+  let add = (a, b) => a + b;
+  assert(add(2, 3) === 5);
+
+  const addBlock = (a, b) => {
+    let res = a + b;
+    return res;
+  };
+  assert(addBlock(5, 4) === 9);
+
+  const makeObject = (a, b) => ({ a, b });
+  const obj = makeObject(33, 44);
+  assert(typeof obj === "object");
+  assert(obj.a === 33);
+  assert(obj.b === 44);
+
+  let returnUndefined = () => {};
+  assert(typeof returnUndefined() === "undefined");
+
+  const makeArray = (a, b) => [a, b];
+  const array = makeArray("3", { foo: 4 });
+  assert(array[0] === "3");
+  assert(array[1].foo === 4);
+
+  let square = x => x * x;
+  assert(square(3) === 9);
+
+  let squareBlock = x => {
+    return x * x;
+  };
+  assert(squareBlock(4) === 16);
+
+  const message = (who => "Hello " + who)("friends!");
+  assert(message === "Hello friends!");
+
+  const sum = ((x, y, z) => x + y + z)(1, 2, 3);
+  assert(sum === 6);
+
+  const product = ((x, y, z) => {
+    let res = x * y * z;
+    return res;
+  })(5, 4, 2);
+  assert(product === 40);
+
+  const half = (x => {
+    return x / 2;
+  })(10);
+  assert(half === 5);
+
+  console.log("PASS");
+} catch {
+  console.log("FAIL");
+}
diff --git a/Libraries/LibJS/Token.h b/Libraries/LibJS/Token.h
index 5ff035c41d..6e81c4fcf1 100644
--- a/Libraries/LibJS/Token.h
+++ b/Libraries/LibJS/Token.h
@@ -34,6 +34,7 @@ namespace JS {
 #define ENUMERATE_JS_TOKENS                           \
     __ENUMERATE_JS_TOKEN(Ampersand)                   \
     __ENUMERATE_JS_TOKEN(AmpersandEquals)             \
+    __ENUMERATE_JS_TOKEN(Arrow)                       \
     __ENUMERATE_JS_TOKEN(Asterisk)                    \
     __ENUMERATE_JS_TOKEN(AsteriskAsteriskEquals)      \
     __ENUMERATE_JS_TOKEN(AsteriskEquals)              \