From c862f085c18f8560c5aa50183cb4fbb9a11656c3 Mon Sep 17 00:00:00 2001
From: Simon Sawert <simon@sawert.se>
Date: Sun, 23 Jun 2024 23:23:40 +0200
Subject: [PATCH] Fix missing expression cuddle check (#154)

---
 testdata/src/default_config/generic_handling.go   | 13 +++++++++++++
 .../src/default_config/generic_handling.go.golden | 15 +++++++++++++++
 wsl.go                                            |  2 +-
 3 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/testdata/src/default_config/generic_handling.go b/testdata/src/default_config/generic_handling.go
index 89490eb..204d71d 100644
--- a/testdata/src/default_config/generic_handling.go
+++ b/testdata/src/default_config/generic_handling.go
@@ -592,3 +592,16 @@ func AnonymousFunc() {
 		return nil // want "return statements should not be cuddled if block has more than two lines"
 	}())
 }
+
+// Issue #153
+func ExpressionForBlock() {
+	for i := 0; i < 4; i++ {
+		fmt.Println()
+	}
+	Up(3) // want "expressions should not be cuddled with blocks"
+
+	if autoheight != 3 {
+		t.Errorf("height should be 3 but is %d", autoheight)
+	}
+	Down(3) // want "expressions should not be cuddled with blocks"
+}
diff --git a/testdata/src/default_config/generic_handling.go.golden b/testdata/src/default_config/generic_handling.go.golden
index 8c6f585..86bd161 100644
--- a/testdata/src/default_config/generic_handling.go.golden
+++ b/testdata/src/default_config/generic_handling.go.golden
@@ -639,3 +639,18 @@ func AnonymousFunc() {
 		return nil // want "return statements should not be cuddled if block has more than two lines"
 	}())
 }
+
+// Issue #153
+func ExpressionForBlock() {
+	for i := 0; i < 4; i++ {
+		fmt.Println()
+	}
+
+	Up(3) // want "expressions should not be cuddled with blocks"
+
+	if autoheight != 3 {
+		t.Errorf("height should be 3 but is %d", autoheight)
+	}
+
+	Down(3) // want "expressions should not be cuddled with blocks"
+}
diff --git a/wsl.go b/wsl.go
index 0b8d8c0..76f4abf 100644
--- a/wsl.go
+++ b/wsl.go
@@ -583,7 +583,7 @@ func (p *processor) parseBlockStatements(statements []ast.Stmt) {
 				}
 
 				p.addWhitespaceBeforeError(t, reasonExpressionCuddledWithDeclOrRet)
-			case *ast.IfStmt, *ast.RangeStmt, *ast.SwitchStmt:
+			case *ast.IfStmt, *ast.RangeStmt, *ast.SwitchStmt, *ast.TypeSwitchStmt, *ast.ForStmt:
 				p.addWhitespaceBeforeError(t, reasonExpressionCuddledWithBlock)
 			}