no-conditional-assignment: rewrite and check conditional expressions

src/rules/noConditionalAssignmentRule.ts

@@ -15,6 +15,7 @@
  * limitations under the License.
  */
 
+import { isAssignmentKind, isNodeKind } from "tsutils";
 import * as ts from "typescript";
 
 import * as Lint from "../index";
@@ -24,7 +25,7 @@ export class Rule extends Lint.Rules.AbstractRule {
     public static metadata: Lint.IRuleMetadata = {
         ruleName: "no-conditional-assignment",
         description: "Disallows any type of assignment in conditionals.",
-        descriptionDetails: "This applies to `do-while`, `for`, `if`, and `while` statements.",
+        descriptionDetails: "This applies to `do-while`, `for`, `if`, and `while` statements and conditional (ternary) expressions.",
         rationale: Lint.Utils.dedent `
             Assignments in conditionals are often typos:
             for example \`if (var1 = var2)\` instead of \`if (var1 == var2)\`.
@@ -40,61 +41,74 @@ export class Rule extends Lint.Rules.AbstractRule {
     public static FAILURE_STRING = "Assignments in conditional expressions are forbidden";
 
     public apply(sourceFile: ts.SourceFile): Lint.RuleFailure[] {
-        const walker = new NoConditionalAssignmentWalker(sourceFile, this.getOptions());
-        return this.applyWithWalker(walker);
+        return this.applyWithFunction(sourceFile, walk);
     }
 }
 
-class NoConditionalAssignmentWalker extends Lint.RuleWalker {
-    private isInConditional = false;
+function walk(ctx: Lint.WalkContext<void>) {
+    let checking = 0;
+    return ts.forEachChild(ctx.sourceFile, cb);
 
-    protected visitIfStatement(node: ts.IfStatement) {
-        this.validateConditionalExpression(node.expression);
-        super.visitIfStatement(node);
-    }
-
-    protected visitWhileStatement(node: ts.WhileStatement) {
-        this.validateConditionalExpression(node.expression);
-        super.visitWhileStatement(node);
-    }
-
-    protected visitDoStatement(node: ts.DoStatement) {
-        this.validateConditionalExpression(node.expression);
-        super.visitDoStatement(node);
-    }
-
-    protected visitForStatement(node: ts.ForStatement) {
-        if (node.condition != null) {
-            this.validateConditionalExpression(node.condition);
+    function cb(node: ts.Node): void {
+        const kind = node.kind;
+        if (!isNodeKind(kind)) {
+            return; // return early for tokens
         }
-        super.visitForStatement(node);
-    }
-
-    protected visitBinaryExpression(expression: ts.BinaryExpression) {
-        if (this.isInConditional) {
-            this.checkForAssignment(expression);
+        switch (kind) {
+            case ts.SyntaxKind.ConditionalExpression:
+                return check((node as ts.ConditionalExpression).condition),
+                       cb((node as ts.ConditionalExpression).whenTrue),
+                       cb((node as ts.ConditionalExpression).whenFalse);
+            case ts.SyntaxKind.IfStatement:
+                return check((node as ts.IfStatement).expression),
+                       cb((node as ts.IfStatement).thenStatement),
+                       maybeCallback(cb, (node as ts.IfStatement).elseStatement);
+            case ts.SyntaxKind.DoStatement:
+            case ts.SyntaxKind.WhileStatement:
+                return check((node as ts.DoStatement | ts.WhileStatement).expression),
+                       cb((node as ts.IterationStatement).statement);
+            case ts.SyntaxKind.ForStatement:
+                return maybeCallback(cb, (node as ts.ForStatement).initializer),
+                       maybeCallback(check, (node as ts.ForStatement).condition),
+                       maybeCallback(cb, (node as ts.ForStatement).incrementor),
+                       cb((node as ts.ForStatement).statement);
         }
-        super.visitBinaryExpression(expression);
-    }
-
-    private validateConditionalExpression(expression: ts.Expression) {
-        this.isInConditional = true;
-        if (expression.kind === ts.SyntaxKind.BinaryExpression) {
-            // check for simple assignment in a conditional, like `if (a = 1) {`
-            this.checkForAssignment(expression as ts.BinaryExpression);
+        if (checking !== 0) {
+            switch (kind) {
+                case ts.SyntaxKind.BinaryExpression:
+                    if (isAssignmentKind((node as ts.BinaryExpression).operatorToken.kind)) {
+                        ctx.addFailureAtNode(node, Rule.FAILURE_STRING);
+                    }
+                    return cb((node as ts.BinaryExpression).left), cb((node as ts.BinaryExpression).right);
+                case ts.SyntaxKind.ParenthesizedExpression:
+                case ts.SyntaxKind.NonNullExpression:
+                case ts.SyntaxKind.AsExpression:
+                case ts.SyntaxKind.TypeAssertionExpression:
+                    return cb((node as ts.AssertionExpression | ts.NonNullExpression | ts.ParenthesizedExpression).expression);
+                case ts.SyntaxKind.PrefixUnaryExpression:
+                    return cb((node as ts.PrefixUnaryExpression).operand);
+                default:
+                    return noCheck(node);
+            }
         }
-        // walk the children of the conditional expression for nested assignments, like `if ((a = 1) && (b == 1)) {`
-        this.walkChildren(expression);
-        this.isInConditional = false;
+        return ts.forEachChild(node, cb);
     }
 
-    private checkForAssignment(expression: ts.BinaryExpression) {
-        if (isAssignmentToken(expression.operatorToken)) {
-            this.addFailureAtNode(expression, Rule.FAILURE_STRING);
-        }
+    function check(node: ts.Node): void {
+        ++checking;
+        cb(node);
+        --checking;
+    }
+    function noCheck(node: ts.Node): void {
+        const old = checking;
+        checking = 0;
+        ts.forEachChild(node, cb);
+        checking = old;
     }
 }
 
-function isAssignmentToken(token: ts.Node) {
-    return token.kind >= ts.SyntaxKind.FirstAssignment && token.kind <= ts.SyntaxKind.LastAssignment;
+function maybeCallback(cb: (node: ts.Node) => void, node?: ts.Node) {
+    if (node !== undefined) {
+        return cb(node);
+    }
 }

test/rules/no-conditional-assignment/test.ts.lint

@@ -15,6 +15,16 @@ for (var x = 8; x == 8; ++x) { }
 for (var x = 8; x == 8; x = 12) { }
 for (;;) { }
 
+while (a ? b : c) {}
+while (--i) {}
+while ((a = b)++) {} // don't bother, it's a compile error anyway
+
+while ((a || (a = b)).foo) {}
+while (a[b = c]) {}
+while ((() => a = b)()) {}
+while (fn(a = b)) {}
+a ? b = a : a = b;
+
 // invalid cases
 if (x = 5) { }
     ~~~~~      [0]
@@ -43,8 +53,27 @@ else if (h || (x <<= 4)) { }
 
 do { } while (x ^= 4) { }
               ~~~~~~      [0]
-while ((a = 5) && ((b == 4) || (c = 3)))
+while ((a = 5) && ((b == 4) || (c = 3))) {}
         ~~~~~                            [0]
                                 ~~~~~    [0]
+while (+(a = b) || -(b = c)) {}
+         ~~~~~ [0]
+                     ~~~~~ [0]
+while (foo() || (a = (b = c)!)) {}
+                 ~~~~~~~~~~~~ [0]
+                      ~~~~~ [0]
+while (a == (b = fn()) as any) {}
+             ~~~~~~~~ [0]
+while (a ? b = a : a = b) {}
+           ~~~~~ [0]
+                   ~~~~~ [0]
+while (!((a = b) ? c : d) || (x = y)) {}
+          ~~~~~ [0]
+                              ~~~~~ [0]
+(a = b) ? c : d;
+ ~~~~~ [0]
+(foo ? a = b : b = a) ? c : d;
+       ~~~~~ [0]
+               ~~~~~ [0]
 
 [0]: Assignments in conditional expressions are forbidden