Adding arrow functions

lib/ast.js

@@ -142,7 +142,7 @@ var AST_SimpleStatement = DEFNODE("SimpleStatement", "body", {
 }, AST_Statement);
 
 function walk_body(node, visitor) {
-    if (node.body instanceof AST_Statement) {
+    if (node.body instanceof AST_Node) {
         node.body._walk(visitor);
     }
     else node.body.forEach(function(stat){

lib/compress.js

@@ -1019,6 +1019,7 @@ merge(Compressor.prototype, {
     });
 
     OPT(AST_Block, function(self, compressor){
+        if (self.body instanceof AST_Node) { return self; }
         self.body = tighten_body(self.body, compressor);
         return self;
     });
@@ -1225,6 +1226,7 @@ merge(Compressor.prototype, {
         var hoist_funs = compressor.option("hoist_funs");
         var hoist_vars = compressor.option("hoist_vars");
         var self = this;
+        if (!(self.body instanceof Array)) { return self; }  // Hoisting makes no sense in an arrow func
         if (hoist_funs || hoist_vars) {
             var dirs = [];
             var hoisted = [];

lib/output.js

@@ -763,6 +763,34 @@ function OutputStream(options) {
         self._do_print(output);
     });
 
+    AST_Arrow.DEFMETHOD("_do_print", function(output){
+        var self = this;
+        var parent = output.parent();
+        var needs_parens = parent instanceof AST_Binary ||
+            parent instanceof AST_Unary ||
+            parent instanceof AST_Call;
+        if (needs_parens) { output.print("(") }
+        if (self.argnames.length === 1 && self.argnames[0] instanceof AST_Symbol) {
+            self.argnames[0].print(output);
+        } else {
+            output.with_parens(function(){
+                self.argnames.forEach(function(arg, i){
+                    if (i) output.comma();
+                    arg.print(output);
+                });
+            });
+        }
+        output.space();
+        output.print('=>');
+        output.space();
+        if (self.body instanceof AST_Node) {
+            this.body.print(output);
+        } else {
+            print_bracketed(this.body, output);
+        }
+        if (needs_parens) { output.print(")") }
+    });
+
     /* -----[ exits ]----- */
     AST_Exit.DEFMETHOD("_do_print", function(output, kind){
         output.print(kind);

lib/parse.js

@@ -510,6 +510,16 @@ function tokenizer($TEXT, filename, html5_comments) {
         return S.regex_allowed ? read_regexp("") : read_operator("/");
     };
 
+    function handle_eq_sign() {
+        next();
+        if (peek() === ">") {
+            next();
+            return token("arrow", "=>");
+        } else {
+            return read_operator("=");
+        }
+    };
+
     function handle_dot() {
         next();
         return is_digit(peek().charCodeAt(0))
@@ -559,6 +569,7 @@ function tokenizer($TEXT, filename, html5_comments) {
           case 34: case 39: return read_string(ch);
           case 46: return handle_dot();
           case 47: return handle_slash();
+          case 61: return handle_eq_sign();
         }
         if (is_digit(code)) return read_num();
         if (PUNC_CHARS(ch)) return token("punc", next());
@@ -975,6 +986,41 @@ function parse($TEXT, options) {
         });
     };
 
+    var arrow_function = function(args) {
+        expect_token("arrow", "=>");
+
+        if (args instanceof AST_SymbolRef) {
+            args = [args];
+        } else if (args instanceof AST_Seq) {
+            args = args.to_array();
+        } else if (args instanceof AST_Node) {
+            croak("Invalid syntax", args.start.line, args.start.col);
+        }
+
+        for (var i = 0; i < args.length; i++) {
+            if (!(args[i] instanceof AST_SymbolRef)) {
+                croak("Invalid parameter for an arrow function", args[i].start.line, args[i].start.col);
+            }
+
+            args[i] = new AST_SymbolFunarg({
+                name: args[i].name,
+                start: args[i].start,
+                end: args[i].end
+            })
+        }
+
+        return new AST_Arrow({
+            argnames: args,
+            body: (function(){
+                if (is("punc", "{")) {
+                    return _function_body();
+                } else {
+                    return expression(true);
+                }
+            })()
+        });
+    };
+
     var function_ = function(ctor) {
         var start = S.token
 
@@ -1486,6 +1532,13 @@ function parse($TEXT, options) {
 
     var maybe_assign = function(no_in) {
         var start = S.token;
+
+        if (start.value == "(" && peek().value == ")") {
+            next();  // (
+            next();  // )
+            return arrow_function([]);
+        }
+
         var left = maybe_conditional(no_in), val = S.token.value;
         if (is("operator") && ASSIGNMENT(val)) {
             if (is_assignable(left)) {
@@ -1500,6 +1553,9 @@ function parse($TEXT, options) {
             }
             croak("Invalid assignment");
         }
+        if (is("arrow")) {
+            return arrow_function(left)
+        }
         return left;
     };
 

lib/transform.js

@@ -166,7 +166,11 @@ TreeTransformer.prototype = new TreeWalker;
     _(AST_Lambda, function(self, tw){
         if (self.name) self.name = self.name.transform(tw);
         self.argnames = do_list(self.argnames, tw);
-        self.body = do_list(self.body, tw);
+        if (self.body instanceof AST_Node) {
+            self.body = self.body.transform(tw);
+        } else {
+            self.body = do_list(self.body, tw);
+        }
     });
 
     _(AST_Call, function(self, tw){

test/compress/harmony.js

@@ -0,0 +1,37 @@
+arrow_functions: {
+    input: {
+        (a) => b;  // 1 args
+        (a, b) => c;  // n args
+        () => b;  // 0 args
+        (a) => (b) => c;  // func returns func returns func
+        (a) => ((b) => c);  // So these parens are dropped
+        () => (b,c) => d;  // func returns func returns func
+        a=>{return b;}
+        a => 'lel';  // Dropping the parens
+    }
+    expect_exact: "a=>b;(a,b)=>c;()=>b;a=>b=>c;a=>b=>c;()=>(b,c)=>d;a=>{return b};a=>\"lel\";"
+}
+
+arrow_function_parens: {
+    input: {
+        something && (() => {});
+    }
+    expect_exact: "something&&(()=>{});"
+}
+arrow_function_parens_2: {
+    input: {
+        (() => null)();
+    }
+    expect_exact: "(()=>null)();"
+}
+
+regression_arrow_functions_and_hoist: {
+    options = {
+        hoist_vars: true,
+        hoist_funs: true
+    }
+    input: {
+        (a) => b;
+    }
+    expect_exact: "a=>b;"
+}