chore: more work done and planning with AI.

pkg/format/body.go

@@ -92,10 +92,11 @@ func formatBodyInner(inner string, st config.Style) string {
 	var b strings.Builder
 
 	// Emit verbatim up to and including DECLARE (keyword-cased).
+	// Normalize CRLF in trivia so the output always uses st.Newline.
 	for i := 0; i <= declareIdx; i++ {
 		t := sig[i]
 		for _, tr := range t.Lead {
-			b.WriteString(tr.Text)
+			b.WriteString(strings.ReplaceAll(tr.Text, "\r\n", nl))
 		}
 		if i == declareIdx {
 			b.WriteString(applyCase(t.Tok.Text, st.KeywordCase))
@@ -122,57 +123,30 @@ func formatDeclareVars(b *strings.Builder, toks []cst.Tok, st config.Style) {
 	nl := st.Newline
 	indent := st.Indent
 	depth := 0
+	var decls [][]cst.Tok
 	var cur []cst.Tok
-	var preComments []string
+	var preCommentSets [][]string
+	var curPreComments []string
 
-	emit := func() {
+	collect := 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)
+		decls = append(decls, cur)
+		preCommentSets = append(preCommentSets, curPreComments)
 		cur = nil
+		curPreComments = 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"))
+					text := strings.TrimRight(strings.ReplaceAll(tr.Text, "\r", ""), " \t")
+					curPreComments = append(curPreComments, text)
 				}
 			}
 		}
-
 		switch t.Tok.Kind {
 		case lexer.LParen, lexer.LBracket:
 			depth++
@@ -181,14 +155,195 @@ func formatDeclareVars(b *strings.Builder, toks []cst.Tok, st config.Style) {
 				depth--
 			}
 		}
-
 		cur = append(cur, t)
-
 		if t.Tok.Kind == lexer.Semicolon && depth == 0 {
-			emit()
+			collect()
+		}
+	}
+	collect()
+
+	// Compute alignment widths when requested.
+	var nameColW, typeColW int
+	if st.PlpgsqlDeclareAlignType || st.PlpgsqlDeclareAlignEq {
+		for _, decl := range decls {
+			if anyComment(decl[1:]) {
+				continue
+			}
+			body := decl
+			if len(body) > 0 && body[len(body)-1].Tok.Kind == lexer.Semicolon {
+				body = body[:len(body)-1]
+			}
+			nw, tw := declareNameTypeWidth(body, st)
+			if nw > nameColW {
+				nameColW = nw
+			}
+			if tw > typeColW {
+				typeColW = tw
+			}
+		}
+	}
+
+	for i, cur := range decls {
+		for _, c := range preCommentSets[i] {
+			b.WriteString(indent)
+			b.WriteString(c)
+			b.WriteString(nl)
+		}
+		// Graceful degradation: mid-declaration comments stay verbatim.
+		if anyComment(cur[1:]) {
+			b.WriteString(indent)
+			b.WriteString(verbatimSpan(cur))
+			b.WriteString(nl)
+			continue
+		}
+
+		body := cur
+		hasSemi := len(body) > 0 && body[len(body)-1].Tok.Kind == lexer.Semicolon
+		if hasSemi {
+			body = body[:len(body)-1]
+		}
+		b.WriteString(indent)
+		if (st.PlpgsqlDeclareAlignType || st.PlpgsqlDeclareAlignEq) && nameColW > 0 {
+			writeDeclareAligned(b, body, st, nameColW, typeColW)
+		} else {
+			for j, t := range body {
+				if j > 0 && needSpace(body[j-1].Tok, t.Tok) {
+					b.WriteByte(' ')
+				}
+				b.WriteString(caseText(t.Tok, st))
+			}
+		}
+		if hasSemi {
+			b.WriteString(";")
+		}
+		b.WriteString(nl)
+	}
+}
+
+// declareNameTypeWidth returns the rendered width of the name and type portions
+// of a DECLARE variable declaration (without the default assignment).
+// Format is: [name type [:= default]] or [name type [DEFAULT default]].
+func declareNameTypeWidth(body []cst.Tok, st config.Style) (nameW, typeW int) {
+	if len(body) < 2 {
+		return 0, 0
+	}
+	// name is always the first token.
+	name := caseText(body[0].Tok, st)
+	nameW = len(name)
+
+	// type runs from body[1] until we hit := / DEFAULT / = at depth 0.
+	var typeTokens []cst.Tok
+	depth := 0
+	for _, t := range body[1:] {
+		switch t.Tok.Kind {
+		case lexer.LParen, lexer.LBracket:
+			depth++
+		case lexer.RParen, lexer.RBracket:
+			if depth > 0 {
+				depth--
+			}
+		}
+		if depth == 0 {
+			low := lowerASCII(t.Tok.Text)
+			if t.Tok.Kind == lexer.Operator && (t.Tok.Text == ":=" || t.Tok.Text == "=") {
+				break
+			}
+			if t.Tok.Kind == lexer.Ident && low == "default" {
+				break
+			}
+		}
+		typeTokens = append(typeTokens, t)
+	}
+	var tb strings.Builder
+	for j, t := range typeTokens {
+		if j > 0 && needSpace(typeTokens[j-1].Tok, t.Tok) {
+			tb.WriteByte(' ')
+		}
+		tb.WriteString(caseText(t.Tok, st))
+	}
+	typeW = len(tb.String())
+	return nameW, typeW
+}
+
+// writeDeclareAligned writes a single DECLARE variable with aligned columns.
+func writeDeclareAligned(b *strings.Builder, body []cst.Tok, st config.Style, nameColW, typeColW int) {
+	if len(body) == 0 {
+		return
+	}
+	name := caseText(body[0].Tok, st)
+	b.WriteString(name)
+	if len(body) == 1 {
+		return
+	}
+
+	// Pad name to nameColW if align_type is requested.
+	if st.PlpgsqlDeclareAlignType {
+		pad := nameColW - len(name)
+		for k := 0; k < pad; k++ {
+			b.WriteByte(' ')
+		}
+	}
+
+	// Collect type tokens.
+	var typeTokens, restTokens []cst.Tok
+	depth := 0
+	pastType := false
+	for _, t := range body[1:] {
+		switch t.Tok.Kind {
+		case lexer.LParen, lexer.LBracket:
+			depth++
+		case lexer.RParen, lexer.RBracket:
+			if depth > 0 {
+				depth--
+			}
+		}
+		if !pastType && depth == 0 {
+			low := lowerASCII(t.Tok.Text)
+			if (t.Tok.Kind == lexer.Operator && (t.Tok.Text == ":=" || t.Tok.Text == "=")) ||
+				(t.Tok.Kind == lexer.Ident && low == "default") {
+				pastType = true
+				restTokens = append(restTokens, t)
+				continue
+			}
+		}
+		if pastType {
+			restTokens = append(restTokens, t)
+		} else {
+			typeTokens = append(typeTokens, t)
+		}
+	}
+
+	var typeStr strings.Builder
+	for j, t := range typeTokens {
+		if j > 0 && needSpace(typeTokens[j-1].Tok, t.Tok) {
+			typeStr.WriteByte(' ')
+		}
+		typeStr.WriteString(caseText(t.Tok, st))
+	}
+	typeRendered := typeStr.String()
+
+	b.WriteByte(' ')
+	b.WriteString(typeRendered)
+
+	if len(restTokens) > 0 {
+		// Pad type to typeColW if align_eq is requested.
+		if st.PlpgsqlDeclareAlignEq {
+			pad := typeColW - len(typeRendered)
+			for k := 0; k < pad; k++ {
+				b.WriteByte(' ')
+			}
+		}
+		for j, t := range restTokens {
+			prev := restTokens[0].Tok
+			if j > 0 {
+				prev = restTokens[j-1].Tok
+			}
+			if j == 0 || needSpace(prev, t.Tok) {
+				b.WriteByte(' ')
+			}
+			b.WriteString(caseText(t.Tok, st))
 		}
 	}
-	emit()
 }
 
 // bline is one logical line within an accumulated statement.
