babel-plugin-__coverage__/lib/index.js

'use strict';
// babel-plugin-__coverage__
//
// This is my first Babel plugin, and I wrote it during the night.
// Therefore, be prepared to see a lot of copypasta and wtf code.

var _babelTemplate = require('babel-template');

var _babelTemplate2 = _interopRequireDefault(_babelTemplate);

var _babelHelperFunctionName = require('babel-helper-function-name');

var _babelHelperFunctionName2 = _interopRequireDefault(_babelHelperFunctionName);

var _fs = require('fs');

function _interopRequireDefault(obj) { return obj && obj.__esModule ? obj : { default: obj }; }

function _toConsumableArray(arr) { if (Array.isArray(arr)) { for (var i = 0, arr2 = Array(arr.length); i < arr.length; i++) { arr2[i] = arr[i]; } return arr2; } else { return Array.from(arr); } }

var coverageTemplate = (0, _babelTemplate2.default)('\n  var GLOBAL = (new Function(\'return this\'))()\n  var COVERAGE = GLOBAL[\'__coverage__\'] || (GLOBAL[\'__coverage__\'] = { })\n  var FILE_COVERAGE = COVERAGE[PATH] = GLOBAL[\'JSON\'].parse(INITIAL)\n');

//
// Takes a relative path and returns a real path.
// Assumes the path name is relative to working directory.
//
function getRealpath(n) {
  try {
    return (0, _fs.realpathSync)(n) || n;
  } catch (e) {
    return n;
  }
}

