diff --git a/README.md b/README.md
index 7a30a7d1..f678d008 100644
--- a/README.md
+++ b/README.md
@@ -12,6 +12,8 @@ Note that this does *not* solve the [halting problem](http://en.wikipedia.org/wi
 
 With loop protection in place, it means that a user can enter the code as follows on JS Bin, and the final `console.log` will still work.
 
+The code is transformed from this:
+
 ```js
 while (true) {
   doSomething();
@@ -20,6 +22,21 @@ while (true) {
 console.log('All finished');
 ```
 
+…to this:
+
+```js
+let i = 0;
+var _LP = Date.now();
+while (true) {
+  if (Date.now() - _LP > 100)
+    break;
+
+  doSomething();
+}
+
+console.log('All finished');
+```
+
 ## Usage
 
 The loop protection is a babel transform, so can be used on the server or in the client.
diff --git a/lib/index.js b/lib/index.js
index 149e47d1..f7f6b3f1 100644
--- a/lib/index.js
+++ b/lib/index.js
@@ -9,7 +9,7 @@ const generateBefore = (t, id) =>
     ),
   ]);
 
-const generateInside = ({ t, id, line, timeout, extra } = {}) => {
+const generateInside = ({ t, id, line, ch, timeout, extra } = {}) => {
   return t.ifStatement(
     t.binaryExpression(
       '>',
@@ -25,11 +25,11 @@ const generateInside = ({ t, id, line, timeout, extra } = {}) => {
     ),
     extra
       ? t.blockStatement([
-          t.expressionStatement(
-            t.callExpression(extra, [t.numericLiteral(line)])
-          ),
-          t.breakStatement(),
-        ])
+        t.expressionStatement(
+          t.callExpression(extra, [t.numericLiteral(line), t.numericLiteral(ch)])
+        ),
+        t.breakStatement(),
+      ])
       : t.breakStatement()
   );
 };
@@ -41,6 +41,7 @@ const protect = (t, timeout, extra) => path => {
     t,
     id,
     line: path.node.loc.start.line,
+    ch: path.node.loc.start.column,
     timeout,
     extra,
   });