@@ -208,12 +363,17 @@ type bline struct {
 //  3. After EXCEPTION the formatter switches to verbatim-indent mode (original
 //     leading whitespace is preserved) to avoid conflicts between styles that
 //     put WHEN at col-0 vs indented.
-//  4. Blank-line counts from the original are preserved.
+//  4. Blank-line counts from the original are preserved (capped by PlpgsqlMaxBlankLines).
 func formatBodyStatements(text string, st config.Style) string {
 	nl := st.Newline
 	normalised := strings.ReplaceAll(text, "\r\n", "\n")
 	rawLines := strings.Split(normalised, "\n")
 
+	maxBlanks := st.PlpgsqlMaxBlankLines
+	if maxBlanks < 0 {
+		maxBlanks = 0
+	}
+
 	var (
 		result        strings.Builder
 		stmt          []bline
@@ -228,7 +388,11 @@ func formatBodyStatements(text string, st config.Style) string {
 		if len(stmt) == 0 {
 			return
 		}
-		for i := 0; i < pendingBlanks; i++ {
+		blanks := pendingBlanks
+		if blanks > maxBlanks {
+			blanks = maxBlanks
+		}
+		for i := 0; i < blanks; i++ {
 			result.WriteString(nl)
 		}
 		pendingBlanks = 0
@@ -259,8 +423,6 @@ func formatBodyStatements(text string, st config.Style) string {
 			}
 			effectiveDepth = blockDepth
 		case "else", "elsif", "elseif":
-			// Emit at one level up; blockDepth unchanged so the body continues
-			// at the same depth (THEN will re-apply depthInc for elsif).
 			if blockDepth > 0 {
 				effectiveDepth = blockDepth - 1
 			}
@@ -271,7 +433,14 @@ func formatBodyStatements(text string, st config.Style) string {
 
 		baseIndent := strings.Repeat(st.Indent, effectiveDepth)
 
-		for i, ll := range stmt {
+		// plpgsql_if_then_newline: when false, THEN stays on the same line as
+		// the condition. When true (default) it's already on its own logical line.
+		stmtLines := stmt
+		if !st.PlpgsqlIfThenNewline && fw == "if" {
+			stmtLines = joinThenToCondition(stmt)
+		}
+
+		for i, ll := range stmtLines {
 			if i == 0 || ll.indent == "" {
 				result.WriteString(baseIndent)
 			} else {
@@ -301,8 +470,6 @@ func formatBodyStatements(text string, st config.Style) string {
 
 		fw := lowerASCII(firstBodyKeyword(stripped))
 		isColZero := indent == ""
-		// Only join when we are at paren-depth 0; content inside parens (e.g.
-		// inside a CTE subquery) should not be merged across lines.
 		joinToPrev := isColZero && parenDepth == 0 && len(stmt) > 0 && !sqlClauseKw[fw]
 
 		if joinToPrev {
@@ -312,7 +479,6 @@ func formatBodyStatements(text string, st config.Style) string {
 			stmt = append(stmt, bline{text: stripped, indent: indent})
 		}
 
-		// Scan tokens to track paren depth and detect flush triggers.
 		var lastD0Kw string
 		for _, tok := range lexer.Lex(stripped) {
 			if tok.IsTrivia() || tok.Kind == lexer.EOF {
@@ -327,6 +493,13 @@ func formatBodyStatements(text string, st config.Style) string {
 				}
 			case lexer.Semicolon:
 				if parenDepth == 0 {
+					// plpgsql_loop_collapse: fold empty FOR … LOOP END LOOP; to one line.
+					if st.PlpgsqlLoopCollapse && len(stmt) > 0 {
+						collapsed, ok := tryCollapseLoop(stmt, st)
+						if ok {
+							stmt = []bline{{text: collapsed, indent: ""}}
+						}
+					}
 					flush()
 				}
 			}
@@ -335,13 +508,9 @@ func formatBodyStatements(text string, st config.Style) string {
 			}
 		}
 
-		// Structural keywords at the end of a line (paren depth 0) trigger a
-		// flush and possibly a block-depth change.
 		if parenDepth == 0 && len(stmt) > 0 {
 			switch lastD0Kw {
 			case "then", "loop", "begin":
-				// ELSIF/ELSEIF headers end with THEN but must NOT increment depth
-				// (blockDepth is already at the right level for the body).
 				fw0 := lowerASCII(firstBodyKeyword(stmt[0].text))
 				if fw0 != "elsif" && fw0 != "elseif" {
 					depthInc = true
@@ -357,6 +526,52 @@ func formatBodyStatements(text string, st config.Style) string {
 	return result.String()
 }
 
+// joinThenToCondition merges a THEN line (on its own bline) into the preceding
+// condition line when plpgsql_if_then_newline is false.
+func joinThenToCondition(lines []bline) []bline {
+	out := make([]bline, 0, len(lines))
+	for i, ll := range lines {
+		if i > 0 && strings.EqualFold(strings.TrimSpace(ll.text), "then") {
+			out[len(out)-1].text = strings.TrimRight(out[len(out)-1].text, " \t") + " THEN"
+		} else {
+			out = append(out, ll)
+		}
+	}
+	return out
+}
+
+// tryCollapseLoop tries to collapse an empty loop body to one line.
+// Detects: FOR … LOOP\n  (empty or only blanks)\nEND LOOP;
+// Returns the collapsed line and true on success.
+func tryCollapseLoop(lines []bline, st config.Style) (string, bool) {
+	if len(lines) < 2 {
+		return "", false
+	}
+	first := strings.TrimSpace(lines[0].text)
+	last := strings.TrimSpace(lines[len(lines)-1].text)
+	firstLow := lowerASCII(first)
+	lastLow := lowerASCII(last)
+
+	// Check last line is END LOOP; or LOOP (for WHILE/FOR empty bodies that end with LOOP).
+	if !strings.HasPrefix(lastLow, "end loop") && lastLow != "end loop;" {
+		return "", false
+	}
+	// Check middle lines are all empty.
+	for _, mid := range lines[1 : len(lines)-1] {
+		if strings.TrimSpace(mid.text) != "" {
+			return "", false
+		}
+	}
+	// Check first line ends with LOOP.
+	if !strings.HasSuffix(firstLow, "loop") {
+		return "", false
+	}
+	_ = st
+	_ = firstLow
+	// Collapse to: <header> END LOOP;
+	return strings.TrimRight(first, " \t") + " " + strings.ToUpper(last), true
+}
+
 // firstBodyKeyword returns the text of the first identifier token in s
 // (lowercased), or "" if the first significant token is not an identifier.
 func firstBodyKeyword(s string) string {

pkg/format/dml.go

@@ -179,9 +179,16 @@ func dmlSegText(seg dmlSeg, st config.Style) string {
 	}
 
 	switch kw {
-	case "select", "set", "returning":
+	case "select", "returning":
 		items := dmlSplitCommas(seg.body)
-		return dmlColList(kwText, items, st)
+		return dmlColListSelect(kwText, items, st)
+	case "set":
+		items := dmlSplitCommas(seg.body)
+		return dmlColListSet(kwText, items, st)
+	case "where":
+		return dmlWhereClause(kwText, seg.body, st)
+	case "join", "left", "right", "inner", "full", "cross", "natural":
+		return dmlJoinClause(kwText, seg.body, st)
 	case "with":
 		return formatWithBody(kwText, seg.body, st)
 	default:
@@ -193,6 +200,97 @@ func dmlSegText(seg dmlSeg, st config.Style) string {
 	}
 }
 
+// dmlJoinClause formats a JOIN clause, applying indent_join when configured.
+func dmlJoinClause(kwText string, body []cst.Tok, st config.Style) string {
+	text := dmlInline(body, st)
+	line := kwText
+	if text != "" {
+		line += " " + text
+	}
+	if !st.IndentJoin {
+		return line
+	}
+	indent := strings.Repeat(st.Indent, st.JoinIndentSize)
+	nl := st.Newline
+	var b strings.Builder
+	for i, part := range strings.Split(line, nl) {
+		if i > 0 {
+			b.WriteString(nl)
+		}
+		b.WriteString(indent)
+		b.WriteString(part)
+	}
+	return b.String()
+}
+
+// dmlWhereClause formats a WHERE clause, splitting AND/OR conditions per
+// the where_wrap and where_and_or_indent settings.
+func dmlWhereClause(kwText string, body []cst.Tok, st config.Style) string {
+	if st.WhereWrap == config.WrapNever {
+		text := dmlInline(body, st)
+		if text == "" {
+			return kwText
+		}
+		return kwText + " " + text
+	}
+
+	// Split at depth-0 AND/OR.
+	conditions := dmlSplitAndOr(body)
+	if len(conditions) <= 1 {
+		text := dmlInline(body, st)
+		if text == "" {
+			return kwText
+		}
+		return kwText + " " + text
+	}
+
+	nl := st.Newline
+	var b strings.Builder
+	b.WriteString(kwText)
+	for i, cond := range conditions {
+		b.WriteString(nl)
+		text := dmlInline(cond, st)
+		if st.WhereAndOrIndent {
+			b.WriteString(st.Indent)
+		}
+		if i == 0 {
+			// First condition: no leading AND/OR
+			b.WriteString("  ") // align with AND/OR token width
+			b.WriteString(text)
+		} else {
+			b.WriteString(text)
+		}
+	}
+	return b.String()
+}
+
+// dmlSplitAndOr splits toks at depth-0 AND/OR tokens, keeping the AND/OR with
+// the following condition.
+func dmlSplitAndOr(toks []cst.Tok) [][]cst.Tok {
+	var result [][]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 depth == 0 && t.Tok.Kind == lexer.Ident {
+			low := lowerASCII(t.Tok.Text)
+			if (low == "and" || low == "or") && i > start {
+				result = append(result, toks[start:i])
+				start = i
+			}
+		}
+	}
+	result = append(result, toks[start:])
+	return result
+}
+
 // formatWithBody formats the body of a WITH clause by splitting CTE definitions
 // at depth-0 commas and formatting the subquery inside each AS (...) block.
 func formatWithBody(kwText string, body []cst.Tok, st config.Style) string {
@@ -370,7 +468,18 @@ func dmlInline(toks []cst.Tok, st config.Style) string {
 		if i > 0 && needSpace(toks[i-1].Tok, t.Tok) {
 			b.WriteByte(' ')
 		}
-		b.WriteString(caseText(t.Tok, st))
+		// Space after comma in calls: func(a, b) vs func(a,b).
+		if st.SpaceAfterCommaInCalls && i > 0 && toks[i-1].Tok.Kind == lexer.Comma {
+			// Only inside parens (caller manages this at depth > 0, but we add space
+			// when the comma is not a clause-level comma — heuristic: always add).
+			b.WriteByte(' ')
+		}
+		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, st))
 	}
 	return b.String()
 }
@@ -405,7 +514,12 @@ func dmlSplitCommas(toks []cst.Tok) [][]cst.Tok {
 // One item: kept on the same line as the keyword.
 // Multiple items: each on its own line with the configured comma style.
 func dmlColList(kwText string, items [][]cst.Tok, st config.Style) string {
-	// Filter out spurious empty items (e.g. trailing comma in source).
+	return dmlColListSelect(kwText, items, st)
+}
+
+// dmlColListSelect formats a SELECT / RETURNING column list with optional
+// align_columns and select_align_as settings.
+func dmlColListSelect(kwText string, items [][]cst.Tok, st config.Style) string {
 	var kept [][]cst.Tok
 	for _, item := range items {
 		if len(item) > 0 {
@@ -426,14 +540,23 @@ func dmlColList(kwText string, items [][]cst.Tok, st config.Style) string {
 		return kwText + " " + body
 	}
 
-	// Multiple items: one per line.
-	first := st.Indent + " " // aligns item text one column past the comma
+	// Render each item text.
+	texts := make([]string, len(items))
+	for i, item := range items {
+		texts[i] = dmlInline(item, st)
+	}
+
+	// align_columns / select_align_as: pad expressions so AS and aliases align.
+	if (st.AlignColumns || st.SelectAlignAs) && len(texts) > 1 {
+		texts = alignSelectItems(texts, st)
+	}
+
+	first := st.Indent + " "
 	cont := st.Indent + ","
 	var b strings.Builder
 	b.WriteString(kwText)
-	for i, item := range items {
+	for i, text := range texts {
 		b.WriteString(nl)
-		text := dmlInline(item, st)
 		if i == 0 || st.Commas != config.CommaLeading {
 			b.WriteString(first)
 			b.WriteString(text)
@@ -447,3 +570,136 @@ func dmlColList(kwText string, items [][]cst.Tok, st config.Style) string {
 	}
 	return b.String()
 }
+
+// dmlColListSet formats an UPDATE SET column list with optional set_align_equal.
+func dmlColListSet(kwText string, items [][]cst.Tok, st config.Style) string {
+	var kept [][]cst.Tok
+	for _, item := range items {
+		if len(item) > 0 {
+			kept = append(kept, item)
+		}
+	}
+	items = kept
+
+	nl := st.Newline
+	switch len(items) {
+	case 0:
+		return kwText
+	case 1:
+		body := dmlInline(items[0], st)
+		if body == "" {
+			return kwText
+		}
+		return kwText + " " + body
+	}
+
+	texts := make([]string, len(items))
+	for i, item := range items {
+		texts[i] = dmlInline(item, st)
+	}
+
+	// set_align_equal: pad lhs so = signs align.
+	if st.SetAlignEqual && len(texts) > 1 {
+		texts = alignSetItems(texts)
+	}
+
+	first := st.Indent + " "
+	cont := st.Indent + ","
+	var b strings.Builder
+	b.WriteString(kwText)
+	for i, text := range texts {
+		b.WriteString(nl)
+		if i == 0 || st.Commas != config.CommaLeading {
+			b.WriteString(first)
+			b.WriteString(text)
+			if st.Commas == config.CommaTrailing && i < len(items)-1 {
+				b.WriteString(",")
+			}
+		} else {
+			b.WriteString(cont)
+			b.WriteString(text)
+		}
+	}
+	return b.String()
+}
+
+// alignSelectItems pads SELECT list item expressions so that AS keywords and
+// alias names align vertically.
+func alignSelectItems(texts []string, st config.Style) []string {
+	// Split each text into (expr, " AS ", alias) or keep as-is.
+	type part struct{ expr, alias string; hasAs bool }
+	parts := make([]part, len(texts))
+	maxExpr := 0
+	for i, t := range texts {
+		// Find " AS " or " as " (case-insensitive).
+		if idx := findAsIndex(t); idx >= 0 {
+			parts[i] = part{expr: t[:idx], alias: t[idx:], hasAs: true}
+			if l := len(t[:idx]); l > maxExpr {
+				maxExpr = l
+			}
+		} else {
+			parts[i] = part{expr: t}
+			if st.AlignColumns {
+				if l := len(t); l > maxExpr {
+					maxExpr = l
+				}
+			}
+		}
+	}
+
+	out := make([]string, len(texts))
+	for i, p := range parts {
+		if !p.hasAs || maxExpr == 0 {
+			out[i] = texts[i]
+			continue
+		}
+		pad := strings.Repeat(" ", maxExpr-len(p.expr))
+		out[i] = p.expr + pad + p.alias
+	}
+	return out
+}
+
+// findAsIndex returns the byte index of " AS " (case-insensitive) in s,
+// or -1 if not present at depth 0.
+func findAsIndex(s string) int {
+	low := lowerASCII(s)
+	// Look for " as " boundary.
+	for i := 0; i < len(low)-3; i++ {
+		if low[i] == ' ' && low[i+1] == 'a' && low[i+2] == 's' && low[i+3] == ' ' {
+			return i + 1 // index of 'a'
+		}
+	}
+	return -1
+}
+
+// alignSetItems pads SET assignment lhs values so that = signs align.
+func alignSetItems(texts []string) []string {
+	maxLhs := 0
+	lhsWidths := make([]int, len(texts))
+	for i, t := range texts {
+		idx := strings.Index(t, " = ")
+		if idx < 0 {
+			idx = strings.Index(t, "=")
+		}
+		if idx >= 0 {
+			lhsWidths[i] = idx
+			if idx > maxLhs {
+				maxLhs = idx
+			}
+		}
+	}
+	if maxLhs == 0 {
+		return texts
+	}
+	out := make([]string, len(texts))
+	for i, t := range texts {
+		if lhsWidths[i] == 0 || lhsWidths[i] == maxLhs {
+			out[i] = t
+			continue
+		}
+		idx := lhsWidths[i]
+		pad := strings.Repeat(" ", maxLhs-idx)
+		out[i] = t[:idx] + pad + t[idx:]
+	}
+	return out
+}

pkg/format/dml_test.go

@@ -226,6 +226,62 @@ func TestDMLIdempotent(t *testing.T) {
 	}
 }
 
+func TestDMLWhereAndOr(t *testing.T) {
+	// where_wrap=always should split AND/OR conditions onto separate lines.
+	src := "select a from t where x = 1 and y = 2 or z = 3;"
+	got := format(src)
+	want := "SELECT a\nFROM t\nWHERE\n    x = 1\n  AND y = 2\n  OR z = 3;\n"
+	if got != want {
+		t.Errorf("where and/or\n--- got ---\n%s\n--- want ---\n%s", got, want)
+	}
+	checkDML(t, "where and/or", got)
+}
+
+func TestDMLIndentJoin(t *testing.T) {
+	st := config.Default()
+	st.IndentJoin = true
+	src := "select a from t join s on t.id = s.id;"
+	got := File(parser.Parse(src), st)
+	want := "SELECT a\nFROM t\n  JOIN s ON t.id = s.id;\n"
+	if got != want {
+		t.Errorf("indent join\n--- got ---\n%s\n--- want ---\n%s", got, want)
+	}
+	// Idempotence with same config.
+	twice := File(parser.Parse(got), st)
+	if twice != got {
+		t.Errorf("indent join not idempotent:\n--- once ---\n%s\n--- twice ---\n%s", got, twice)
+	}
+}
+
+func TestDMLSetAlignEqual(t *testing.T) {
+	st := config.Default()
+	st.SetAlignEqual = true
+	src := "update t set a = 1, bb = 2, ccc = 3 where id = 1;"
+	got := File(parser.Parse(src), st)
+	// All = signs should align.
+	if got == "" {
+		t.Error("empty output")
+	}
+	// Idempotence.
+	twice := File(parser.Parse(got), st)
+	if twice != got {
+		t.Errorf("set_align_equal not idempotent:\n--- once ---\n%s\n--- twice ---\n%s", got, twice)
+	}
+}
+
+func TestDMLAlignParamTypes(t *testing.T) {
+	src := "create function f(in p_name text, in p_long_name integer, out p_result boolean) returns void language sql as $$ select 1 $$;"
+	got := format(src)
+	// p_name and p_long_name should have aligned types.
+	if got == "" {
+		t.Error("empty output")
+	}
+	twice := format(got)
+	if twice != got {
+		t.Errorf("align_param_types not idempotent:\n--- once ---\n%s\n--- twice ---\n%s", got, twice)
+	}
+}
+
 func TestCorpusUnaffectedByDML(t *testing.T) {
 	// Verify the corpus (which contains only CREATE FUNCTION) is not affected
 	// by the new DML formatting path.

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

pkg/format/format_test.go

@@ -24,9 +24,9 @@ func TestFormatHeaderGolden(t *testing.T) {
 
 	want := "--select * from dropall('resolvespec_login');\n" +
 		"CREATE OR REPLACE FUNCTION resolvespec_login(\n" +
-		"   INOUT p_data jsonb\n" +
+		"   INOUT p_data    jsonb\n" +
 		"  ,OUT p_success boolean\n" +
-		"  ,OUT p_error text\n" +
+		"  ,OUT p_error   text\n" +
 		")\n" +
 		"LANGUAGE plpgsql\n" +
 		"VOLATILE\n" +
@@ -71,6 +71,42 @@ func TestFormatBodyBroken(t *testing.T) {
 	}
 }
 
+func TestFormatMmProcBroken(t *testing.T) {
+	dir := filepath.Join("..", "..", "testdata", "corpus")
+	brokenData, err := os.ReadFile(filepath.Join(dir, "test_mm_proc_broken.pgsql"))
+	if err != nil {
+		t.Skipf("no test_mm_proc_broken.pgsql: %v", err)
+	}
+	goldenData, err := os.ReadFile(filepath.Join(dir, "test_mm_proc.pgsql"))
+	if err != nil {
+		t.Skipf("no test_mm_proc.pgsql: %v", err)
+	}
+
+	got := format(string(brokenData))
+	want := string(goldenData)
+	if got != want {
+		// Find and report the first differing line.
+		gotLines := strings.Split(got, "\n")
+		wantLines := strings.Split(want, "\n")
+		for i := 0; i < len(gotLines) && i < len(wantLines); i++ {
+			if gotLines[i] != wantLines[i] {
+				t.Errorf("format(test_mm_proc_broken) != test_mm_proc.pgsql at line %d\n  got:  %q\n  want: %q", i+1, gotLines[i], wantLines[i])
+				break
+			}
+		}
+		if len(gotLines) != len(wantLines) {
+			t.Errorf("format(test_mm_proc_broken): got %d lines, want %d lines", len(gotLines), len(wantLines))
+		}
+	}
+	twice := format(got)
+	if twice != got {
+		t.Errorf("format(test_mm_proc_broken) is not idempotent")
+	}
+	if !semanticallyEqual(string(brokenData), got) {
+		t.Errorf("format(test_mm_proc_broken) changed semantics")
+	}
+}
+
 func TestCorpusIdempotentAndSafe(t *testing.T) {
 	dir := filepath.Join("..", "..", "testdata", "corpus")
 	entries, err := os.ReadDir(dir)
@@ -79,7 +115,7 @@ func TestCorpusIdempotentAndSafe(t *testing.T) {
 	}
 	var seen int
 	for _, e := range entries {
-		if e.IsDir() || !strings.HasSuffix(e.Name(), ".pgsql") {
+		if e.IsDir() || !strings.HasSuffix(e.Name(), ".pgsql") || strings.HasSuffix(e.Name(), "_broken.pgsql") {
 			continue
 		}
 		seen++

pkg/format/keywords.go

@@ -51,5 +51,23 @@ func words(s string) map[string]bool {
 	return m
 }
 
-func isKeyword(lower string) bool  { return keywords[lower] }
-func isTypeName(lower string) bool { return typeNames[lower] }
+// builtinFunctions are built-in function names controlled by BuiltinCase.
+var builtinFunctions = words(`
+abs age array_agg array_length array_lower array_ndims array_upper
+bit_length btrim cardinality ceil ceiling char_length character_length
+chr clock_timestamp coalesce concat concat_ws count
+currval decode div encode exp extract floor
+generate_series greatest initcap jsonb_agg jsonb_object_agg
+justify_days justify_hours justify_interval
+lastval least length lower lpad ltrim
+max md5 min mod now nullif
+overlay pg_sleep position power quote_ident quote_literal
+random regexp_match regexp_matches regexp_replace replace reverse round rpad rtrim
+setval split_part sqrt string_agg strpos substr substring sum
+to_char to_date to_json to_jsonb to_number to_timestamp to_tsvector
+translate trim trunc unnest upper width_bucket
+`)
+
+func isKeyword(lower string) bool       { return keywords[lower] }
+func isTypeName(lower string) bool      { return typeNames[lower] }
+func isBuiltinFunc(lower string) bool   { return builtinFunctions[lower] }