allow `collapse_vars` across conditional branches (#2867)

2 files changed, 213 insertions, 38 deletions

diff --git a/lib/compress.js b/lib/compress.js
index 77636cb2..6f9d64f9 100644
--- a/lib/compress.js
+++ b/lib/compress.js
@@ -952,32 +952,11 @@ merge(Compressor.prototype, {
             var stat_index = statements.length;
             var scanner = new TreeTransformer(function(node, descend) {
                 if (abort) return node;
-                // Scan case expressions first in a switch statement
-                if (node instanceof AST_Switch) {
-                    if (!hit) {
-                        if (node !== hit_stack[hit_index]) return node;
-                        hit_index++;
-                    }
-                    node.expression = node.expression.transform(scanner);
-                    for (var i = 0, len = node.body.length; !abort && i < len; i++) {
-                        var branch = node.body[i];
-                        if (branch instanceof AST_Case) {
-                            if (!hit) {
-                                if (branch !== hit_stack[hit_index]) continue;
-                                hit_index++;
-                            }
-                            branch.expression = branch.expression.transform(scanner);
-                            if (side_effects || !replace_all) break;
-                        }
-                    }
-                    abort = true;
-                    return node;
-                }
                 // Skip nodes before `candidate` as quickly as possible
                 if (!hit) {
                     if (node !== hit_stack[hit_index]) return node;
                     hit_index++;
-                    if (hit_index < hit_stack.length) return;
+                    if (hit_index < hit_stack.length) return handle_custom_scan_order(node);
                     hit = true;
                     stop_after = find_stop(node, 0);
                     if (stop_after === node) abort = true;
@@ -997,10 +976,21 @@ merge(Compressor.prototype, {
                     abort = true;
                     return node;
                 }
+                // Stop only if candidate is found within conditional branches
+                if (!stop_if_hit && (side_effects || !replace_all)
+                    && (parent instanceof AST_Binary && lazy_op(parent.operator) && parent.left !== node
+                        || parent instanceof AST_Conditional && parent.condition !== node
+                        || parent instanceof AST_If && parent.condition !== node)) {
+                    stop_if_hit = parent;
+                }
                 // Replace variable with assignment when found
                 if (can_replace
                     && !(node instanceof AST_SymbolDeclaration)
                     && lhs.equivalent_to(node)) {
+                    if (stop_if_hit) {
+                        abort = true;
+                        return node;
+                    }
                     if (is_lhs(node, parent)) {
                         if (value_def) replaced++;
                         return node;
@@ -1056,18 +1046,15 @@ merge(Compressor.prototype, {
                     || (sym = is_lhs(node.left, node))
                         && (sym instanceof AST_PropAccess || sym.name in lvalues)
                     || may_throw
-                        && (in_try ? node.has_side_effects(compressor) : side_effects_external(node))
-                    || (side_effects || !replace_all)
-                        && (parent instanceof AST_Binary && lazy_op(parent.operator)
-                            || parent instanceof AST_Conditional
-                            || parent instanceof AST_If)) {
+                        && (in_try ? node.has_side_effects(compressor) : side_effects_external(node))) {
                     stop_after = node;
                     if (node instanceof AST_Scope) abort = true;
                 }
-                // Skip (non-executed) functions
-                if (node instanceof AST_Scope) return node;
+                return handle_custom_scan_order(node);
             }, function(node) {
-                if (!abort && stop_after === node) abort = true;
+                if (abort) return;
+                if (stop_after === node) abort = true;
+                if (stop_if_hit === node) stop_if_hit = null;
             });
             var multi_replacer = new TreeTransformer(function(node) {
                 if (abort) return node;
@@ -1106,6 +1093,7 @@ merge(Compressor.prototype, {
                     var candidate = hit_stack[hit_stack.length - 1];
                     var value_def = null;
                     var stop_after = null;
+                    var stop_if_hit = null;
                     var lhs = get_lhs(candidate);
                     if (!lhs || is_lhs_read_only(lhs) || lhs.has_side_effects(compressor)) continue;
                     // Locate symbols which may execute code outside of scanning range
@@ -1149,6 +1137,28 @@ merge(Compressor.prototype, {
                 }
             }
 
+            function handle_custom_scan_order(node) {
+                // Skip (non-executed) functions
+                if (node instanceof AST_Scope) return node;
+                // Scan case expressions first in a switch statement
+                if (node instanceof AST_Switch) {
+                    node.expression = node.expression.transform(scanner);
+                    for (var i = 0, len = node.body.length; !abort && i < len; i++) {
+                        var branch = node.body[i];
+                        if (branch instanceof AST_Case) {
+                            if (!hit) {
+                                if (branch !== hit_stack[hit_index]) continue;
+                                hit_index++;
+                            }
+                            branch.expression = branch.expression.transform(scanner);
+                            if (side_effects || !replace_all) break;
+                        }
+                    }
+                    abort = true;
+                    return node;
+                }
+            }
+
             function extract_args() {
                 var iife, fn = compressor.self();
                 if (fn instanceof AST_Function
@@ -1265,18 +1275,49 @@ merge(Compressor.prototype, {
                 hit_stack.pop();
             }
 
-            function find_stop(node, level) {
+            function find_stop(node, level, write_only) {
                 var parent = scanner.parent(level);
-                if (parent instanceof AST_Binary) return node;
+                if (parent instanceof AST_Assign) {
+                    if (write_only
+                        && !(parent.left instanceof AST_PropAccess
+                            || parent.left.name in lvalues)) {
+                        return find_stop(parent, level + 1, write_only);
+                    }
+                    return node;
+                }
+                if (parent instanceof AST_Binary) {
+                    if (write_only && (!lazy_op(parent.operator) || parent.left === node)) {
+                        return find_stop(parent, level + 1, write_only);
+                    }
+                    return node;
+                }
                 if (parent instanceof AST_Call) return node;
                 if (parent instanceof AST_Case) return node;
-                if (parent instanceof AST_Conditional) return node;
-                if (parent instanceof AST_Definitions) return find_stop(parent, level + 1);
-                if (parent instanceof AST_Exit) return node;
-                if (parent instanceof AST_If) return node;
+                if (parent instanceof AST_Conditional) {
+                    if (write_only && parent.condition === node) {
+                        return find_stop(parent, level + 1, write_only);
+                    }
+                    return node;
+                }
+                if (parent instanceof AST_Definitions) {
+                    return find_stop(parent, level + 1, true);
+                }
+                if (parent instanceof AST_Exit) {
+                    return write_only ? find_stop(parent, level + 1, write_only) : node;
+                }
+                if (parent instanceof AST_If) {
+                    if (write_only && parent.condition === node) {
+                        return find_stop(parent, level + 1, write_only);
+                    }
+                    return node;
+                }
                 if (parent instanceof AST_IterationStatement) return node;
-                if (parent instanceof AST_Sequence) return find_stop(parent, level + 1);
-                if (parent instanceof AST_SimpleStatement) return find_stop(parent, level + 1);
+                if (parent instanceof AST_Sequence) {
+                    return find_stop(parent, level + 1, parent.tail_node() !== node);
+                }
+                if (parent instanceof AST_SimpleStatement) {
+                    return find_stop(parent, level + 1, true);
+                }
                 if (parent instanceof AST_Switch) return node;
                 if (parent instanceof AST_VarDef) return node;
                 return null;
diff --git a/test/compress/collapse_vars.js b/test/compress/collapse_vars.js
index 948389b1..a2571c24 100644
--- a/test/compress/collapse_vars.js
+++ b/test/compress/collapse_vars.js
@@ -4249,3 +4249,137 @@ issue_2858: {
     }
     expect_stdout: "undefined"
 }
+
+cond_branch_1: {
+    options = {
+        collapse_vars: true,
+        sequences: true,
+        unused: true,
+    }
+    input: {
+        function f1(b, c) {
+            var log = console.log;
+            var a = ++c;
+            if (b) b++;
+            log(a, b);
+        }
+        function f2(b, c) {
+            var log = console.log;
+            var a = ++c;
+            b && b++;
+            log(a, b);
+        }
+        function f3(b, c) {
+            var log = console.log;
+            var a = ++c;
+            b ? b++ : b--;
+            log(a, b);
+        }
+        f1(1, 2);
+        f2(3, 4);
+        f3(5, 6);
+    }
+    expect: {
+        function f1(b, c) {
+            var log = console.log;
+            if (b) b++;
+            log(++c, b);
+        }
+        function f2(b, c) {
+            var log = console.log;
+            b && b++,
+            log(++c, b);
+        }
+        function f3(b, c) {
+            var log = console.log;
+            b ? b++ : b--,
+            log(++c, b);
+        }
+        f1(1, 2),
+        f2(3, 4),
+        f3(5, 6);
+    }
+    expect_stdout: [
+        "3 2",
+        "5 4",
+        "7 6",
+    ]
+}
+
+cond_branch_2: {
+    options = {
+        collapse_vars: true,
+        sequences: true,
+        unused: true,
+    }
+    input: {
+        function f1(b, c) {
+            var log = console.log;
+            var a = ++c;
+            if (b) b += a;
+            log(a, b);
+        }
+        function f2(b, c) {
+            var log = console.log;
+            var a = ++c;
+            b && (b += a);
+            log(a, b);
+        }
+        function f3(b, c) {
+            var log = console.log;
+            var a = ++c;
+            b ? b += a : b--;
+            log(a, b);
+        }
+        f1(1, 2);
+        f2(3, 4);
+        f3(5, 6);
+    }
+    expect: {
+        function f1(b, c) {
+            var log = console.log;
+            var a = ++c;
+            if (b) b += a;
+            log(a, b);
+        }
+        function f2(b, c) {
+            var log = console.log;
+            var a = ++c;
+            b && (b += a),
+            log(a, b);
+        }
+        function f3(b, c) {
+            var log = console.log;
+            var a = ++c;
+            b ? b += a : b--,
+            log(a, b);
+        }
+        f1(1, 2),
+        f2(3, 4),
+        f3(5, 6);
+    }
+    expect_stdout: [
+        "3 4",
+        "5 8",
+        "7 12",
+    ]
+}
+
+cond_branch_switch: {
+    options = {
+        collapse_vars: true,
+    }
+    input: {
+        var c = 0;
+        if (c = 1 + c, 0) switch (c = 1 + c) {
+        }
+        console.log(c);
+    }
+    expect: {
+        var c = 0;
+        if (c = 1 + c, 0) switch (c = 1 + c) {
+        }
+        console.log(c);
+    }
+    expect_stdout: "1"
+}