messing around with asts

counter.js

module.exports = {
  counter: {},
  count: function count(key) {
    if (this.counter[key] === undefined) this.counter[key] = 0
    this.counter[key]++
  },
  output: function output() {
    for (var key in this.counter) {
      console.log(key + ' was called ' + this.counter[key] + ' times')
    }
  }
}

main.js

var Module = require('module')
var es = require('esprima')
var gen = require('escodegen')
var oldCompile = Module.prototype._compile
var counter = require('./counter')

Module.prototype._compile = function wrappedCompile(content, filename) {
  try {
    var p = es.parse(content)
    injectCode(p)
    p.body = es.parse('var counter = require("./counter")').body.concat(p.body)
    content = gen.generate(p)
  } catch (e) {
    console.log('borked things', e.message)
  }
  return oldCompile.bind(this)(content, filename)
}


function injectCode(tokens) {
  if (!tokens || !tokens.body) return

  var seenBlocks = []
  function generateReturnStatement(fromFunction) {
    if (fromFunction === null) return []
    var returnInject = 'console.log("' + fromFunction.name + ' returned")'
    return es.parse(returnInject).body
  }

  var toRelax = Array.prototype.concat([], tokens.body)
  var currentToken
  var currentTokenIsBlockOwner
  var lastFunctionSeen
  var lastBlock

  while(toRelax.length !== 0) {
    currentToken = toRelax.pop()
    currentTokenIsBlockOwner = false
    lastBlock = seenBlocks[seenBlocks.length - 1]


    switch(currentToken.type) {
      case 'FunctionDeclaration':
      case 'FunctionExpression':
      case 'ArrowFunctionExpression':
        injectFunctionBody(currentToken)
        toRelax.push(currentToken.body)
        currentTokenIsBlockOwner = true
        seenBlocks.push({
          isFunction: true,
          leftToRelax: currentToken.body.body.length,
          name: (currentToken.id && currentToken.id.name) || 'anonymous'
        })
        break
      case 'ObjectExpression':
        toRelax = toRelax.concat(currentToken.properties)
        break
      case 'Property':
        toRelax.push(currentToken.value)
        break
      case 'ExpressionStatement':
        toRelax.push(currentToken.expression)
        break
      case 'AssignmentExpression':
        toRelax.push(currentToken.left)
        toRelax.push(currentToken.right)
        break
      case 'BlockStatement':
        var returnStatement = generateReturnStatement(lastFunctionSeen)
        var returnSeen = false
        for (var i = 0; i < currentToken.body.length; i++) {
          var statement = currentToken.body[i]
          if (statement.type === 'ReturnStatement') {
            returnSeen = true
            currentToken.body = spliceIntoArrayAt(currentToken.body, returnStatement, i)
            i += returnStatement.length
          }
          toRelax.push(statement)
        }
        if (!returnSeen && lastBlock.isFunction) {
          currentToken.body = currentToken.body.concat(returnStatement)
        }
        break
      case 'VariableDeclaration':
        toRelax = toRelax.concat(currentToken.declarations)
        break
      case 'VariableDeclarator':
        toRelax.push(currentToken.init)
        break
      case 'WhileStatement':
      case 'ForStatement':
        seenBlocks.push({
          isFunction: false,
          leftToRelax: currentToken.body.body.length
        })

        currentTokenIsBlockOwner = true
        toRelax.push(currentToken.body)
        break
      case 'IfStatement':
        seenBlocks.push({
          isFunction: false,
          leftToRelax: currentToken.consequent.body.length
        })

        if (currentToken.alternate !== null) {
          seenBlocks.push({
            isFunction: false,
            leftToRelax: currentToken.alternate.body.length
          })
        }

        currentTokenIsBlockOwner = true
      case 'ConditionalExpression':
        toRelax.push(currentToken.consequent)
        if (currentToken.alternate !== null) {
          toRelax.push(currentToken.alternate)
        }
        break
    }

    lastFunctionSeen = null;
    for (var i = seenBlocks.length - 1; i >= 0; --i) {
      var block = seenBlocks[i]
      if (block.isFunction) {
        lastFunctionSeen = block
        break
      }
    }

    if (!currentTokenIsBlockOwner && seenBlocks.length !== 0) {
      while (lastBlock.leftToRelax === 0) {
        lastBlock = seenBlocks.pop()
        lastBlock.leftToRelax--
      }
    }

  }
}

function spliceIntoArrayAt(first, second, i) {
  var result = []
  for (var j = 0; j < first.length; j++) {
    if (j === i) {
      for (var k = 0; k < second.length; k++) {
        result.push(second[k])
      }
    }
    result.push(first[j])
  }
  return result
}

function injectFunctionBody(token) {
    injectedCode = 'counter.count("' + ((token.id && token.id.name) || 'anonymous') + '")'
    token.body.body = es.parse(injectedCode).body.concat(token.body.body)
}

var asdf = require('./test.js')
setInterval(asdf.asdf, 100)
setInterval(asdf, 100)
setInterval(counter.output.bind(counter), 3000)

test.js

module.exports = (a) => {return a}
module.exports.asdf = asdf
function asdf() {
  function nested() {
  }
  nested()
  nested()
  var q = false ? nested : function test () {}
  function branchTest() {
    if (true) {
    } else {
    }
    for (var i = 0; i < 10; i++) {
      if (i === 9) {
        return
      }
    }
  }
  q()
  branchTest()
  return
}

var tacos = {
  burrito: function burritos() {
  }
}

for (var i = 0; i < 10; i++) {
}
while(i < 9) {
  function tests() {return}
}
setTimeout(tacos.burrito, 1000)j