PgTidy/pkg/format/safety.go

package format

import (
	"fmt"
	"strings"

	"git.warky.dev/wdevs/pgtidy/pkg/config"
	"git.warky.dev/wdevs/pgtidy/pkg/lexer"
	"git.warky.dev/wdevs/pgtidy/pkg/parser"
)

// VerifySafe runs every safety invariant against a formatting result before it
// is written to disk or returned to an editor. src is the original input, out
// the formatted output, and st the style out was produced with. It returns nil
// when out is safe to emit, otherwise an error naming the invariant that failed.
//
// The checks, in order of cost:
//
//  1. Semantic equivalence — the non-trivia (code) token stream is unchanged:
//     identifiers/keywords compare case-insensitively, everything else exactly,
//     recursing into dollar-quoted bodies. Comments and whitespace are trivia
//     and are deliberately ignored here.
//  2. Comment preservation — every -- and /* */ comment in src reappears in out,
//     in the same order, with the same content (ignoring only trailing
//     whitespace and CRLF/LF). The formatter may move or re-indent a comment but
//     must never drop, merge, split, or reword one.
//  3. Structural balance — the ( ) [ ] and BEGIN/CASE/IF/LOOP…END nesting
//     profile of out matches src's, counting only real code tokens (anything
//     inside a comment or a string/dollar-quoted literal is ignored).
//  4. Idempotence — formatting out again yields out unchanged.
//
// Any failure means the formatter has a bug: the caller must keep the original
// source and never emit out.
func VerifySafe(src, out string, st config.Style) error {
	if !SemanticallyEqual(src, out) {
		return fmt.Errorf("code token stream changed")
	}
	if err := CommentsPreserved(src, out); err != nil {
		return err
	}
	if err := StructurallyBalanced(src, out); err != nil {
		return err
	}
	if reformatted := File(parser.Parse(out), st); reformatted != out {
		return fmt.Errorf("output is not idempotent (a second format pass would change it)")
	}
	return nil
}

// SemanticallyEqual reports whether a and b have the same non-trivia token
// stream, i.e. formatting may only ever change whitespace/comment trivia and
// layout — it must never add, remove, or alter a token of actual code.
// Unquoted identifiers/keywords compare case-insensitively (casing is a
// style choice); everything else (strings, numbers, operators, punctuation)
// must match exactly. Dollar-quoted body tokens are compared recursively so
// that independent body reformatting doesn't trigger a false failure.
//
// The CLI and LSP must call this (via VerifySafe) before ever writing or
// emitting formatted output: if it returns false, the formatter has a bug and
// the original source must be kept, never the (corrupting) formatted output.
func SemanticallyEqual(a, b string) bool {
	ta := significantTokens(a)
	tb := significantTokens(b)
	if len(ta) != len(tb) {
		return false
	}
	for i := range ta {
		if ta[i].Kind != tb[i].Kind {
			return false
		}
		switch ta[i].Kind {
		case lexer.Ident:
			if !strings.EqualFold(ta[i].Text, tb[i].Text) {
				return false
			}
		case lexer.String, lexer.EscapeString, lexer.BitString, lexer.HexString, lexer.UnicodeString:
			// A CRLF vs LF difference inside a multi-line string literal is a
			// line-ending normalisation, not a change of code content — the
			// formatter always re-emits layout with st.Newline. Compare the
			// literal modulo \r\n ↔ \n.
			if normNL(ta[i].Text) != normNL(tb[i].Text) {
				return false
			}
		case lexer.DollarString:
			_, innerA, _, okA := splitDollarQuote(ta[i].Text)
			_, innerB, _, okB := splitDollarQuote(tb[i].Text)
			if okA != okB || (okA && !SemanticallyEqual(innerA, innerB)) {
				return false
			}
		default:
			if ta[i].Text != tb[i].Text {
				return false
			}
		}
	}
	return true
}

// CommentsPreserved reports whether every comment in a survives into b with its
// text intact. Comments are compared in document order; each is reduced to its
// sequence of non-blank text lines (line endings normalised, every line trimmed
// of surrounding whitespace, blank lines dropped) so that the formatter is free
// to move or re-indent a comment but can never drop, merge, split, reorder, or
// reword one. Comments inside dollar-quoted bodies are included (the bodies are
// lexed recursively). A non-nil error describes the first divergence.
func CommentsPreserved(a, b string) error {
	ca := comments(a)
	cb := comments(b)
	if len(ca) != len(cb) {
		return fmt.Errorf("comment count changed: input has %d, output has %d", len(ca), len(cb))
	}
	for i := range ca {
		if ca[i] != cb[i] {
			return fmt.Errorf("comment %d/%d changed:\n  input:  %q\n  output: %q", i+1, len(ca), ca[i], cb[i])
		}
	}
	return nil
}

