diff --git a/lib/index.js b/lib/index.js
index 5e4d3c8..c7a5a48 100644
--- a/lib/index.js
+++ b/lib/index.js
@@ -1,69 +1,66 @@
 'use strict'
 
 var acorn = require('acorn')
-var walk = require('walk-ast')
 var MagicString = require('magic-string')
+var walk = require('estree-walker').walk
 
 function optimizeJs (jsString, opts) {
   opts = opts || {}
   var ast = acorn.parse(jsString)
   var magicString = new MagicString(jsString)
 
-  function walkIt (node) {
-    if (node.type === 'FunctionExpression') {
-      handleFunctionExpression(node)
-    }
-  }
-
-  function handleFunctionExpression (node) {
-    var prePreChar = jsString.charAt(node.start - 2)
-    var preChar = jsString.charAt(node.start - 1)
-    var postChar = jsString.charAt(node.end)
-    var postPostChar = jsString.charAt(node.end + 1)
-
-    // assuming this node is an argument to a function, return true if it itself
-    // is already padded with parentheses
-    function isPaddedArgument (node) {
-      var idx = node.parentNode.arguments.indexOf(node)
-      if (idx === 0) { // first arg
-        if (prePreChar === '(' && preChar === '(' && postChar === ')') { // already padded
-          return true
-        }
-      } else if (idx === node.parentNode.arguments.length - 1) { // last arg
-        if (preChar === '(' && postChar === ')' && postPostChar === ')') { // already padded
-          return true
-        }
-      } else { // middle arg
-        if (preChar === '(' && postChar === ')') { // already padded
-          return true
+  walk(ast, {
+    enter: function (node, parent) {
+      // assuming this node is an argument to a function, return true if it itself
+      // is already padded with parentheses
+      function isPaddedArgument (node) {
+        var idx = parent.arguments.indexOf(node)
+        if (idx === 0) { // first arg
+          if (prePreChar === '(' && preChar === '(' && postChar === ')') { // already padded
+            return true
+          }
+        } else if (idx === parent.arguments.length - 1) { // last arg
+          if (preChar === '(' && postChar === ')' && postPostChar === ')') { // already padded
+            return true
+          }
+        } else { // middle arg
+          if (preChar === '(' && postChar === ')') { // already padded
+            return true
+          }
         }
+        return false
       }
-      return false
-    }
 
-    if (node.parentNode && node.parentNode.type === 'CallExpression') {
-      // this function is getting called itself or
-      // it is getting passed in to another call expression
-      // the else statement is strictly never hit, but I think the code is easier to read this way
-      /* istanbul ignore else */
-      if (node.parentNode.arguments.length && node.parentNode.arguments.indexOf(node) !== -1) {
-        // function passed in to another function. these are almost _always_ executed, e.g.
-        // UMD bundles, Browserify bundles, Node-style errbacks, Promise then()s and catch()s, etc.
-        if (!isPaddedArgument(node)) { // don't double-pad
-          magicString = magicString.insertLeft(node.start, '(')
-            .insertRight(node.end, ')')
-        }
-      } else if (node.parentNode.callee === node) {
-        // this function is getting immediately invoked, e.g. function(){}()
-        if (preChar !== '(') {
-          magicString.insertLeft(node.start, '(')
-            .insertRight(node.end, ')')
+      if (node.type === 'FunctionExpression') {
+        var prePreChar = jsString.charAt(node.start - 2)
+        var preChar = jsString.charAt(node.start - 1)
+        var postChar = jsString.charAt(node.end)
+        var postPostChar = jsString.charAt(node.end + 1)
+
+        if (parent && parent.type === 'CallExpression') {
+          // this function is getting called itself or
+          // it is getting passed in to another call expression
+          // the else statement is strictly never hit, but I think the code is easier to read this way
+          /* istanbul ignore else */
+          if (parent.arguments && parent.arguments.indexOf(node) !== -1) {
+            // function passed in to another function. these are almost _always_ executed, e.g.
+            // UMD bundles, Browserify bundles, Node-style errbacks, Promise then()s and catch()s, etc.
+            if (!isPaddedArgument(node)) { // don't double-pad
+              magicString = magicString.insertLeft(node.start, '(')
+                .insertRight(node.end, ')')
+            }
+          } else if (parent.callee === node) {
+            // this function is getting immediately invoked, e.g. function(){}()
+            if (preChar !== '(') {
+              magicString.insertLeft(node.start, '(')
+                .insertRight(node.end, ')')
+            }
+          }
         }
       }
     }
-  }
+  })
 
-  walk(ast, walkIt)
   var out = magicString.toString()
   if (opts.sourceMap) {
     out += '\n//# sourceMappingURL=' + magicString.generateMap().toUrl()
diff --git a/package.json b/package.json
index 413a26c..587df57 100644
--- a/package.json
+++ b/package.json
@@ -23,8 +23,8 @@
   "dependencies": {
     "acorn": "^3.3.0",
     "concat-stream": "^1.5.1",
+    "estree-walker": "^0.3.0",
     "magic-string": "^0.16.0",
-    "walk-ast": "^0.0.2",
     "yargs": "^4.8.1"
   },
   "devDependencies": {
diff --git a/test/cases/arrow-function-no-error/input.js b/test/cases/arrow-function-no-error/input.js
new file mode 100644
index 0000000..b7e90bd
--- /dev/null
+++ b/test/cases/arrow-function-no-error/input.js
@@ -0,0 +1,7 @@
+// examples taken from https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Functions/Arrow_functions
+var func = x => x * x;
+var func = (x, y) => { return x + y; };
+
+setInterval(() => {
+  this.age++;
+}, 1000);
\ No newline at end of file
diff --git a/test/cases/arrow-function-no-error/output.js b/test/cases/arrow-function-no-error/output.js
new file mode 100644
index 0000000..b7e90bd
--- /dev/null
+++ b/test/cases/arrow-function-no-error/output.js
@@ -0,0 +1,7 @@
+// examples taken from https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Functions/Arrow_functions
+var func = x => x * x;
+var func = (x, y) => { return x + y; };
+
+setInterval(() => {
+  this.age++;
+}, 1000);
\ No newline at end of file