quickfix: add check for turning if-else chains into tagged switches

dominikh · dominikh · commit 7ad909facfe7 · 2021-01-22T17:36:50.000+01:00
Closes dominikhgh-6
diff --git a/quickfix/analysis.go b/quickfix/analysis.go
@@ -20,4 +20,8 @@ var Analyzers = lint.InitializeAnalyzers(Docs, map[string]*analysis.Analyzer{
 		Run:      CheckTaglessSwitch,
 		Requires: []*analysis.Analyzer{inspect.Analyzer},
 	},
+	"QF1003": {
+		Run:      CheckIfElseToSwitch,
+		Requires: []*analysis.Analyzer{inspect.Analyzer},
+	},
 })
diff --git a/quickfix/doc.go b/quickfix/doc.go
@@ -39,4 +39,30 @@ After:
 `,
 		Since: "Unreleased",
 	},
+	"QF1003": {
+		Title: "Convert if/else-if chain to tagged switch",
+		Text: ` A series of if/else-if checks comparing the same variable against values can be replaced with a tagged switch.
+
+Before:
+
+	if x == 1 || x == 2 {
+		...
+	} else if x == 3 {
+		...
+	} else {
+		...
+	}
+
+After:
+
+	switch x {
+	case 1, 2:
+		...
+	case 3:
+		...
+	default:
+		...
+	}`,
+		Since: "Unreleased",
+	},
 }
diff --git a/quickfix/lint.go b/quickfix/lint.go
@@ -404,3 +404,115 @@ func CheckTaglessSwitch(pass *analysis.Pass) (interface{}, error) {
 	code.Preorder(pass, fn, (*ast.SwitchStmt)(nil))
 	return nil, nil
 }
+
+func CheckIfElseToSwitch(pass *analysis.Pass) (interface{}, error) {
+	fn := func(node ast.Node, stack []ast.Node) {
+		if _, ok := stack[len(stack)-2].(*ast.IfStmt); ok {
+			// this if statement is part of an if-else chain
+			return
+		}
+		ifstmt := node.(*ast.IfStmt)
+
+		m := map[ast.Expr][]*ast.BinaryExpr{}
+		for item := ifstmt; item != nil; {
+			if item.Init != nil {
+				return
+			}
+			if item.Body == nil {
+				return
+			}
+
+			skip := false
+			ast.Inspect(item.Body, func(node ast.Node) bool {
+				if branch, ok := node.(*ast.BranchStmt); ok && branch.Tok != token.GOTO {
+					skip = true
+					return false
+				}
+				return true
+			})
+			if skip {
+				return
+			}
+
+			var pairs []*ast.BinaryExpr
+			if !findSwitchPairs(pass, item.Cond, &pairs) {
+				return
+			}
+			m[item.Cond] = pairs
+			switch els := item.Else.(type) {
+			case *ast.IfStmt:
+				item = els
+			case *ast.BlockStmt, nil:
+				item = nil
+			default:
+				panic(fmt.Sprintf("unreachable: %T", els))
+			}
+		}
+
+		var x ast.Expr
+		for _, pair := range m {
+			if len(pair) == 0 {
+				continue
+			}
+			if x == nil {
+				x = pair[0].X
+			} else {
+				if !astutil.Equal(x, pair[0].X) {
+					return
+				}
+			}
+		}
+		if x == nil {
+			// shouldn't happen
+			return
+		}
+
+		// We require at least two 'if' to make this suggestion, to
+		// avoid clutter in the editor.
+		if len(m) < 2 {
+			return
+		}
+
+		var edits []analysis.TextEdit
+		for item := ifstmt; item != nil; {
+			var end token.Pos
+			if item.Else != nil {
+				end = item.Else.Pos()
+			} else {
+				// delete up to but not including the closing brace.
+				end = item.Body.Rbrace
+			}
+
+			var conds []string
+			for _, cond := range m[item.Cond] {
+				y := cond.Y
+				if p, ok := y.(*ast.ParenExpr); ok {
+					y = p.X
+				}
+				conds = append(conds, report.Render(pass, y))
+			}
+			sconds := strings.Join(conds, ", ")
+			edits = append(edits,
+				edit.ReplaceWithString(pass.Fset, edit.Range{item.If, item.Body.Lbrace + 1}, "case "+sconds+":"),
+				edit.Delete(edit.Range{item.Body.Rbrace, end}))
+
+			switch els := item.Else.(type) {
+			case *ast.IfStmt:
+				item = els
+			case *ast.BlockStmt:
+				edits = append(edits, edit.ReplaceWithString(pass.Fset, edit.Range{els.Lbrace, els.Lbrace + 1}, "default:"))
+				item = nil
+			case nil:
+				item = nil
+			default:
+				panic(fmt.Sprintf("unreachable: %T", els))
+			}
+		}
+		// FIXME this forces the first case to begin in column 0. try to fix the indentation
+		edits = append(edits, edit.ReplaceWithString(pass.Fset, edit.Range{ifstmt.If, ifstmt.If}, fmt.Sprintf("switch %s {\n", report.Render(pass, x))))
+		report.Report(pass, ifstmt, fmt.Sprintf("could use tagged switch on %s", report.Render(pass, x)),
+			report.Fixes(edit.Fix("Replace with tagged switch", edits...)))
+	}
+	code.PreorderStack(pass, fn, (*ast.IfStmt)(nil))
+	return nil, nil
+}
diff --git a/quickfix/lint_test.go b/quickfix/lint_test.go
@@ -11,6 +11,7 @@ func TestAll(t *testing.T) {
 		"QF1000": {{Dir: "CheckStringsIndexByte"}},
 		"QF1001": {{Dir: "CheckDeMorgan"}},
 		"QF1002": {{Dir: "CheckTaglessSwitch"}},
+		"QF1003": {{Dir: "CheckIfElseToSwitch"}},
 	}
 
 	testutil.Run(t, Analyzers, checks)
diff --git a/quickfix/testdata/src/CheckIfElseToSwitch/CheckIfElseToSwitch.go b/quickfix/testdata/src/CheckIfElseToSwitch/CheckIfElseToSwitch.go
@@ -0,0 +1,45 @@
+package pkg
+
+func fn() {
+	var x, y int
+	var z []int
+	var a bool
+
+	if x == 1 || x == 2 { // want `could use tagged switch on x`
+	} else if x == 3 {
+	}
+
+	if x == 1 || x == 2 { // want `could use tagged switch on x`
+	} else if x == 3 {
+	} else {
+	}
+
+	if x == 1 || x == 2 {
+	} else if y == 3 {
+	} else {
+	}
+
+	if a == (x == y) { // want `could use tagged switch on a`
+	} else if a == (x != y) {
+	}
+
+	if z[0] == 1 || z[0] == 2 { // want `could use tagged switch on z\[0\]`
+	} else if z[0] == 3 {
+	}
+
+	for {
+		if x == 1 || x == 2 { // want `could use tagged switch on x`
+		} else if x == 3 {
+		}
+	}
+
+	for {
+		if x == 1 || x == 2 {
+		} else if x == 3 {
+			break
+		}
+	}
+
+	if x == 1 || x == 2 {
+	}
+}
diff --git a/quickfix/testdata/src/CheckIfElseToSwitch/CheckIfElseToSwitch.go.golden b/quickfix/testdata/src/CheckIfElseToSwitch/CheckIfElseToSwitch.go.golden
@@ -0,0 +1,50 @@
+package pkg
+
+func fn() {
+	var x, y int
+	var z []int
+	var a bool
+
+	switch x {
+	case 1, 2: // want `could use tagged switch on x`
+	case 3:
+	}
+
+	switch x {
+	case 1, 2: // want `could use tagged switch on x`
+	case 3:
+	default:
+	}
+
+	if x == 1 || x == 2 {
+	} else if y == 3 {
+	} else {
+	}
+
+	switch a {
+	case x == y: // want `could use tagged switch on a`
+	case x != y:
+	}
+
+	switch z[0] {
+	case 1, 2: // want `could use tagged switch on z\[0\]`
+	case 3:
+	}
+
+	for {
+		switch x {
+		case 1, 2: // want `could use tagged switch on x`
+		case 3:
+		}
+	}
+
+	for {
+		if x == 1 || x == 2 {
+		} else if x == 3 {
+			break
+		}
+	}
+
+	if x == 1 || x == 2 {
+	}
+}

Original file line number	Diff line number	Diff line change
`@@ -11,6 +11,7 @@ func TestAll(t *testing.T) {`
`11`	`11`	`"QF1000": {{Dir: "CheckStringsIndexByte"}},`
`12`	`12`	`"QF1001": {{Dir: "CheckDeMorgan"}},`
`13`	`13`	`"QF1002": {{Dir: "CheckTaglessSwitch"}},`
	`14`	`+ "QF1003": {{Dir: "CheckIfElseToSwitch"}},`
`14`	`15`	`}`
`15`	`16`
`16`	`17`	`testutil.Run(t, Analyzers, checks)`