Fix all bugs for the grouping of decl

Author: bombsimon

go.mod

@@ -2,7 +2,10 @@ module github.com/bombsimon/wsl/v3
 
 go 1.19
 
-require golang.org/x/tools v0.6.0
+require (
+	github.com/dave/dst v0.27.2
+	golang.org/x/tools v0.6.0
+)
 
 require (
 	golang.org/x/mod v0.8.0 // indirect

go.sum

@@ -1,3 +1,7 @@
+github.com/dave/dst v0.27.2 h1:4Y5VFTkhGLC1oddtNwuxxe36pnyLxMFXT51FOzH8Ekc=
+github.com/dave/dst v0.27.2/go.mod h1:jHh6EOibnHgcUW3WjKHisiooEkYwqpHLBSX1iOBhEyc=
+github.com/dave/jennifer v1.5.0 h1:HmgPN93bVDpkQyYbqhCHj5QlgvUkvEOzMyEvKLgCRrg=
+github.com/sergi/go-diff v1.2.0 h1:XU+rvMAioB0UC3q1MFrIQy4Vo5/4VsRDQQXHsEya6xQ=
 golang.org/x/mod v0.8.0 h1:LUYupSeNrTNCGzR/hVBk2NHZO4hXcVaW1k4Qx7rjPx8=
 golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
 golang.org/x/sync v0.1.0 h1:wsuoTGHzEhffawBOhz5CYhcrV4IdKZbEyZjBMuTp12o=

testdata/src/default_config/var_decl.go

@@ -37,4 +37,16 @@ func fn() {
 		Addr:        ":8081",
 		ReadTimeout: 2 * time.Second,
 	}
+
+	var a1 = 1 // want "declarations should never be cuddled"
+	var s1 = http.Server{
+		Addr:        ":8080", // important comment
+		ReadTimeout: time.Second,
+	}
+
+	var s2 = http.Server{ // want "declarations should never be cuddled"
+		Addr:        ":8080", // important comment
+		ReadTimeout: time.Second,
+	}
+	var a2 = 1 // also comment
 }

testdata/src/default_config/var_decl.go.golden

