Repository: bombsimon/wsl
Pull request #175 (2025-06-19 21:46:36)

Improve fixer for `DeclStmt`
New tracking issue for #121 after #122 but now with the v5 codebase. 

First iteration bails out if we have any comments but as long as no comments are related to statement we can group them safely.

- [x] Group any decl statement without comments
- [ ] ~~Group any decl statement with comments~~


coveralls: 
[![Coverage Status](https://coveralls.io/builds/74243983/badge)](https://coveralls.io/builds/74243983)

coverage: 94.163% (-0.2%) from 94.404%
when pulling **9f9e3fa8c95e95a33c8b37e9a897db6906a38f93 on decl-fixer-2**
into **dacbaad127129f3ebac0c0f612753bd4e5bd32b4 on main**.


Patch:
diff --git a/CHECKS.md b/CHECKS.md
index 87aa024..e630a83 100644
--- a/CHECKS.md
+++ b/CHECKS.md
@@ -121,9 +121,8 @@ together they should be declared in the same group with parenthesis into a
 single statement. The benefit of this is that it also aligns the declaration or
 assignment increasing readability.
 
-> **NOTE** The fixer can't do smart adjustments and currently only add
-> whitespaces.
-> This is tracked in [#21](https://github.com/bombsimon/wsl/issues/121)
+> **NOTE** The fixer can't do smart adjustments if there are comments on the
+> same line as the declaration.
 
 <table>
 <thead><tr><th>Bad</th><th>Good</th></tr></thead>
diff --git a/analyzer.go b/analyzer.go
index db2b852..fd9d054 100644
--- a/analyzer.go
+++ b/analyzer.go
@@ -117,7 +117,7 @@ func (wa *wslAnalyzer) run(pass *analysis.Pass) (any, error) {
 				textEdits = append(textEdits, analysis.TextEdit{
 					Pos:     f.fixRangeStart,
 					End:     f.fixRangeEnd,
-					NewText: []byte("\n"),
+					NewText: f.fix,
 				})
 			}
 
diff --git a/cursor.go b/cursor.go
index a4ad021..0d27582 100644
--- a/cursor.go
+++ b/cursor.go
@@ -25,6 +25,17 @@ func (c *Cursor) SetChecker(ct CheckType) {
 	c.checkType = ct
 }
 
+func (c *Cursor) NextNode() ast.Node {
+	defer c.Save()()
+
+	var nextNode ast.Node
+	if c.Next() {
+		nextNode = c.Stmt()
+	}
+
+	return nextNode
+}
+
 func (c *Cursor) Next() bool {
 	if c.currentIdx >= len(c.statements)-1 {
 		return false
diff --git a/wsl.go b/wsl.go
index 1fb1957..4f99e84 100644
--- a/wsl.go
+++ b/wsl.go
@@ -1,8 +1,10 @@
 package wsl
 
 import (
+	"bytes"
 	"fmt"
 	"go/ast"
+	"go/format"
 	"go/token"
 	"go/types"
 	"slices"
@@ -19,6 +21,7 @@ const (
 type fixRange struct {
 	fixRangeStart token.Pos
 	fixRangeEnd   token.Pos
+	fix           []byte
 }
 
 type issue struct {
@@ -30,20 +33,22 @@ type issue struct {
 }
 
 type WSL struct {
-	file     *ast.File
-	fset     *token.FileSet
-	typeInfo *types.Info
-	issues   map[token.Pos]issue
-	config   *Configuration
+	file         *ast.File
+	fset         *token.FileSet
+	typeInfo     *types.Info
+	issues       map[token.Pos]issue
+	config       *Configuration
+	groupedDecls map[token.Pos]struct{}
 }
 
 func New(file *ast.File, pass *analysis.Pass, cfg *Configuration) *WSL {
 	return &WSL{
-		fset:     pass.Fset,
-		file:     file,
-		typeInfo: pass.TypesInfo,
-		issues:   make(map[token.Pos]issue),
-		config:   cfg,
+		fset:         pass.Fset,
+		file:         file,
+		typeInfo:     pass.TypesInfo,
+		issues:       make(map[token.Pos]issue),
+		config:       cfg,
+		groupedDecls: make(map[token.Pos]struct{}),
 	}
 }
 
@@ -219,16 +224,15 @@ func (w *WSL) checkCuddlingMaxAllowed(
 	cursor *Cursor,
 	maxAllowedStatements int,
 ) {
-	var previousNode ast.Node
-
-	resetCursor := cursor.Save()
+	previousNode := cursor.PreviousNode()
 
-	if cursor.Previous() {
-		previousNode = cursor.Stmt()
+	if previousNode != nil {
+		if _, ok := w.groupedDecls[previousNode.End()]; ok {
+			w.addErrorTooManyStatements(cursor.Stmt().Pos(), cursor.checkType)
+			return
+		}
 	}
 
-	resetCursor()
-
 	numStmtsAbove := w.numberOfStatementsAbove(cursor)
 	previousIdents := identsFromNode(previousNode, true)
 
@@ -478,6 +482,12 @@ func (w *WSL) checkDeclStmt(stmt *ast.DeclStmt, cursor *Cursor) {
 		return
 	}
 
+	// Try to do smart grouping and if we succeed return, otherwise do
+	// line-by-line fixing.
+	if w.maybeGroupDecl(stmt, cursor) {
+		return
+	}
+
 	w.addErrorNeverAllow(stmt.Pos(), cursor.checkType)
 }
 
@@ -917,6 +927,81 @@ func (w *WSL) checkTrailingNewline(body *ast.BlockStmt) {
 	}
 }
 
+func (w *WSL) maybeGroupDecl(stmt *ast.DeclStmt, cursor *Cursor) bool {
+	firstNode := asGenDeclWithValueSpecs(cursor.PreviousNode())
+	if firstNode == nil {
+		return false
+	}
+
+	currentNode := asGenDeclWithValueSpecs(stmt)
+	if currentNode == nil {
+		return false
+	}
+
+	// Both are not same type, e.g. `const` or `var`
+	if firstNode.Tok != currentNode.Tok {
+		return false
+	}
+
+	group := &ast.GenDecl{
+		Tok:    firstNode.Tok,
+		Lparen: 1,
+		Specs:  firstNode.Specs,
+	}
+
+	group.Specs = append(group.Specs, currentNode.Specs...)
+
+	reportNodes := []ast.Node{currentNode}
+	lastNode := currentNode
+
+	for {
+		nextPeeked := cursor.NextNode()
+		if nextPeeked == nil {
+			break
+		}
+
+		if w.lineFor(lastNode.End()) < w.lineFor(nextPeeked.Pos())-1 {
+			break
+		}
+
+		nextNode := asGenDeclWithValueSpecs(nextPeeked)
+		if nextNode == nil {
+			break
+		}
+
+		if nextNode.Tok != firstNode.Tok {
+			break
+		}
+
+		cursor.Next()
+
+		group.Specs = append(group.Specs, nextNode.Specs...)
+		reportNodes = append(reportNodes, nextNode)
+		lastNode = nextNode
+	}
+
+	var buf bytes.Buffer
+	if err := format.Node(&buf, token.NewFileSet(), group); err != nil {
+		return false
+	}
+
+	// We add a diagnostic to every subsequent statement to properly represent
+	// the violations. Duplicate fixes for the same range is fine.
+	for _, n := range reportNodes {
+		w.groupedDecls[n.End()] = struct{}{}
+
+		w.addErrorWithMessageAndFix(
+			n.Pos(),
+			firstNode.Pos(),
+			lastNode.End(),
+			fmt.Sprintf("%s (never cuddle %s)", messageMissingWhitespaceAbove, CheckDecl),
+			buf.Bytes(),
+		)
+	}
+
+	return true
+}
+
 func (w *WSL) maybeCheckBlock(
 	node ast.Node,
 	blockStmt *ast.BlockStmt,
@@ -1032,6 +1117,10 @@ func (w *WSL) addError(report, start, end token.Pos, message string, ct CheckTyp
 }
 
 func (w *WSL) addErrorWithMessage(report, start, end token.Pos, message string) {
+	w.addErrorWithMessageAndFix(report, start, end, message, []byte("\n"))
+}
+
+func (w *WSL) addErrorWithMessageAndFix(report, start, end token.Pos, message string, fix []byte) {
 	iss, ok := w.issues[report]
 	if !ok {
 		iss = issue{
@@ -1043,11 +1132,43 @@ func (w *WSL) addErrorWithMessage(report, start, end token.Pos, message string)
 	iss.fixRanges = append(iss.fixRanges, fixRange{
 		fixRangeStart: start,
 		fixRangeEnd:   end,
+		fix:           fix,
 	})
 
 	w.issues[report] = iss
 }
 
+func asGenDeclWithValueSpecs(n ast.Node) *ast.GenDecl {
+	decl, ok := n.(*ast.DeclStmt)
+	if !ok {
+		return nil
+	}
+
+	genDecl, ok := decl.Decl.(*ast.GenDecl)
+	if !ok {
+		return nil
+	}
+
+	for _, spec := range genDecl.Specs {
+		// We only care about value specs and not type specs or import
+		// specs. We will never see any import specs but type specs we just
+		// separate with an empty line as usual.
+		valueSpec, ok := spec.(*ast.ValueSpec)
+		if !ok {
+			return nil
+		}
+
+		// It's very hard to get comments right in the ast and with the current
+		// way the ast package works we simply don't support grouping at all if
+		// there are any comments related to the node.
+		if valueSpec.Doc != nil || valueSpec.Comment != nil {
+			return nil
+		}
+	}
+
+	return genDecl
+}
+
 func hasIntersection(a, b ast.Node) bool {
 	return len(nodeIdentIntersection(a, b)) > 0
 }


Test patch:
diff --git a/testdata/src/default_config/decl/decl.go b/testdata/src/default_config/decl/decl.go
index 80c5ead..99866ae 100644
--- a/testdata/src/default_config/decl/decl.go
+++ b/testdata/src/default_config/decl/decl.go
@@ -1,30 +1,139 @@
 package testpkg
 
 func fn1() {
+ 	// want +2 `missing whitespace above this line \(never cuddle decl\)`
 	var a = 1
-	var b = 2 // want `missing whitespace above this line \(never cuddle decl\)`
+	var b = 2
 
+ 	// want +2 `missing whitespace above this line \(never cuddle decl\)`
 	const c = 3
-	const d = 4 // want `missing whitespace above this line \(never cuddle decl\)`
+	const d = 4
 
+	// want +3 `missing whitespace above this line \(never cuddle decl\)`
+	// want +3 `missing whitespace above this line \(invalid statement above assign\)`
 	e := 5
-	var f = 6 // want `missing whitespace above this line \(never cuddle decl\)`
-	g := 7 // want `missing whitespace above this line \(invalid statement above assign\)`
+	var f = 6
+	g := 7
+
+	// want +2 `missing whitespace above this line \(never cuddle decl\)`
+	var a = 1
+	var b = a
 }
 
 func fn2() {
-	var a = 1
-	var b = a // want `missing whitespace above this line \(never cuddle decl\)`
+	var x = func() { // want +1 `unnecessary whitespace \(leading-whitespace\)`
+
+		return 1
+	}()
 }
 
 func fn3() {
-	a := 1
-	var b = a // want `missing whitespace above this line \(never cuddle decl\)`
+	// want +4 `missing whitespace above this line \(never cuddle decl\)`
+	// want +4 `missing whitespace above this line \(never cuddle decl\)`
+	// want +4 `missing whitespace above this line \(never cuddle decl\)`
+	var a = 1
+	var b = 2
+	var c = 3
+	var d = 4
 }
 
 func fn4() {
-	var x = func() { // want +1 `unnecessary whitespace \(leading-whitespace\)`
+	// want +5 `missing whitespace above this line \(never cuddle decl\)`
+	var (
+		a = 1
+		b = 2
+	)
+	var (
+		c = 3
+		d = 4
+	)
+}
 
+func fn5() {
+	// want +7 `missing whitespace above this line \(never cuddle decl\)`
+	// want +7 `missing whitespace above this line \(never cuddle decl\)`
+	// want +9 `missing whitespace above this line \(never cuddle decl\)`
+	var (
+		a = 1
+		b = 2
+	)
+	var x int
+	var z = func() int {
 		return 1
-	}()
+	}
+	var (
+		c = 3
+		d = 4
+	)
+}
+
+func fn6() {
+	// want +5 `missing whitespace above this line \(never cuddle decl\)`
+	// want +5 `missing whitespace above this line \(never cuddle decl\)`
+	// want +5 `missing whitespace above this line \(never cuddle decl\)`
+	// want +5 `missing whitespace above this line \(never cuddle decl\)`
+	var a = 1
+	var b = 2
+	var c = 3 // test
+	var d = 4
+	var e = 5
+}
+
+func fn7() {
+	// want +2 `missing whitespace above this line \(never cuddle decl\)`
+	var a = 1
+	var b = 2
+
+	var c = 3
+
+	// want +2 `missing whitespace above this line \(never cuddle decl\)`
+	var d = 4
+	var e = 5
+}
+
+func fn8() {
+ 	// want +3 `missing whitespace above this line \(never cuddle decl\)`
+	// Comment above
+	var g = 7
+	var h = 8
+	// Comment after
+}
+
+func fn9() {
+	type S1 struct {
+		N int
+	}
+	type S2 struct { // want `missing whitespace above this line \(never cuddle decl\)`
+		N int
+	}
+}
+
+func fn10() {
+ 	// want +5 `missing whitespace above this line \(never cuddle decl\)`
+ 	// want +5 `missing whitespace above this line \(never cuddle decl\)`
+ 	// want +5 `missing whitespace above this line \(never cuddle decl\)`
+	a := 1
+	b := 2
+	var c int
+	var d = "string"
+	var e string = "string"
+	f := 3 // want `missing whitespace above this line \(invalid statement above assign\)`
+}
+
+func fn11() {
+ 	// want +2 `missing whitespace above this line \(never cuddle decl\)`
+	var a int
+	var b int
+	if b > 0 { // want `missing whitespace above this line \(too many statements above if\)`
+		_ = 1
+	}
+}
+
+func fn12() {
+	// want +2 `missing whitespace above this line \(never cuddle decl\)`
+	var a int
+	var b int // Not grouped due to this comment
+	if b > 0 {
+		_ = 1
+	}
 }
diff --git a/testdata/src/default_config/decl/decl.go.golden b/testdata/src/default_config/decl/decl.go.golden
index c6601d6..fd5e95b 100644
--- a/testdata/src/default_config/decl/decl.go.golden
+++ b/testdata/src/default_config/decl/decl.go.golden
@@ -1,35 +1,165 @@
 package testpkg
 
 func fn1() {
-	var a = 1
+ 	// want +2 `missing whitespace above this line \(never cuddle decl\)`
+	var (
+		a = 1
+		b = 2
+	)
 
-	var b = 2 // want `missing whitespace above this line \(never cuddle decl\)`
-
-	const c = 3
-
-	const d = 4 // want `missing whitespace above this line \(never cuddle decl\)`
+ 	// want +2 `missing whitespace above this line \(never cuddle decl\)`
+	const (
+		c = 3
+		d = 4
+	)
 
+	// want +3 `missing whitespace above this line \(never cuddle decl\)`
+	// want +3 `missing whitespace above this line \(invalid statement above assign\)`
 	e := 5
 
-	var f = 6 // want `missing whitespace above this line \(never cuddle decl\)`
+	var f = 6
+
+	g := 7
 
-	g := 7 // want `missing whitespace above this line \(invalid statement above assign\)`
+	// want +2 `missing whitespace above this line \(never cuddle decl\)`
+	var (
+		a = 1
+		b = a
+	)
 }
 
 func fn2() {
-	var a = 1
-
-	var b = a // want `missing whitespace above this line \(never cuddle decl\)`
+	var x = func() { // want +1 `unnecessary whitespace \(leading-whitespace\)`
+		return 1
+	}()
 }
 
 func fn3() {
+	// want +4 `missing whitespace above this line \(never cuddle decl\)`
+	// want +4 `missing whitespace above this line \(never cuddle decl\)`
+	// want +4 `missing whitespace above this line \(never cuddle decl\)`
+	var (
+		a = 1
+		b = 2
+		c = 3
+		d = 4
+	)
+}
+
+func fn4() {
+	// want +5 `missing whitespace above this line \(never cuddle decl\)`
+	var (
+		a = 1
+		b = 2
+		c = 3
+		d = 4
+	)
+}
+
+func fn5() {
+	// want +7 `missing whitespace above this line \(never cuddle decl\)`
+	// want +7 `missing whitespace above this line \(never cuddle decl\)`
+	// want +9 `missing whitespace above this line \(never cuddle decl\)`
+	var (
+		a = 1
+		b = 2
+		x int
+		z = func() int {
+			return 1
+		}
+		c = 3
+		d = 4
+	)
+}
+
+func fn6() {
+	// want +5 `missing whitespace above this line \(never cuddle decl\)`
+	// want +5 `missing whitespace above this line \(never cuddle decl\)`
+	// want +5 `missing whitespace above this line \(never cuddle decl\)`
+	// want +5 `missing whitespace above this line \(never cuddle decl\)`
+	var (
+		a = 1
+		b = 2
+	)
+
+	var c = 3 // test
+
+	var (
+		d = 4
+		e = 5
+	)
+}
+
+func fn7() {
+	// want +2 `missing whitespace above this line \(never cuddle decl\)`
+	var (
+		a = 1
+		b = 2
+	)
+
+	var c = 3
+
+	// want +2 `missing whitespace above this line \(never cuddle decl\)`
+	var (
+		d = 4
+		e = 5
+	)
+}
+
+func fn8() {
+ 	// want +3 `missing whitespace above this line \(never cuddle decl\)`
+	// Comment above
+	var (
+		g = 7
+		h = 8
+	)
+	// Comment after
+}
+
+func fn9() {
+	type S1 struct {
+		N int
+	}
+
+	type S2 struct { // want `missing whitespace above this line \(never cuddle decl\)`
+		N int
+	}
+}
+
+func fn10() {
+ 	// want +5 `missing whitespace above this line \(never cuddle decl\)`
+ 	// want +5 `missing whitespace above this line \(never cuddle decl\)`
+ 	// want +5 `missing whitespace above this line \(never cuddle decl\)`
 	a := 1
+	b := 2
 
-	var b = a // want `missing whitespace above this line \(never cuddle decl\)`
+	var (
+		c int
+		d        = "string"
+		e string = "string"
+	)
+
+	f := 3 // want `missing whitespace above this line \(invalid statement above assign\)`
 }
 
-func fn4() {
-	var x = func() { // want +1 `unnecessary whitespace \(leading-whitespace\)`
-		return 1
-	}()
+func fn11() {
+ 	// want +2 `missing whitespace above this line \(never cuddle decl\)`
+	var (
+		a int
+		b int
+	)
+
+	if b > 0 { // want `missing whitespace above this line \(too many statements above if\)`
+		_ = 1
+	}
+}
+
+func fn12() {
+ 	// want +2 `missing whitespace above this line \(never cuddle decl\)`
+	var a int
+
+	var b int // Not grouped due to this comment
+	if b > 0 {
+		_ = 1
+	}
 }