arrow-functions-with-scope-runtime.js

global.__$scope = new WeakMap();
function __$withScope(f, s) {
  global.__$scope.set(f, s);
  return f;
}

arrow-functions-with-scope.js

module.exports = function (babel) {
  var t = babel.types;
  return {
    visitor: {
      ArrowFunctionExpression: function ArrowFunctionExpression(path, state) {
        path.arrowFunctionToShadowed()
        var includedReference = {};
        var references = []
        var afScope = path.scope;
        path.traverse({Identifier: function(path) {
            if (path.parent.type === 'MemberExpression' && path.parent.property === path.node) return;

            var id = path.node.name
            if (includedReference[id] !== undefined) return

            var s = path.scope
            while (true) {
                if (id in s.bindings) return
                if (s === afScope) break
                s = s.parent
            }
            references.push(t.objectProperty(t.identifier(id), t.identifier(id)))
            includedReference[id] = true
        }})
        var passedScope = t.objectExpression(references)

        var replacement = t.callExpression(t.identifier("__$withScope"), [path.node, passedScope])
        path.replaceWith(replacement)
      }
    }
  }
}

test.in.js

function test() {
    var s = 1;
    var t = 2;
    var arr = [1,2,3,4];
    var res = arr.map(el => arg => {
        var x = s + arg + arr.indexOf(el)
        return el + x + 1 + t
    })
    return res;
}

test.out.js

function test() {
    var s = 1;
    var t = 2;
    var arr = [1, 2, 3, 4];
    var res = arr.map(__$withScope(function (el) {
        return __$withScope(function (arg) {
            var x = s + arg + arr.indexOf(el);
            return el + x + 1 + t;
        }, {
            s: s,
            arr: arr,
            el: el,
            t: t
        });
    }, {
        s: s,
        arr: arr,
        t: t
    }));
    return res;
}