gasm-sdk/lsp/handlers.go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: BSD-3-Clause

package lsp

import (
	"fmt"
	"sort"
	"strings"
	"unicode"

	"sourcedock.dev/petrbalvin/gasm-devkit/arch"
	"sourcedock.dev/petrbalvin/gasm-devkit/ast"
	"sourcedock.dev/petrbalvin/gasm-devkit/format"
	"sourcedock.dev/petrbalvin/gasm-devkit/lexer"
	"sourcedock.dev/petrbalvin/gasm-devkit/parser"
	"sourcedock.dev/petrbalvin/gasm-devkit/token"
)

// textflagMacros are the flag names defined by textflag.h; they are highlighted
// as macros and offered as completions after a TEXT/GLOBL directive.
var textflagMacros = map[string]bool{
	"NOPROFILE": true, "DUPOK": true, "NOSPLIT": true, "RODATA": true,
	"NOPTR": true, "WRAPPER": true, "NEEDCTXT": true, "TOPFRAME": true,
	"LEAF": true, "ABI0": true, "REFLECTDATA": true,
}

// completion builds the completion list for a document.
func (s *Server) completion(p completionParams) []CompletionItem {
	a := arch.ForArch(arch.FromFilename(uriPath(p.TextDocument.URI)))
	tab := a
	items := []CompletionItem{
		{Label: "TEXT", Kind: ciKeyword, Detail: "define a function"},
		{Label: "DATA", Kind: ciKeyword, Detail: "initialise a data symbol"},
		{Label: "GLOBL", Kind: ciKeyword, Detail: "declare a global symbol"},
	}
	for name := range textflagMacros {
		items = append(items, CompletionItem{Label: name, Kind: ciKeyword, Detail: "textflag.h flag"})
	}
	for name, desc := range map[string]string{
		"FP": "frame pointer (arguments/results)", "SP": "stack pointer",
		"SB": "static base (globals)", "PC": "program counter",
	} {
		items = append(items, CompletionItem{Label: name, Kind: ciConstant, Detail: desc})
	}
	for _, in := range tab.Instructions() {
		items = append(items, CompletionItem{
			Label: in.Name, Kind: ciFunction, Detail: in.Summary, Documentation: in.Summary,
		})
	}
	for _, r := range tab.Registers() {
		kind := ciVariable
		if r.Class == arch.Vector || r.Class == arch.Mask || r.Class == arch.Float || r.Class == arch.VecARM {
			kind = ciClass
		}
		items = append(items, CompletionItem{Label: r.Name, Kind: kind, Detail: r.Desc})
	}
	// Local labels defined in the document.
	if f, _ := parser.Parse("", s.docs[p.TextDocument.URI]); f != nil {
		for _, name := range labelNames(f) {
			items = append(items, CompletionItem{Label: name, Kind: ciModule, Detail: "local label"})
		}
	}
	sort.Slice(items, func(i, j int) bool { return items[i].Label < items[j].Label })
	return items
}

// hover returns documentation for the symbol under the cursor.
func (s *Server) hover(p hoverParams) *Hover {
	text := s.docs[p.TextDocument.URI]
	word, rng := wordAt(text, p.Position)
	if word == "" {
		return nil
	}
	a := arch.ForArch(arch.FromFilename(uriPath(p.TextDocument.URI)))

	var md string
	if in, ok := a.Lookup(word); ok {
		md = "**" + in.Name + "** — " + in.Summary
	} else if r, ok := a.Register(word); ok {
		md = "**" + r.Name + "** — " + r.Class.String() + " register. " + r.Desc
	} else if desc, ok := arch.PseudoRegDesc(word); ok {
		md = "**" + strings.ToUpper(word) + "** — pseudo-register. " + desc
	} else if textflagMacros[strings.ToUpper(word)] {
		md = "**" + strings.ToUpper(word) + "** — textflag.h flag"
	} else {
		return nil
	}
	return &Hover{
		Contents: markupContent{Kind: "markdown", Value: md},
		Range:    rng,
	}
}

// definition returns the location of the label definition for a label reference.
func (s *Server) definition(p definitionParams) []Location {
	text := s.docs[p.TextDocument.URI]
	word, _ := wordAt(text, p.Position)
	if word == "" {
		return nil
	}

	// Parse the document to find label definitions.
	f, errs := parser.Parse(uriPath(p.TextDocument.URI), text)
	if f == nil || len(errs) > 0 {
		return nil
	}

	// Find the label definition.
	for _, d := range f.Decls {
		if t, ok := d.(*ast.Text); ok {
			for _, stmt := range t.Body {
				if lbl, ok := stmt.(*ast.Label); ok {
					if lbl.Name.Text == word {
						return []Location{{
							URI: p.TextDocument.URI,
							Range: Range{
								Start: Position{Line: lbl.Name.Pos.Line - 1, Character: lbl.Name.Pos.Column - 1},
								End:   Position{Line: lbl.Name.Pos.Line - 1, Character: lbl.Name.Pos.Column - 1 + len(word)},
							},
						}}
					}
				}
			}
		}
	}
	return nil
}

