chore: more work done and planning with AI.

pkg/format/format.go

@@ -53,9 +53,12 @@ func (p *printer) writeItem(n cst.Node) {
 	case *cst.CreateFunction:
 		p.writeCreateFunction(v)
 	case *cst.Raw:
-		if isDMLStart(v.Toks) {
+		switch {
+		case isDMLStart(v.Toks):
 			p.b.WriteString(formatDML(v.Toks, p.st))
-		} else {
+		case isDoBlock(v.Toks):
+			p.b.WriteString(formatDoBlock(v.Toks, p.st))
+		default:
 			p.b.WriteString(verbatimSpan(v.Toks))
 		}
 	default:
@@ -63,11 +66,64 @@ func (p *printer) writeItem(n cst.Node) {
 	}
 }
 
+// isDoBlock reports whether toks is a DO $$ ... $$ statement.
+func isDoBlock(toks []cst.Tok) bool {
+	for _, t := range toks {
+		if t.Tok.Kind == lexer.Ident {
+			return lowerASCII(t.Tok.Text) == "do"
+		}
+		if !t.Tok.IsTrivia() {
+			return false
+		}
+	}
+	return false
+}
+
+// formatDoBlock formats a DO $$ ... $$ block by applying formatBody to the
+// dollar-quoted string and emitting DO + newline + formatted body.
+func formatDoBlock(toks []cst.Tok, st config.Style) string {
+	// Find the DO keyword, the dollar-string body, and the optional semicolon.
+	var doTok, bodyTok *cst.Tok
+	hasSemi := false
+	for i := range toks {
+		t := &toks[i]
+		if t.Tok.IsTrivia() || t.Tok.Kind == lexer.EOF {
+			continue
+		}
+		low := lowerASCII(t.Tok.Text)
+		if t.Tok.Kind == lexer.Ident && low == "do" && doTok == nil {
+			doTok = t
+			continue
+		}
+		if doTok != nil && t.Tok.Kind == lexer.DollarString && bodyTok == nil {
+			bodyTok = t
+			continue
+		}
+		if t.Tok.Kind == lexer.Semicolon {
+			hasSemi = true
+		}
+	}
+	if doTok == nil || bodyTok == nil {
+		return verbatimSpan(toks)
+	}
+
+	nl := st.Newline
+	var b strings.Builder
+	b.WriteString(applyCase(doTok.Tok.Text, st.KeywordCase))
+	b.WriteString(nl)
+	b.WriteString(formatBody(bodyTok.Tok.Text, st))
+	if hasSemi {
+		b.WriteString(";")
+	}
+	return b.String()
+}
+
 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.
+	// We still format the body dollar-string independently since it is self-contained.
 	if headerHasComments(cf) {
-		p.b.WriteString(verbatimSpan(cst.Tokens(cf)))
+		p.b.WriteString(verbatimSpanFormatBody(cst.Tokens(cf), cf.Body, p.st))
 		return
 	}
 
@@ -80,11 +136,22 @@ func (p *printer) writeCreateFunction(cf *cst.CreateFunction) {
 	}
 	p.b.WriteString("(")
 
+	// Build formatted param texts first so we can measure widths.
+	paramTexts := make([]string, len(cf.Params))
+	for i, param := range cf.Params {
+		paramTexts[i] = p.inline(param.Toks)
+	}
+
+	// align_param_types: pad param names so type columns align.
+	if p.st.AlignParamTypes && len(cf.Params) > 1 {
+		paramTexts = alignParamTypes(paramTexts)
+	}
+
 	first := p.st.Indent + " " // align item text one column past the comma
 	cont := p.st.Indent
-	for i, param := range cf.Params {
+	for i, text := range paramTexts {
+		param := cf.Params[i]
 		p.nl()
-		text := p.inline(param.Toks)
 		if i == 0 || p.st.Commas != config.CommaLeading {
 			p.b.WriteString(first)
 			p.b.WriteString(text)
@@ -96,6 +163,16 @@ func (p *printer) writeCreateFunction(cf *cst.CreateFunction) {
 			p.b.WriteString(",")
 			p.b.WriteString(text)
 		}
+		// Emit trailing inline comment from the separator (e.g. --description after param).
+		if param.Sep != nil {
+			for _, tr := range param.Sep.Lead {
+				if tr.Kind == lexer.LineComment || tr.Kind == lexer.BlockComment {
+					p.b.WriteByte(' ')
+					p.b.WriteString(strings.TrimRight(tr.Text, " \t"))
+					break
+				}
+			}
+		}
 	}
 	p.nl()
 	p.b.WriteString(")")
@@ -105,7 +182,12 @@ func (p *printer) writeCreateFunction(cf *cst.CreateFunction) {
 		p.b.WriteString(p.inline(clause))
 	}
 	if cf.As != nil {
-		p.nl()
+		// routine_as_wrap: when false, AS stays on the same line as the last option.
+		if p.st.RoutineAsWrap {
+			p.nl()
+		} else {
+			p.b.WriteByte(' ')
+		}
 		p.b.WriteString(p.inline([]cst.Tok{{Tok: cf.As.Tok}}))
 	}
 	if cf.Body != nil {
@@ -139,7 +221,12 @@ func (p *printer) inline(toks []cst.Tok) string {
 		if i > 0 && needSpace(toks[i-1].Tok, t.Tok) {
 			b.WriteByte(' ')
 		}
-		b.WriteString(caseText(t.Tok, p.st))
+		var prev lexer.Token
+		if i > 0 {
+			prev = toks[i-1].Tok
+		}
+		nextIsLParen := i+1 < len(toks) && toks[i+1].Tok.Kind == lexer.LParen
+		b.WriteString(caseTextCtx(t.Tok, prev, nextIsLParen, p.st))
 	}
 	return b.String()
 }
@@ -147,7 +234,8 @@ func (p *printer) inline(toks []cst.Tok) 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"))
+			text := strings.TrimRight(strings.ReplaceAll(tr.Text, "\r", ""), " \t")
+			p.b.WriteString(text)
 			p.nl()
 		}
 	}
@@ -157,7 +245,7 @@ func (p *printer) trailingComments(lead cst.Trivia) {
 	cs := commentsOf(lead)
 	for _, c := range cs {
 		p.nl()
-		p.b.WriteString(strings.TrimRight(c.Text, " \t"))
+		p.b.WriteString(strings.TrimRight(strings.ReplaceAll(c.Text, "\r", ""), " \t"))
 	}
 }
 
@@ -211,8 +299,14 @@ func needSpace(a, b lexer.Token) bool {
 }
 
 func caseText(t lexer.Token, st config.Style) string {
+	return caseTextCtx(t, lexer.Token{}, false, st)
+}
+
+// caseTextCtx applies casing with context: prev is the preceding significant
+// token, nextIsLParen indicates the next significant token is '('.
+func caseTextCtx(t lexer.Token, prev lexer.Token, nextIsLParen bool, st config.Style) string {
 	if t.Kind != lexer.Ident {
-		return t.Text // only unquoted words are re-cased
+		return t.Text
 	}
 	low := lowerASCII(t.Text)
 	switch {
@@ -220,6 +314,10 @@ func caseText(t lexer.Token, st config.Style) string {
 		return applyCase(t.Text, st.TypeCase)
 	case isKeyword(low):
 		return applyCase(t.Text, st.KeywordCase)
+	case nextIsLParen && isBuiltinFunc(low):
+		return applyCase(t.Text, st.BuiltinCase)
+	case prev.Kind == lexer.Ident && lowerASCII(prev.Text) == "as":
+		return applyCase(t.Text, st.AliasCase)
 	default:
 		return applyCase(t.Text, st.IdentCase)
 	}
@@ -239,18 +337,25 @@ func applyCase(s string, c config.Case) string {
 // --- 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.
+// the formatter cannot confidently relocate. Param separator (Sep) comments
+// are excluded — those are trailing inline comments on param lines that the
+// formatter emits explicitly after each param text. The first token's leading
+// trivia and the body token are also excluded.
 func headerHasComments(cf *cst.CreateFunction) bool {
+	// Build a set of Sep token offsets so we can skip them.
+	sepOffsets := make(map[int]bool, len(cf.Params))
+	for _, p := range cf.Params {
+		if p.Sep != nil {
+			sepOffsets[p.Sep.Tok.Off] = true
+		}
+	}
 	all := cst.Tokens(cf)
 	for i, t := range all {
-		if i == 0 {
+		if i == 0 || t.Tok.Kind == lexer.Semicolon {
 			continue
 		}
-		if t.Tok.Kind == lexer.Semicolon {
-			continue
+		if sepOffsets[t.Tok.Off] {
+			continue // Sep comments handled separately
 		}
 		if hasComment(t) {
 			return true
@@ -302,6 +407,27 @@ func verbatimSpan(toks []cst.Tok) string {
 	return b.String()
 }
 
+// verbatimSpanFormatBody emits toks verbatim but replaces bodyTok's text with
+// formatBody output. Used when the function header has comments we cannot
+// safely relocate but the body can still be independently formatted.
+// If bodyTok is nil the function is identical to verbatimSpan.
+func verbatimSpanFormatBody(toks []cst.Tok, bodyTok *cst.Tok, st config.Style) string {
+	var b strings.Builder
+	for i, t := range toks {
+		if i > 0 {
+			for _, tr := range t.Lead {
+				b.WriteString(tr.Text)
+			}
+		}
+		if bodyTok != nil && t.Tok.Kind == lexer.DollarString && t.Tok.Off == bodyTok.Tok.Off {
+			b.WriteString(formatBody(t.Tok.Text, st))
+		} else {
+			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 {
@@ -331,3 +457,66 @@ func lowerASCII(s string) string {
 	}
 	return string(b)
 }
+
+// alignParamTypes pads param names so the type column aligns across all params.
+// Expected format per param: "[mode] name type [DEFAULT expr]".
+// Mode keywords (IN/OUT/INOUT/VARIADIC) are detected and skipped.
+// Params without a type are passed through unchanged.
+func alignParamTypes(params []string) []string {
+	type pp struct{ mode, name, rest string }
+	parsed := make([]pp, len(params))
+	maxNameW := 0
+
+	modeKws := map[string]bool{"in": true, "out": true, "inout": true, "variadic": true}
+
+	for i, s := range params {
+		fields := strings.Fields(s)
+		if len(fields) < 2 {
+			parsed[i].rest = s
+			continue
+		}
+		nameIdx := 0
+		if modeKws[lowerASCII(fields[0])] {
+			nameIdx = 1
+		}
+		if nameIdx >= len(fields) || nameIdx+1 >= len(fields) {
+			// No type field — keep verbatim.
+			parsed[i].rest = s
+			continue
+		}
+		if nameIdx > 0 {
+			parsed[i].mode = fields[0]
+		}
+		parsed[i].name = fields[nameIdx]
+		parsed[i].rest = strings.Join(fields[nameIdx+1:], " ")
+		if len(parsed[i].name) > maxNameW {
+			maxNameW = len(parsed[i].name)
+		}
+	}
+
+	if maxNameW == 0 {
+		return params
+	}
+
+	out := make([]string, len(params))
+	for i, p := range parsed {
+		if p.name == "" {
+			out[i] = params[i]
+			continue
+		}
+		var b strings.Builder
+		if p.mode != "" {
+			b.WriteString(p.mode)
+			b.WriteByte(' ')
+		}
+		b.WriteString(p.name)
+		// Pad name to (maxNameW+1) so the type column starts at a consistent offset.
+		pad := maxNameW + 1 - len(p.name)
+		for k := 0; k < pad; k++ {
+			b.WriteByte(' ')
+		}
+		b.WriteString(p.rest)
+		out[i] = b.String()
+	}
+	return out
+}