     (x) -> '"' + x + '"'
   ]
 
-  originalSource = code.split('\n').map (line) -> sourceMappings.reduce(((src, f) -> f(src)), line)
-  originalSource = originalSource.slice(0, -1) if _.last(originalSource) == '""' # useless trimming - just to keep the semantics the same as for jscoverage
-
-  output = []
-  output.push "/* automatically generated by jscov - do not edit */"
-  output.push "if (typeof _$jscoverage === 'undefined') _$jscoverage = {};"
-  output.push "if (! _$jscoverage['#{filename}']) {"
-  output.push "  _$jscoverage['#{filename}'] = [];"
-  trackedLines.forEach (line) ->
-    output.push "  _$jscoverage['#{filename}'][#{line}] = 0;"
-  output.push "}"
-  output.push escodegen.generate(ast, { indent: "  " })
-  output.push "_$jscoverage['#{filename}'].source = [" + originalSource.join(",") + "];"
-
-  output.join('\n')
+  (originalCode, coveredCode, filename, trackedLines) ->
+
+    originalSource = originalCode.split('\n').map (line) -> sourceMappings.reduce(((src, f) -> f(src)), line)
+
+    # useless trimming - just to keep the semantics the same as for jscoverage
+    originalSource = originalSource.slice(0, -1) if _.last(originalSource) == '""'
+
+    output = []
+    output.push "/* automatically generated by jscov - do not edit */"
+    output.push "if (typeof _$jscoverage === 'undefined') _$jscoverage = {};"
+    output.push "if (! _$jscoverage['#{filename}']) {"
+    output.push "  _$jscoverage['#{filename}'] = [];"
+    trackedLines.forEach (line) ->
+      output.push "  _$jscoverage['#{filename}'][#{line}] = 0;"
+    output.push "}"
+    output.push coveredCode
+    output.push "_$jscoverage['#{filename}'].source = [" + originalSource.join(",") + "];"
+
+    output.join('\n')
+
+
+exports.rewriteSource = (code, filename) ->
+
+  injectList = []
+
+  ast = esprima.parse(code, { loc: true })
+
+  formatTree(ast)
+
+  escodegen.traverse ast,
+    enter: (node) ->
+      if node.type in ['BlockStatement', 'Program']
+        node.body = _.flatten node.body.map (x) ->
+          if x.type == 'FunctionDeclaration'
+            injectList.push(x.body.loc.start.line)
+            [inject(x.body, filename), x]
+          else
+            injectList.push(x.loc.start.line)
+            [inject(x, filename), x]
+      if node.type == 'SwitchCase'
+        node.consequent = _.flatten node.consequent.map (x) ->
+          injectList.push(x.loc.start.line)
+          [inject(x, filename), x]
+      if node.type == 'IfStatement'
+        ['consequent', 'alternate'].forEach (src) ->
+          if node[src]? && node[src].type != 'BlockStatement'
+            node[src] =
+              type: 'BlockStatement'
+              body: [node[src]]
+      if node.type in ['ForInStatement', 'ForStatement', 'WhileStatement', 'WithStatement', 'DoWhileStatement'] && node.body? && node.body.type != 'BlockStatement'
+        node.body =
+          type: 'BlockStatement'
+          body: [node.body]
+
+  trackedLines = _.sortBy(_.unique(injectList), _.identity)
+
+  writeFile(code, escodegen.generate(ast, { indent: "  " }), filename, trackedLines)
 
 
 exports.rewriteFolder = (source, target, options, callback) ->

wrench = require 'wrench'
coffee = require 'coffee-script'

reservedWords = [
  'break'
  'case'
  'catch'
  'continue'
  'debugger'
  'default'
  'delete'
  'do'
  'else'
  'finally'
  'for'
  'function'
  'if'
  'in'
  'instanceof'
  'new'
  'return'
  'switch'
  'this'
  'throw'
  'try'
  'typeof'
  'var'
  'void'
  'while'
  'with'

  # These are not keyword according to JavaScript, but JSCoverage treats them as if they were.
  # Just follow suit...
  'throws'
]

isValidIdentifier = (name) ->
  name? && (name.toString().match(/^[_a-zA-Z][_a-zA-Z0-9]*$/) || name.toString().match(/^[1-9][0-9]*$/)) && reservedWords.indexOf(name) == -1

exports.rewriteSource = (code, filename) ->

  injectList = []

  inject = (x) ->
    injectList.push(x.loc.start.line)

    type: 'ExpressionStatement'
    expression:
      type: 'UpdateExpression'
      operator: '++'
      prefix: false
      argument:
        type: 'MemberExpression'
        computed: true
        property:
          type: 'Literal'
          value: x.loc.start.line
        object:
          type: 'MemberExpression'
          computed: true
          object:
            type: 'Identifier'
            name: '_$jscoverage'
          property:
            type: 'Literal'
            value: filename


  replaceNode = (node, newVal) ->
    props = Object.getOwnPropertyNames(node)
    props.forEach (prop) -> delete node[prop]

    Object.keys(newVal).forEach (prop) ->
      node[prop] = newVal[prop]

  makeLiteral = (node, value) ->
    if typeof value == 'number' && value < 0
      replaceNode(node, {
        type: 'UnaryExpression'
        operator: '-'
        argument: {
          type: 'Literal'
          value: -value
        }
      })
    else
      node.type = 'Literal'
      node.value = value


  ast = esprima.parse(code, { loc: true })

  # formatting (no difference in result, just here to give it exactly the same semantics as JSCoverage)
  # this block should be part of the comparisons in the tests; not the source
  format = true

            node.computed = false
            node.property = { type: 'Identifier', name: node.property.value }
        if node.type == 'BinaryExpression' && node.left.type == 'Literal' && node.right.type == 'Literal'
          if typeof node.left.value == 'string' && typeof node.right.value == 'string' && node.operator == '+'
            makeLiteral(node, node.left.value + node.right.value)
            format = true
          if typeof node.left.value == 'number' && typeof node.right.value == 'number' && node.operator in ['+', '-', '*', '%', '/', '<<', '>>', '>>>']
            makeLiteral(node, evalBinaryExpression(node.left.value, node.operator, node.right.value))
            format = true
        if node.type == 'UnaryExpression' && node.argument.type == 'Literal'
          if node.operator == '!'
            makeLiteral(node, !node.argument.value)
            format = true
          if node.operator == "~" && typeof node.argument.value == 'number'
            makeLiteral(node, ~node.argument.value)
            format = true
        if node.type == 'ConditionalExpression' && node.test.type == 'Literal'
          if typeof node.test.value == 'string' || typeof node.test.value == 'number' || typeof node.test.value == 'boolean'
            if node.test.value
              replaceNode(node, node.consequent)
            else
              replaceNode(node, node.alternate)  


inject = (x, filename) ->
  type: 'ExpressionStatement'
  expression:
    type: 'UpdateExpression'
    operator: '++'
    prefix: false
    argument:
      type: 'MemberExpression'
      computed: true
      property:
        type: 'Literal'
        value: x.loc.start.line
      object:
        type: 'MemberExpression'
        computed: true
        object:
          type: 'Identifier'
          name: '_$jscoverage'
        property:
          type: 'Literal'
          value: filename


writeFile = do ->
  sourceMappings = [
    (x) -> x.replace(/&/g, '&amp;')
    (x) -> x.replace(/</g, '&lt;')