// references returns all locations where the symbol under the cursor appears.
func (s *Server) references(p referenceParams) []Location {
	text := s.docs[p.TextDocument.URI]
	word, _ := wordAt(text, p.Position)
	if word == "" {
		return nil
	}
	f, errs := parser.Parse(uriPath(p.TextDocument.URI), text)
	if f == nil || len(errs) > 0 {
		return nil
	}
	uri := p.TextDocument.URI
	var out []Location
	for _, d := range f.Decls {
		t, ok := d.(*ast.Text)
		if !ok {
			continue
		}
		// Include the definition if requested.
		if p.Context.IncludeDeclaration {
			if t.Name.Name == word {
				out = append(out, Location{
					URI:   uri,
					Range: symRange(t.Name),
				})
			}
		}
		for _, stmt := range t.Body {
			switch st := stmt.(type) {
			case *ast.Label:
				if st.Name.Text == word {
					out = append(out, Location{
						URI:   uri,
						Range: tokenRange(st.Name),
					})
				}
			case *ast.Instr:
				for _, op := range st.Operands {
					if op.Addr.Sym != nil && op.Addr.Sym.Name == word {
						out = append(out, Location{
							URI: uri,
							Range: Range{
								Start: Position{Line: op.Pos.Line - 1, Character: op.Pos.Column - 1},
								End:   Position{Line: op.Pos.Line - 1, Character: op.Pos.Column - 1 + runeLen(word)},
							},
						})
					}
				}
			}
		}
	}
	return out
}

// rename renames a label across the document.
func (s *Server) rename(p renameParams) *WorkspaceEdit {
	text := s.docs[p.TextDocument.URI]
	word, _ := wordAt(text, p.Position)
	if word == "" || word == p.NewName {
		return nil
	}
	refs := s.references(referenceParams{
		TextDocument: p.TextDocument,
		Position:     p.Position,
		Context:      referenceContext{IncludeDeclaration: true},
	})
	if len(refs) == 0 {
		return nil
	}
	edits := make([]TextEdit, 0, len(refs))
	for _, loc := range refs {
		edits = append(edits, TextEdit{Range: loc.Range, NewText: p.NewName})
	}
	return &WorkspaceEdit{
		Changes: map[string][]TextEdit{p.TextDocument.URI: edits},
	}
}

// documentFormatting returns a single edit that replaces the whole document
// with its formatted version (canonical indentation and spacing).
func (s *Server) documentFormatting(p documentFormattingParams) []TextEdit {
	text := s.docs[p.TextDocument.URI]
	if text == "" {
		return nil
	}
	formatted := format.Source(uriPath(p.TextDocument.URI), text)
	if formatted == text {
		return nil
	}
	lines := strings.Split(text, "\n")
	endLine := len(lines) - 1
	endChar := 0
	if endLine >= 0 {
		endChar = len([]rune(lines[endLine]))
	}
	return []TextEdit{{
		Range:   Range{Start: Position{Line: 0, Character: 0}, End: Position{Line: endLine, Character: endChar}},
		NewText: formatted,
	}}
}

// inlayHints returns hints for frame sizes and argument areas.
func (s *Server) inlayHints(p inlayHintParams) []InlayHint {
	text := s.docs[p.TextDocument.URI]
	f, _ := parser.Parse(uriPath(p.TextDocument.URI), text)
	if f == nil {
		return nil
	}
	var out []InlayHint
	for _, d := range f.Decls {
		t, ok := d.(*ast.Text)
		if !ok {
			continue
		}
		// Hint after the args size: show frame size.
		if t.Frame != nil && t.Frame.Imm.HasVal && t.Args != nil && t.Args.Imm.HasVal {
			// Place hint after the args operand using its raw text length.
			col := t.Args.Pos.Column - 1 + len(t.Args.Raw)
			out = append(out, InlayHint{
				Position: Position{Line: t.Args.Pos.Line - 1, Character: col},
				Label:    fmt.Sprintf(" frame=%d", t.Frame.Imm.Val),
				Kind:     inlayHintTypeParameter,
				Tooltip:  fmt.Sprintf("local frame size: %d bytes", t.Frame.Imm.Val),
			})
		}
	}
	return out
}

// formatFile produces a canonical formatting of a parsed file.
// documentSymbols returns functions and their labels, plus global symbols.
func (s *Server) documentSymbols(p documentSymbolParams) []DocumentSymbol {
	text := s.docs[p.TextDocument.URI]
	f, _ := parser.Parse(uriPath(p.TextDocument.URI), text)
	if f == nil {
		return nil
	}
	var out []DocumentSymbol
	for _, d := range f.Decls {
		switch dd := d.(type) {
		case *ast.Text:
			sym := DocumentSymbol{
				Name:           dd.Name.Name,
				Detail:         "TEXT " + strings.Join(dd.Flags, " "),
				Kind:           symFunction,
				Range:          textRange(dd),
				SelectionRange: symRange(dd.Name),
			}
			for _, st := range dd.Body {
				if l, ok := st.(*ast.Label); ok {
					sym.Children = append(sym.Children, DocumentSymbol{
						Name:           l.Name.Text,
						Kind:           symVariable,
						Range:          tokenRange(l.Name),
						SelectionRange: tokenRange(l.Name),
					})
				}
			}
			out = append(out, sym)
		case *ast.Globl:
			out = append(out, DocumentSymbol{
				Name: dd.Name.Name, Detail: "GLOBL", Kind: symConstant,
				Range: symRange(dd.Name), SelectionRange: symRange(dd.Name),
			})
		case *ast.Data:
			out = append(out, DocumentSymbol{
				Name: dd.Name.Name, Detail: "DATA", Kind: symConstant,
				Range: symRange(dd.Name), SelectionRange: symRange(dd.Name),
			})
		}
	}
	return out
}

// semTok is one classified token before delta encoding.
type semTok struct {
	line, char, length, typ int
}

// semanticTokens encodes syntax highlighting as LSP semantic tokens.
func (s *Server) semanticTokens(p semanticTokensParams) SemanticTokens {
	text := s.docs[p.TextDocument.URI]
	a := arch.ForArch(arch.FromFilename(uriPath(p.TextDocument.URI)))
	f, _ := parser.Parse("", text)
	labels := map[string]bool{}
	for _, name := range labelNames(f) {
		labels[name] = true
	}

	toks := lexer.Tokenize(text)
	lines := groupLines(toks)

	var encoded []semTok
	for _, line := range lines {
		encoded = append(encoded, classifyLine(line, a, labels)...)
	}

	return SemanticTokens{Data: deltaEncode(encoded)}
}

// classifyLine assigns a semantic token type to each significant token on a line.
func classifyLine(line []token.Token, a *arch.Table, labels map[string]bool) []semTok {
	if len(line) == 0 {
		return nil
	}
	var out []semTok
	first := firstSignificant(line)
	if first < 0 {
		return nil
	}

	isDirective := line[first].Kind == token.Ident &&
		(line[first].Text == "TEXT" || line[first].Text == "DATA" || line[first].Text == "GLOBL")
	isLabel := line[first].Kind == token.Ident && first+1 < len(line) &&
		line[first+1].Kind == token.Colon
	isInstr := !isDirective && !isLabel && line[first].Kind == token.Ident

	mnemonicDone := false
	for i, t := range line {
		typ := -1
		switch t.Kind {
		case token.Comment:
			typ = stComment
		case token.Number:
			typ = stNumber
		case token.String, token.Rune:
			typ = stString
		case token.Hash:
			typ = stMacro
		case token.Ident:
			typ = classifyIdent(line, i, first, t.Text, a, labels,
				isDirective, isLabel, isInstr, &mnemonicDone)
		case token.Colon, token.Comma, token.LParen, token.RParen,
			token.Plus, token.Minus, token.Star, token.Slash, token.Dollar,
			token.LAngle, token.RAngle, token.LShift, token.RShift, token.Arrow, token.At:
			typ = stOperator
		}
		if typ >= 0 {
			out = append(out, semTok{
				line:   t.Pos.Line - 1,
				char:   t.Pos.Column - 1,
				length: runeLen(t.Text),
				typ:    typ,
			})
		}
	}
	return out
}

// classifyIdent decides the semantic type of an identifier token.
func classifyIdent(line []token.Token, i, first int, text string, a *arch.Table,
	labels map[string]bool, isDirective, isLabel, isInstr bool, mnemonicDone *bool) int {

	upper := strings.ToUpper(text)
	switch {
	case isDirective && i == first:
		return stKeyword
	case isDirective && textflagMacros[upper]:
		return stMacro
	case isLabel && i == first:
		return stNamespace
	case arch.IsPseudoReg(text):
		return stProperty
	case labels[text]:
		return stNamespace
	}
	if r, ok := a.Register(text); ok {
		switch r.Class {
		case arch.Vector, arch.Mask, arch.Float, arch.VecARM:
			return stType
		default:
			return stVariable
		}
	}
	if isInstr && i == first && !*mnemonicDone {
		*mnemonicDone = true
		return stFunction
	}
	// Argument/symbol names and anything else.
	return stVariable
}

