fix(tools/spxls): resolve issues with LS functionality

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

tools/spxls/internal/server/compile.go

@@ -56,6 +56,12 @@ type compileResult struct {
 	// mainPkg is the main package.
 	mainPkg *types.Package
 
+	// mainPkgGameType is the Game type in the main package.
+	mainPkgGameType *types.Named
+
+	// mainPkgSpriteTypes stores sprite types in the main package.
+	mainPkgSpriteTypes []*types.Named
+
 	// mainASTPkg is the main package AST.
 	mainASTPkg *gopast.Package
 
@@ -80,9 +86,6 @@ type compileResult struct {
 	// process.
 	typeInfo *goptypesutil.Info
 
-	// spxSpriteNames stores spx sprite names.
-	spxSpriteNames []string
-
 	// spxResourceSet is the set of spx resources.
 	spxResourceSet SpxResourceSet
 
@@ -342,6 +345,29 @@ func (r *compileResult) selectorTypeNameForIdent(ident *gopast.Ident) string {
 	return ""
 }
 
+// isGameOrSpriteThisPtr reports whether the given object is the "this" pointer
+// of Game or a Sprite type method in the main package.
+func (r *compileResult) isGameOrSpriteThisPtr(obj types.Object) bool {
+	if obj == nil || obj.Pkg() == nil || obj.Pkg().Path() != "main" || obj.Name() != "this" {
+		return false
+	}
+
+	v, ok := obj.(*types.Var)
+	if !ok {
+		return false
+	}
+	ptr, ok := v.Type().(*types.Pointer)
+	if !ok {
+		return false
+	}
+
+	named, ok := ptr.Elem().(*types.Named)
+	if !ok {
+		return false
+	}
+	return named == r.mainPkgGameType || slices.Contains(r.mainPkgSpriteTypes, named)
+}
+
 // isDefinedInFirstVarBlock reports whether the given object is defined in the
 // first var block of an AST file.
 func (r *compileResult) isDefinedInFirstVarBlock(obj types.Object) bool {
@@ -591,7 +617,10 @@ func (s *Server) compileUncached(snapshot *vfs.MapFS, spxFiles []string) (*compi
 		diagnostics:                   make(map[DocumentURI][]Diagnostic, len(spxFiles)),
 	}
 
-	gpfs := vfs.NewGopParserFS(snapshot)
+	var (
+		gpfs        = vfs.NewGopParserFS(snapshot)
+		spriteNames = make([]string, 0, len(spxFiles)-1)
+	)
 	for _, spxFile := range spxFiles {
 		documentURI := s.toDocumentURI(spxFile)
 		result.diagnostics[documentURI] = []Diagnostic{}
@@ -644,7 +673,7 @@ func (s *Server) compileUncached(snapshot *vfs.MapFS, spxFiles []string) (*compi
 			if spxFileBaseName := path.Base(spxFile); spxFileBaseName == "main.spx" {
 				result.mainSpxFile = spxFile
 			} else {
-				result.spxSpriteNames = append(result.spxSpriteNames, strings.TrimSuffix(spxFileBaseName, ".spx"))
+				spriteNames = append(spriteNames, strings.TrimSuffix(spxFileBaseName, ".spx"))
 			}
 
 			for _, decl := range astFile.Decls {
@@ -707,10 +736,21 @@ func (s *Server) compileUncached(snapshot *vfs.MapFS, spxFiles []string) (*compi
 		},
 		nil,
 		result.typeInfo,
-	).Files(nil, gopASTFileMapToSlice(result.mainASTPkg.Files)); err != nil {
+	).Files(nil, slices.Collect(maps.Values(result.mainASTPkg.Files))); err != nil {
 		// Errors should be handled by the type checker.
 	}
 
+	if obj := result.mainPkg.Scope().Lookup("Game"); obj != nil {
+		result.mainPkgGameType = obj.Type().(*types.Named)
+	}
+
+	result.mainPkgSpriteTypes = make([]*types.Named, 0, len(spriteNames))
+	for _, spxSpriteName := range spriteNames {
+		if obj := result.mainPkg.Scope().Lookup(spxSpriteName); obj != nil {
+			result.mainPkgSpriteTypes = append(result.mainPkgSpriteTypes, obj.Type().(*types.Named))
+		}
+	}
+
 	s.inspectForSpxResourceSet(snapshot, result)
 	s.inspectForSpxResourceAutoBindingsAndRefsAtDecls(result)
 	s.inspectForSpxResourceRefs(result)
@@ -819,11 +859,11 @@ func (s *Server) inspectForSpxResourceAutoBindingsAndRefsAtDecls(result *compile
 		case spxSpriteTypeFullName:
 			isSpxSpriteResourceAutoBinding = true
 		default:
-			for _, spxSpriteName := range result.spxSpriteNames {
-				if objTypeName != "main."+spxSpriteName {
+			for _, spxSpriteType := range result.mainPkgSpriteTypes {
+				if objType != spxSpriteType {
 					continue
 				}
-				if ident.Name != spxSpriteName {
+				if ident.Name != spxSpriteType.Obj().Name() {
 					result.addDiagnostics(documentURI, Diagnostic{
 						Severity: SeverityError,
 						Range:    identRange,

tools/spxls/internal/server/diagnostic_test.go

@@ -725,4 +725,30 @@ run "assets", {Title: "My Game"}
 			assert.Empty(t, fullReport.Items)
 		}
 	})
+
+	t.Run("NoTypeSpriteVarDeclaration", func(t *testing.T) {
+		s := New(newMapFSWithoutModTime(map[string][]byte{
+			"main.spx": []byte(`// A spx game.
+
+var (
+	MySprite
+)
+
+run "assets", {Title: "My Game"}
+`),
+			"MySprite.spx":                       []byte(``),
+			"assets/index.json":                  []byte(`{}`),
+			"assets/sprites/MySprite/index.json": []byte(`{}`),
+		}), nil)
+
+		report, err := s.workspaceDiagnostic(&WorkspaceDiagnosticParams{})
+		require.NoError(t, err)
+		require.NotNil(t, report)
+		assert.Len(t, report.Items, 2)
+		for _, item := range report.Items {
+			fullReport := item.Value.(WorkspaceFullDocumentDiagnosticReport)
+			assert.Equal(t, string(DiagnosticFull), fullReport.Kind)
+			assert.Empty(t, fullReport.Items)
+		}
+	})
 }

tools/spxls/internal/server/format.go

@@ -9,6 +9,7 @@ import (
 	gopast "github.com/goplus/gop/ast"
 	gopfmt "github.com/goplus/gop/format"
 	gopparser "github.com/goplus/gop/parser"
+	"github.com/goplus/gop/parser/fsx/memfs"
 	goptoken "github.com/goplus/gop/token"
 )
 
@@ -58,25 +59,32 @@ func (s *Server) textDocumentFormatting(params *DocumentFormattingParams) ([]Tex
 func formatSpx(src []byte) ([]byte, error) {
 	// Parse the source into AST.
 	fset := goptoken.NewFileSet()
-	f, err := gopparser.ParseFile(fset, "main.spx", src, gopparser.ParseComments|gopparser.ParseGoPlusClass)
+	astFile, err := gopparser.ParseFSEntry(fset, memfs.SingleFile("", "main.spx", string(src)), "main.spx", nil, gopparser.Config{
+		Mode: gopparser.ParseComments | gopparser.AllErrors | gopparser.ParseGoPlusClass,
+	})
 	if err != nil {
 		return nil, err
 	}
 
-	gopast.SortImports(fset, f)
+	// Sort import statements first.
+	gopast.SortImports(fset, astFile)
 
-	// Find all var blocks and function declarations.
+	// Collect all declarations.
 	var (
-		varBlocks  []*gopast.GenDecl
-		funcDecls  []gopast.Decl
-		otherDecls []gopast.Decl
+		importDecls []gopast.Decl
+		varBlocks   []*gopast.GenDecl
+		funcDecls   []gopast.Decl
+		otherDecls  []gopast.Decl
 	)
-	for _, decl := range f.Decls {
+	for _, decl := range astFile.Decls {
 		switch d := decl.(type) {
 		case *gopast.GenDecl:
-			if d.Tok == goptoken.VAR {
+			switch d.Tok {
+			case goptoken.IMPORT:
+				importDecls = append(importDecls, d)
+			case goptoken.VAR:
 				varBlocks = append(varBlocks, d)
-			} else {
+			default:
 				otherDecls = append(otherDecls, d)
 			}
 		case *gopast.FuncDecl:
@@ -86,10 +94,11 @@ func formatSpx(src []byte) ([]byte, error) {
 		}
 	}
 
-	// Reorder declarations: vars -> funcs -> others.
+	// Reorder declarations: imports -> 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(f.Decls))
+	newDecls := make([]gopast.Decl, 0, len(astFile.Decls))
+	newDecls = append(newDecls, importDecls...)
 	if len(varBlocks) > 0 {
 		// Merge multiple var blocks into a single one.
 		firstVarBlock := varBlocks[0]
@@ -112,11 +121,11 @@ func formatSpx(src []byte) ([]byte, error) {
 	}
 	newDecls = append(newDecls, funcDecls...)
 	newDecls = append(newDecls, otherDecls...)
-	f.Decls = newDecls
+	astFile.Decls = newDecls
 
 	// Format the modified AST.
 	var buf bytes.Buffer
-	if err := gopfmt.Node(&buf, fset, f); err != nil {
+	if err := gopfmt.Node(&buf, fset, astFile); err != nil {
 		return nil, err
 	}
 	return buf.Bytes(), nil

tools/spxls/internal/server/format_test.go

@@ -179,9 +179,27 @@ var (
 
 run "assets", {Title: "My Game"}
 `),
-			"MySprite.spx":                       []byte(``),
-			"assets/index.json":                  []byte(`{}`),
-			"assets/sprites/MySprite/index.json": []byte(`{}`),
+		}), nil)
+		params := &DocumentFormattingParams{
+			TextDocument: TextDocumentIdentifier{URI: "file:///main.spx"},
+		}
+
+		edits, err := s.textDocumentFormatting(params)
+		require.NoError(t, err)
+		require.Nil(t, edits)
+	})
+
+	t.Run("WithImportStmt", func(t *testing.T) {
+		s := New(newMapFSWithoutModTime(map[string][]byte{
+			"main.spx": []byte(`// A spx game.
+import "math"
+
+onClick => {
+	println math.floor(2.5)
+}
+
+run "assets", {Title: "My Game"}
+`),
 		}), nil)
 		params := &DocumentFormattingParams{
 			TextDocument: TextDocumentIdentifier{URI: "file:///main.spx"},

tools/spxls/internal/server/rename.go

@@ -27,7 +27,7 @@ func (s *Server) textDocumentPrepareRename(params *PrepareRenameParams) (*Range,
 		return nil, nil
 	}
 	obj := result.typeInfo.ObjectOf(ident)
-	if !isRenameableObject(obj) {
+	if !isRenameableObject(obj) || result.isGameOrSpriteThisPtr(obj) {
 		return nil, nil
 	}
 
@@ -63,7 +63,7 @@ func (s *Server) textDocumentRename(params *RenameParams) (*WorkspaceEdit, error
 	}
 
 	obj := result.typeInfo.ObjectOf(result.identAtASTFilePosition(astFile, params.Position))
-	if !isRenameableObject(obj) {
+	if !isRenameableObject(obj) || result.isGameOrSpriteThisPtr(obj) {
 		return nil, nil
 	}
 

tools/spxls/internal/server/rename_test.go

@@ -61,6 +61,42 @@ onStart => {
 		require.NoError(t, err)
 		require.Nil(t, range3)
 	})
+
+	t.Run("MainGameAndSpriteThisPtrs", func(t *testing.T) {
+		s := New(newMapFSWithoutModTime(map[string][]byte{
+			"main.spx": []byte(`
+onClick => {
+	_ = this
+}
+run "assets", {Title: "My Game"}
+`),
+			"MySprite.spx": []byte(`
+onClick => {
+	_ = this
+}
+`),
+			"assets/index.json":                  []byte(`{}`),
+			"assets/sprites/MySprite/index.json": []byte(`{}`),
+		}), nil)
+
+		range1, err := s.textDocumentPrepareRename(&PrepareRenameParams{
+			TextDocumentPositionParams: TextDocumentPositionParams{
+				TextDocument: TextDocumentIdentifier{URI: "file:///main.spx"},
+				Position:     Position{Line: 2, Character: 5},
+			},
+		})
+		require.NoError(t, err)
+		require.Nil(t, range1)
+
+		range2, err := s.textDocumentPrepareRename(&PrepareRenameParams{
+			TextDocumentPositionParams: TextDocumentPositionParams{
+				TextDocument: TextDocumentIdentifier{URI: "file:///MySprite.spx"},
+				Position:     Position{Line: 2, Character: 5},
+			},
+		})
+		require.NoError(t, err)
+		require.Nil(t, range2)
+	})
 }
 
 func TestServerTextDocumentRename(t *testing.T) {
@@ -218,6 +254,40 @@ onStart => {
 			NewText: "NewSprite",
 		})
 	})
+
+	t.Run("MainGameAndSpriteThisPtrs", func(t *testing.T) {
+		s := New(newMapFSWithoutModTime(map[string][]byte{
+			"main.spx": []byte(`
+onClick => {
+	_ = this
+}
+run "assets", {Title: "My Game"}
+`),
+			"MySprite.spx": []byte(`
+onClick => {
+	_ = this
+}
+`),
+			"assets/index.json":                  []byte(`{}`),
+			"assets/sprites/MySprite/index.json": []byte(`{}`),
+		}), nil)
+
+		mainSpxWorkspaceEdit, err := s.textDocumentRename(&RenameParams{
+			TextDocument: TextDocumentIdentifier{URI: "file:///main.spx"},
+			Position:     Position{Line: 2, Character: 5},
+			NewName:      "that",
+		})
+		require.NoError(t, err)
+		require.Nil(t, mainSpxWorkspaceEdit)
+
+		mySpriteSpxWorkspaceEdit, err := s.textDocumentRename(&RenameParams{
+			TextDocument: TextDocumentIdentifier{URI: "file:///MySprite.spx"},
+			Position:     Position{Line: 2, Character: 5},
+			NewName:      "that",
+		})
+		require.NoError(t, err)
+		require.Nil(t, mySpriteSpxWorkspaceEdit)
+	})
 }
 
 func TestServerSpxRenameBackdropResource(t *testing.T) {

tools/spxls/internal/server/spx_definition.go

@@ -260,15 +260,15 @@ var GetSpxPkg = sync.OnceValue(func() *types.Package {
 })
 
 // GetSpxGameType returns the [spx.Game] type.
-var GetSpxGameType = sync.OnceValue(func() types.Type {
+var GetSpxGameType = sync.OnceValue(func() *types.Named {
 	spxPkg := GetSpxPkg()
-	return spxPkg.Scope().Lookup("Game").Type()
+	return spxPkg.Scope().Lookup("Game").Type().(*types.Named)
 })
 
 // GetSpxSpriteImplType returns the [spx.SpriteImpl] type.
-var GetSpxSpriteImplType = sync.OnceValue(func() types.Type {
+var GetSpxSpriteImplType = sync.OnceValue(func() *types.Named {
 	spxPkg := GetSpxPkg()
-	return spxPkg.Scope().Lookup("SpriteImpl").Type()
+	return spxPkg.Scope().Lookup("SpriteImpl").Type().(*types.Named)
 })
 
 // GetSpxPkgDefinitions returns the spx definitions for the spx package.

tools/spxls/internal/server/util.go

@@ -31,15 +31,6 @@ func listSpxFiles(rootFS fs.ReadDirFS) ([]string, error) {
 	return files, nil
 }
 
-// gopASTFileMapToSlice converts a map of [gopast.File] to a slice of [gopast.File].
-func gopASTFileMapToSlice(fileMap map[string]*gopast.File) []*gopast.File {
-	files := make([]*gopast.File, 0, len(fileMap))
-	for _, file := range fileMap {
-		files = append(files, file)
-	}
-	return files
-}
-
 // unwrapPointerType returns the underlying type of t. For pointer types, it
 // returns the element type that the pointer points to. For non-pointer types,
 // it returns the type unchanged.