fix corner case in `booleans` (#3659)

fixes #3658

2 files changed, 108 insertions, 35 deletions

diff --git a/lib/ast.js b/lib/ast.js
index 30885905..11d1f361 100644
--- a/lib/ast.js
+++ b/lib/ast.js
@@ -969,6 +969,7 @@ TreeWalker.prototype = {
                 || p instanceof AST_DWLoop && p.condition === self
                 || p instanceof AST_For && p.condition === self
                 || p instanceof AST_If && p.condition === self
+                || p instanceof AST_Return && p.in_bool
                 || p instanceof AST_UnaryPrefix && p.operator == "!" && p.expression === self) {
                 return true;
             }
diff --git a/lib/compress.js b/lib/compress.js
index 4252c58e..76ad3996 100644
--- a/lib/compress.js
+++ b/lib/compress.js
@@ -209,6 +209,7 @@ merge(Compressor.prototype, {
         if (is_scope) {
             node.hoist_properties(this);
             node.hoist_declarations(this);
+            node.process_boolean_returns(this);
         }
         // Before https://github.com/mishoo/UglifyJS2/pull/1602 AST_Node.optimize()
         // would call AST_Node.transform() if a different instance of AST_Node is
@@ -591,15 +592,22 @@ merge(Compressor.prototype, {
         def(AST_Call, function(tw, descend) {
             tw.find_parent(AST_Scope).may_call_this();
             var exp = this.expression;
-            if (!(exp instanceof AST_SymbolRef)) return;
-            var def = exp.definition();
-            if (tw.in_boolean_context()) def.bool_fn++;
-            if (!(def.fixed instanceof AST_Defun)) return;
-            var defun = mark_defun(tw, def);
-            if (!defun) return;
-            descend();
-            defun.walk(tw);
-            return true;
+            if (exp instanceof AST_SymbolRef) {
+                var def = exp.definition();
+                if (this.TYPE == "Call" && tw.in_boolean_context()) def.bool_fn++;
+                if (!(def.fixed instanceof AST_Defun)) return;
+                var defun = mark_defun(tw, def);
+                if (!defun) return;
+                descend();
+                defun.walk(tw);
+                return true;
+            } else if (this.TYPE == "Call"
+                && exp instanceof AST_Assign
+                && exp.operator == "="
+                && exp.left instanceof AST_SymbolRef
+                && tw.in_boolean_context()) {
+                exp.left.definition().bool_fn++;
+            }
         });
         def(AST_Case, function(tw) {
             push(tw);
@@ -4344,6 +4352,96 @@ merge(Compressor.prototype, {
         self.body = dirs.concat(hoisted, self.body);
     });
 
+    function scan_local_returns(fn, transform) {
+        fn.walk(new TreeWalker(function(node) {
+            if (node instanceof AST_Return) {
+                transform(node);
+                return true;
+            }
+            if (node instanceof AST_Scope && node !== fn) return true;
+        }));
+    }
+
+    function map_bool_returns(fn) {
+        var map = Object.create(null);
+        scan_local_returns(fn, function(node) {
+            var value = node.value;
+            if (value) value = value.tail_node();
+            if (value instanceof AST_SymbolRef) {
+                var id = value.definition().id;
+                map[id] = (map[id] || 0) + 1;
+            }
+        });
+        return map;
+    }
+
+    function all_bool(def, bool_returns, compressor) {
+        return def.bool_fn + (bool_returns[def.id] || 0) === def.references.length
+            && !compressor.exposed(def);
+    }
+
+    function process_boolean_returns(fn, compressor) {
+        scan_local_returns(fn, function(node) {
+            node.in_bool = true;
+            var value = node.value;
+            if (value) {
+                var ev = value.is_truthy() || value.tail_node().evaluate(compressor);
+                if (!ev) {
+                    value = value.drop_side_effect_free(compressor);
+                    if (node.value !== value) node.value = value ? make_sequence(node.value, [
+                        value,
+                        make_node(AST_Number, node.value, {
+                            value: 0
+                        })
+                    ]) : null;
+                } else if (ev && !(ev instanceof AST_Node)) {
+                    value = value.drop_side_effect_free(compressor);
+                    if (node.value !== value) node.value = value ? make_sequence(node.value, [
+                        value,
+                        make_node(AST_Number, node.value, {
+                            value: 1
+                        })
+                    ]) : make_node(AST_Number, node.value, {
+                        value: 1
+                    });
+                }
+            }
+        });
+    }
+
+    AST_Scope.DEFMETHOD("process_boolean_returns", noop);
+    AST_Defun.DEFMETHOD("process_boolean_returns", function(compressor) {
+        if (!compressor.option("booleans")) return;
+        var bool_returns = map_bool_returns(this);
+        if (!all_bool(this.name.definition(), bool_returns, compressor)) return;
+        process_boolean_returns(this, compressor);
+    });
+    AST_Function.DEFMETHOD("process_boolean_returns", function(compressor) {
+        if (!compressor.option("booleans")) return;
+        var bool_returns = map_bool_returns(this);
+        if (this.name && !all_bool(this.name.definition(), bool_returns, compressor)) return;
+        var parent = compressor.parent();
+        if (parent instanceof AST_Assign) {
+            if (parent.operator != "=") return;
+            var sym = parent.left;
+            if (!(sym instanceof AST_SymbolRef)) return;
+            if (!all_bool(sym.definition(), bool_returns, compressor)) return;
+        } else if (parent instanceof AST_Call && parent.expression !== this) {
+            var exp = parent.expression;
+            if (exp instanceof AST_SymbolRef) exp = exp.fixed_value();
+            if (!(exp instanceof AST_Lambda)) return;
+            if (exp.uses_arguments || exp.pinned()) return;
+            var sym = exp.argnames[parent.args.indexOf(this)];
+            if (sym && !all_bool(sym.definition(), bool_returns, compressor)) return;
+        } else if (parent.TYPE == "Call") {
+            compressor.pop();
+            var in_bool = compressor.in_boolean_context();
+            compressor.push(this);
+            if (!in_bool) return;
+        } else return;
+        process_boolean_returns(this, compressor);
+    });
+
     AST_Scope.DEFMETHOD("var_names", function() {
         var var_names = this._var_names;
         if (!var_names) {
@@ -6686,32 +6784,6 @@ merge(Compressor.prototype, {
         if (compressor.option("reduce_vars") && is_lhs(compressor.self(), parent) !== compressor.self()) {
             var def = self.definition();
             var fixed = self.fixed_value();
-            if (compressor.option("booleans") && def.bool_fn === def.references.length && fixed instanceof AST_Lambda) {
-                def.bool_fn = null;
-                fixed.process_expression(false, function(node) {
-                    if (!node.value) return node;
-                    var value = node.value.is_truthy() || node.value.tail_node().evaluate(compressor);
-                    if (!value) {
-                        value = node.value.drop_side_effect_free(compressor);
-                        node.value = value ? make_sequence(node.value, [
-                            value,
-                            make_node(AST_Number, node.value, {
-                                value: 0
-                            })
-                        ]) : null;
-                    } else if (value && !(value instanceof AST_Node)) {
-                        var num = make_node(AST_Number, node.value, {
-                            value: 1
-                        });
-                        value = node.value.drop_side_effect_free(compressor);
-                        node.value = value ? make_sequence(node.value, [
-                            value,
-                            num
-                        ]) : num;
-                    }
-                    return node;
-                });
-            }
             var single_use = def.single_use && !(parent instanceof AST_Call && parent.is_expr_pure(compressor));
             if (single_use && fixed instanceof AST_Lambda) {
                 if (def.scope !== self.scope