@@ -60,10 +61,19 @@ module.exports = (timeout = 100, extra = null) => {
     extra = `() => console.error("${string.replace(/"/g, '\\"')}")`;
   }
   return ({ types: t, transform }) => {
-    const callback = extra
-      ? transform(extra).ast.program.body[0].expression
+    const node = extra
+      ? transform(extra).ast.program.body[0]
       : null;
 
+    // console.log(node && node.type)
+
+    let callback = null;
+    if (t.isExpressionStatement(node)) {
+      callback = node.expression;
+    } else if (t.isFunctionDeclaration(node)) {
+      callback = t.functionExpression(null, node.params, node.body);
+    }
+
     return {
       visitor: {
         WhileStatement: protect(t, timeout, callback),
diff --git a/package-lock.json b/package-lock.json
index 90018647..33692306 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,6 +1,6 @@
 {
   "name": "loop-protect",
-  "version": "2.0.0-beta.1",
+  "version": "2.1.0",
   "lockfileVersion": 1,
   "requires": true,
   "dependencies": {
diff --git a/package.json b/package.json
index f3917c70..6fa97980 100644
--- a/package.json
+++ b/package.json
@@ -3,7 +3,7 @@
   "name": "loop-protect",
   "description": "Prevent infinite loops in dynamically eval'd JavaScript.",
   "main": "lib/",
-  "version": "2.1.0",
+  "version": "2.1.2",
   "homepage": "https://github.com/jsbin/loop-protect",
   "repository": {
     "type": "git",
diff --git a/test/callback.test.js b/test/callback.test.js
new file mode 100644
index 00000000..7f21bfc9
--- /dev/null
+++ b/test/callback.test.js
@@ -0,0 +1,57 @@
+/* eslint-env node, jest */
+const Babel = require('babel-standalone');
+
+const code = `let i = 0; while (true) { i++; }; done(i)`;
+
+let done = jest.fn();
+
+beforeEach(() => {
+  done = jest.fn();
+});
+
+const transform = id => code => Babel.transform(new Function(code).toString(), {
+  plugins: [id],
+}).code; // eslint-disable-line no-new-func
+
+const run = code => {
+  // console.log(code);
+  eval(`(${code})()`); // eslint-disable-line no-eval
+};
+
+test('no callback', () => {
+  const id = 'lp1';
+  Babel.registerPlugin(id, require('../lib')(100));
+  const after = transform(id)(code);
+  run(after);
+  expect(done).toBeCalledWith(expect.any(Number));
+});
+
+test('anonymous callback', () => {
+  const id = 'lp2';
+  Babel.registerPlugin(id, require('../lib')(100, line => done(`line: ${line}`)));
+  const after = transform(id)(code);
+  run(after);
+  expect(done).toHaveBeenCalledWith('line: 2');
+});
+
+test('arrow function callback', () => {
+  const id = 'lp3';
+  const callback = line => done(`lp3: ${line}`);
+
+  Babel.registerPlugin(id, require('../lib')(100, callback));
+  const after = transform(id)(code);
+  run(after);
+  expect(done).toHaveBeenCalledWith(`${id}: 2`);
+});
+
+test('named function callback', () => {
+  const id = 'lp4';
+  function callback(line, ch) {
+    done(`lp4: ${line}`);
+  }
+
+  Babel.registerPlugin(id, require('../lib')(100, callback));
+  const after = transform(id)(code);
+  run(after);
+  expect(done).toHaveBeenCalledWith(`${id}: 2`);
+});
diff --git a/test/loop-protect.test.js b/test/loop-protect.test.js
index 553af949..cf79ab5b 100644
--- a/test/loop-protect.test.js
+++ b/test/loop-protect.test.js
@@ -69,8 +69,8 @@ const loopProtect = code =>
   }).code; // eslint-disable-line no-new-func
 const run = code => eval(`(${code})()`); // eslint-disable-line no-eval
 
-describe('loop', function() {
-  beforeEach(function() {
+describe('loop', function () {
+  beforeEach(function () {
     spy = sinon.spy(run);
   });
 
@@ -139,7 +139,7 @@ describe('loop', function() {
     assert(run(compiled) === true);
   });
 
-  it('should ignore comments', function() {
+  it('should ignore comments', function () {
     var c = code.ignorecomments;
     var compiled = loopProtect(c);
     // console.log('\n---------\n' + c + '\n---------\n' + compiled);
@@ -147,7 +147,7 @@ describe('loop', function() {
     assert(result === true);
   });
 
-  it('should rewrite for loops', function() {
+  it('should rewrite for loops', function () {
     var c = code.simplefor;
     var compiled = loopProtect(c);
     assert(compiled !== c);
@@ -161,7 +161,7 @@ describe('loop', function() {
     assert(result === 10);
   });
 
-  it('should handle one liner for with an inline function', function() {
+  it('should handle one liner for with an inline function', function () {
     var c = code.onelineforinline;
     var compiled = loopProtect(c);
     assert(compiled !== c);
@@ -169,7 +169,7 @@ describe('loop', function() {
     assert(result === true, 'value is ' + result);
   });
 
-  it('should rewrite one line for loops', function() {
+  it('should rewrite one line for loops', function () {
     var c = code.onelinefor;
     var compiled = loopProtect(c);
     assert(compiled !== c);
@@ -186,7 +186,7 @@ describe('loop', function() {
     assert(result === 0);
   });
 
-  it('should rewrite one line while loops', function() {
+  it('should rewrite one line while loops', function () {
     var c = code.onelinewhile2;
     var compiled = loopProtect(c);
     assert(compiled !== c);
@@ -195,7 +195,7 @@ describe('loop', function() {
     assert(result === undefined);
   });
 
-  it('should protect infinite while', function() {
+  it('should protect infinite while', function () {
     var c = code.whiletrue;
     var compiled = loopProtect(c);
 
@@ -203,14 +203,14 @@ describe('loop', function() {
     assert(spy(compiled) === true);
   });
 
-  it('should protect infinite for', function() {
+  it('should protect infinite for', function () {
     var c = code.irl1;
     var compiled = loopProtect(c);
     assert(compiled !== c);
     // assert(spy(compiled) === 0);
   });
 
-  it('should allow nested loops to run', function() {
+  it('should allow nested loops to run', function () {
     var c = code.irl2;
     var compiled = loopProtect(c);
     var r = run(compiled);
@@ -218,7 +218,7 @@ describe('loop', function() {
     expect(r).toBe(60000);
   });
 
-  it('should rewrite loops when curlies are on the next line', function() {
+  it('should rewrite loops when curlies are on the next line', function () {
     var c = code.dirtybraces;
     var compiled = loopProtect(c);
     var r = spy(compiled);
@@ -226,7 +226,7 @@ describe('loop', function() {
     assert(r === 10000, r);
   });
 
-  it('should find one liners on multiple lines', function() {
+  it('should find one liners on multiple lines', function () {
     var c = code.onelinenewliner;
     var compiled = loopProtect(c);
     var r = spy(compiled);
@@ -238,21 +238,21 @@ describe('loop', function() {
     assert(r === 10000, 'return value does not match 10000: ' + r);
   });
 
-  it('should handle brackets inside of loop conditionals', function() {
+  it('should handle brackets inside of loop conditionals', function () {
     var c = code.brackets;
     var compiled = loopProtect(c);
     assert(compiled !== c);
     assert(spy(compiled) === 11);
   });
 
-  it('should not corrupt multi-line (on more than one line) loops', function() {
+  it('should not corrupt multi-line (on more than one line) loops', function () {
     var c = code.lotolines;
     var compiled = loopProtect(c);
     assert(compiled !== c);
     assert(spy(compiled) === 8);
   });
 
-  it('should protect do loops', function() {
+  it('should protect do loops', function () {
     var c = code.dowhile;
     var compiled = loopProtect(c);
     assert(compiled !== c);
@@ -274,12 +274,12 @@ describe('loop', function() {
   });
 });
 
-describe('labels', function() {
-  beforeEach(function() {
+describe('labels', function () {
+  beforeEach(function () {
     spy = sinon.spy(run);
   });
 
-  it('should handle continue statements and gotos', function() {
+  it('should handle continue statements and gotos', function () {
     var c = code.continues;
     var compiled = loopProtect(c);
     assert(spy(compiled) === 10);
@@ -289,13 +289,13 @@ describe('labels', function() {
     assert(spy(compiled) === 2);
   });
 
-  it('should handle labels with comments', function() {
+  it('should handle labels with comments', function () {
     var c = code.labelWithComment;
     var compiled = loopProtect(c);
     assert(spy(compiled) === 10);
   });
 
-  it('should handle things that *look* like labels', function() {
+  it('should handle things that *look* like labels', function () {
     var c = code.notlabels2;
     var compiled = loopProtect(c);
     assert(compiled !== c);
@@ -315,7 +315,7 @@ describe('labels', function() {
     // assert(result === 10, 'actual ' + result);
   });
 
-  it('should handle if statement without {}', function() {
+  it('should handle if statement without {}', function () {
     var c = code.loopbehindif;
     var compiled = loopProtect(c);
     assert(compiled !== c);