Use TreeWalker for more accurate @const results and update tests

2 files changed, 73 insertions, 8 deletions

diff --git a/lib/scope.js b/lib/scope.js
index 22fb150d..144ae48e 100644
--- a/lib/scope.js
+++ b/lib/scope.js
@@ -358,13 +358,31 @@ AST_Symbol.DEFMETHOD("global", function(){
 });
 
 AST_Symbol.DEFMETHOD("has_const_pragma", function() {
-    var token = this.scope.body[0] && this.scope.body[0].start;
-    var comments = token && token.comments_before;
-    if (comments && comments.length > 0) {
-        var last = comments[comments.length - 1];
-        return /@const/.test(last.value);
-    }
-    return false;
+    var symbol = this;
+    var symbol_has_pragma = false;
+    var pragma_found = false;
+    var found_symbol = false;
+    // Walk the current scope, looking for a comment with the @const pragma.
+    // If it exists, mark a bool that will remain true only for the next iteration.
+    // If the next iteration is this symbol, then we return true.
+    // Otherwise we stop descending and get out of here.
+    var tw = new TreeWalker(function(node, descend){
+        // This is our symbol. Was the pragma before this?
+        if (node.name === symbol) {
+            found_symbol = true;
+            symbol_has_pragma = pragma_found;
+        }
+
+        // Look for the /** @const */ pragma
+        var comments_before = node.start && node.start.comments_before;
+        var lastComment = comments_before && comments_before[comments_before.length - 1];
+        pragma_found = lastComment && /@const/.test(lastComment.value);
+
+        // no need to descend after finding our node
+        return found_symbol;
+    });
+    this.scope.walk(tw);
+    return symbol_has_pragma;
 })
 
 AST_Toplevel.DEFMETHOD("_default_mangler_options", function(options){
diff --git a/test/compress/dead-code.js b/test/compress/dead-code.js
index f79b04de..8aad336c 100644
--- a/test/compress/dead-code.js
+++ b/test/compress/dead-code.js
@@ -97,6 +97,7 @@ dead_code_const_declaration: {
         evaluate     : true

     };

     input: {

+        var unused;

         const CONST_FOO = false;

         if (CONST_FOO) {

             console.log("unreachable");

@@ -105,6 +106,7 @@ dead_code_const_declaration: {
         }

     }

     expect: {

+        var unused;

         const CONST_FOO = !1;

         var moo;

         function bar() {}

@@ -120,7 +122,8 @@ dead_code_const_annotation: {
         evaluate     : true

     };

     input: {

-        /** @const*/ var CONST_FOO_ANN = false;

+        var unused;

+        /** @const */ var CONST_FOO_ANN = false;

         if (CONST_FOO_ANN) {

             console.log("unreachable");

             var moo;

@@ -128,8 +131,52 @@ dead_code_const_annotation: {
         }

     }

     expect: {

+        var unused;

         var CONST_FOO_ANN = !1;

         var moo;

         function bar() {}

     }

 }

+

+dead_code_const_annotation_complex_scope: {

+    options = {

+        dead_code    : true,

+        loops        : true,

+        booleans     : true,

+        conditionals : true,

+        evaluate     : true

+    };

+    input: {

+        var unused_var;

+        /** @const */ var test = 'test';

+        /** @const */ var CONST_FOO_ANN = false;

+        var unused_var_2;

+        if (CONST_FOO_ANN) {

+            console.log("unreachable");

+            var moo;

+            function bar() {}

+        }

+        if (test === 'test') {

+            var beef = 'good';

+            /** @const */ var meat = 'beef';

+            var pork = 'bad';

+            if (meat === 'pork') {

+                console.log('also unreachable');

+            } else if (pork === 'good') {

+                console.log('reached, not const');

+            }

+        }

+    }

+    expect: {

+        var unused_var;

+        var test = 'test';

+        var CONST_FOO_ANN = !1;

+        var unused_var_2;

+        var moo;

+        function bar() {}

+        var beef = 'good';

+        var meat = 'beef';

+        var pork = 'bad';

+        'good' === pork && console.log('reached, not const');

+    }

+}