// comments returns the normalised text of every -- and /* */ comment in src, in
// order, descending into dollar-quoted bodies.
func comments(src string) []string {
	var out []string
	for _, t := range lexer.Lex(src) {
		switch t.Kind {
		case lexer.LineComment, lexer.BlockComment:
			out = append(out, normComment(t.Text))
		case lexer.DollarString:
			if _, inner, _, ok := splitDollarQuote(t.Text); ok {
				out = append(out, comments(inner)...)
			}
		}
	}
	return out
}

// normComment canonicalises a comment token to its content — the ordered list of
// non-blank text lines, each stripped of surrounding whitespace, joined with LF.
// Line endings and indentation are layout, not content, so they are discarded;
// dropping or rewording an actual line of comment text still shows up.
func normComment(s string) string {
	s = strings.ReplaceAll(s, "\r\n", "\n")
	s = strings.ReplaceAll(s, "\r", "\n")
	var lines []string
	for _, ln := range strings.Split(s, "\n") {
		if ln = strings.TrimSpace(ln); ln != "" {
			lines = append(lines, ln)
		}
	}
	return strings.Join(lines, "\n")
}

// StructurallyBalanced reports whether a and b have the same delimiter and block
// nesting profile: identical counts of ( ) [ ] and of the PL/pgSQL block
// keywords BEGIN / CASE / IF / LOOP / END (and the compound END IF / END LOOP /
// END CASE), plus an identical running paren/bracket depth trace. Only real code
// tokens are counted — anything inside a -- or /* */ comment is trivia and is
// skipped, and string / dollar-quoted literals are opaque single tokens whose
// contents never register (dollar-quoted bodies are recursed into separately).
//
// Given SemanticallyEqual, this is defence in depth: an independent re-count
// with different code that catches a structural token slipping through a bug in
// the token-stream comparison (e.g. its dollar-quote or CRLF handling), and it
// pins down *where* the structure broke.
func StructurallyBalanced(a, b string) error {
	pa := structureProfile(a)
	pb := structureProfile(b)
	if pa.parenDepthTrace != pb.parenDepthTrace {
		return fmt.Errorf("parenthesis/bracket nesting changed")
	}
	for _, k := range structureKeys {
		if pa.counts[k] != pb.counts[k] {
			return fmt.Errorf("structural token %q count changed: input %d, output %d", k, pa.counts[k], pb.counts[k])
		}
	}
	return nil
}

var structureKeys = []string{"(", ")", "[", "]", "begin", "case", "if", "loop", "end", "end if", "end loop", "end case"}

type structProfile struct {
	counts map[string]int
	// parenDepthTrace is the sequence of running ( ) [ ] depths after each
	// bracket token, joined with commas — a compact fingerprint of the nesting
	// shape that diverges as soon as an open/close is added, dropped, or moved.
	parenDepthTrace string
}

func structureProfile(src string) structProfile {
	p := structProfile{counts: map[string]int{}}
	var trace strings.Builder
	depth := 0
	toks := significantTokens(src) // trivia (comments/whitespace) already excluded
	for i := 0; i < len(toks); i++ {
		t := toks[i]
		switch t.Kind {
		case lexer.LParen:
			p.counts["("]++
			depth++
			fmt.Fprintf(&trace, "%d,", depth)
		case lexer.RParen:
			p.counts[")"]++
			depth--
			fmt.Fprintf(&trace, "%d,", depth)
		case lexer.LBracket:
			p.counts["["]++
			depth++
			fmt.Fprintf(&trace, "%d,", depth)
		case lexer.RBracket:
			p.counts["]"]++
			depth--
			fmt.Fprintf(&trace, "%d,", depth)
		case lexer.Ident:
			switch lowerASCII(t.Text) {
			case "begin", "case", "if", "loop":
				p.counts[lowerASCII(t.Text)]++
			case "end":
				p.counts["end"]++
				if i+1 < len(toks) && toks[i+1].Kind == lexer.Ident {
					switch lowerASCII(toks[i+1].Text) {
					case "if":
						p.counts["end if"]++
					case "loop":
						p.counts["end loop"]++
					case "case":
						p.counts["end case"]++
					}
				}
			}
		case lexer.DollarString:
			if _, inner, _, ok := splitDollarQuote(t.Text); ok {
				sub := structureProfile(inner)
				for _, k := range structureKeys {
					p.counts[k] += sub.counts[k]
				}
				trace.WriteString("[" + sub.parenDepthTrace + "]")
			}
		}
	}
	p.parenDepthTrace = trace.String()
	return p
}

// normNL collapses CRLF to LF so string literals compare independent of the
// source file's line-ending convention.
func normNL(s string) string { return strings.ReplaceAll(s, "\r\n", "\n") }

// significantTokens lexes src and returns its tokens excluding EOF and trivia
// (whitespace/comments).
func significantTokens(src string) []lexer.Token {
	var out []lexer.Token
	for _, t := range lexer.Lex(src) {
		if t.Kind == lexer.EOF || t.IsTrivia() {
			continue
		}
		out = append(out, t)
	}
	return out
}