snapshot of https://github.com/sqlc-dev/sqlc.git (main, 2026-09-16; blobs > 131072B omitted)

internal/sql/ast/comment.go

@@ -0,0 +1,235 @@
+package ast
+
+import (
+	"sort"
+
+	"github.com/sqlc-dev/sqlc/internal/sql/format"
+)
+
+// File is a parsed query file: its statements together with the comments
+// the parser's lexer saw. Engines whose parsers surface their trivia — the
+// lexical channel of whitespace and comments the grammar never sees, in
+// Roslyn's terminology — return it from ParseFile, so the formatter gets
+// statements and comments from one lexer pass.
+type File struct {
+	Stmts    []Statement
+	Comments []Comment
+}
+
+// Comment is a single SQL comment, positioned by byte offsets into the
+// source the statement was parsed from — the same coordinates the engine
+// parsers stamp on nodes.
+type Comment struct {
+	// Text is the comment as written, marker included ("-- x", "/* x */").
+	Text  string
+	Start int
+	End   int
+	// OwnLine reports that only blank space preceded the comment on its
+	// line, so it leads the code after it rather than trailing the code
+	// before it.
+	OwnLine bool
+}
+
+// Line reports whether the comment runs to the end of its line (-- or #),
+// so nothing may be printed after it on the same line.
+func (c Comment) Line() bool {
+	return len(c.Text) < 2 || c.Text[0] != '/' || c.Text[1] != '*'
+}
+
+// commentRec is one attached comment: the anchor nodes on either side of it
+// in print order, and whether it trails the code before it (same source
+// line) or leads the code after it (its own line).
+//
+// The next anchor is where the printer emits the comment; the prev anchor
+// is the node the comment belongs to when the tree is edited — a trailing
+// comment travels with the node it annotates, the way dave/dst attaches
+// decorations in the Go ecosystem.
+type commentRec struct {
+	c          Comment
+	prev, next Node
+	trailing   bool
+}
+
+// CommentTable holds a statement's comments attached to its nodes. It is
+// built once by AttachComments, before any printing (or, later, editing):
+// each comment is classified against the statement's anchor nodes by source
+// position and line, and from then on positions are never consulted again —
+// the printer emits comments by node identity, which is what lets an edited
+// or synthetic tree print its comments correctly.
+type CommentTable struct {
+	recs []commentRec
+	// byNext indexes recs by their next anchor, in source order; end holds
+	// the comments with no anchor after them, which trail the statement.
+	byNext map[Node][]int
+	end    []int
+	taken  []bool
+	nTaken int
+	// breaks marks the emission points where the author broke the line:
+	// the printer keeps those breaks, gofmt-style, instead of imposing a
+	// line width.
+	breaks map[Node]bool
+}
+
+// breakAt reports whether the author broke the line at this emission point.
+func (t *CommentTable) breakAt(n Node) bool {
+	return t != nil && t.breaks[n]
+}
+
+// Exhausted reports whether every attached comment was printed.
+func (t *CommentTable) Exhausted() bool {
+	return t == nil || t.nTaken == len(t.recs)
+}
+
+// take returns and consumes the comments anchored to n, in order.
+func (t *CommentTable) take(n Node) []commentRec {
+	if t == nil || n == nil {
+		return nil
+	}
+	idxs := t.byNext[n]
+	if len(idxs) == 0 {
+		return nil
+	}
+	out := make([]commentRec, 0, len(idxs))
+	for _, i := range idxs {
+		if t.taken[i] {
+			continue
+		}
+		t.taken[i] = true
+		t.nTaken++
+		out = append(out, t.recs[i])
+	}
+	return out
+}
+
+// takeRemaining returns and consumes every comment not yet printed; the
+// statement is over, so everything left trails it.
+func (t *CommentTable) takeRemaining() []commentRec {
+	if t == nil {
+		return nil
+	}
+	var out []commentRec
+	for i := range t.recs {
+		if t.taken[i] {
+			continue
+		}
+		t.taken[i] = true
+		t.nTaken++
+		out = append(out, t.recs[i])
+	}
+	return out
+}
+
+// AttachComments classifies a statement's interior comments against its
+// nodes, producing the table the printer (and any future rewriting tool)
+// works from. Placement follows gofmt's rules, decided here once from
+// source positions and lines: a comment on the same line as the code before
+// it trails that code; any other comment leads the first node printed after
+// it; a comment after the last node trails the statement.
+func AttachComments(raw *RawStmt, d format.Dialect, comments []Comment, src string) *CommentTable {
+	anchors := collectAnchors(raw, d)
+	lines := []int{0}
+	for i := 0; i < len(src); i++ {
+		if src[i] == '\n' {
+			lines = append(lines, i+1)
+		}
+	}
+	lineOf := func(pos int) int {
+		return sort.SearchInts(lines, pos+1) - 1
+	}
+
+	table := &CommentTable{byNext: make(map[Node][]int, len(anchors))}
+	for _, c := range comments {
+		// prev: the printed node with the greatest position before the
+		// comment (for the trailing/leading call). next: the emission point
+		// for the first printed node after the comment — the node itself,
+		// or, when boundary markers immediately precede it in print order,
+		// the earliest of those markers, which is where the printer will
+		// look for this comment first.
+		var prev, next Node
+		prevPos := -1
+		for _, a := range anchors {
+			if !a.marker && a.pos < c.Start && a.pos > prevPos {
+				prev, prevPos = a.node, a.pos
+			}
+		}
+		for i, a := range anchors {
+			if a.marker || a.pos <= c.Start {
+				continue
+			}
+			j := i
+			for j > 0 && anchors[j-1].marker {
+				j--
+			}
+			next = anchors[j].node
+			break
+		}
+		rec := commentRec{
+			c:        c,
+			prev:     prev,
+			next:     next,
+			trailing: !c.OwnLine && prev != nil && lineOf(prevPos) == lineOf(c.Start),
+		}
+		table.recs = append(table.recs, rec)
+		i := len(table.recs) - 1
+		if next == nil {
+			table.end = append(table.end, i)
+		} else {
+			table.byNext[next] = append(table.byNext[next], i)
+		}
+	}
+	table.taken = make([]bool, len(table.recs))
+
+	// Record where the author broke the line: an emission point whose
+	// neighbouring printed nodes sit on different source lines. The printer
+	// keeps these breaks (and only imposes new ones for comments), so the
+	// statement's shape stays the author's.
+	table.breaks = make(map[Node]bool)
+	for i, a := range anchors {
+		if !a.marker {
+			continue
+		}
+		prevPos, nextPos := -1, -1
+		for j := i - 1; j >= 0; j-- {
+			if !anchors[j].marker {
+				prevPos = anchors[j].pos
+				break
+			}
+		}
+		for j := i + 1; j < len(anchors); j++ {
+			if !anchors[j].marker {
+				nextPos = anchors[j].pos
+				break
+			}
+		}
+		if prevPos >= 0 && nextPos >= 0 && lineOf(prevPos) != lineOf(nextPos) {
+			table.breaks[a.node] = true
+		}
+	}
+	return table
+}
+
+type anchor struct {
+	node Node
+	pos  int
+	// marker anchors are emission points (beforeClause, list boundaries)
+	// rather than printed nodes; they carry no position of their own.
+	marker bool
+}
+
+// collectAnchors renders the statement once, flat, recording every
+// positioned node in the order the printer visits them. Print order is what
+// comment emission is defined against, so classifying against it keeps the
+// attach-time decision and the print-time emission point identical.
+func collectAnchors(n Node, d format.Dialect) (out []anchor) {
+	defer func() {
+		// A formatter panic here surfaces later, on the real print; anchors
+		// collected so far still place most comments.
+		recover()
+	}()
+	tb := NewTrackedBuffer()
+	tb.anchors = &out
+	if ft, ok := n.(nodeFormatter); ok {
+		ft.Format(tb, d)
+	}
+	return out
+}