@@ -14,15 +14,12 @@ func fn() {
 	)
 
 	var (
-		c = 1 // want "declarations should never be cuddled"
-
+		c             = 1 // want "declarations should never be cuddled"
 		d             = 1 // like this
 		extendedAlign = 33
-
-		x = 1
-		y = 2
-
-		afsa = 1
+		x             = 1
+		y             = 2
+		afsa          = 1
 	)
 
 	if true {
@@ -35,7 +32,7 @@ func fn() {
 	)
 
 	var (
-		s = http.Server{
+		s = http.Server{ // want "declarations should never be cuddled"
 			Addr:        ":8080", // important comment
 			ReadTimeout: time.Second,
 		}
@@ -44,4 +41,20 @@ func fn() {
 			ReadTimeout: 2 * time.Second,
 		}
 	)
+
+	var (
+		a1 = 1 // want "declarations should never be cuddled"
+		s1 = http.Server{
+			Addr:        ":8080", // important comment
+			ReadTimeout: time.Second,
+		}
+	)
+
+	var (
+		s2 = http.Server{ // want "declarations should never be cuddled"
+			Addr:        ":8080", // important comment
+			ReadTimeout: time.Second,
+		}
+		a2 = 1 // also comment
+	)
 }

testdata/src/wip/wip.go

@@ -1,25 +1 @@
 package testpkg
-
-func fn() {
-    var a = 1 // want "declarations should never be cuddled"
-    var baba  = 2 // what about
-    var xx  = 2 // Some comments?
-    var yy  = 2
-
-    var c = 1 // want "declarations should never be cuddled"
-	var (
-		d             = 1 // like this
-		extendedAlign = 33
-	)
-	var (
-		x = 1
-		y = 2
-	)
-	var afsa = 1
-
-	if true {
-		//
-	}
-	var x = 1
-	var z = 1
-}

testdata/src/wip/wip.go.golden

@@ -1,22 +1 @@
 package testpkg
-
-func fn() {
-	var (
-		a    = 1 // want "declarations should never be cuddled"
-		baba = 2 // what about
-		xx   = 2 // Some comments?
-		yy   = 2
-	)
-
-	var (
-		c = 1 // want "declarations should never be cuddled"
-
-		d             = 1 // like this
-		extendedAlign = 33
-
-		x = 1
-		y = 2
-
-		afsa = 1
-	)
-}

wsl.go

@@ -4,11 +4,13 @@ import (
 	"bytes"
 	"fmt"
 	"go/ast"
-	"go/printer"
 	"go/token"
 	"reflect"
 	"sort"
 	"strings"
+
+	"github.com/dave/dst"
+	"github.com/dave/dst/decorator"
 )
 
 // Error reason strings.
@@ -194,20 +196,29 @@ type result struct {
 // processor is the type that keeps track of the file and fileset and holds the
 // results from parsing the AST.
 type processor struct {
-	config   *Configuration
-	file     *ast.File
-	fileSet  *token.FileSet
-	Result   map[token.Pos]result
-	Warnings []string
+	config    *Configuration
+	file      *ast.File
+	fileSet   *token.FileSet
+	decorator *decorator.Decorator
+	Result    map[token.Pos]result
+	Warnings  []string
 }
 
 // newProcessorWithConfig will create a Processor with the passed configuration.
 func newProcessorWithConfig(file *ast.File, fileSet *token.FileSet, cfg *Configuration) *processor {
+	dec := decorator.NewDecorator(fileSet)
+
+	_, err := dec.DecorateFile(file)
+	if err != nil {
+		panic(err)
+	}
+
 	return &processor{
-		config:  cfg,
-		file:    file,
-		fileSet: fileSet,
-		Result:  make(map[token.Pos]result),
+		config:    cfg,
+		file:      file,
+		fileSet:   fileSet,
+		decorator: dec,
+		Result:    make(map[token.Pos]result),
 	}
 }
 
@@ -562,6 +573,28 @@ func (p *processor) parseBlockStatements(statements []ast.Stmt) {
 			// and look forward until we no longer have a cuddled decl.
 			var endNode ast.Node = previousDeclNode
 
+			// There's a difference on where `dst` decorates decl stmts based on
+			// if it's a single value spec or if it's a grouped var block or any
+			// kind of node with multiple child nodes.
+			// If we have a single `var a = 1` we need to move the decoration
+			// from the *dst.DeclStmt to the *dst.ValueSpec which holds the
+			// assigned value. This will make sure we keep the comment when
+			// combining all the specs.
+			moveCommentsIfGrouped := func(node *ast.DeclStmt, specs []dst.Spec) {
+				if len(specs) != 1 {
+					return
+				}
+
+				dds, _ := p.decorator.Dst.Nodes[node].(*dst.DeclStmt)
+
+				if spec, ok := specs[0].(*dst.ValueSpec); ok {
+					if spec.Decs.NodeDecs.Start == nil && spec.Decs.NodeDecs.End == nil {
+						spec.Decs.NodeDecs = dds.Decs.NodeDecs
+						dds.Decs.NodeDecs = dst.NodeDecs{}
+					}
+				}
+			}
+
 			seenDeclStmtIdx = i
 			for seenDeclStmtIdx < len(statements) {
 				// If next statement is not a decl there's nothing more to
@@ -586,22 +619,68 @@ func (p *processor) parseBlockStatements(statements []ast.Stmt) {
 					break
 				}
 
-				previousGenDecl.Specs = append(previousGenDecl.Specs, nextGenDecl.Specs...)
-				previousDeclNode.Decl = previousGenDecl
+				pd, _ := p.decorator.Dst.Nodes[previousGenDecl].(*dst.GenDecl)
+				nd, _ := p.decorator.Dst.Nodes[nextGenDecl].(*dst.GenDecl)
+
+				// We must set Rparen to true on the first node to indicate that
+				// it's a multiline block and that the comment should end with
+				// the assignment, not the closing parenthesis.
+				pd.Rparen = true
+
+				moveCommentsIfGrouped(previousDeclNode, pd.Specs)
+				moveCommentsIfGrouped(nextStatement, nd.Specs)
+
+				pd.Specs = append(pd.Specs, nd.Specs...)
+
+				pn, _ := p.decorator.Dst.Nodes[previousDeclNode].(*dst.DeclStmt)
+				pn.Decl = pd
 
 				endNode = nextStatement
 				seenDeclStmtIdx++
 			}
 
+			// More hacks to ensure new text spans the whole way. If we have a
+			// comment on our end node that starts after the node ends it's a
+			// trailing comment so we need to span over it wit our new text.
+			//
+			//   var a = 1
+			//   var b = 2 // comment
+			//   -------------------^ Need to end here
+			commentMap := ast.NewCommentMap(p.fileSet, endNode, p.file.Comments)
+			if cm, ok := commentMap[endNode]; ok {
+				lastComment := cm[len(cm)-1]
+				if lastComment.Pos() > endNode.End() {
+					endNode = lastComment
+				}
+			}
+
+			decl, ok := p.decorator.Dst.Nodes[previousStatement].(*dst.DeclStmt)
+			if !ok {
+				p.addWarning("failed to get node", stmt.Pos(), stmt)
+				continue
+			}
+
+			dummyFile := &dst.File{
+				Name: dst.NewIdent("dummy"),
+				Decls: []dst.Decl{
+					decl.Decl,
+				},
+			}
+
 			b := bytes.Buffer{}
-			_ = printer.Fprint(&b, p.fileSet, previousDeclNode)
+			decorator.Fprint(&b, dummyFile)
+
+			stripLen := len("package dummy\n\n")
+			newText := b.Bytes()
+			newText = newText[stripLen:]
+			newText = newText[:len(newText)-1]
 
 			p.addErrorRangeWithFix(
 				previousDeclNode.Pos(),
 				previousDeclNode.Pos(),
 				endNode.End(),
 				reasonNeverCuddleDeclare,
-				b.Bytes(),
+				newText,
 			)
 
 		case *ast.ExprStmt: