From e6d79b5431dd1e013ca0eba2bd61767659509436 Mon Sep 17 00:00:00 2001
From: Simon Sawert <simon@sawert.se>
Date: Mon, 6 Jan 2025 22:52:45 +0100
Subject: [PATCH] Improved decl

---
 testdata/src/default_config/decl/decl.go      |  5 +++
 .../src/default_config/decl/decl.go.golden    |  6 +++
 wsl.go                                        | 37 +++++++++++++------
 3 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/testdata/src/default_config/decl/decl.go b/testdata/src/default_config/decl/decl.go
index 0402968..9f30b37 100644
--- a/testdata/src/default_config/decl/decl.go
+++ b/testdata/src/default_config/decl/decl.go
@@ -10,3 +10,8 @@ func fn1() {
 	_ = a
 	_ = b
 }
+
+func fn2() {
+	var a = 1
+	var b = a // want "missing whitespace decreases readability"
+}
diff --git a/testdata/src/default_config/decl/decl.go.golden b/testdata/src/default_config/decl/decl.go.golden
index a778706..8402723 100644
--- a/testdata/src/default_config/decl/decl.go.golden
+++ b/testdata/src/default_config/decl/decl.go.golden
@@ -12,3 +12,9 @@ func fn1() {
 	_ = a
 	_ = b
 }
+
+func fn2() {
+	var a = 1
+
+	var b = a // want "missing whitespace decreases readability"
+}
diff --git a/wsl.go b/wsl.go
index 3c9b534..1c8b3e1 100644
--- a/wsl.go
+++ b/wsl.go
@@ -108,8 +108,20 @@ func (w *WSL) Run() {
 }
 
 func (w *WSL) CheckCuddling(stmt ast.Node, cursor *Cursor, maxAllowedStatements int) {
-	reset := cursor.Save()
-	defer reset()
+	w.checkCuddlingWithDecl(stmt, cursor, maxAllowedStatements, true)
+}
+
+func (w *WSL) CheckCuddlingNoDecl(stmt ast.Node, cursor *Cursor, maxAllowedStatements int) {
+	w.checkCuddlingWithDecl(stmt, cursor, maxAllowedStatements, false)
+}
+
+func (w *WSL) checkCuddlingWithDecl(
+	stmt ast.Node,
+	cursor *Cursor,
+	maxAllowedStatements int,
+	declIsValid bool,
+) {
+	defer cursor.Save()()
 
 	currentIdents := allIdents(cursor.Stmt())
 	previousIdents := []*ast.Ident{}
@@ -137,6 +149,10 @@ func (w *WSL) CheckCuddling(stmt ast.Node, cursor *Cursor, maxAllowedStatements
 	_, prevIsIncDec := previousNode.(*ast.IncDecStmt)
 	_, currIsDefer := stmt.(*ast.DeferStmt)
 
+	if !declIsValid {
+		prevIsDecl = false
+	}
+
 	// We're cuddled but not with an assign, declare or defer statement which is
 	// never allowed.
 	if !prevIsAssign && !prevIsDecl && !currIsDefer && !prevIsIncDec {
@@ -402,12 +418,7 @@ func (w *WSL) CheckBranch(stmt *ast.BranchStmt, cursor *Cursor) {
 		return
 	}
 
-	w.addError(
-		stmt.Pos(),
-		stmt.Pos(),
-		stmt.Pos(),
-		MessageAddWhitespace,
-	)
+	w.addError(stmt.Pos(), stmt.Pos(), stmt.Pos(), MessageAddWhitespace)
 }
 
 func (w *WSL) CheckDecl(stmt *ast.DeclStmt, cursor *Cursor) {
@@ -419,7 +430,7 @@ func (w *WSL) CheckDecl(stmt *ast.DeclStmt, cursor *Cursor) {
 		return
 	}
 
-	w.CheckCuddling(stmt, cursor, 1)
+	w.CheckCuddlingNoDecl(stmt, cursor, 1)
 }
 
 func (w *WSL) CheckBlock(block *ast.BlockStmt) {
@@ -428,7 +439,6 @@ func (w *WSL) CheckBlock(block *ast.BlockStmt) {
 
 	cursor := NewCursor(-1, block.List)
 	for cursor.Next() {
-		// fmt.Printf("%d: %T\n", cursor.currentIdx, cursor.Stmt())
 		w.CheckStmt(cursor.Stmt(), cursor)
 	}
 }
@@ -603,8 +613,7 @@ func (w *WSL) CheckExpr(expr ast.Expr, cursor *Cursor) {
 // numberOfStatementsAbove will find out how many lines above the cursor's
 // current statement there is without any newlines between.
 func (w *WSL) numberOfStatementsAbove(cursor *Cursor) int {
-	reset := cursor.Save()
-	defer reset()
+	defer cursor.Save()()
 
 	statementsWithoutNewlines := 0
 	currentStmtStartLine := w.lineFor(cursor.Stmt().Pos())
@@ -785,6 +794,10 @@ func allIdents(node ast.Node) []*ast.Ident {
 		for _, name := range n.Names {
 			idents = append(idents, allIdents(name)...)
 		}
+
+		for _, value := range n.Values {
+			idents = append(idents, allIdents(value)...)
+		}
 	case *ast.AssignStmt:
 		// TODO: For TypeSwitchStatements, this can be a false positive by
 		// allowing shadowing and "tricking" usage;