refactor(format): subqueries and CTE lists on Doc-IR; fix CTE commas with trailing style

docs/todo.md

@@ -167,7 +167,7 @@ Legend: ✅ done · 🚧 in progress · ⬜ not started
   function call at the token level — formatted without space (known limitation).
 - Note: SQL keywords inside PL/pgSQL function bodies remain lowercase (matching
   the corpus golden files); casing is applied only to top-level DML.
-- ✅ Doc-IR core (`pkg/format/doc.go`: Text/Line/SoftLine/Group/Indent/IfBreak + `Render`) and `line_width` (default 120, 0 = unlimited). First consumer: `where_wrap: when_long` (`dmlWhereWhenLong`). `select_wrap` (default `always`) and `join_wrap` (default `never`) also take `always|when_long|never`; `when_long` keeps the SELECT list / `JOIN … ON` on one line when it fits `line_width`, otherwise uses the existing one-column-per-line / ON-per-line layout. Items containing a subquery or wrapped CASE always use the broken layout. _Still TODO: move the remaining DML builders (subqueries, CTEs, VALUES) onto Doc._
+- ✅ Doc-IR core (`pkg/format/doc.go`: Text/Line/SoftLine/Group/Indent/IfBreak + `Render`) and `line_width` (default 120, 0 = unlimited). First consumer: `where_wrap: when_long` (`dmlWhereWhenLong`). `select_wrap` (default `always`) and `join_wrap` (default `never`) also take `always|when_long|never`; `when_long` keeps the SELECT list / `JOIN … ON` on one line when it fits `line_width`, otherwise uses the existing one-column-per-line / ON-per-line layout. Items containing a subquery or wrapped CASE always use the broken layout. Subquery wrapping and multi-CTE lists are on the Doc-IR too (`HardLine`, `Lines`; CTE lists reuse `dmlCommaList`, which fixes dropped commas with `commas: trailing`). VALUES rows already share `dmlCommaList`.
 - ✅ LSP range formatting (already implemented; this note was stale). LSP also has hover (diagnostic rule + message), documentSymbol (functions/procedures), willSaveWaitUntil (format-on-save) and token-wide diagnostic ranges — see `docs/lsp-status.md`.
 
 ## ✅ Config expansion — DataGrip settings parity

pkg/format/dml.go

@@ -338,7 +338,6 @@ func dmlSplitAndOr(toks []cst.Tok) [][]cst.Tok {
 // 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 {
-	nl := st.Newline
 
 	cteDefs := dmlSplitCommas(body)
 	// Filter spurious empty items.
@@ -357,37 +356,11 @@ func formatWithBody(kwText string, body []cst.Tok, st config.Style) string {
 		return kwText + " " + formatCTEDef(cteDefs[0], st)
 	default:
 		// Multiple CTEs: one per line with the configured comma style.
-		first := st.Indent + " "
+		texts := make([]string, len(cteDefs))
-		cont := st.Indent + ","
-		contPad := strings.Repeat(" ", len(cont)) // same width as cont, no comma
-
-		var b strings.Builder
-		b.WriteString(kwText)
 		for i, cteDef := range cteDefs {
-			b.WriteString(nl)
+			texts[i] = formatCTEDef(cteDef, st)
-
-			var headPfx, tailPfx string
-			if i == 0 || st.Commas != config.CommaLeading {
-				headPfx = first
-				tailPfx = first
-			} else {
-				headPfx = cont
-				tailPfx = contPad
-			}
-
-			cteText := formatCTEDef(cteDef, st)
-			cteLines := strings.Split(cteText, nl)
-			for j, line := range cteLines {
-				if j > 0 {
-					b.WriteString(nl)
-					b.WriteString(tailPfx)
-				} else {
-					b.WriteString(headPfx)
-				}
-				b.WriteString(line)
-			}
 		}
-		return b.String()
+		return dmlCommaList(kwText, texts, st)
 	}
 }
 
