PgTidy/pkg/format/body.go

package format

import (
	"strings"

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

// formatBody applies house-style formatting to a PL/pgSQL dollar-quoted body
// token. Currently only the DECLARE section is formatted; the rest is verbatim.
func formatBody(bodyText string, st config.Style) string {
	open, inner, close, ok := splitDollarQuote(bodyText)
	if !ok {
		return bodyText
	}
	return open + formatBodyInner(inner, st) + close
}

// splitDollarQuote splits a dollar-quoted token (e.g. "$$...\n$$" or
// "$S$...$S$") into (open tag, inner text, close tag). The open and close tags
// are the same string; the last occurrence in s is taken as the close tag.
func splitDollarQuote(s string) (open, inner, close string, ok bool) {
	if len(s) == 0 || s[0] != '$' {
		return
	}
	end := strings.Index(s[1:], "$")
	if end < 0 {
		return
	}
	openLen := end + 2
	open = s[:openLen]
	closeStart := strings.LastIndex(s, open)
	if closeStart < openLen {
		return
	}
	inner = s[openLen:closeStart]
	close = s[closeStart:]
	ok = true
	return
}

func formatBodyInner(inner string, st config.Style) string {
	sig, _ := cst.Attach(lexer.Lex(inner))
	nl := st.Newline

	// Find DECLARE at depth 0.
	declareIdx := -1
	for i, t := range sig {
		if t.Tok.Kind == lexer.Ident && lowerASCII(t.Tok.Text) == "declare" {
			declareIdx = i
			break
		}
	}
	if declareIdx < 0 {
		return inner
	}

	// Find BEGIN at depth 0 after DECLARE.
	beginIdx := -1
	depth := 0
	for i := declareIdx + 1; i < len(sig); i++ {
		switch sig[i].Tok.Kind {
		case lexer.LParen, lexer.LBracket:
			depth++
		case lexer.RParen, lexer.RBracket:
			if depth > 0 {
				depth--
			}
		}
		if depth == 0 && sig[i].Tok.Kind == lexer.Ident && lowerASCII(sig[i].Tok.Text) == "begin" {
			beginIdx = i
			break
		}
	}
	if beginIdx < 0 {
		return inner
	}

	var b strings.Builder

	// Emit verbatim up to and including DECLARE (keyword-cased).
	for i := 0; i <= declareIdx; i++ {
		t := sig[i]
		for _, tr := range t.Lead {
			b.WriteString(tr.Text)
		}
		if i == declareIdx {
			b.WriteString(applyCase(t.Tok.Text, st.KeywordCase))
		} else {
			b.WriteString(t.Tok.Text)
		}
	}
	b.WriteString(nl)

	// Format each variable declaration in the DECLARE section.
	formatDeclareVars(&b, sig[declareIdx+1:beginIdx], st)

	// Emit BEGIN and everything after it verbatim from the original source.
	// The newline before BEGIN is supplied by the last declaration's line end.
	b.WriteString(inner[sig[beginIdx].Tok.Off:])

	return b.String()
}

// formatDeclareVars writes each variable declaration as a single indented line.
// Comments in the leading trivia of a declaration's first token are preserved
// on their own lines. If a declaration carries mid-body comments it is emitted
// verbatim to avoid losing them.
func formatDeclareVars(b *strings.Builder, toks []cst.Tok, st config.Style) {
	nl := st.Newline
	indent := st.Indent
	depth := 0
	var cur []cst.Tok
	var preComments []string

	emit := func() {
		if len(cur) == 0 {
			return
		}
		for _, c := range preComments {
			b.WriteString(indent)
			b.WriteString(c)
			b.WriteString(nl)
		}
		preComments = nil

		// Graceful degradation: mid-declaration comments stay verbatim.
		if anyComment(cur[1:]) {
			b.WriteString(indent)
			b.WriteString(verbatimSpan(cur))
			b.WriteString(nl)
			cur = nil
			return
		}

		body := cur
		hasSemi := len(body) > 0 && body[len(body)-1].Tok.Kind == lexer.Semicolon
		if hasSemi {
			body = body[:len(body)-1]
		}
		b.WriteString(indent)
		for i, t := range body {
			if i > 0 && needSpace(body[i-1].Tok, t.Tok) {
				b.WriteByte(' ')
			}
			b.WriteString(caseText(t.Tok, st))
		}
		if hasSemi {
			b.WriteString(";")
		}
		b.WriteString(nl)
		cur = nil
	}

	for _, t := range toks {
		if len(cur) == 0 {
			for _, tr := range t.Lead {
				if tr.Kind == lexer.LineComment || tr.Kind == lexer.BlockComment {
					preComments = append(preComments, strings.TrimRight(tr.Text, " \t"))
				}
			}
		}

		switch t.Tok.Kind {
		case lexer.LParen, lexer.LBracket:
			depth++
		case lexer.RParen, lexer.RBracket:
			if depth > 0 {
				depth--
			}
		}

		cur = append(cur, t)

		if t.Tok.Kind == lexer.Semicolon && depth == 0 {
			emit()
		}
	}
	emit()
}