From 5b2216b34804a618e478e71250d0a49a9d02e0e5 Mon Sep 17 00:00:00 2001
From: Ingvar Stepanyan <me@rreverser.com>
Date: Fri, 6 Feb 2015 13:55:51 +0200
Subject: [PATCH] Add tail recursion optimization.

As per ES6, VMs should perform tail call optimization and prevent growth of call stack.
This adds tail call optimization for recursion case (when function has explicit name and calls itself in `return`).
Cross-function optimization is not currently performed as it's more complicated and requires value tracking.
---
 .../transformers/es6/tail-call.js             | 215 ++++++++++++++++++
 lib/6to5/transformation/transformers/index.js |   2 +
 lib/6to5/types/index.js                       |   2 +-
 .../es6-tail-call/call-apply/actual.js        |   7 +
 .../es6-tail-call/call-apply/expected.js      |  27 +++
 .../es6-tail-call/expressions/actual.js       |   3 +
 .../es6-tail-call/expressions/expected.js     |  30 +++
 .../es6-tail-call/recursion/actual.js         |   8 +
 .../es6-tail-call/recursion/expected.js       |  22 ++
 .../es6-tail-call/try-catch/actual.js         |  39 ++++
 .../es6-tail-call/try-catch/expected.js       |  65 ++++++
 11 files changed, 419 insertions(+), 1 deletion(-)
 create mode 100644 lib/6to5/transformation/transformers/es6/tail-call.js
 create mode 100644 test/fixtures/transformation/es6-tail-call/call-apply/actual.js
 create mode 100644 test/fixtures/transformation/es6-tail-call/call-apply/expected.js
 create mode 100644 test/fixtures/transformation/es6-tail-call/expressions/actual.js
 create mode 100644 test/fixtures/transformation/es6-tail-call/expressions/expected.js
 create mode 100644 test/fixtures/transformation/es6-tail-call/recursion/actual.js
 create mode 100644 test/fixtures/transformation/es6-tail-call/recursion/expected.js
 create mode 100644 test/fixtures/transformation/es6-tail-call/try-catch/actual.js
 create mode 100644 test/fixtures/transformation/es6-tail-call/try-catch/expected.js

