PgTidy/pkg/format/format.go

// Package format renders a cst.File back to source text in the configured
// house style.
//
// Scope (V1): CREATE FUNCTION/PROCEDURE headers are laid out to house style
// (one parameter per line with leading commas, each option clause on its own
// line, AS/$$ on their own lines). The PL/pgSQL body and any statement the
// parser left as cst.Raw are emitted verbatim — this upholds the safety
// invariants while body formatting is built out (task #4).
//
// Guarantees: formatting changes only trivia/layout (never literal/identifier
// semantics), and is idempotent.
package format

import (
	"strings"

	"github.com/hein/pgtidy/pkg/config"
	"github.com/hein/pgtidy/pkg/cst"
	"github.com/hein/pgtidy/pkg/lexer"
)

// File formats a parsed file with the given style.
func File(f *cst.File, st config.Style) string {
	p := &printer{st: st}
	for i, item := range f.Items {
		toks := cst.Tokens(item)
		if len(toks) == 0 {
			continue
		}
		lead := toks[0].Lead
		if i > 0 {
			p.nl()
			if hasBlankLine(lead) {
				p.nl()
			}
		}
		p.leadingComments(lead)
		p.writeItem(item)
	}
	p.trailingComments(f.Trailing)
	return ensureTrailingNewline(p.b.String(), st.Newline)
}

type printer struct {
	st config.Style
	b  strings.Builder
}

func (p *printer) nl() { p.b.WriteString(p.st.Newline) }

func (p *printer) writeItem(n cst.Node) {
	switch v := n.(type) {
	case *cst.CreateFunction:
		p.writeCreateFunction(v)
	default:
		p.b.WriteString(verbatimSpan(cst.Tokens(n)))
	}
}

func (p *printer) writeCreateFunction(cf *cst.CreateFunction) {
	// Safety: if the header carries comments we cannot confidently relocate,
	// emit the whole statement verbatim rather than risk dropping them.
	if headerHasComments(cf) {
		p.b.WriteString(verbatimSpan(cst.Tokens(cf)))
		return
	}

	head := p.inline(cf.Head)
	name := p.inline(cf.Name)
	p.b.WriteString(head)
	if name != "" {
		p.b.WriteString(" ")
		p.b.WriteString(name)
	}
	p.b.WriteString("(")

	first := p.st.Indent + " " // align item text one column past the comma
	cont := p.st.Indent
	for i, param := range cf.Params {
		p.nl()
		text := p.inline(param.Toks)
		if i == 0 || p.st.Commas != config.CommaLeading {
			p.b.WriteString(first)
			p.b.WriteString(text)
			if p.st.Commas == config.CommaTrailing && i < len(cf.Params)-1 {
				p.b.WriteString(",")
			}
		} else {
			p.b.WriteString(cont)
			p.b.WriteString(",")
			p.b.WriteString(text)
		}
	}
	p.nl()
	p.b.WriteString(")")

	for _, clause := range cf.Options {
		p.nl()
		p.b.WriteString(p.inline(clause))
	}
	if cf.As != nil {
		p.nl()
		p.b.WriteString(p.inline([]cst.Tok{{Tok: cf.As.Tok}}))
	}
	if cf.Body != nil {
		p.nl()
		p.b.WriteString(cf.Body.Tok.Text) // body emitted verbatim (formatted later)
	}
	for _, clause := range cf.Tail {
		p.nl()
		p.b.WriteString(p.inline(clause))
	}
	if cf.Semi != nil {
		p.b.WriteString(";")
	}
}

// inline renders a run of tokens on one line, applying spacing and casing.
// If the run contains comment trivia it is emitted verbatim to avoid losing
// or misplacing the comments.
func (p *printer) inline(toks []cst.Tok) string {
	if len(toks) == 0 {
		return ""
	}
	// A comment on the first token's lead is the span's leading/separation
	// comment, handled by the caller and never emitted here, so it does not
	// force verbatim. Internal comments (on later tokens) do.
	if anyComment(toks[1:]) {
		return verbatimSpan(toks)
	}
	var b strings.Builder
	for i, t := range toks {
		if i > 0 && needSpace(toks[i-1].Tok, t.Tok) {
			b.WriteByte(' ')
		}
		b.WriteString(caseText(t.Tok, p.st))
	}
	return b.String()
}