module.exports = function (_ref) {
  var t = _ref.types;

  //
  // Return the immediate data structure local to a file.
  //
  function getData(context) {
    var path = getRealpath(context.file.opts.filename);
    //
    // XXX: Is it OK to mutate `context.file`? I don’t know but it works!
    //
    return context.file.__coverage__data || (context.file.__coverage__data = {
      //
      // Initial data that will be added in front of generated source code
      //
      base: {
        path: path,
        s: {},
        b: {},
        f: {},
        statementMap: {},
        fnMap: {},
        branchMap: {}
      },
      //
      // The counter that generates the next ID for each statement type.
      nextId: {
        s: 1,
        b: 1,
        f: 1
      }
    });
  }

  //
  // Turns a `SourceLocation` into a plain object.
  //
  function locToObject(loc) {
    return {
      start: {
        line: loc.start.line,
        column: loc.start.column
      },
      end: {
        line: loc.end.line,
        column: loc.end.column
      }
    };
  }

  //
  // Generates an AST representing an expression that will increment the
  // code coverage counter.
  //
  function increase(context, type, id, index) {
    var wrap = index != null
    // If `index` present, turn `x` into `x[index]`.
    ? function (x) {
      return t.memberExpression(x, t.numericLiteral(index), true);
    } : function (x) {
      return x;
    };
    return t.unaryExpression('++', wrap(t.memberExpression(t.memberExpression(getData(context).id, t.identifier(type)), t.stringLiteral(id), true)));
  }

  //
  // Adds coverage traking expression to a path.
  //
  // - If it’s a statement (`a`), turns into `++coverage; a`.
  // - If it’s an expression (`x`), turns into `(++coverage, x)`.
  //
  function instrument(path, increment) {
    if (path.isStatement()) {
      path.insertBefore(t.expressionStatement(increment));
    } else if (path.isExpression()) {
      path.replaceWith(t.sequenceExpression([increment, path.node]));
    } else {
      throw new Error('wtf? I can’t cover a ' + path.node.type + '!!!!??');
    }
  }

  //
  // Adds coverage to any statement.
  //
  function instrumentStatement(context, path) {
    var node = path.node;

    // Don’t cover code generated by Babel.
    if (!node.loc) return;

    // Make sure we don’t cover already instrumented code (only applies to statements).
    // XXX: Hacky node mutation again. PRs welcome!
    if (node.__coverage__instrumented) return;
    node.__coverage__instrumented = true;

    var data = getData(context);
    var id = String(data.nextId.s++);
    data.base.s[id] = 0;
    data.base.statementMap[id] = locToObject(node.loc);
    instrument(path, increase(context, 's', id));
  }

  //
  // Returns the next branch ID and adds the information to `branchMap` object.
  //
  function nextBranchId(context, line, type, locations) {
    var data = getData(context);
    var id = String(data.nextId.b++);
    data.base.b[id] = locations.map(function () {
      return 0;
    });
    data.base.branchMap[id] = { line: line, type: type, locations: locations.map(locToObject) };
    return id;
  }

  //
  // `a` => `++coverage; a` For most common type of statements.
  //
  function coverStatement(path) {
    instrumentStatement(this, path);
  }

  //
  // `var x = 1` => `var x = (++coverage, 1)`
  //
  function coverVariableDeclarator(path) {
    if (!path.node.init) return;
    instrumentStatement(this, path.get('init'));
  }

  //
  // Adds branch coverage to `if` statements.
  //
  function coverIfStatement(path) {
    if (!path.node.loc) return;
    instrumentStatement(this, path);
    if (!path.get('consequent').node) path.set('consequent', t.emptyStatement());
    if (!path.get('alternate').node) path.set('alternate', t.emptyStatement());
    var node = path.node;
    var loc1 = node.consequent.loc || node.loc;
    var loc2 = node.alternate.loc || loc1;
    var id = nextBranchId(this, node.loc.start.line, 'if', [loc1, loc2]);
    instrument(path.get('consequent'), increase(this, 'b', id, 0));
    instrument(path.get('alternate'), increase(this, 'b', id, 1));
  }

  //
  // Adds branch coverage to `switch` statements.
  //
  function coverSwitchStatement(path) {
    if (!path.node.loc) return;
    instrumentStatement(this, path);
    var validCases = path.get('cases').filter(function (p) {
      return p.node.loc;
    });
    var id = nextBranchId(this, path.node.loc.start.line, 'switch', validCases.map(function (p) {
      return p.node.loc;
    }));
    var index = 0;
    var _iteratorNormalCompletion = true;
    var _didIteratorError = false;
    var _iteratorError = undefined;

    try {
      for (var _iterator = validCases[Symbol.iterator](), _step; !(_iteratorNormalCompletion = (_step = _iterator.next()).done); _iteratorNormalCompletion = true) {
        var p = _step.value;

        if (p.node.test) {
          instrumentStatement(this, p.get('test'));
        }
        p.node.consequent.unshift(increase(this, 'b', id, index++));
      }
    } catch (err) {
      _didIteratorError = true;
      _iteratorError = err;
    } finally {
      try {
        if (!_iteratorNormalCompletion && _iterator.return) {
          _iterator.return();
        }
      } finally {
        if (_didIteratorError) {
          throw _iteratorError;
        }
      }
    }
  }

  //
  // `for (;; x)` => `for (;; ++coverage, x)`.
  // Because the increment may be stopped in the first iteration due to `break`.
  //
  function coverForStatement(path) {
    instrumentStatement(this, path);
    if (path.get('update').node) {
      instrumentStatement(this, path.get('update'));
    }
  }

  //
  // Covers a function.
  //
  function coverFunction(path) {
    if (!path.node.loc) return;
    var node = path.node;
    var data = getData(this);
    var id = String(data.nextId.f++);
    var nameOf = function nameOf(namedNode) {
      return namedNode && namedNode.id && namedNode.id.name || null;
    };
    data.base.f[id] = 0;
    data.base.fnMap[id] = {
      name: nameOf((0, _babelHelperFunctionName2.default)(path)), // I love Babel!
      line: node.loc.start.line,
      loc: locToObject(node.loc)
    };
    var increment = increase(this, 'f', id);
    var body = path.get('body');
    if (body.isBlockStatement()) {
      body.node.body.unshift(t.expressionStatement(increment));
    } else if (body.isExpression()) {
      body.replaceWith(t.sequenceExpression([increment, body.node]));
    } else {
      throw new Error('wtf?? Can’t cover function with ' + body.node.type);
    }
  }

  //
  // `a ? b : c` => `a ? (++coverage, b) : (++coverage, c)`.
  // Also adds branch coverage.
  //
  function coverConditionalExpression(path) {
    if (!path.node.loc) return;
    var node = path.node;
    var loc1 = node.consequent.loc || node.loc;
    var loc2 = node.alternate.loc || loc1;
    var id = nextBranchId(this, node.loc.start.line, 'cond-expr', [loc1, loc2]);
    instrumentStatement(this, path.get('consequent'));
    instrumentStatement(this, path.get('alternate'));
    instrument(path.get('consequent'), increase(this, 'b', id, 0));
    instrument(path.get('alternate'), increase(this, 'b', id, 1));
  }

  //
  // `a || b` => `a || (++coverage, b)`. Required due to short circuiting.
  // Also adds branch coverage.
  //
  function coverLogicalExpression(path) {
    if (!path.node.loc) return;
    var node = path.node;
    var loc1 = node.left.loc || node.loc;
    var loc2 = node.right.loc || loc1;
    var id = nextBranchId(this, node.loc.start.line, 'binary-expr', [loc1, loc2]);
    // left is always evaluated!!
    instrumentStatement(this, path.get('right'));
    instrument(path.get('left'), increase(this, 'b', id, 0));
    instrument(path.get('right'), increase(this, 'b', id, 1));
  }

  return {
    visitor: {
      //
      // Shamelessly copied from istanbul.
      //
      ExpressionStatement: coverStatement,
      BreakStatement: coverStatement,
      ContinueStatement: coverStatement,
      DebuggerStatement: coverStatement,
      ReturnStatement: coverStatement,
      ThrowStatement: coverStatement,
      TryStatement: coverStatement,
      VariableDeclarator: coverVariableDeclarator,
      IfStatement: coverIfStatement,
      ForStatement: coverForStatement,
      ForInStatement: coverStatement,
      ForOfStatement: coverStatement,
      WhileStatement: coverStatement,
      DoWhileStatement: coverStatement,
      SwitchStatement: coverSwitchStatement,
      ArrowFunctionExpression: coverFunction,
      FunctionExpression: coverFunction,
      FunctionDeclaration: coverFunction,
      LabeledStatement: coverStatement,
      ConditionalExpression: coverConditionalExpression,
      LogicalExpression: coverLogicalExpression,

      Program: {
        enter: function enter(path) {
          // Save the variable name used for tracking coverage.
          getData(this).id = path.scope.generateUidIdentifier('__coverage__file');
        },
        exit: function exit(path) {
          var _path$node$body;

          // Prepends the coverage runtime.
          var realPath = getRealpath(this.file.opts.filename);
          (_path$node$body = path.node.body).unshift.apply(_path$node$body, _toConsumableArray(coverageTemplate({
            GLOBAL: path.scope.generateUidIdentifier('__coverage__global'),
            COVERAGE: path.scope.generateUidIdentifier('__coverage__object'),
            FILE_COVERAGE: getData(this).id,
            PATH: t.stringLiteral(realPath),
            INITIAL: t.stringLiteral(JSON.stringify(getData(this).base))
          })));
        }
      }
    }
  };
};