PgTidy/pkg/cst/cst.go

// Package cst defines PgTidy's concrete syntax tree.
//
// The CST is lossless: every lexer token (significant tokens and the trivia —
// whitespace/comments — attached to them) is reachable, so File.Source()
// reproduces the original input byte-for-byte. The parser builds structured
// nodes only where it is confident; everything else is captured verbatim in a
// Raw node. This makes "graceful degradation" a property of the data model
// rather than something the printer must remember to do.
package cst

import (
	"strings"

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

// Trivia is a run of whitespace/comment tokens preceding a significant token.
type Trivia []lexer.Token

// Tok is a significant (non-trivia) token together with the trivia that
// immediately precedes it in the source.
type Tok struct {
	Lead Trivia
	Tok  lexer.Token
}

// Text returns the significant token's text.
func (t Tok) Text() string { return t.Tok.Text }

// Is reports whether the token is an unquoted word equal (case-insensitively)
// to kw.
func (t Tok) Is(kw string) bool {
	return t.Tok.Kind == lexer.Ident && strings.EqualFold(t.Tok.Text, kw)
}

// Comments returns just the comment tokens from the leading trivia.
func (t Tok) Comments() []lexer.Token {
	var out []lexer.Token
	for _, tr := range t.Lead {
		if tr.Kind == lexer.LineComment || tr.Kind == lexer.BlockComment {
			out = append(out, tr)
		}
	}
	return out
}

// Attach converts a lossless lexer stream into significant tokens, each
// carrying its leading trivia, plus any trailing trivia before EOF.
func Attach(toks []lexer.Token) (sig []Tok, trailing Trivia) {
	var lead Trivia
	for _, t := range toks {
		if t.Kind == lexer.EOF {
			trailing = lead
			break
		}
		if t.IsTrivia() {
			lead = append(lead, t)
			continue
		}
		sig = append(sig, Tok{Lead: lead, Tok: t})
		lead = nil
	}
	return sig, trailing
}

// Node is any element of a File.
type Node interface {
	appendTokens(*[]Tok)
}

// Tokens returns all significant tokens of a node in source order.
func Tokens(n Node) []Tok {
	var t []Tok
	n.appendTokens(&t)
	return t
}

// File is the whole parsed input.
type File struct {
	Items    []Node
	Trailing Trivia // trivia after the last significant token
}

// Source reconstructs the original input from the tree. With a correct parse
// this is byte-for-byte identical to the lexer input.
func (f *File) Source() string {
	var toks []Tok
	for _, it := range f.Items {
		it.appendTokens(&toks)
	}
	var b strings.Builder
	for _, t := range toks {
		for _, tr := range t.Lead {
			b.WriteString(tr.Text)
		}
		b.WriteString(t.Tok.Text)
	}
	for _, tr := range f.Trailing {
		b.WriteString(tr.Text)
	}
	return b.String()
}

// Raw is an unstructured statement: a verbatim run of significant tokens. This
// is the graceful-degradation fallback for anything the parser does not (yet)
// structure.
type Raw struct {
	Toks []Tok
}

func (r *Raw) appendTokens(out *[]Tok) { *out = append(*out, r.Toks...) }

// Param is one entry in a function/procedure parameter list, with the comma
// that separates it from the next entry (nil on the last parameter).
type Param struct {
	Toks []Tok
	Sep  *Tok // trailing comma, nil if last
}

// CreateFunction is a parsed CREATE [OR REPLACE] FUNCTION|PROCEDURE statement.
//
// Field order matches source order; appendTokens emits them contiguously so the
// node round-trips exactly.
type CreateFunction struct {
	Head    []Tok   // CREATE [OR REPLACE] FUNCTION|PROCEDURE
	Name    []Tok   // name (possibly schema-qualified)
	LParen  Tok     //
	Params  []Param // parameter list (may be empty)
	RParen  Tok     //
	Options [][]Tok // option clauses before the body (RETURNS/LANGUAGE/VOLATILE/...)
	As      *Tok    // the AS keyword, if present
	Body    *Tok    // the body token (dollar-quoted string or string), if present
	Tail    [][]Tok // option clauses after the body (rare)
	Semi    *Tok    // terminating semicolon, if present
}

// IsProcedure reports whether this is a PROCEDURE (vs FUNCTION).
func (c *CreateFunction) IsProcedure() bool {
	for _, t := range c.Head {
		if t.Is("procedure") {
			return true
		}
	}
	return false
}

func (c *CreateFunction) appendTokens(out *[]Tok) {
	*out = append(*out, c.Head...)
	*out = append(*out, c.Name...)
	*out = append(*out, c.LParen)
	for _, p := range c.Params {
		*out = append(*out, p.Toks...)
		if p.Sep != nil {
			*out = append(*out, *p.Sep)
		}
	}
	*out = append(*out, c.RParen)
	for _, cl := range c.Options {
		*out = append(*out, cl...)
	}
	if c.As != nil {
		*out = append(*out, *c.As)
	}
	if c.Body != nil {
		*out = append(*out, *c.Body)
	}
	for _, cl := range c.Tail {
		*out = append(*out, cl...)
	}
	if c.Semi != nil {
		*out = append(*out, *c.Semi)
	}
}