func (p *printer) leadingComments(lead cst.Trivia) {
	for _, tr := range lead {
		if tr.Kind == lexer.LineComment || tr.Kind == lexer.BlockComment {
			p.b.WriteString(strings.TrimRight(tr.Text, " \t"))
			p.nl()
		}
	}
}

func (p *printer) trailingComments(lead cst.Trivia) {
	cs := commentsOf(lead)
	for _, c := range cs {
		p.nl()
		p.b.WriteString(strings.TrimRight(c.Text, " \t"))
	}
}

// --- spacing & casing ---

// tightOps are operators printed without surrounding spaces.
var tightOps = map[string]bool{"::": true, ":": true, "->": true, "->>": true}

func needSpace(a, b lexer.Token) bool {
	// No space after.
	switch a.Kind {
	case lexer.LParen, lexer.LBracket, lexer.Dot:
		return false
	case lexer.Operator:
		if tightOps[a.Text] {
			return false
		}
	}
	// No space before.
	switch b.Kind {
	case lexer.RParen, lexer.RBracket, lexer.Comma, lexer.Semicolon, lexer.Dot:
		return false
	case lexer.LParen:
		switch a.Kind {
		case lexer.Ident, lexer.QuotedIdent, lexer.RParen, lexer.RBracket, lexer.Param:
			return false // function call / type modifier
		}
	case lexer.Operator:
		if tightOps[b.Text] {
			return false
		}
	}
	return true
}

func caseText(t lexer.Token, st config.Style) string {
	if t.Kind != lexer.Ident {
		return t.Text // only unquoted words are re-cased
	}
	low := lowerASCII(t.Text)
	switch {
	case isTypeName(low):
		return applyCase(t.Text, st.TypeCase)
	case isKeyword(low):
		return applyCase(t.Text, st.KeywordCase)
	default:
		return applyCase(t.Text, st.IdentCase)
	}
}

func applyCase(s string, c config.Case) string {
	switch c {
	case config.CaseUpper:
		return strings.ToUpper(s)
	case config.CaseLower:
		return strings.ToLower(s)
	default:
		return s
	}
}

// --- helpers ---

// headerHasComments reports whether the function header carries comment trivia
// the formatter cannot confidently relocate. The first token's leading trivia
// is excluded: that is the statement's leading comment, which File() emits
// separately. The body token's own text is excluded too (it is emitted
// verbatim), but a comment in front of the body is caught.
func headerHasComments(cf *cst.CreateFunction) bool {
	all := cst.Tokens(cf)
	for i, t := range all {
		if i == 0 {
			continue
		}
		if t.Tok.Kind == lexer.Semicolon {
			continue
		}
		if hasComment(t) {
			return true
		}
	}
	return false
}

func anyComment(toks []cst.Tok) bool {
	for _, t := range toks {
		if hasComment(t) {
			return true
		}
	}
	return false
}

func hasComment(t cst.Tok) bool {
	for _, tr := range t.Lead {
		if tr.Kind == lexer.LineComment || tr.Kind == lexer.BlockComment {
			return true
		}
	}
	return false
}

func commentsOf(lead cst.Trivia) []lexer.Token {
	var out []lexer.Token
	for _, tr := range lead {
		if tr.Kind == lexer.LineComment || tr.Kind == lexer.BlockComment {
			out = append(out, tr)
		}
	}
	return out
}

// verbatimSpan emits tokens exactly as in source, excluding the leading trivia
// of the first token (separation is controlled by the caller).
func verbatimSpan(toks []cst.Tok) string {
	var b strings.Builder
	for i, t := range toks {
		if i > 0 {
			for _, tr := range t.Lead {
				b.WriteString(tr.Text)
			}
		}
		b.WriteString(t.Tok.Text)
	}
	return b.String()
}

// hasBlankLine reports whether leading whitespace trivia contains a blank line
// (two or more newlines), indicating the author wanted statements separated.
func hasBlankLine(lead cst.Trivia) bool {
	n := 0
	for _, tr := range lead {
		if tr.Kind == lexer.Whitespace {
			n += strings.Count(tr.Text, "\n")
		}
	}
	return n >= 2
}

func ensureTrailingNewline(s, nl string) string {
	s = strings.TrimRight(s, "\n")
	if s == "" {
		return s
	}
	return s + nl
}

func lowerASCII(s string) string {
	b := []byte(s)
	for i, c := range b {
		if c >= 'A' && c <= 'Z' {
			b[i] = c + 32
		}
	}
	return string(b)
}