extend fuzzy RHS folding (#3006)

- `a = []; if (1) x();` => `if (a = []) x();`

3 files changed, 138 insertions, 58 deletions

diff --git a/lib/ast.js b/lib/ast.js
index a23dece6..ded3bbfb 100644
--- a/lib/ast.js
+++ b/lib/ast.js
@@ -916,5 +916,25 @@ TreeWalker.prototype = {
                 || node instanceof AST_Break && x instanceof AST_Switch)
                 return x;
         }
+    },
+    in_boolean_context: function() {
+        var self = this.self();
+        for (var i = 0, p; p = this.parent(i); i++) {
+            if (p instanceof AST_SimpleStatement
+                || p instanceof AST_Conditional && p.condition === self
+                || p instanceof AST_DWLoop && p.condition === self
+                || p instanceof AST_For && p.condition === self
+                || p instanceof AST_If && p.condition === self
+                || p instanceof AST_UnaryPrefix && p.operator == "!" && p.expression === self) {
+                return true;
+            }
+            if (p instanceof AST_Binary && (p.operator == "&&" || p.operator == "||")
+                || p instanceof AST_Conditional
+                || p.tail_node() === self) {
+                self = p;
+            } else {
+                return false;
+            }
+        }
     }
 };
diff --git a/lib/compress.js b/lib/compress.js
index 39220d68..f9fe3c63 100644
--- a/lib/compress.js
+++ b/lib/compress.js
@@ -147,27 +147,6 @@ merge(Compressor.prototype, {
                 return true;
         return false;
     },
-    in_boolean_context: function() {
-        if (!this.option("booleans")) return false;
-        var self = this.self();
-        for (var i = 0, p; p = this.parent(i); i++) {
-            if (p instanceof AST_SimpleStatement
-                || p instanceof AST_Conditional && p.condition === self
-                || p instanceof AST_DWLoop && p.condition === self
-                || p instanceof AST_For && p.condition === self
-                || p instanceof AST_If && p.condition === self
-                || p instanceof AST_UnaryPrefix && p.operator == "!" && p.expression === self) {
-                return true;
-            }
-            if (p instanceof AST_Binary && (p.operator == "&&" || p.operator == "||")
-                || p instanceof AST_Conditional
-                || p.tail_node() === self) {
-                self = p;
-            } else {
-                return false;
-            }
-        }
-    },
     compress: function(node) {
         if (this.option("expression")) {
             node.process_expression(true);
@@ -980,7 +959,7 @@ merge(Compressor.prototype, {
             var args;
             var candidates = [];
             var stat_index = statements.length;
-            var scanner = new TreeTransformer(function(node, descend) {
+            var scanner = new TreeTransformer(function(node) {
                 if (abort) return node;
                 // Skip nodes before `candidate` as quickly as possible
                 if (!hit) {
@@ -1019,7 +998,7 @@ merge(Compressor.prototype, {
                 if (can_replace
                     && !(node instanceof AST_SymbolDeclaration)
                     && (scan_lhs && (hit_lhs = lhs.equivalent_to(node))
-                        || scan_rhs && (hit_rhs = rhs.equivalent_to(node)))) {
+                        || scan_rhs && (hit_rhs = scan_rhs(node, this)))) {
                     if (stop_if_hit && (hit_rhs || !lhs_local || !replace_all)) {
                         abort = true;
                         return node;
@@ -1388,12 +1367,16 @@ merge(Compressor.prototype, {
             }
 
             function foldable(expr) {
-                if (expr.is_constant()) return true;
-                if (expr instanceof AST_Array) return false;
-                if (expr instanceof AST_Function) return false;
-                if (expr instanceof AST_Object) return false;
-                if (expr instanceof AST_RegExp) return false;
-                if (expr instanceof AST_Symbol) return true;
+                if (expr instanceof AST_SymbolRef) {
+                    var value = expr.evaluate(compressor);
+                    if (value === expr) return rhs_exact_match;
+                    return rhs_fuzzy_match(value, rhs_exact_match);
+                }
+                if (expr instanceof AST_This) return rhs_exact_match;
+                if (expr.is_truthy()) return rhs_fuzzy_match(true, return_false);
+                if (expr.is_constant()) {
+                    return rhs_fuzzy_match(expr.evaluate(compressor), rhs_exact_match);
+                }
                 if (!(lhs instanceof AST_SymbolRef)) return false;
                 if (expr.has_side_effects(compressor)) return false;
                 var circular;
@@ -1404,7 +1387,25 @@ merge(Compressor.prototype, {
                         circular = true;
                     }
                 }));
-                return !circular;
+                return !circular && rhs_exact_match;
+            }
+
+            function rhs_exact_match(node) {
+                return rhs.equivalent_to(node);
+            }
+
+            function rhs_fuzzy_match(value, fallback) {
+                return function(node, tw) {
+                    if (tw.in_boolean_context()) {
+                        if (value && node.is_truthy() && !node.has_side_effects(compressor)) {
+                            return true;
+                        }
+                        if (node.is_constant()) {
+                            return !node.evaluate(compressor) == !value;
+                        }
+                    }
+                    return fallback(node);
+                };
             }
 
             function get_lvalues(expr) {
@@ -5089,7 +5090,7 @@ merge(Compressor.prototype, {
             }
             break;
         }
-        if (self.operator == "+" && compressor.in_boolean_context()) {
+        if (compressor.option("booleans") && self.operator == "+" && compressor.in_boolean_context()) {
             var ll = self.left.evaluate(compressor);
             var rr = self.right.evaluate(compressor);
             if (ll && typeof ll == "string") {
@@ -5154,7 +5155,7 @@ merge(Compressor.prototype, {
                 }
                 var rr = self.right.evaluate(compressor);
                 if (!rr) {
-                    if (compressor.in_boolean_context()) {
+                    if (compressor.option("booleans") && compressor.in_boolean_context()) {
                         compressor.warn("Boolean && always false [{file}:{line},{col}]", self.start);
                         return make_sequence(self, [
                             self.left,
@@ -5163,7 +5164,8 @@ merge(Compressor.prototype, {
                     } else self.falsy = true;
                 } else if (!(rr instanceof AST_Node)) {
                     var parent = compressor.parent();
-                    if (parent.operator == "&&" && parent.left === compressor.self() || compressor.in_boolean_context()) {
+                    if (parent.operator == "&&" && parent.left === compressor.self()
+                        || compressor.option("booleans") && compressor.in_boolean_context()) {
                         compressor.warn("Dropping side-effect-free && [{file}:{line},{col}]", self.start);
                         return self.left.optimize(compressor);
                     }
@@ -5190,12 +5192,13 @@ merge(Compressor.prototype, {
                 var rr = self.right.evaluate(compressor);
                 if (!rr) {
                     var parent = compressor.parent();
-                    if (parent.operator == "||" && parent.left === compressor.self() || compressor.in_boolean_context()) {
+                    if (parent.operator == "||" && parent.left === compressor.self()
+                        || compressor.option("booleans") && compressor.in_boolean_context()) {
                         compressor.warn("Dropping side-effect-free || [{file}:{line},{col}]", self.start);
                         return self.left.optimize(compressor);
                     }
                 } else if (!(rr instanceof AST_Node)) {
-                    if (compressor.in_boolean_context()) {
+                    if (compressor.option("booleans") && compressor.in_boolean_context()) {
                         compressor.warn("Boolean || always true [{file}:{line},{col}]", self.start);
                         return make_sequence(self, [
                             self.left,
@@ -5833,7 +5836,7 @@ merge(Compressor.prototype, {
                 right: alternative
             }).optimize(compressor);
         }
-        var in_bool = compressor.in_boolean_context();
+        var in_bool = compressor.option("booleans") && compressor.in_boolean_context();
         if (is_true(self.consequent)) {
             if (is_false(self.alternative)) {
                 // c ? true : false ---> !!c
@@ -5931,32 +5934,29 @@ merge(Compressor.prototype, {
     });
 
     OPT(AST_Boolean, function(self, compressor){
+        if (!compressor.option("booleans")) return self;
         if (compressor.in_boolean_context()) return make_node(AST_Number, self, {
             value: +self.value
         });
-        if (compressor.option("booleans")) {
-            var p = compressor.parent();
-            if (p instanceof AST_Binary && (p.operator == "=="
-                                            || p.operator == "!=")) {
-                compressor.warn("Non-strict equality against boolean: {operator} {value} [{file}:{line},{col}]", {
-                    operator : p.operator,
-                    value    : self.value,
-                    file     : p.start.file,
-                    line     : p.start.line,
-                    col      : p.start.col,
-                });
-                return make_node(AST_Number, self, {
-                    value: +self.value
-                });
-            }
-            return make_node(AST_UnaryPrefix, self, {
-                operator: "!",
-                expression: make_node(AST_Number, self, {
-                    value: 1 - self.value
-                })
+        var p = compressor.parent();
+        if (p instanceof AST_Binary && (p.operator == "==" || p.operator == "!=")) {
+            compressor.warn("Non-strict equality against boolean: {operator} {value} [{file}:{line},{col}]", {
+                operator : p.operator,
+                value    : self.value,
+                file     : p.start.file,
+                line     : p.start.line,
+                col      : p.start.col,
+            });
+            return make_node(AST_Number, self, {
+                value: +self.value
             });
         }
-        return self;
+        return make_node(AST_UnaryPrefix, self, {
+            operator: "!",
+            expression: make_node(AST_Number, self, {
+                value: 1 - self.value
+            })
+        });
     });
 
     OPT(AST_Sub, function(self, compressor){
diff --git a/test/compress/collapse_vars.js b/test/compress/collapse_vars.js
index 8b0f033a..886a7564 100644
--- a/test/compress/collapse_vars.js
+++ b/test/compress/collapse_vars.js
@@ -4976,7 +4976,7 @@ collapse_rhs_array: {
     expect_stdout: "false false false"
 }
 
-collapse_rhs_boolean: {
+collapse_rhs_boolean_1: {
     options = {
         collapse_vars: true,
     }
@@ -5001,6 +5001,66 @@ collapse_rhs_boolean: {
     expect_stdout: "true true true"
 }
 
+collapse_rhs_boolean_2: {
+    options = {
+        collapse_vars: true,
+    }
+    input: {
+        var a;
+        (function f1() {
+            a = function() {};
+            if (/foo/)
+                console.log(typeof a);
+        })();
+        console.log(function f2() {
+            a = [];
+            return !1;
+        }());
+    }
+    expect: {
+        var a;
+        (function f1() {
+            if (a = function() {})
+                console.log(typeof a);
+        })();
+        console.log(function f2() {
+            return !(a = []);
+        }());
+    }
+    expect_stdout: [
+        "function",
+        "false",
+    ]
+}
+
+collapse_rhs_boolean_3: {
+    options = {
+        booleans: true,
+        collapse_vars: true,
+        conditionals: true,
+    }
+    input: {
+        var a, f, g, h, i, n, s, t, x, y;
+        if (x()) {
+            n = a;
+        } else if (y()) {
+            n = f();
+        } else if (s) {
+            i = false;
+            n = g(true);
+        } else if (t) {
+            i = false;
+            n = h(true);
+        } else {
+            n = [];
+        }
+    }
+    expect: {
+        var a, f, g, h, i, n, s, t, x, y;
+        n = x() ? a : y() ? f() : s ? g(!(i = !1)) : t ? h(!(i = !1)) : [];
+    }
+}
+
 collapse_rhs_function: {
     options = {
         collapse_vars: true,