PgTidy/pkg/parser/parser.go

// Package parser builds a cst.File from PostgreSQL source.
//
// The parser is intentionally partial: it splits input into statements and
// structures the constructs the formatter currently understands (today:
// CREATE FUNCTION/PROCEDURE). Anything else is preserved verbatim as a cst.Raw
// node. This guarantees the parser never loses or corrupts input — see the
// round-trip test.
package parser

import (
	"git.warky.dev/wdevs/pgtidy/pkg/cst"
	"git.warky.dev/wdevs/pgtidy/pkg/lexer"
)

// Parse lexes and parses src into a lossless cst.File.
func Parse(src string) *cst.File {
	sig, trailing := cst.Attach(lexer.Lex(src))
	f := &cst.File{Trailing: trailing}
	for _, stmt := range splitStatements(sig) {
		f.Items = append(f.Items, parseStatement(stmt))
	}
	return f
}

// splitStatements breaks the significant-token stream into statements at
// top-level (paren-depth 0) semicolons. The terminating semicolon is included
// in the statement it ends.
func splitStatements(sig []cst.Tok) [][]cst.Tok {
	var stmts [][]cst.Tok
	var cur []cst.Tok
	depth := 0
	for _, t := range sig {
		cur = append(cur, t)
		switch t.Tok.Kind {
		case lexer.LParen:
			depth++
		case lexer.RParen:
			if depth > 0 {
				depth--
			}
		case lexer.Semicolon:
			if depth == 0 {
				stmts = append(stmts, cur)
				cur = nil
			}
		}
	}
	if len(cur) > 0 {
		stmts = append(stmts, cur)
	}
	return stmts
}

func parseStatement(stmt []cst.Tok) cst.Node {
	if isCreateFunction(stmt) {
		if cf, ok := parseCreateFunction(stmt); ok {
			return cf
		}
	}
	return &cst.Raw{Toks: stmt}
}

// isCreateFunction reports whether stmt begins with CREATE ... FUNCTION|PROCEDURE.
func isCreateFunction(stmt []cst.Tok) bool {
	if len(stmt) == 0 || !stmt[0].Is("create") {
		return false
	}
	for i := 1; i < len(stmt) && i < 4; i++ {
		if stmt[i].Is("function") || stmt[i].Is("procedure") {
			return true
		}
	}
	return false
}

// parseCreateFunction structures a CREATE FUNCTION/PROCEDURE statement. It
// returns ok=false (so the caller falls back to Raw) if the shape is not what
// it expects.
func parseCreateFunction(stmt []cst.Tok) (*cst.CreateFunction, bool) {
	cf := &cst.CreateFunction{}
	i := 0

	// Head: CREATE [OR REPLACE] (FUNCTION|PROCEDURE)
	if i >= len(stmt) || !stmt[i].Is("create") {
		return nil, false
	}
	i++
	if i+1 < len(stmt) && stmt[i].Is("or") && stmt[i+1].Is("replace") {
		i += 2
	}
	if i >= len(stmt) || !(stmt[i].Is("function") || stmt[i].Is("procedure")) {
		return nil, false
	}
	i++
	cf.Head = stmt[:i]

	// Name: everything up to the opening '(' of the parameter list.
	nameStart := i
	for i < len(stmt) && stmt[i].Tok.Kind != lexer.LParen {
		i++
	}
	if i >= len(stmt) {
		return nil, false // no parameter list
	}
	cf.Name = stmt[nameStart:i]
	cf.LParen = stmt[i]
	i++

	// Parameters: up to the matching ')'.
	params, rparen, ok := parseParamList(stmt, i)
	if !ok {
		return nil, false
	}
	cf.Params = params
	cf.RParen = stmt[rparen]
	i = rparen + 1

	// Trailing semicolon (handle now so it isn't swept into options/tail).
	end := len(stmt)
	if end > 0 && stmt[end-1].Tok.Kind == lexer.Semicolon {
		cf.Semi = &stmt[end-1]
		end--
	}

	// Options up to AS; then the body; then any tail options.
	asIdx := indexOfKeyword(stmt, i, end, "as")
	if asIdx < 0 {
		// No AS clause (e.g. RETURN / BEGIN ATOMIC bodies). Keep options
		// structured but leave the body unstructured.
		cf.Options = splitClauses(stmt[i:end])
		return cf, true
	}
	cf.Options = splitClauses(stmt[i:asIdx])
	cf.As = &stmt[asIdx]
	bi := asIdx + 1
	if bi < end && isBodyToken(stmt[bi]) {
		cf.Body = &stmt[bi]
		bi++
	}
	if bi < end {
		cf.Tail = splitClauses(stmt[bi:end])
	}
	return cf, true
}

