enhance `merge_vars` (#4105)

2 files changed, 85 insertions, 17 deletions

diff --git a/lib/compress.js b/lib/compress.js
index b605dda7..51f65268 100644
--- a/lib/compress.js
+++ b/lib/compress.js
@@ -4309,7 +4309,7 @@ merge(Compressor.prototype, {
 
     AST_Scope.DEFMETHOD("merge_variables", function(compressor) {
         if (!compressor.option("merge_vars")) return;
-        var self = this, segment = self;
+        var self = this, segment;
         var first = [], last = [], index = 0;
         var references = Object.create(null);
         var prev = Object.create(null);
@@ -4334,8 +4334,9 @@ merge(Compressor.prototype, {
             if (node instanceof AST_Conditional) {
                 node.condition.walk(tw);
                 var save = segment;
-                segment = node;
+                segment = node.consequent;
                 node.consequent.walk(tw);
+                segment = node.alternative;
                 node.alternative.walk(tw);
                 segment = save;
                 return true;
@@ -4364,7 +4365,10 @@ merge(Compressor.prototype, {
                 var save = segment;
                 segment = node;
                 node.body.walk(tw);
-                if (node.alternative) node.alternative.walk(tw);
+                if (node.alternative) {
+                    segment = node.alternative;
+                    node.alternative.walk(tw);
+                }
                 segment = save;
                 return true;
             }
@@ -4388,10 +4392,21 @@ merge(Compressor.prototype, {
             }
             if (node instanceof AST_Switch) {
                 node.expression.walk(tw);
-                var save = segment;
-                segment = node;
+                var save = segment, first = true;
                 node.body.forEach(function(branch) {
-                    branch.walk(tw);
+                    if (branch instanceof AST_Default) return;
+                    if (first) {
+                        first = false;
+                    } else {
+                        segment = branch.expression;
+                    }
+                    branch.expression.walk(tw);
+                });
+                node.body.forEach(function(branch) {
+                    segment = branch;
+                    branch.body.forEach(function(stat) {
+                        stat.walk(tw);
+                    });
                 });
                 segment = save;
                 return true;
@@ -4410,7 +4425,10 @@ merge(Compressor.prototype, {
                 node.body.forEach(function(branch) {
                     branch.walk(tw);
                 });
-                if (node.bcatch) node.bcatch.walk(tw);
+                if (node.bcatch) {
+                    segment = node.bcatch;
+                    node.bcatch.walk(tw);
+                }
                 segment = save;
                 if (node.bfinally) node.bfinally.walk(tw);
                 return true;
@@ -4437,11 +4455,16 @@ merge(Compressor.prototype, {
             if (!(def.id in prev)) continue;
             if (!references[def.id]) continue;
             while (def.id in merged) def = merged[def.id];
+            var skipped = [];
             do {
                 var tail = last.pop();
                 if (!tail) continue;
                 if (tail.index > head.index) continue;
                 if (!references[tail.definition.id]) continue;
+                if (references[def.id].segment !== references[tail.definition.id].segment) {
+                    skipped.unshift(tail);
+                    continue;
+                }
                 var orig = [], refs = [];
                 references[tail.definition.id].forEach(function(sym) {
                     push(sym);
@@ -4456,6 +4479,7 @@ merge(Compressor.prototype, {
                 merged[tail.definition.id] = def;
                 break;
             } while (last.length);
+            if (skipped.length) last = last.concat(skipped);
         }
 
         function read(def) {
@@ -4468,21 +4492,28 @@ merge(Compressor.prototype, {
 
         function mark(sym, write_only) {
             var def = sym.definition();
-            if (segment !== self) references[def.id] = false;
             if (def.id in references) {
-                if (!references[def.id]) return;
-                references[def.id].push(sym);
+                var refs = references[def.id];
+                if (!refs) return;
+                if (refs.segment !== segment) return references[def.id] = false;
+                refs.push(sym);
                 if (def.id in prev) last[prev[def.id]] = null;
                 read(def);
-            } else if (compressor.exposed(def) || self.variables.get(def.name) !== def) {
+            } else if (self.variables.get(def.name) !== def || compressor.exposed(def)) {
                 references[def.id] = false;
             } else {
-                references[def.id] = [ sym ];
-                if (!write_only) return read(def);
-                first.push({
-                    index: index++,
-                    definition: def,
-                });
+                var refs = [ sym ];
+                references[def.id] = refs;
+                if (write_only) {
+                    refs.segment = segment;
+                    first.push({
+                        index: index++,
+                        definition: def,
+                    });
+                } else {
+                    refs.segment = self;
+                    read(def);
+                }
             }
         }
 
diff --git a/test/compress/merge_vars.js b/test/compress/merge_vars.js
index 6923fe73..4c5f6dac 100644
--- a/test/compress/merge_vars.js
+++ b/test/compress/merge_vars.js
@@ -76,6 +76,43 @@ merge_toplevel: {
     ]
 }
 
+segment: {
+    options = {
+        merge_vars: true,
+        toplevel: true,
+    }
+    input: {
+        var a = "foo";
+        console.log(a);
+        for (var c, i = 0; i < 1; i++) {
+            var b = "bar";
+            console.log(b);
+            c = "baz";
+            console.log(c);
+        }
+        var d = "moo";
+        console.log(d);
+    }
+    expect: {
+        var d = "foo";
+        console.log(d);
+        for (var c, i = 0; i < 1; i++) {
+            var c = "bar";
+            console.log(c);
+            c = "baz";
+            console.log(c);
+        }
+        var d = "moo";
+        console.log(d);
+    }
+    expect_stdout: [
+        "foo",
+        "bar",
+        "baz",
+        "moo",
+    ]
+}
+
 init_scope_vars: {
     options = {
         merge_vars: true,