// deltaEncode converts absolute token positions to the LSP relative encoding.
func deltaEncode(toks []semTok) []int {
	data := make([]int, 0, len(toks)*5)
	prevLine, prevChar := 0, 0
	for _, t := range toks {
		dLine := t.line - prevLine
		dChar := t.char
		if dLine == 0 {
			dChar = t.char - prevChar
		}
		data = append(data, dLine, dChar, t.length, t.typ, 0)
		prevLine, prevChar = t.line, t.char
	}
	return data
}

// --- shared helpers ---------------------------------------------------------

// wordAt extracts the identifier surrounding pos and its range.
func wordAt(text string, pos Position) (string, Range) {
	lines := strings.Split(text, "\n")
	if pos.Line < 0 || pos.Line >= len(lines) {
		return "", Range{}
	}
	runes := []rune(lines[pos.Line])
	col := pos.Character
	if col < 0 || col > len(runes) {
		return "", Range{}
	}
	isWord := func(r rune) bool {
		return r == '_' || r == '\u00B7' || unicode.IsLetter(r) || unicode.IsDigit(r)
	}
	start, end := col, col
	for start > 0 && isWord(runes[start-1]) {
		start--
	}
	for end < len(runes) && isWord(runes[end]) {
		end++
	}
	if start == end {
		return "", Range{}
	}
	rng := Range{
		Start: Position{Line: pos.Line, Character: start},
		End:   Position{Line: pos.Line, Character: end},
	}
	return string(runes[start:end]), rng
}

// labelNames collects every label defined in a file.
func labelNames(f *ast.File) []string {
	if f == nil {
		return nil
	}
	seen := map[string]bool{}
	var out []string
	collect := func(body []ast.Stmt) {
		for _, st := range body {
			if l, ok := st.(*ast.Label); ok && !seen[l.Name.Text] {
				seen[l.Name.Text] = true
				out = append(out, l.Name.Text)
			}
		}
	}
	for _, d := range f.Decls {
		if t, ok := d.(*ast.Text); ok {
			collect(t.Body)
		}
	}
	collect(f.Orphans)
	return out
}

// groupLines splits a token stream into lines, keeping Newline boundaries but
// dropping the Newline and EOF tokens themselves.
func groupLines(toks []token.Token) [][]token.Token {
	var lines [][]token.Token
	var cur []token.Token
	for _, t := range toks {
		if t.Kind == token.EOF {
			break
		}
		if t.Kind == token.Newline {
			lines = append(lines, cur)
			cur = nil
			continue
		}
		cur = append(cur, t)
	}
	if len(cur) > 0 {
		lines = append(lines, cur)
	}
	return lines
}

func firstSignificant(line []token.Token) int {
	for i, t := range line {
		if t.Kind != token.Comment {
			return i
		}
	}
	return -1
}

func runeLen(s string) int { return len([]rune(s)) }

// symRange builds a range covering a symbol from its position and raw text.
func symRange(sym *ast.Symbol) Range {
	start := Position{Line: sym.Pos.Line - 1, Character: sym.Pos.Column - 1}
	end := start
	end.Character += runeLen(sym.Name)
	return Range{Start: start, End: end}
}

// tokenRange builds a range covering one token.
func tokenRange(t token.Token) Range {
	return Range{
		Start: Position{Line: t.Pos.Line - 1, Character: t.Pos.Column - 1},
		End:   Position{Line: t.End.Line - 1, Character: t.End.Column - 1},
	}
}

// textRange spans a TEXT function from its keyword to the end of its body.
func textRange(t *ast.Text) Range {
	start := Position{Line: t.Keyword.Pos.Line - 1, Character: t.Keyword.Pos.Column - 1}
	end := start
	end.Character += runeLen(t.Keyword.Text)
	if n := len(t.Body); n > 0 {
		last := t.Body[n-1]
		if in, ok := last.(*ast.Instr); ok {
			end = Position{Line: in.Mnemonic.End.Line - 1, Character: in.Mnemonic.End.Column - 1}
		} else {
			lp := last.Pos()
			end = Position{Line: lp.Line - 1, Character: lp.Column - 1}
		}
	}
	return Range{Start: start, End: end}
}