// parseParamList parses comma-separated parameters starting at index `start`
// (the token after '('). It returns the parameters, the index of the matching
// ')', and ok.
func parseParamList(stmt []cst.Tok, start int) ([]cst.Param, int, bool) {
	var params []cst.Param
	depth := 0
	itemStart := start
	i := start
	flush := func(endExclusive int, sep *cst.Tok) {
		toks := stmt[itemStart:endExclusive]
		if len(toks) == 0 && sep == nil {
			return // empty parameter list ()
		}
		params = append(params, cst.Param{Toks: toks, Sep: sep})
	}
	for ; i < len(stmt); i++ {
		switch stmt[i].Tok.Kind {
		case lexer.LParen, lexer.LBracket:
			depth++
		case lexer.RParen:
			if depth == 0 {
				// Closing the parameter list.
				flush(i, nil)
				return params, i, true
			}
			depth--
		case lexer.RBracket:
			if depth > 0 {
				depth--
			}
		case lexer.Comma:
			if depth == 0 {
				sep := &stmt[i]
				flush(i, sep)
				itemStart = i + 1
			}
		}
	}
	return nil, 0, false // unbalanced
}

// splitClauses groups a run of option tokens into clauses, each beginning at a
// recognized clause keyword (at paren depth 0). Tokens before the first
// recognized keyword, if any, are attached to a leading clause so nothing is
// dropped.
func splitClauses(toks []cst.Tok) [][]cst.Tok {
	if len(toks) == 0 {
		return nil
	}
	var clauses [][]cst.Tok
	depth := 0
	start := 0
	for i, t := range toks {
		switch t.Tok.Kind {
		case lexer.LParen, lexer.LBracket:
			depth++
		case lexer.RParen, lexer.RBracket:
			if depth > 0 {
				depth--
			}
		}
		if i > start && depth == 0 && isClauseStarter(t) {
			clauses = append(clauses, toks[start:i])
			start = i
		}
	}
	clauses = append(clauses, toks[start:])
	return clauses
}

var clauseStarters = map[string]bool{
	"returns": true, "language": true, "transform": true, "window": true,
	"immutable": true, "stable": true, "volatile": true, "leakproof": true,
	"not": true, "called": true, "strict": true, "external": true,
	"security": true, "parallel": true, "cost": true, "rows": true,
	"support": true, "set": true, "as": true,
}

func isClauseStarter(t cst.Tok) bool {
	if t.Tok.Kind != lexer.Ident {
		return false
	}
	return clauseStarters[lowerASCII(t.Tok.Text)]
}

func isBodyToken(t cst.Tok) bool {
	return t.Tok.Kind == lexer.DollarString || t.Tok.Kind == lexer.String
}

// indexOfKeyword returns the index in [from,to) of the first top-level (paren
// depth 0) token equal to kw, or -1.
func indexOfKeyword(stmt []cst.Tok, from, to int, kw string) int {
	depth := 0
	for i := from; i < to; i++ {
		switch stmt[i].Tok.Kind {
		case lexer.LParen, lexer.LBracket:
			depth++
		case lexer.RParen, lexer.RBracket:
			if depth > 0 {
				depth--
			}
		}
		if depth == 0 && stmt[i].Is(kw) {
			return i
		}
	}
	return -1
}

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)
}