From 9e69ee9cf56e70ed6f81cbd8280c14b68e4505c3 Mon Sep 17 00:00:00 2001
From: brettlangdon <brett@blangdon.com>
Date: Mon, 26 Jan 2015 08:06:20 -0500
Subject: [PATCH] rewrite traversal code to not duplicate

---
 lib/index.js | 66 ++++++++++++++++++++++++++++------------------------
 1 file changed, 36 insertions(+), 30 deletions(-)

diff --git a/lib/index.js b/lib/index.js
index d3efcb4..b4662d6 100644
--- a/lib/index.js
+++ b/lib/index.js
@@ -7,53 +7,59 @@ var mapName = function(elm){
     return elm.name;
 };
 
-var findReturns = function(expr){
-    var returns = []
+var traverseForType = function(type, expr, callback){
+    if(!expr || !expr.type){
+        return;
+    }
+
     if(expr.type === 'BlockStatement'){
         expr.body.forEach(function(expr){
-            if(expr.type === 'ReturnStatement'){
-                if(expr.argument.type === 'Literal'){
-                    returns.push(expr.argument.value);
+            if(expr.type === type){
+                if(expr.argument.type === 'NewExpression'){
+                    callback(expr.argument.callee.name);
+                } else if(expr.argument.type === 'Literal'){
+                    callback(expr.argument.raw);
+                } else if(expr.argument.type === 'Identifier'){
+                    callback(expr.argument.name);
                 } else if(expr.argument.type === 'ObjectExpression'){
-                    returns.push('Object');
+                    returns.push('[Object]');
                 } else if(expr.argument.type === 'ArrayExpression'){
-                    returns.push('Array');
+                    returns.push('[Array]');
                 } else if(expr.argument.type === 'FunctionExpression'){
-                    returns.push('Function');
-                } else if(expr.argument.type === 'Identifier'){
-                    returns.push(expr.argument.name);
+                    returns.push('[Function]');
                 }
-            } else if(expr.type === 'IfStatement'){
-                returns = returns.concat(findReturns(expr));
+            } else {
+                traverseForType(type, expr, callback);
             }
         });
     } else if(expr.type === 'IfStatement'){
-        returns = returns.concat(findReturns(expr.consequent));
+        traverseForType(type, expr.consequent, callback)
         if(expr.alternate){
-            returns = returns.concat(findReturns(expr.alternate));
+            traverseForType(type, expr.alternate, callback);
         }
+    } else if(expr.type === 'TryStatement'){
+        traverseForType(type, expr.block, callback);
+        expr.handlers.forEach(function(expr){
+            traverseForType(type, expr.body, callback);
+        });
+    } else if(expr.type === 'WhileStatement'){
+        traverseForType(type, expr.body, callback);
     }
+}
+
+var findReturns = function(expr){
+    var returns = []
+    traverseForType('ReturnStatement', expr, function(name){
+        returns.push(name);
+    });
     return returns;
 };
 
 var findRaises = function(expr){
     var raises = []
-    if(expr.type === 'BlockStatement'){
-        expr.body.forEach(function(expr){
-            if(expr.type === 'ThrowStatement'){
-                if(expr.argument.type === 'NewExpression'){
-                    raises.push(expr.argument.callee.name);
-                }
-            } else if(expr.type === 'IfStatement'){
-                raises = raises.concat(findRaises(expr));
-            }
-        });
-    } else if(expr.type === 'IfStatement'){
-        raises = raises.concat(findRaises(expr.consequent));
-        if(expr.alternate){
-            raises = raises.concat(findRaises(expr.alternate));
-        }
-    }
+    traverseForType('ThrowStatement', expr, function(name){
+        raises.push(name);
+    });
     return raises;
 };