Repository: bombsimon/wsl (go)

Issue:
Case arms should be considered for intersecting variables
Most statements that has a block as a condition, e.g. `if [condition]` or `for range [condition]`. For `switch` statements we check what we're switching on, e.g. `switch [condition]` but I think we should also consider the arms the same way.

Today, this is not allowed:

```go
myVariable := someFn()
switch [condition without myVariable] {
case isGood(myVariable):
    // ...
case isBad(myVariable):
    // ...
}
```

Since they're so closely tied together and in the same context, I think it would make sense to allow this. By doing so, we would also support the same pattern with anonymous `switch`:

```go
myVariable := 42
switch {
case myVariable < 10:
    // ...
case myVariable == 10:
   // ...
case myBariable > 10:
   // ...
}
```

This would also apply for ~~type switches and~~ `select`:

```go
ctx := someCtx()
select {
case ctx.Done():
    // ...
case time.After(time.Minute):
    // ...
}
```

Pull request description:
Allow idents from case arms when cuddling
Closes #171

Patch:
diff --git a/CHECKS.md b/CHECKS.md
index a26f209..87aa024 100644
--- a/CHECKS.md
+++ b/CHECKS.md
@@ -764,6 +764,8 @@ func Fn() int {
 
 ## `select`
 
+Identifiers used in case arms of select statements are allowed to be cuddled.
+
 > Configurable via `allow-first-in-block` to allow cuddling if the variable is
 > used _first_ in the block (enabled by default).
 >
@@ -790,6 +792,14 @@ case <-stop:
 </td><td valign="top">
 
 ```go
+stop := make(chan struct{})
+select {
+case <-time.After(time.Second):
+    // ...
+case <-stop:
+    // ...
+}
+
 x := 0
 
 select {
@@ -864,6 +874,10 @@ b := 1
 
 ## `switch`
 
+In addition to checking the switch condition, switch statements also checks
+identifiers in all case arms. If a variable is used in one or more of the case
+arms it's allowed to be cuddled.
+
 > Configurable via `allow-first-in-block` to allow cuddling if the variable is
 > used _first_ in the block (enabled by default).
 >
@@ -900,6 +914,22 @@ case 2:
 </td><td valign="top">
 
 ```go
+n := 1
+switch n {
+case 1:
+    // ...
+case 2:
+    // ...
+}
+
+n := 1
+switch {
+case n < 1:
+    // ...
+case n > 1:
+    // ...
+}
+
 x := 0
 
 switch y {
diff --git a/wsl.go b/wsl.go
index cd60f14..177aef3 100644
--- a/wsl.go
+++ b/wsl.go
@@ -193,22 +193,29 @@ func (w *WSL) checkBody(body []ast.Stmt) {
 	}
 }
 
-func (w *WSL) checkCuddlingBlock(stmt ast.Node, blockList []ast.Stmt, cursor *Cursor, maxAllowedStatements int) {
+func (w *WSL) checkCuddlingBlock(
+	stmt ast.Node,
+	blockList []ast.Stmt,
+	allowedIdents []*ast.Ident,
+	cursor *Cursor,
+	maxAllowedStatements int,
+) {
 	var firstBlockStmt ast.Node
 	if len(blockList) > 0 {
 		firstBlockStmt = blockList[0]
 	}
 
-	w.checkCuddlingMaxAllowed(stmt, firstBlockStmt, cursor, maxAllowedStatements)
+	w.checkCuddlingMaxAllowed(stmt, firstBlockStmt, allowedIdents, cursor, maxAllowedStatements)
 }
 
 func (w *WSL) checkCuddling(stmt ast.Node, cursor *Cursor, maxAllowedStatements int) {
-	w.checkCuddlingMaxAllowed(stmt, nil, cursor, maxAllowedStatements)
+	w.checkCuddlingMaxAllowed(stmt, nil, []*ast.Ident{}, cursor, maxAllowedStatements)
 }
 
 func (w *WSL) checkCuddlingMaxAllowed(
 	stmt ast.Node,
 	firstBlockStmt ast.Node,
+	allowedIdents []*ast.Ident,
 	cursor *Cursor,
 	maxAllowedStatements int,
 ) {
@@ -285,6 +292,10 @@ func (w *WSL) checkCuddlingMaxAllowed(
 		return
 	}
 
+	if checkIntersection(allowedIdents) {
+		return
+	}
+
 	intersects := identIntersection(currentIdents, previousIdents)
 	if len(intersects) > 0 {
 		return
@@ -641,7 +652,7 @@ func (w *WSL) checkIf(stmt *ast.IfStmt, cursor *Cursor, isElse bool) {
 
 	if _, ok := w.config.Checks[CheckIf]; !isElse && ok {
 		cursor.SetChecker(CheckIf)
-		w.checkCuddlingBlock(stmt, stmt.Body.List, cursor, 1)
+		w.checkCuddlingBlock(stmt, stmt.Body.List, []*ast.Ident{}, cursor, 1)
 	}
 }
 
@@ -739,7 +750,7 @@ func (w *WSL) checkSend(stmt *ast.SendStmt, cursor *Cursor) {
 		return true
 	})
 
-	w.checkCuddlingBlock(stmt, stmts, cursor, 1)
+	w.checkCuddlingBlock(stmt, stmts, []*ast.Ident{}, cursor, 1)
 }
 
 func (w *WSL) checkSwitch(stmt *ast.SwitchStmt, cursor *Cursor) {
@@ -917,12 +928,18 @@ func (w *WSL) maybeCheckBlock(
 	if _, ok := w.config.Checks[check]; ok {
 		cursor.SetChecker(check)
 
-		blockList := []ast.Stmt{}
-		if check != CheckSwitch && check != CheckTypeSwitch {
+		var (
+			blockList     []ast.Stmt
+			allowedIdents []*ast.Ident
+		)
+
+		if check != CheckSwitch && check != CheckTypeSwitch && check != CheckSelect {
 			blockList = blockStmt.List
+		} else {
+			allowedIdents = identsFromCaseArms(node)
 		}
 
-		w.checkCuddlingBlock(node, blockList, cursor, 1)
+		w.checkCuddlingBlock(node, blockList, allowedIdents, cursor, 1)
 	}
 }
 
@@ -1085,3 +1102,48 @@ func identsFromNode(node ast.Node, skipBlock bool) []*ast.Ident {
 
 	return idents
 }
+
+func identsFromCaseArms(node ast.Node) []*ast.Ident {
+	var (
+		idents []*ast.Ident
+		nodes  []ast.Stmt
+		seen   = map[string]struct{}{}
+
+		addUnseen = func(node ast.Node) {
+			for _, ident := range identsFromNode(node, true) {
+				if _, ok := seen[ident.Name]; ok {
+					continue
+				}
+
+				seen[ident.Name] = struct{}{}
+				idents = append(idents, ident)
+			}
+		}
+	)
+
+	switch v := node.(type) {
+	case *ast.SwitchStmt:
+		nodes = v.Body.List
+	case *ast.TypeSwitchStmt:
+		nodes = v.Body.List
+	case *ast.SelectStmt:
+		nodes = v.Body.List
+	default:
+		return idents
+	}
+
+	for _, node := range nodes {
+		switch n := node.(type) {
+		case *ast.CommClause:
+			addUnseen(n.Comm)
+		case *ast.CaseClause:
+			for _, n := range n.List {
+				addUnseen(n)
+			}
+		default:
+			continue
+		}
+	}
+
+	return idents
+}


Test patch:
diff --git a/testdata/src/default_config/select/select.go b/testdata/src/default_config/select/select.go
index c8d8844..63d7c27 100644
--- a/testdata/src/default_config/select/select.go
+++ b/testdata/src/default_config/select/select.go
@@ -26,3 +26,32 @@ func fn1(ctx context.Context, ch1 chan struct{}) {
 
 	_ = x
 }
+
+func fn2(ch1 chan struct{}) {
+	ch := make(chan struct{})
+	select {
+	case <-ch1:
+		_ = 1
+	case <-ch:
+		_ = 1
+	}
+
+	ctx := getCtx()
+	select {
+	case <-ch1:
+		_ = 1
+	case <-ctx.Done():
+		_ = 1
+	}
+
+	ch2 := make(chan struct{})
+	ch3 := make(chan struct{}) // want `missing whitespace above this line \(too many statements above select\)`
+	select {
+	case <-ch1:
+		_ = 1
+	case <-ch2:
+		_ = 1
+	case <-ch3:
+		_ = 1
+	}
+}
diff --git a/testdata/src/default_config/select/select.go.golden b/testdata/src/default_config/select/select.go.golden
index 6090756..b01226f 100644
--- a/testdata/src/default_config/select/select.go.golden
+++ b/testdata/src/default_config/select/select.go.golden
@@ -27,3 +27,33 @@ func fn1(ctx context.Context, ch1 chan struct{}) {
 
 	_ = x
 }
+
+func fn2(ch1 chan struct{}) {
+	ch := make(chan struct{})
+	select {
+	case <-ch1:
+		_ = 1
+	case <-ch:
+		_ = 1
+	}
+
+	ctx := getCtx()
+	select {
+	case <-ch1:
+		_ = 1
+	case <-ctx.Done():
+		_ = 1
+	}
+
+	ch2 := make(chan struct{})
+
+	ch3 := make(chan struct{}) // want `missing whitespace above this line \(too many statements above select\)`
+	select {
+	case <-ch1:
+		_ = 1
+	case <-ch2:
+		_ = 1
+	case <-ch3:
+		_ = 1
+	}
+}
diff --git a/testdata/src/default_config/switch/switch.go b/testdata/src/default_config/switch/switch.go
index 3ae5dc3..a3c1225 100644
--- a/testdata/src/default_config/switch/switch.go
+++ b/testdata/src/default_config/switch/switch.go
@@ -55,10 +55,31 @@ func fn4() {
 
 func fn5() {
 	timeout := 10
-	switch { // want `missing whitespace above this line \(no shared variables above switch\)`
+	switch {
 	case timeout > 5:
 		timeout = 100
 	case timeout == 10:
 		timeout = 0
 	}
+
+	timeout2 := 20
+	timeout3 := 30 // want `missing whitespace above this line \(too many statements above switch\)`
+	switch {
+	case timeout2 > 5:
+		timeout = 100
+	case timeout3 == 10:
+		timeout = 0
+	}
+}
+
+func fn6(a, b int) {
+	n := 1
+	switch { // want `missing whitespace above this line \(no shared variables above switch\)`
+	case a > b:
+		_ = n
+	case b > a:
+		_ = n
+	default:
+		_ = n
+	}
 }
diff --git a/testdata/src/default_config/switch/switch.go.golden b/testdata/src/default_config/switch/switch.go.golden
index e84d6d8..0d3fca3 100644
--- a/testdata/src/default_config/switch/switch.go.golden
+++ b/testdata/src/default_config/switch/switch.go.golden
@@ -58,11 +58,33 @@ func fn4() {
 
 func fn5() {
 	timeout := 10
-
-	switch { // want `missing whitespace above this line \(no shared variables above switch\)`
+	switch {
 	case timeout > 5:
 		timeout = 100
 	case timeout == 10:
 		timeout = 0
 	}
+
+	timeout2 := 20
+
+	timeout3 := 30 // want `missing whitespace above this line \(too many statements above switch\)`
+	switch {
+	case timeout2 > 5:
+		timeout = 100
+	case timeout3 == 10:
+		timeout = 0
+	}
+}
+
+func fn6(a, b int) {
+	n := 1
+
+	switch { // want `missing whitespace above this line \(no shared variables above switch\)`
+	case a > b:
+		_ = n
+	case b > a:
+		_ = n
+	default:
+		_ = n
+	}
 }