fix(tools/spxls): preserve doc comments after formatting

Fixes #1269

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

tools/spxls/internal/server/compile.go

@@ -124,6 +124,30 @@ type astFileLine struct {
 	line    int
 }
 
+// newCompileResult creates a new [compileResult].
+func newCompileResult() *compileResult {
+	return &compileResult{
+		fset:                      goptoken.NewFileSet(),
+		mainPkg:                   types.NewPackage("main", "main"),
+		mainASTPkg:                &gopast.Package{Name: "main", Files: make(map[string]*gopast.File)},
+		mainASTPkgSpecToGenDecl:   make(map[gopast.Spec]*gopast.GenDecl),
+		mainASTPkgIdentToFuncDecl: make(map[*gopast.Ident]*gopast.FuncDecl),
+		firstVarBlocks:            make(map[*gopast.File]*gopast.GenDecl),
+		typeInfo: &goptypesutil.Info{
+			Types:      make(map[gopast.Expr]types.TypeAndValue),
+			Defs:       make(map[*gopast.Ident]types.Object),
+			Uses:       make(map[*gopast.Ident]types.Object),
+			Implicits:  make(map[gopast.Node]types.Object),
+			Selections: make(map[*gopast.SelectorExpr]*types.Selection),
+			Scopes:     make(map[gopast.Node]*types.Scope),
+		},
+		spxSoundResourceAutoBindings:  make(map[types.Object]struct{}),
+		spxSpriteResourceAutoBindings: make(map[types.Object]struct{}),
+		diagnostics:                   make(map[DocumentURI][]Diagnostic),
+		documentURIs:                  make(map[string]DocumentURI),
+	}
+}
+
 // isInFset reports whether the given position exists in the file set.
 func (r *compileResult) isInFset(pos goptoken.Pos) bool {
 	return r.fset.File(pos) != nil
@@ -597,7 +621,8 @@ type compileCache struct {
 	spxFileModTimes map[string]time.Time
 }
 
-// compile compiles spx source files and returns compile result.
+// compile compiles spx source files and returns compile result. It uses cached
+// result if available.
 func (s *Server) compile() (*compileResult, error) {
 	snapshot := s.workspaceRootFS.Snapshot()
 	spxFiles, err := listSpxFiles(snapshot)
@@ -635,8 +660,8 @@ func (s *Server) compile() (*compileResult, error) {
 		}
 	}
 
-	// Compile uncached if cache is not used.
-	result, err := s.compileUncached(snapshot, spxFiles)
+	// Compile at the given snapshot if cache is not used.
+	result, err := s.compileAt(snapshot)
 	if err != nil {
 		return nil, err
 	}
@@ -658,30 +683,19 @@ func (s *Server) compile() (*compileResult, error) {
 	return result, nil
 }
 
-// compileUncached compiles spx source files without using cache.
-func (s *Server) compileUncached(snapshot *vfs.MapFS, spxFiles []string) (*compileResult, error) {
-	result := &compileResult{
-		fset:                      goptoken.NewFileSet(),
-		mainPkg:                   types.NewPackage("main", "main"),
-		mainASTPkg:                &gopast.Package{Name: "main", Files: make(map[string]*gopast.File)},
-		mainASTPkgSpecToGenDecl:   make(map[gopast.Spec]*gopast.GenDecl),
-		mainASTPkgIdentToFuncDecl: make(map[*gopast.Ident]*gopast.FuncDecl),
-		firstVarBlocks:            make(map[*gopast.File]*gopast.GenDecl),
-		typeInfo: &goptypesutil.Info{
-			Types:      make(map[gopast.Expr]types.TypeAndValue),
-			Defs:       make(map[*gopast.Ident]types.Object),
-			Uses:       make(map[*gopast.Ident]types.Object),
-			Implicits:  make(map[gopast.Node]types.Object),
-			Selections: make(map[*gopast.SelectorExpr]*types.Selection),
-			Scopes:     make(map[gopast.Node]*types.Scope),
-		},
-		spxSoundResourceAutoBindings:  make(map[types.Object]struct{}),
-		spxSpriteResourceAutoBindings: make(map[types.Object]struct{}),
-		diagnostics:                   make(map[DocumentURI][]Diagnostic, len(spxFiles)),
-		documentURIs:                  make(map[string]DocumentURI, len(spxFiles)),
+// compileAt compiles spx source files at the given snapshot and returns the
+// compile result.
+func (s *Server) compileAt(snapshot *vfs.MapFS) (*compileResult, error) {
+	spxFiles, err := listSpxFiles(snapshot)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get spx files: %w", err)
+	}
+	if len(spxFiles) == 0 {
+		return nil, errNoMainSpxFile
 	}
 
 	var (
+		result      = newCompileResult()
 		gpfs        = vfs.NewGopParserFS(snapshot)
 		spriteNames = make([]string, 0, len(spxFiles)-1)
 	)

tools/spxls/internal/server/format.go

@@ -50,8 +50,8 @@ func (s *Server) textDocumentFormatting(params *DocumentFormattingParams) ([]Tex
 
 // formatSpx formats an spx source file. If no change is needed, it returns `(nil, nil, nil)`.
 func (s *Server) formatSpx(spxFileName string) (formatted, original []byte, err error) {
-	// Parse the source into AST.
-	compileResult, err := s.compile()
+	// NOTE: Avoid using cached results as we need to modify AST nodes.
+	compileResult, err := s.compileAt(s.workspaceRootFS.Snapshot())
 	if err != nil {
 		return nil, nil, err
 	}
@@ -61,6 +61,25 @@ func (s *Server) formatSpx(spxFileName string) (formatted, original []byte, err
 	}
 	original = astFile.Code
 
+	// Find the first syntax error position if any.
+	var errorPos goptoken.Pos
+	gopast.Inspect(astFile, func(node gopast.Node) bool {
+		switch node := node.(type) {
+		case *gopast.BadExpr, *gopast.BadStmt, *gopast.BadDecl:
+			if !errorPos.IsValid() || node.Pos() < errorPos {
+				errorPos = node.Pos()
+			}
+		}
+		return true
+	})
+	if errorPos.IsValid() &&
+		astFile.ShadowEntry != nil &&
+		astFile.ShadowEntry.Pos().IsValid() &&
+		errorPos > astFile.ShadowEntry.Pos() &&
+		errorPos <= astFile.ShadowEntry.End() {
+		errorPos = astFile.ShadowEntry.Pos()
+	}
+
 	eliminateUnusedLambdaParams(compileResult, astFile)
 
 	fset := compileResult.fset
@@ -69,72 +88,182 @@ func (s *Server) formatSpx(spxFileName string) (formatted, original []byte, err
 
 	// Collect all declarations.
 	var (
-		importDecls []gopast.Decl
-		typeDecls   []gopast.Decl
-		constDecls  []gopast.Decl
-		varBlocks   []*gopast.GenDecl
-		funcDecls   []gopast.Decl
-		otherDecls  []gopast.Decl
+		importDecls       []gopast.Decl
+		typeDecls         []gopast.Decl
+		constDecls        []gopast.Decl
+		varBlocks         []*gopast.GenDecl
+		funcDecls         []gopast.Decl
+		overloadFuncDecls []gopast.Decl
+		otherDecls        []gopast.Decl
+		processedComments = make(map[*gopast.CommentGroup]struct{})
 	)
 	for _, decl := range astFile.Decls {
-		switch d := decl.(type) {
+		// Skip the declaration if it appears after the error position.
+		if errorPos.IsValid() && errorPos <= decl.Pos() {
+			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 d.Tok {
+			switch decl.Tok {
 			case goptoken.IMPORT:
-				importDecls = append(importDecls, d)
+				importDecls = append(importDecls, decl)
 			case goptoken.TYPE:
-				typeDecls = append(typeDecls, d)
+				typeDecls = append(typeDecls, decl)
 			case goptoken.CONST:
-				constDecls = append(constDecls, d)
+				constDecls = append(constDecls, decl)
 			case goptoken.VAR:
-				varBlocks = append(varBlocks, d)
+				varBlocks = append(varBlocks, decl)
 			default:
-				otherDecls = append(otherDecls, d)
+				otherDecls = append(otherDecls, decl)
 			}
 		case *gopast.FuncDecl:
-			funcDecls = append(funcDecls, d)
+			if decl.Shadow {
+				continue
+			}
+			funcDecls = append(funcDecls, decl)
+		case *gopast.OverloadFuncDecl:
+			overloadFuncDecls = append(overloadFuncDecls, decl)
 		default:
-			otherDecls = append(otherDecls, d)
+			otherDecls = append(otherDecls, decl)
 		}
 	}
 
 	// Reorder declarations: imports -> types -> consts -> vars -> funcs -> others.
 	//
 	// See https://github.com/goplus/builder/issues/591 and https://github.com/goplus/builder/issues/752.
-	newDecls := make([]gopast.Decl, 0, len(astFile.Decls))
-	newDecls = append(newDecls, importDecls...)
-	newDecls = append(newDecls, typeDecls...)
-	newDecls = append(newDecls, constDecls...)
+	sortedDecls := make([]gopast.Decl, 0, len(astFile.Decls))
+	sortedDecls = append(sortedDecls, importDecls...)
+	sortedDecls = append(sortedDecls, typeDecls...)
+	sortedDecls = append(sortedDecls, constDecls...)
 	if len(varBlocks) > 0 {
 		// Merge multiple var blocks into a single one.
 		firstVarBlock := varBlocks[0]
 		firstVarBlock.Lparen = firstVarBlock.Pos()
+		firstVarBlock.Rparen = varBlocks[len(varBlocks)-1].End()
 		if len(varBlocks) > 1 {
-			firstVarBlock.Rparen = varBlocks[len(varBlocks)-1].End()
 			for _, varBlock := range varBlocks[1:] {
-				for _, spec := range varBlock.Specs {
+				for i, spec := range varBlock.Specs {
 					valueSpec, ok := spec.(*gopast.ValueSpec)
 					if !ok {
 						return nil, nil, fmt.Errorf("unexpected non-value spec in var block: %T", spec)
 					}
+
+					// If this is the first spec in the var block, associate the var block's doc with it.
+					if i == 0 && varBlock.Doc != nil && len(varBlock.Doc.List) > 0 {
+						if valueSpec.Doc == nil {
+							valueSpec.Doc = varBlock.Doc
+						} else {
+							// Merge block doc with existing spec doc.
+							mergedList := append(varBlock.Doc.List, valueSpec.Doc.List...)
+							valueSpec.Doc = &gopast.CommentGroup{List: mergedList}
+						}
+					}
+
 					firstVarBlock.Specs = append(firstVarBlock.Specs, valueSpec)
 				}
 			}
-		} else {
-			firstVarBlock.Rparen = firstVarBlock.End()
 		}
-		newDecls = append(newDecls, firstVarBlock)
+		sortedDecls = append(sortedDecls, firstVarBlock)
 	}
-	newDecls = append(newDecls, funcDecls...)
-	newDecls = append(newDecls, otherDecls...)
-	astFile.Decls = newDecls
+	sortedDecls = append(sortedDecls, funcDecls...)
+	sortedDecls = append(sortedDecls, overloadFuncDecls...)
+	sortedDecls = append(sortedDecls, otherDecls...)
 
 	// Format the modified AST.
-	var buf bytes.Buffer
-	if err := gopfmt.Node(&buf, fset, astFile); err != nil {
-		return nil, nil, err
+	var formattedBuf bytes.Buffer
+
+	// Handle sorted declarations and mixed floating comments.
+	var lastPos goptoken.Pos
+	for _, decl := range sortedDecls {
+		// Add floating comments that appear before this declaration.
+		for _, cg := range astFile.Comments {
+			if errorPos.IsValid() && errorPos <= cg.Pos() {
+				continue
+			}
+			if _, ok := processedComments[cg]; ok {
+				continue
+			}
+			if cg.Pos() >= decl.Pos() {
+				continue
+			}
+			if lastPos != 0 && cg.Pos() < lastPos {
+				continue
+			}
+			if formattedBuf.Len() > 0 {
+				formattedBuf.WriteByte('\n')
+			}
+			for _, c := range cg.List {
+				formattedBuf.WriteString(c.Text)
+				formattedBuf.WriteByte('\n')
+			}
+			processedComments[cg] = struct{}{}
+		}
+
+		// Add the declaration.
+		if formattedBuf.Len() > 0 {
+			formattedBuf.WriteByte('\n')
+		}
+		if err := gopfmt.Node(&formattedBuf, fset, decl); err != nil {
+			return nil, nil, err
+		}
+		formattedBuf.WriteByte('\n')
+		lastPos = decl.End()
+	}
+
+	// Handle remaining floating comments that appear after all declarations.
+	for _, cg := range astFile.Comments {
+		if errorPos.IsValid() && errorPos <= cg.Pos() {
+			continue
+		}
+		if _, ok := processedComments[cg]; ok {
+			continue
+		}
+		if formattedBuf.Len() > 0 {
+			formattedBuf.WriteByte('\n')
+		}
+		for _, c := range cg.List {
+			formattedBuf.WriteString(c.Text)
+			formattedBuf.WriteByte('\n')
+		}
+	}
+
+	// Add the shadow entry of the AST file if it exists and not empty.
+	if astFile.ShadowEntry != nil &&
+		astFile.ShadowEntry.Pos().IsValid() &&
+		astFile.ShadowEntry.Pos() != errorPos &&
+		len(astFile.ShadowEntry.Body.List) > 0 {
+		if formattedBuf.Len() > 0 {
+			formattedBuf.WriteByte('\n')
+		}
+		if err := gopfmt.Node(&formattedBuf, fset, astFile.ShadowEntry.Body.List); err != nil {
+			return nil, nil, err
+		}
+		formattedBuf.WriteByte('\n')
+	}
+
+	// Add the original code after the error position if any.
+	if errorPos.IsValid() {
+		offset := fset.Position(errorPos).Offset
+		formattedBuf.Write(original[offset:])
+	}
+
+	formatted = formattedBuf.Bytes()
+	if formatted == nil {
+		formatted = []byte{}
 	}
-	return buf.Bytes(), original, nil
+	return formatted, original, nil
 }
 
 // eliminateUnusedLambdaParams eliminates useless lambda parameter declarations.