fix(tools/spxls): improve declaration end position handling in formatting (#1341)

- Fixed declaration end position handling to account for trailing
comments.
- Ensured correct trailing newlines after declarations and comments.
- Fixed the exclusion of non-floating comments.

Signed-off-by: Aofei Sheng <[email protected]>

Author: aofei

tools/spxls/internal/server/format.go

@@ -6,6 +6,7 @@ import (
 	"go/types"
 	"io/fs"
 	"path"
+	"slices"
 	"time"
 
 	"github.com/goplus/builder/tools/spxls/internal/vfs"
@@ -175,21 +176,10 @@ func (s *Server) formatSpxDecls(snapshot *vfs.MapFS, spxFile string) ([]byte, er
 	)
 	for _, decl := range astFile.Decls {
 		// Skip the declaration if it appears after the error position.
-		if errorPos.IsValid() && errorPos <= decl.Pos() {
+		if errorPos.IsValid() && decl.Pos() >= errorPos {
 			continue
 		}
 
-		// Pre-process all comments within the declaration to exclude
-		// them from floating comments.
-		gopast.Inspect(decl, func(node gopast.Node) bool {
-			cg, ok := node.(*gopast.CommentGroup)
-			if !ok {
-				return true
-			}
-			processedComments[cg] = struct{}{}
-			return true
-		})
-
 		switch decl := decl.(type) {
 		case *gopast.GenDecl:
 			switch decl.Tok {
@@ -214,6 +204,23 @@ func (s *Server) formatSpxDecls(snapshot *vfs.MapFS, spxFile string) ([]byte, er
 		default:
 			otherDecls = append(otherDecls, decl)
 		}
+
+		// Pre-process all comments within the declaration to exclude
+		// them from floating comments.
+		if doc := getDeclDoc(decl); doc != nil {
+			processedComments[doc] = struct{}{}
+		}
+		startLine := compileResult.fset.Position(decl.Pos()).Line
+		endLine := compileResult.fset.Position(decl.End()).Line
+		for _, cg := range astFile.Comments {
+			if _, ok := processedComments[cg]; ok {
+				continue
+			}
+			cgStartLine := compileResult.fset.Position(cg.Pos()).Line
+			if cgStartLine >= startLine && cgStartLine <= endLine {
+				processedComments[cg] = struct{}{}
+			}
+		}
 	}
 
 	// Reorder declarations: imports -> types -> consts -> vars -> funcs -> others.
@@ -259,30 +266,51 @@ func (s *Server) formatSpxDecls(snapshot *vfs.MapFS, spxFile string) ([]byte, er
 
 	// Format the sorted declarations.
 	formattedBuf := bytes.NewBuffer(make([]byte, 0, len(astFile.Code)))
+	ensureTrailingNewlines := func(count int) {
+		if formattedBuf.Len() == 0 {
+			return
+		}
+		for _, b := range slices.Backward(formattedBuf.Bytes()) {
+			if b != '\n' {
+				break
+			}
+			count--
+		}
+		for range count {
+			formattedBuf.WriteByte('\n')
+		}
+	}
 
 	// Handle declarations and floating comments in order of their position.
-	var lastDeclEnd goptoken.Pos
 	processDecl := func(decl gopast.Decl) error {
 		startPos := decl.Pos()
 		if doc := getDeclDoc(decl); doc != nil {
 			startPos = doc.Pos()
 		}
 
-		if formattedBuf.Len() > 0 {
-			formattedBuf.WriteByte('\n')
+		endPos := decl.End()
+		endLine := compileResult.fset.Position(endPos).Line
+		for _, cg := range astFile.Comments {
+			if compileResult.fset.Position(cg.Pos()).Line != endLine {
+				continue
+			}
+			if cg.Pos() > endPos {
+				endPos = cg.End()
+				break
+			}
 		}
+
+		ensureTrailingNewlines(2)
 		if genDecl, ok := decl.(*gopast.GenDecl); ok && genDecl.Tok == goptoken.VAR {
 			if err := gopfmt.Node(formattedBuf, compileResult.fset, decl); err != nil {
 				return err
 			}
 		} else {
 			start := compileResult.fset.Position(startPos).Offset
-			end := compileResult.fset.Position(decl.End()).Offset
+			end := compileResult.fset.Position(endPos).Offset
 			formattedBuf.Write(astFile.Code[start:end])
 		}
-		formattedBuf.WriteByte('\n')
-
-		lastDeclEnd = decl.End()
+		ensureTrailingNewlines(1)
 		return nil
 	}
 
@@ -294,10 +322,9 @@ func (s *Server) formatSpxDecls(snapshot *vfs.MapFS, spxFile string) ([]byte, er
 		}
 		processedComments[cg] = struct{}{}
 
-		if errorPos.IsValid() && errorPos <= cg.Pos() {
-			continue
-		}
-		if lastDeclEnd.IsValid() && cg.Pos() < lastDeclEnd {
+		// Skip the comment if it appears after the error position or
+		// shadow entry position.
+		if errorPos.IsValid() && cg.Pos() >= errorPos {
 			continue
 		}
 		if shadowEntryPos.IsValid() && cg.Pos() >= shadowEntryPos {
@@ -307,9 +334,6 @@ func (s *Server) formatSpxDecls(snapshot *vfs.MapFS, spxFile string) ([]byte, er
 		// Process declarations that should come before this comment.
 		for ; declIndex < len(sortedDecls); declIndex++ {
 			decl := sortedDecls[declIndex]
-			if errorPos.IsValid() && errorPos <= decl.Pos() {
-				continue
-			}
 
 			startPos := decl.Pos()
 			if doc := getDeclDoc(decl); doc != nil {
@@ -325,35 +349,27 @@ func (s *Server) formatSpxDecls(snapshot *vfs.MapFS, spxFile string) ([]byte, er
 		}
 
 		// Add the floating comment.
-		if formattedBuf.Len() > 0 {
-			formattedBuf.WriteByte('\n')
-		}
+		ensureTrailingNewlines(2)
 		start := compileResult.fset.Position(cg.Pos()).Offset
 		end := compileResult.fset.Position(cg.End()).Offset
 		formattedBuf.Write(astFile.Code[start:end])
-		formattedBuf.WriteByte('\n')
+		ensureTrailingNewlines(1)
 	}
 
 	// Process remaining declarations before shadow entry.
 	for ; declIndex < len(sortedDecls); declIndex++ {
-		decl := sortedDecls[declIndex]
-		if errorPos.IsValid() && errorPos <= decl.Pos() {
-			continue
-		}
-		if err := processDecl(decl); err != nil {
+		if err := processDecl(sortedDecls[declIndex]); err != nil {
 			return nil, err
 		}
 	}
 
 	// Add the shadow entry if it exists and not empty.
 	if shadowEntryPos.IsValid() {
-		if formattedBuf.Len() > 0 {
-			formattedBuf.WriteByte('\n')
-		}
+		ensureTrailingNewlines(2)
 		start := compileResult.fset.Position(shadowEntryPos).Offset
 		end := compileResult.fset.Position(astFile.ShadowEntry.End()).Offset
 		formattedBuf.Write(astFile.Code[start:end])
-		formattedBuf.WriteByte('\n')
+		ensureTrailingNewlines(1)
 	}
 
 	formatted := formattedBuf.Bytes()

tools/spxls/internal/server/format_test.go

@@ -363,7 +363,9 @@ var a int
 // floating comment2
 
 // comment for func test
-func test() {}
+func test() {
+	// comment inside func test
+}
 
 // floating comment3
 
@@ -383,7 +385,7 @@ const b = "123"
 		assert.Contains(t, edits, TextEdit{
 			Range: Range{
 				Start: Position{Line: 0, Character: 0},
-				End:   Position{Line: 18, Character: 0},
+				End:   Position{Line: 20, Character: 0},
 			},
 			NewText: `import "fmt"
 
@@ -402,9 +404,47 @@ var (
 )
 
 // comment for func test
-func test() {}
+func test() {
+	// comment inside func test
+}
 
 // floating comment4
+`,
+		})
+	})
+
+	t.Run("WithTrailingComments", func(t *testing.T) {
+		s := New(newMapFSWithoutModTime(map[string][]byte{
+			"main.spx": []byte(`import "fmt" // trailing comment for import "fmt"
+
+const foo = "bar" // trailing comment for const foo
+
+var a int // trailing comment for var a
+
+func test() {} // trailing comment for func test
+`),
+		}), nil)
+		params := &DocumentFormattingParams{
+			TextDocument: TextDocumentIdentifier{URI: "file:///main.spx"},
+		}
+
+		edits, err := s.textDocumentFormatting(params)
+		require.NoError(t, err)
+		require.Len(t, edits, 1)
+		assert.Contains(t, edits, TextEdit{
+			Range: Range{
+				Start: Position{Line: 0, Character: 0},
+				End:   Position{Line: 7, Character: 0},
+			},
+			NewText: `import "fmt" // trailing comment for import "fmt"
+
+const foo = "bar" // trailing comment for const foo
+
+var (
+	a int
+) // trailing comment for var a
+
+func test() {} // trailing comment for func test
 `,
 		})
 	})