diff --git a/lib/6to5/transformation/transformers/es6/tail-call.js b/lib/6to5/transformation/transformers/es6/tail-call.js
new file mode 100644
index 0000000000..b9af8e209a
--- /dev/null
+++ b/lib/6to5/transformation/transformers/es6/tail-call.js
@@ -0,0 +1,215 @@
+"use strict";
+
+var t = require("../../../types");
+
+function returnBlock(expr) {
+  return t.blockStatement([t.returnStatement(expr)]);
+}
+
+function transformExpression(node, scope, state) {
+  if (!node) return;
+
+  return (function subTransform(node) {
+    switch (node.type) {
+      case "ConditionalExpression":
+        var callConsequent = subTransform(node.consequent);
+        var callAlternate = subTransform(node.alternate);
+        if (!callConsequent && !callAlternate) {
+          return;
+        }
+        // if ternary operator had tail recursion in value, convert to optimized if-statement
+        node.type = "IfStatement";
+        node.consequent = callConsequent ? t.toBlock(callConsequent) : returnBlock(node.consequent);
+        if (callAlternate) {
+          node.alternate = t.isIfStatement(callAlternate) ? callAlternate : t.toBlock(callAlternate);
+        } else {
+          node.alternate = returnBlock(node.alternate);
+        }
+        return node;
+
+      case "LogicalExpression":
+        // only call in right-value of can be optimized
+        var callRight = subTransform(node.right);
+        if (!callRight) {
+          return;
+        }
+        // cache left value as it might have side-effects
+        var leftId = state.getLeftId();
+        var testExpr = t.assignmentExpression(
+          "=",
+          leftId,
+          node.left
+        );
+        if (node.operator === "&&") {
+          testExpr = t.unaryExpression("!", testExpr);
+        }
+        return [t.ifStatement(testExpr, returnBlock(leftId))].concat(callRight);
+
+      case "SequenceExpression":
+        var seq = node.expressions;
+        // only last element can be optimized
+        var lastCall = subTransform(seq[seq.length - 1]);
+        if (!lastCall) {
+          return;
+        }
+        // remove converted expression from sequence
+        // and convert to regular expression if needed
+        if (--seq.length === 1) {
+          node = seq[0];
+        }
+        return [t.expressionStatement(node)].concat(lastCall);
+
+      case "CallExpression":
+        var callee = node.callee, prop, thisBinding, args;
+
+        if (t.isMemberExpression(callee, { computed: false }) &&
+            t.isIdentifier(prop = callee.property)) {
+          switch (prop.name) {
+            case "call":
+              args = t.arrayExpression(node.arguments.slice(1));
+              break;
+            case "apply":
+              args = node.arguments[1] || t.identifier("undefined");
+              break;
+            default:
+              return;
+          }
+          thisBinding = node.arguments[0];
+          callee = callee.object;
+        }
+
+        // only tail recursion can be optimized as for now
+        if (!t.isIdentifier(callee) || !scope.bindingEquals(callee.name, state.ownerId)) {
+          return;
+        }
+
+        state.hasTailRecursion = true;
+
+        return [
+          t.expressionStatement(t.assignmentExpression(
+            "=",
+            state.getArgumentsId(),
+            args || t.arrayExpression(node.arguments)
+          )),
+          t.expressionStatement(t.assignmentExpression(
+            "=",
+            state.getThisId(),
+            thisBinding || t.identifier("undefined")
+          )),
+          t.returnStatement(t.assignmentExpression(
+            "=",
+            state.getShouldContinueId(),
+            t.literal(true)
+          ))
+        ];
+    }
+  })(node);
+}
+
+var functionChildrenVisitor = {
+  enter: function (node, parent, scope, state) {
+    if (t.isReturnStatement(node)) {
+      // prevent entrance by current visitor
+      this.skip();
+      // transform return argument into statement if
+      // it contains tail recursion
+      return transformExpression(node.argument, scope, state);
+    } else if (t.isFunction(node)) {
+      return this.skip();
+    } else if (t.isTryStatement(parent)) {
+      if (node === parent.block) {
+        return this.skip();
+      } else if (node === parent.finalizer) {
+        return;
+      } else {
+        if (parent.finalizer) {
+          this.skip();
+        }
+        return;
+      }
+    }
+  }
+};
+
+var functionVisitor = {
+  enter: function (node, parent, scope, state) {
+    // traverse all child nodes of this function and find `arguments` and `this`
+    scope.traverse(node, functionChildrenVisitor, state);
+
+    return this.skip();
+  }
+};
+
+exports.FunctionDeclaration =
+exports.FunctionExpression = function (node, parent, scope) {
+  // only tail recursion can be optimized as for now,
+  // so we can skip anonymous functions entirely
+  var ownerId = node.id;
+  if (!ownerId) return;
+
+  var argumentsId, thisId, shouldContinueId, leftId;
+
+  var state = {
+    hasTailRecursion: false,
+    ownerId: ownerId,
+    getArgumentsId: function () {
+      return argumentsId = argumentsId || scope.generateUidIdentifier("arguments");
+    },
+    getThisId: function () {
+      return thisId = thisId || scope.generateUidIdentifier("this");
+    },
+    getShouldContinueId: function () {
+      return shouldContinueId = shouldContinueId || scope.generateUidIdentifier("shouldContinue");
+    },
+    getLeftId: function () {
+      return leftId = leftId || scope.generateUidIdentifier("left");
+    }
+  };
+
+  // traverse the function and look for tail recursion
+  scope.traverse(node, functionVisitor, state);
+
+  if (!state.hasTailRecursion) return;
+
+  var block = t.ensureBlock(node);
+
+  if (leftId) {
+    block.body.unshift(t.variableDeclaration("var", [
+      t.variableDeclarator(leftId)
+    ]));
+  }
+
+  var resultId = scope.generateUidIdentifier("result");
+  state.getShouldContinueId();
+
+  node.body = t.blockStatement([
+    t.variableDeclaration("var", [
+      t.variableDeclarator(argumentsId, t.identifier("arguments")),
+      t.variableDeclarator(thisId, t.thisExpression()),
+      t.variableDeclarator(shouldContinueId),
+      t.variableDeclarator(resultId)
+    ]),
+    t.doWhileStatement(t.blockStatement([
+      t.expressionStatement(t.assignmentExpression(
+        "=",
+        shouldContinueId,
+        t.literal(false)
+      )),
+      t.expressionStatement(t.assignmentExpression(
+        "=",
+        resultId,
+        t.callExpression(
+          t.memberExpression(
+            t.functionExpression(null, node.params, block),
+            t.identifier("apply"),
+            false
+          ),
+          [thisId, argumentsId]
+        )
+      ))
+    ]), shouldContinueId),
+    t.returnStatement(resultId)
+  ]);
+
+  node.params = [];
+};
diff --git a/lib/6to5/transformation/transformers/index.js b/lib/6to5/transformation/transformers/index.js
index f778ae0c8e..ab5aff2ee9 100644
--- a/lib/6to5/transformation/transformers/index.js
+++ b/lib/6to5/transformation/transformers/index.js
@@ -52,6 +52,8 @@ module.exports = {
   // needs to be after `es6.blockScoping` due to needing `letReferences` set on blocks
   "es6.blockScopingTDZ":                   require("./es6/block-scoping-tdz"),
 
+  "es6.tailCall":                          require("./es6/tail-call"),
+
   "es6.parameters.default":                require("./es6/parameters.default"),
   "es6.parameters.rest":                   require("./es6/parameters.rest"),
 
diff --git a/lib/6to5/types/index.js b/lib/6to5/types/index.js
index d58fbf1372..d5301328cd 100644
--- a/lib/6to5/types/index.js
+++ b/lib/6to5/types/index.js
@@ -395,7 +395,7 @@ t.toIdentifier = function (name) {
 
 t.ensureBlock = function (node, key) {
   key = key || "body";
-  node[key] = t.toBlock(node[key], node);
+  return node[key] = t.toBlock(node[key], node);
 };
 
 /**
diff --git a/test/fixtures/transformation/es6-tail-call/call-apply/actual.js b/test/fixtures/transformation/es6-tail-call/call-apply/actual.js
new file mode 100644
index 0000000000..2bb42c6c81
--- /dev/null
+++ b/test/fixtures/transformation/es6-tail-call/call-apply/actual.js
@@ -0,0 +1,7 @@
+(function f(n) {
+  if (n <= 0) {
+    console.log(this, arguments);
+    return "foo";
+  }
+  return Math.random() > 0.5 ? f.call(this, n - 1) : f.apply(this, [n - 1]);
+})(1e6) === "foo";
diff --git a/test/fixtures/transformation/es6-tail-call/call-apply/expected.js b/test/fixtures/transformation/es6-tail-call/call-apply/expected.js
new file mode 100644
index 0000000000..806974a174
--- /dev/null
+++ b/test/fixtures/transformation/es6-tail-call/call-apply/expected.js
@@ -0,0 +1,27 @@
+"use strict";
+
+(function f() {
+  var _arguments = arguments,
+      _this = this,
+      _shouldContinue,
+      _result;
+  do {
+    _shouldContinue = false;
+    _result = (function (n) {
+      if (n <= 0) {
+        console.log(this, arguments);
+        return "foo";
+      }
+      if (Math.random() > 0.5) {
+        _arguments = [n - 1];
+        _this = this;
+        return _shouldContinue = true;
+      } else {
+        _arguments = [n - 1];
+        _this = this;
+        return _shouldContinue = true;
+      }
+    }).apply(_this, _arguments);
+  } while (_shouldContinue);
+  return _result;
+})(1000000) === "foo";
diff --git a/test/fixtures/transformation/es6-tail-call/expressions/actual.js b/test/fixtures/transformation/es6-tail-call/expressions/actual.js
new file mode 100644
index 0000000000..824fdf80d5
--- /dev/null
+++ b/test/fixtures/transformation/es6-tail-call/expressions/actual.js
@@ -0,0 +1,3 @@
+(function f(n) {
+  return n <= 0 ? "foo" : (doSmth(), getTrueValue() && (getFalseValue() || f(n - 1)));
+})(1e6, true) === "foo";
diff --git a/test/fixtures/transformation/es6-tail-call/expressions/expected.js b/test/fixtures/transformation/es6-tail-call/expressions/expected.js
new file mode 100644
index 0000000000..fe2b8c5d97
--- /dev/null
+++ b/test/fixtures/transformation/es6-tail-call/expressions/expected.js
@@ -0,0 +1,30 @@
+"use strict";
+
+(function f() {
+  var _arguments = arguments,
+      _this = this,
+      _shouldContinue,
+      _result;
+  do {
+    _shouldContinue = false;
+    _result = (function (n) {
+      var _left;
+      if (n <= 0) {
+        return "foo";
+      } else {
+        doSmth();
+
+        if (!(_left = getTrueValue())) {
+          return _left;
+        }
+        if (_left = getFalseValue()) {
+          return _left;
+        }
+        _arguments = [n - 1];
+        _this = undefined;
+        return _shouldContinue = true;
+      }
+    }).apply(_this, _arguments);
+  } while (_shouldContinue);
+  return _result;
+})(1000000, true) === "foo";
diff --git a/test/fixtures/transformation/es6-tail-call/recursion/actual.js b/test/fixtures/transformation/es6-tail-call/recursion/actual.js
new file mode 100644
index 0000000000..9d2b656565
--- /dev/null
+++ b/test/fixtures/transformation/es6-tail-call/recursion/actual.js
@@ -0,0 +1,8 @@
+(function f(n, /* should be undefined after first pass */ m) {
+  if (n <= 0) {
+    return "foo";
+  }
+  // Should be clean (undefined) on each pass
+  var local;
+  return f(n - 1);
+})(1e6, true) === "foo";
diff --git a/test/fixtures/transformation/es6-tail-call/recursion/expected.js b/test/fixtures/transformation/es6-tail-call/recursion/expected.js
new file mode 100644
index 0000000000..d2831ed807
--- /dev/null
+++ b/test/fixtures/transformation/es6-tail-call/recursion/expected.js
@@ -0,0 +1,22 @@
+"use strict";
+
+(function f() {
+  var _arguments = arguments,
+      _this = this,
+      _shouldContinue,
+      _result;
+  do {
+    _shouldContinue = false;
+    _result = (function (n, /* should be undefined after first pass */m) {
+      if (n <= 0) {
+        return "foo";
+      }
+      // Should be clean (undefined) on each pass
+      var local;
+      _arguments = [n - 1];
+      _this = undefined;
+      return _shouldContinue = true;
+    }).apply(_this, _arguments);
+  } while (_shouldContinue);
+  return _result;
+})(1000000, true) === "foo";
diff --git a/test/fixtures/transformation/es6-tail-call/try-catch/actual.js b/test/fixtures/transformation/es6-tail-call/try-catch/actual.js
new file mode 100644
index 0000000000..0a6be74326
--- /dev/null
+++ b/test/fixtures/transformation/es6-tail-call/try-catch/actual.js
@@ -0,0 +1,39 @@
+(function f(n) {
+  if (n <= 0) {
+    return "foo";
+  }
+  try {
+    return f(n - 1);
+  } catch (e) {}
+})(1e6) === "foo";
+
+(function f(n) {
+  if (n <= 0) {
+    return "foo";
+  }
+  try {
+    throw new Error();
+  } catch (e) {
+    return f(n - 1);
+  }
+})(1e6) === "foo";
+
+(function f(n) {
+  if (n <= 0) {
+    return "foo";
+  }
+  try {
+    throw new Error();
+  } catch (e) {
+    return f(n - 1);
+  } finally {}
+})(1e6) === "foo";
+
+(function f(n) {
+  if (n <= 0) {
+    return "foo";
+  }
+  try {} finally {
+    return f(n - 1);
+  }
+})(1e6) === "foo";
diff --git a/test/fixtures/transformation/es6-tail-call/try-catch/expected.js b/test/fixtures/transformation/es6-tail-call/try-catch/expected.js
new file mode 100644
index 0000000000..0896197b35
--- /dev/null
+++ b/test/fixtures/transformation/es6-tail-call/try-catch/expected.js
@@ -0,0 +1,65 @@
+"use strict";
+
+(function f(n) {
+  if (n <= 0) {
+    return "foo";
+  }
+  try {
+    return f(n - 1);
+  } catch (e) {}
+})(1000000) === "foo";
+
+(function f() {
+  var _arguments = arguments,
+      _this = this,
+      _shouldContinue,
+      _result;
+  do {
+    _shouldContinue = false;
+    _result = (function (n) {
+      if (n <= 0) {
+        return "foo";
+      }
+      try {
+        throw new Error();
+      } catch (e) {
+        _arguments = [n - 1];
+        _this = undefined;
+        return _shouldContinue = true;
+      }
+    }).apply(_this, _arguments);
+  } while (_shouldContinue);
+  return _result;
+})(1000000) === "foo";
+
+(function f(n) {
+  if (n <= 0) {
+    return "foo";
+  }
+  try {
+    throw new Error();
+  } catch (e) {
+    return f(n - 1);
+  } finally {}
+})(1000000) === "foo";
+
+(function f() {
+  var _arguments = arguments,
+      _this = this,
+      _shouldContinue,
+      _result;
+  do {
+    _shouldContinue = false;
+    _result = (function (n) {
+      if (n <= 0) {
+        return "foo";
+      }
+      try {} finally {
+        _arguments = [n - 1];
+        _this = undefined;
+        return _shouldContinue = true;
+      }
+    }).apply(_this, _arguments);
+  } while (_shouldContinue);
+  return _result;
+})(1000000) === "foo";