@@ -515,26 +488,21 @@ func dmlWrapSubquery(inner []cst.Tok, st config.Style) string {
 	if sub == "" {
 		return "()"
 	}
-	lines := strings.Split(sub, nl)
+	// The first line stays on the "(" line unless subquery_content says
-
+	// new_line; every later line is indented one level.
-	var b strings.Builder
+	var body []Doc
-	b.WriteString("(")
+	for i, line := range strings.Split(sub, nl) {
-	for i, line := range lines {
+		if i > 0 || st.SubqueryContent == config.PlacementNewLine {
-		if i == 0 && st.SubqueryContent != config.PlacementNewLine {
+			body = append(body, HardLine())
-			b.WriteString(line)
-			continue
 		}
-		b.WriteString(nl)
+		body = append(body, Text(line))
-		if line != "" {
-			b.WriteString(st.Indent)
-		}
-		b.WriteString(line)
 	}
+	doc := []Doc{Text("("), Indent(Concat(body...))}
 	if st.SubqueryClosing == config.PlacementNewLine {
-		b.WriteString(nl)
+		doc = append(doc, HardLine())
 	}
-	b.WriteString(")")
+	doc = append(doc, Text(")"))
-	return b.String()
+	return Render(Concat(doc...), 0, st.Indent, nl)
 }
 
 // litNL stands in for a newline inside a multi-line string literal while DML

pkg/format/dml_test.go

@@ -710,3 +710,21 @@ func TestSelectWrapBrokenHonoursCommaStyle(t *testing.T) {
 		t.Errorf("leading commas expected when broken:\n%s", got)
 	}
 }
+
+func TestMultiCTETrailingCommas(t *testing.T) {
+	src := "with a as (select 1), b as (select 2) select * from a, b"
+	st := config.Default()
+	st.Commas = config.CommaTrailing
+	got := File(parser.Parse(src), st)
+	if !strings.Contains(got, "a AS (") || !strings.Contains(got, "),\n") {
+		t.Errorf("trailing comma between CTEs expected:\n%s", got)
+	}
+	if err := VerifySafe(src, got, st); err != nil {
+		t.Errorf("safety gate: %v\n%s", err, got)
+	}
+	st.Commas = config.CommaLeading
+	got = File(parser.Parse(src), st)
+	if !strings.Contains(got, ",b AS (") {
+		t.Errorf("leading comma expected:\n%s", got)
+	}
+}

pkg/format/doc.go

@@ -13,7 +13,7 @@ type Doc interface{ isDoc() }
 
 type (
 	docText   string
-	docLine   struct{ soft bool }
+	docLine   struct{ soft, hard bool }
 	docConcat []Doc
 	docIndent struct{ d Doc }
 	docGroup  struct{ d Doc }
@@ -37,6 +37,24 @@ func Line() Doc { return docLine{} }
 // SoftLine is nothing when flat and a newline (plus indentation) when broken.
 func SoftLine() Doc { return docLine{soft: true} }
 
+// HardLine is always a newline (plus indentation), even inside a flat group; a
+// group containing one can never be printed flat.
+func HardLine() Doc { return docLine{hard: true} }
+
+// Lines turns multi-line text into Text pieces joined by HardLines so that each
+// line picks up the surrounding Indent. Empty lines carry no indentation.
+func Lines(s, nl string) Doc {
+	parts := strings.Split(s, nl)
+	ds := make([]Doc, 0, 2*len(parts))
+	for i, p := range parts {
+		if i > 0 {
+			ds = append(ds, HardLine())
+		}
+		ds = append(ds, Text(p))
+	}
+	return Concat(ds...)
+}
+
 // Concat joins docs in order.
 func Concat(ds ...Doc) Doc { return docConcat(ds) }
 
@@ -59,12 +77,17 @@ type docCmd struct {
 func Render(d Doc, width int, unit, nl string) string {
 	var b strings.Builder
 	col := 0
+	pending := "" // indentation owed to the next non-empty text
 	stack := []docCmd{{"", false, d}}
 	for len(stack) > 0 {
 		c := stack[len(stack)-1]
 		stack = stack[:len(stack)-1]
 		switch v := c.d.(type) {
 		case docText:
+			if len(v) > 0 {
+				b.WriteString(pending)
+				pending = ""
+			}
 			b.WriteString(string(v))
 			if i := strings.LastIndexByte(string(v), '\n'); i >= 0 {
 				col = len(v) - i - 1
@@ -84,7 +107,7 @@ func Render(d Doc, width int, unit, nl string) string {
 				stack = append(stack, docCmd{c.indent, c.flat, v.broken})
 			}
 		case docLine:
-			if c.flat {
+			if c.flat && !v.hard {
 				if !v.soft {
 					b.WriteByte(' ')
 					col++
@@ -92,7 +115,7 @@ func Render(d Doc, width int, unit, nl string) string {
 				break
 			}
 			b.WriteString(nl)
-			b.WriteString(c.indent)
+			pending = c.indent
 			col = len(c.indent)
 		case docGroup:
 			flat := c.flat || width <= 0 || fitsFlat(v.d, width-col, c.indent, unit, stack)
@@ -146,6 +169,9 @@ func fitsFlat(d Doc, rem int, indent, unit string, rest []docCmd) bool {
 				work = append(work, item{v.broken, it.flat})
 			}
 		case docLine:
+			if v.hard {
+				return false
+			}
 			if !it.flat {
 				return true // a real line break ends the measured run
 			}

pkg/format/doc_test.go

@@ -38,3 +38,10 @@ func TestDocTrailingTextCountsTowardFit(t *testing.T) {
 		t.Errorf("got %q want %q", got, want)
 	}
 }
+
+func TestDocHardLineForcesBreakAndSkipsBlankIndent(t *testing.T) {
+	d := Group(Concat(Text("("), Indent(Concat(HardLine(), Lines("a\n\nb", "\n"))), HardLine(), Text(")")))
+	if got, want := Render(d, 0, "  ", "\n"), "(\n  a\n\n  b\n)"; got != want {
+		t.Errorf("got %q want %q", got, want)
+	}
+}