feat: optional extension to early-return rule (#1133)

RULES_DESCRIPTIONS.md

@@ -346,12 +346,13 @@ if !cond {
 _Configuration_: ([]string) rule flags. Available flags are:
 
 * _preserveScope_: do not suggest refactorings that would increase variable scope
+* _allowJump_: suggest a new jump (`return`, `continue` or `break` statement) if it could unnest multiple statements. By default, only relocation of _existing_ jumps (i.e. from the `else` clause) are suggested.
 
 Example:
 
 ```toml
 [rule.early-return]
-  arguments = ["preserveScope"]
+  arguments = ["preserveScope", "allowJump"]
 ```
 
 ## empty-block

internal/ifelse/args.go

@@ -4,8 +4,15 @@ package ifelse
 // that would enlarge variable scope
 const PreserveScope = "preserveScope"
 
+// AllowJump is a configuration argument that permits early-return to
+// suggest introducing a new jump (return, continue, etc) statement
+// to reduce nesting. By default, suggestions only bring existing jumps
+// earlier.
+const AllowJump = "allowJump"
+
 // Args contains arguments common to the early-return, indent-error-flow
-// and superfluous-else rules (currently just preserveScope)
+// and superfluous-else rules
 type Args struct {
 	PreserveScope bool
+	AllowJump     bool
 }

internal/ifelse/branch.go

@@ -9,8 +9,8 @@ import (
 // Branch contains information about a branch within an if-else chain.
 type Branch struct {
 	BranchKind
-	Call          // The function called at the end for kind Panic or Exit.
-	HasDecls bool // The branch has one or more declarations (at the top level block)
+	Call  // The function called at the end for kind Panic or Exit.
+	block []ast.Stmt
 }
 
 // BlockBranch gets the Branch of an ast.BlockStmt.
@@ -21,7 +21,7 @@ func BlockBranch(block *ast.BlockStmt) Branch {
 	}
 
 	branch := StmtBranch(block.List[blockLen-1])
-	branch.HasDecls = hasDecls(block)
+	branch.block = block.List
 	return branch
 }
 
@@ -61,10 +61,22 @@ func StmtBranch(stmt ast.Stmt) Branch {
 // String returns a brief string representation
 func (b Branch) String() string {
 	switch b.BranchKind {
+	case Empty:
+		return "{ }"
+	case Regular:
+		return "{ ... }"
 	case Panic, Exit:
-		return fmt.Sprintf("... %v()", b.Call)
+		if b.IsShort() {
+			return fmt.Sprintf("{ %v() }", b.Call)
+		} else {
+			return fmt.Sprintf("{ ... %v() }", b.Call)
+		}
 	default:
-		return b.BranchKind.String()
+		if b.IsShort() {
+			return fmt.Sprintf("{ %v }", b.BranchKind)
+		} else {
+			return fmt.Sprintf("{ ... %v }", b.BranchKind)
+		}
 	}
 }
 
@@ -78,8 +90,9 @@ func (b Branch) LongString() string {
 	}
 }
 
-func hasDecls(block *ast.BlockStmt) bool {
-	for _, stmt := range block.List {
+// HasDecls returns whether the branch has any top-level declarations
+func (b Branch) HasDecls() bool {
+	for _, stmt := range b.block {
 		switch stmt := stmt.(type) {
 		case *ast.DeclStmt:
 			return true
@@ -91,3 +104,32 @@ func hasDecls(block *ast.BlockStmt) bool {
 	}
 	return false
 }
+
+// IsShort returns whether the branch is empty or consists of a single statement
+func (b Branch) IsShort() bool {
+	switch len(b.block) {
+	case 0:
+		return true
+	case 1:
+	default:
+		return false
+	}
+	switch b.block[0].(type) {
+	case *ast.BlockStmt:
+		return false
+	case *ast.IfStmt:
+		return false
+	case *ast.SwitchStmt:
+		return false
+	case *ast.TypeSwitchStmt:
+		return false
+	case *ast.SelectStmt:
+		return false
+	case *ast.ForStmt:
+		return false
+	case *ast.RangeStmt:
+		return false
+	default:
+		return true
+	}
+}

internal/ifelse/branch_kind.go

@@ -58,19 +58,19 @@ func (k BranchKind) String() string {
 	case Empty:
 		return ""
 	case Regular:
-		return "..."
+		return ""
 	case Return:
-		return "... return"
+		return "return"
 	case Continue:
-		return "... continue"
+		return "continue"
 	case Break:
-		return "... break"
+		return "break"
 	case Goto:
-		return "... goto"
+		return "goto"
 	case Panic:
-		return "... panic()"
+		return "panic()"
 	case Exit:
-		return "... os.Exit()"
+		return "os.Exit()"
 	default:
 		panic("invalid kind")
 	}

internal/ifelse/chain.go

@@ -2,9 +2,11 @@ package ifelse
 
 // Chain contains information about an if-else chain.
 type Chain struct {
-	If                   Branch // what happens at the end of the "if" block
-	Else                 Branch // what happens at the end of the "else" block
-	HasInitializer       bool   // is there an "if"-initializer somewhere in the chain?
-	HasPriorNonDeviating bool   // is there a prior "if" block that does NOT deviate control flow?
-	AtBlockEnd           bool   // whether the chain is placed at the end of the surrounding block
+	If                   Branch     // what happens at the end of the "if" block
+	HasElse              bool       // is there an "else" block?
+	Else                 Branch     // what happens at the end of the "else" block
+	HasInitializer       bool       // is there an "if"-initializer somewhere in the chain?
+	HasPriorNonDeviating bool       // is there a prior "if" block that does NOT deviate control flow?
+	AtBlockEnd           bool       // whether the chain is placed at the end of the surrounding block
+	BlockEndKind         BranchKind // control flow at end of surrounding block (e.g. "return" for function body)
 }

internal/ifelse/rule.go

@@ -31,8 +31,11 @@ type Rule interface {
 func Apply(rule Rule, node ast.Node, target Target, args lint.Arguments) []lint.Failure {
 	v := &visitor{rule: rule, target: target}
 	for _, arg := range args {
-		if arg == PreserveScope {
+		switch arg {
+		case PreserveScope:
 			v.args.PreserveScope = true
+		case AllowJump:
+			v.args.AllowJump = true
 		}
 	}
 	ast.Walk(v, node)
@@ -47,59 +50,90 @@ type visitor struct {
 }
 
 func (v *visitor) Visit(node ast.Node) ast.Visitor {
-	block, ok := node.(*ast.BlockStmt)
-	if !ok {
+	switch stmt := node.(type) {
+	case *ast.FuncDecl:
+		if stmt.Body != nil {
+			v.visitBlock(stmt.Body.List, Return)
+		}
+	case *ast.FuncLit:
+		v.visitBlock(stmt.Body.List, Return)
+	case *ast.ForStmt:
+		v.visitBlock(stmt.Body.List, Continue)
+	case *ast.RangeStmt:
+		v.visitBlock(stmt.Body.List, Continue)
+	case *ast.CaseClause:
+		v.visitBlock(stmt.Body, Break)
+	case *ast.BlockStmt:
+		v.visitBlock(stmt.List, Regular)
+	default:
 		return v
 	}
+	return nil
+}
 
-	for i, stmt := range block.List {
-		if ifStmt, ok := stmt.(*ast.IfStmt); ok {
-			v.visitChain(ifStmt, Chain{AtBlockEnd: i == len(block.List)-1})
+func (v *visitor) visitBlock(stmts []ast.Stmt, endKind BranchKind) {
+	for i, stmt := range stmts {
+		ifStmt, ok := stmt.(*ast.IfStmt)
+		if !ok {
+			ast.Walk(v, stmt)
 			continue
 		}
-		ast.Walk(v, stmt)
+		var chain Chain
+		if i == len(stmts)-1 {
+			chain.AtBlockEnd = true
+			chain.BlockEndKind = endKind
+		}
+		v.visitIf(ifStmt, chain)
 	}
-	return nil
 }
 
-func (v *visitor) visitChain(ifStmt *ast.IfStmt, chain Chain) {
+func (v *visitor) visitIf(ifStmt *ast.IfStmt, chain Chain) {
 	// look for other if-else chains nested inside this if { } block
-	ast.Walk(v, ifStmt.Body)
-
-	if ifStmt.Else == nil {
-		// no else branch
-		return
-	}
+	v.visitBlock(ifStmt.Body.List, chain.BlockEndKind)
 
 	if as, ok := ifStmt.Init.(*ast.AssignStmt); ok && as.Tok == token.DEFINE {
 		chain.HasInitializer = true
 	}
 	chain.If = BlockBranch(ifStmt.Body)
 
+	if ifStmt.Else == nil {
+		if v.args.AllowJump {
+			v.checkRule(ifStmt, chain)
+		}
+		return
+	}
+
 	switch elseBlock := ifStmt.Else.(type) {
 	case *ast.IfStmt:
 		if !chain.If.Deviates() {
 			chain.HasPriorNonDeviating = true
 		}
-		v.visitChain(elseBlock, chain)
+		v.visitIf(elseBlock, chain)
 	case *ast.BlockStmt:
 		// look for other if-else chains nested inside this else { } block
-		ast.Walk(v, elseBlock)
+		v.visitBlock(elseBlock.List, chain.BlockEndKind)
 
+		chain.HasElse = true
 		chain.Else = BlockBranch(elseBlock)
-		if failMsg := v.rule.CheckIfElse(chain, v.args); failMsg != "" {
-			if chain.HasInitializer {
-				// if statement has a := initializer, so we might need to move the assignment
-				// onto its own line in case the body references it
-				failMsg += " (move short variable declaration to its own line if necessary)"
-			}
-			v.failures = append(v.failures, lint.Failure{
-				Confidence: 1,
-				Node:       v.target.node(ifStmt),
-				Failure:    failMsg,
-			})
-		}
+		v.checkRule(ifStmt, chain)
 	default:
-		panic("invalid node type for else")
+		panic("unexpected node type for else")
+	}
+}
+
+func (v *visitor) checkRule(ifStmt *ast.IfStmt, chain Chain) {
+	failMsg := v.rule.CheckIfElse(chain, v.args)
+	if failMsg == "" {
+		return
+	}
+	if chain.HasInitializer {
+		// if statement has a := initializer, so we might need to move the assignment
+		// onto its own line in case the body references it
+		failMsg += " (move short variable declaration to its own line if necessary)"
 	}
+	v.failures = append(v.failures, lint.Failure{
+		Confidence: 1,
+		Node:       v.target.node(ifStmt),
+		Failure:    failMsg,
+	})
 }

rule/early_return.go

@@ -23,8 +23,14 @@ func (*EarlyReturnRule) Name() string {
 
 // CheckIfElse evaluates the rule against an ifelse.Chain and returns a failure message if applicable.
 func (*EarlyReturnRule) CheckIfElse(chain ifelse.Chain, args ifelse.Args) string {
-	if !chain.Else.Deviates() {
-		// this rule only applies if the else-block deviates control flow
+	if chain.HasElse {
+		if !chain.Else.BranchKind.Deviates() {
+			// this rule only applies if the else-block deviates control flow
+			return ""
+		}
+	} else if !args.AllowJump || !chain.AtBlockEnd || !chain.BlockEndKind.Deviates() || chain.If.IsShort() {
+		// this kind of refactor requires introducing a new indented "return", "continue" or "break" statement,
+		// so ignore unless we are able to outdent multiple statements in exchange.
 		return ""
 	}
 
@@ -34,18 +40,22 @@ func (*EarlyReturnRule) CheckIfElse(chain ifelse.Chain, args ifelse.Args) string
 		return ""
 	}
 
-	if chain.If.Deviates() {
+	if chain.HasElse && chain.If.Deviates() {
 		// avoid overlapping with superfluous-else
 		return ""
 	}
 
-	if args.PreserveScope && !chain.AtBlockEnd && (chain.HasInitializer || chain.If.HasDecls) {
+	if args.PreserveScope && !chain.AtBlockEnd && (chain.HasInitializer || chain.If.HasDecls()) {
 		// avoid increasing variable scope
 		return ""
 	}
 
+	if !chain.HasElse {
+		return fmt.Sprintf("if c { ... } can be rewritten if !c { %v } ... to reduce nesting", chain.BlockEndKind)
+	}
+
 	if chain.If.IsEmpty() {
-		return fmt.Sprintf("if c { } else { %[1]v } can be simplified to if !c { %[1]v }", chain.Else)
+		return fmt.Sprintf("if c { } else %[1]v can be simplified to if !c %[1]v", chain.Else)
 	}
-	return fmt.Sprintf("if c { ... } else { %[1]v } can be simplified to if !c { %[1]v } ...", chain.Else)
+	return fmt.Sprintf("if c { ... } else %[1]v can be simplified to if !c %[1]v ...", chain.Else)
 }

rule/indent_error_flow.go

@@ -20,6 +20,10 @@ func (*IndentErrorFlowRule) Name() string {
 
 // CheckIfElse evaluates the rule against an ifelse.Chain and returns a failure message if applicable.
 func (*IndentErrorFlowRule) CheckIfElse(chain ifelse.Chain, args ifelse.Args) string {
+	if !chain.HasElse {
+		return ""
+	}
+
 	if !chain.If.Deviates() {
 		// this rule only applies if the if-block deviates control flow
 		return ""
@@ -36,7 +40,7 @@ func (*IndentErrorFlowRule) CheckIfElse(chain ifelse.Chain, args ifelse.Args) st
 		return ""
 	}
 
-	if args.PreserveScope && !chain.AtBlockEnd && (chain.HasInitializer || chain.Else.HasDecls) {
+	if args.PreserveScope && !chain.AtBlockEnd && (chain.HasInitializer || chain.Else.HasDecls()) {
 		// avoid increasing variable scope
 		return ""
 	}

rule/superfluous_else.go

@@ -22,6 +22,10 @@ func (*SuperfluousElseRule) Name() string {
 
 // CheckIfElse evaluates the rule against an ifelse.Chain and returns a failure message if applicable.
 func (*SuperfluousElseRule) CheckIfElse(chain ifelse.Chain, args ifelse.Args) string {
+	if !chain.HasElse {
+		return ""
+	}
+
 	if !chain.If.Deviates() {
 		// this rule only applies if the if-block deviates control flow
 		return ""
@@ -38,7 +42,7 @@ func (*SuperfluousElseRule) CheckIfElse(chain ifelse.Chain, args ifelse.Args) st
 		return ""
 	}
 
-	if args.PreserveScope && !chain.AtBlockEnd && (chain.HasInitializer || chain.Else.HasDecls) {
+	if args.PreserveScope && !chain.AtBlockEnd && (chain.HasInitializer || chain.Else.HasDecls()) {
 		// avoid increasing variable scope
 		return ""
 	}

test/early_return_test.go

@@ -12,4 +12,5 @@ import (
 func TestEarlyReturn(t *testing.T) {
 	testRule(t, "early_return", &rule.EarlyReturnRule{})
 	testRule(t, "early_return_scope", &rule.EarlyReturnRule{}, &lint.RuleConfig{Arguments: []any{ifelse.PreserveScope}})
+	testRule(t, "early_return_jump", &rule.EarlyReturnRule{}, &lint.RuleConfig{Arguments: []any{ifelse.AllowJump}})
 }