@@ -2,7 +2,8 @@ package ast
22
33import (
44 "sort"
5- "strings"
5+
6+ "github.com/sqlc-dev/sqlc/internal/sql/format"
67)
78
89// File is a parsed query file: its statements together with the comments
@@ -32,59 +33,168 @@ type Comment struct {
3233// Line reports whether the comment runs to the end of its line (-- or #),
3334// so nothing may be printed after it on the same line.
3435func (c Comment ) Line () bool {
35- return ! strings . HasPrefix (c .Text , "/*" )
36+ return len (c .Text ) < 2 || c . Text [ 0 ] != '/' || c . Text [ 1 ] != '*'
3637}
3738
38- // CommentSet carries a statement's interior comments through the printer,
39- // gofmt-style: the printer holds a cursor over the comment list and flushes
40- // every comment positioned before the node it is about to print.
41- type CommentSet struct {
42- comments []Comment
43- // lineStarts holds the byte offset of the first byte of every line of
44- // the source, sqlc's stand-in for go/token.FileSet: it turns a byte
45- // offset into a line number, which is what placement decisions compare.
46- lineStarts []int
39+ // commentRec is one attached comment: the anchor nodes on either side of it
40+ // in print order, and whether it trails the code before it (same source
41+ // line) or leads the code after it (its own line).
42+ //
43+ // The next anchor is where the printer emits the comment; the prev anchor
44+ // is the node the comment belongs to when the tree is edited — a trailing
45+ // comment travels with the node it annotates, the way dave/dst attaches
46+ // decorations in the Go ecosystem.
47+ type commentRec struct {
48+ c Comment
49+ prev , next Node
50+ trailing bool
51+ }
4752
48- next int // index of the next unprinted comment
49- cursor int // greatest node position printed so far
53+ // CommentTable holds a statement's comments attached to its nodes. It is
54+ // built once by AttachComments, before any printing (or, later, editing):
55+ // each comment is classified against the statement's anchor nodes by source
56+ // position and line, and from then on positions are never consulted again —
57+ // the printer emits comments by node identity, which is what lets an edited
58+ // or synthetic tree print its comments correctly.
59+ type CommentTable struct {
60+ recs []commentRec
61+ // byNext indexes recs by their next anchor, in source order; end holds
62+ // the comments with no anchor after them, which trail the statement.
63+ byNext map [Node ][]int
64+ end []int
65+ taken []bool
66+ nTaken int
5067}
5168
52- // NewCommentSet prepares comments for printing against source text. Only
53- // comments inside the printed statement should be included; comments above
54- // the statement and after its terminator are the caller's to keep.
55- func NewCommentSet (comments []Comment , src string ) * CommentSet {
56- cs := & CommentSet {comments : comments , lineStarts : []int {0 }}
57- for i := 0 ; i < len (src ); i ++ {
58- if src [i ] == '\n' {
59- cs .lineStarts = append (cs .lineStarts , i + 1 )
69+ // Exhausted reports whether every attached comment was printed.
70+ func (t * CommentTable ) Exhausted () bool {
71+ return t == nil || t .nTaken == len (t .recs )
72+ }
73+
74+ // take returns and consumes the comments anchored to n, in order.
75+ func (t * CommentTable ) take (n Node ) []commentRec {
76+ if t == nil || n == nil {
77+ return nil
78+ }
79+ idxs := t .byNext [n ]
80+ if len (idxs ) == 0 {
81+ return nil
82+ }
83+ out := make ([]commentRec , 0 , len (idxs ))
84+ for _ , i := range idxs {
85+ if t .taken [i ] {
86+ continue
6087 }
88+ t .taken [i ] = true
89+ t .nTaken ++
90+ out = append (out , t .recs [i ])
6191 }
62- sort .SliceStable (cs .comments , func (i , j int ) bool {
63- return cs .comments [i ].Start < cs .comments [j ].Start
64- })
65- return cs
92+ return out
6693}
6794
68- func (cs * CommentSet ) lineOf (pos int ) int {
69- return sort .SearchInts (cs .lineStarts , pos + 1 ) - 1
95+ // takeRemaining returns and consumes every comment not yet printed; the
96+ // statement is over, so everything left trails it.
97+ func (t * CommentTable ) takeRemaining () []commentRec {
98+ if t == nil {
99+ return nil
100+ }
101+ var out []commentRec
102+ for i := range t .recs {
103+ if t .taken [i ] {
104+ continue
105+ }
106+ t .taken [i ] = true
107+ t .nTaken ++
108+ out = append (out , t .recs [i ])
109+ }
110+ return out
70111}
71112
72- // advance moves the cursor forward to pos; the cursor never moves back.
73- func (cs * CommentSet ) advance (pos int ) {
74- if pos > cs .cursor {
75- cs .cursor = pos
113+ // AttachComments classifies a statement's interior comments against its
114+ // nodes, producing the table the printer (and any future rewriting tool)
115+ // works from. Placement follows gofmt's rules, decided here once from
116+ // source positions and lines: a comment on the same line as the code before
117+ // it trails that code; any other comment leads the first node printed after
118+ // it; a comment after the last node trails the statement.
119+ func AttachComments (raw * RawStmt , d format.Dialect , comments []Comment , src string ) * CommentTable {
120+ anchors := collectAnchors (raw , d )
121+ lines := []int {0 }
122+ for i := 0 ; i < len (src ); i ++ {
123+ if src [i ] == '\n' {
124+ lines = append (lines , i + 1 )
125+ }
126+ }
127+ lineOf := func (pos int ) int {
128+ return sort .SearchInts (lines , pos + 1 ) - 1
129+ }
130+
131+ table := & CommentTable {byNext : make (map [Node ][]int , len (anchors ))}
132+ for _ , c := range comments {
133+ // prev: the printed node with the greatest position before the
134+ // comment (for the trailing/leading call). next: the emission point
135+ // for the first printed node after the comment — the node itself,
136+ // or, when boundary markers immediately precede it in print order,
137+ // the earliest of those markers, which is where the printer will
138+ // look for this comment first.
139+ var prev , next Node
140+ prevPos := - 1
141+ for _ , a := range anchors {
142+ if ! a .marker && a .pos < c .Start && a .pos > prevPos {
143+ prev , prevPos = a .node , a .pos
144+ }
145+ }
146+ for i , a := range anchors {
147+ if a .marker || a .pos <= c .Start {
148+ continue
149+ }
150+ j := i
151+ for j > 0 && anchors [j - 1 ].marker {
152+ j --
153+ }
154+ next = anchors [j ].node
155+ break
156+ }
157+ rec := commentRec {
158+ c : c ,
159+ prev : prev ,
160+ next : next ,
161+ trailing : ! c .OwnLine && prev != nil && lineOf (prevPos ) == lineOf (c .Start ),
162+ }
163+ table .recs = append (table .recs , rec )
164+ i := len (table .recs ) - 1
165+ if next == nil {
166+ table .end = append (table .end , i )
167+ } else {
168+ table .byNext [next ] = append (table .byNext [next ], i )
169+ }
76170 }
171+ table .taken = make ([]bool , len (table .recs ))
172+ return table
77173}
78174
79- // pending returns the next unprinted comment, if any.
80- func ( cs * CommentSet ) pending () ( Comment , bool ) {
81- if cs . next >= len ( cs . comments ) {
82- return Comment {}, false
83- }
84- return cs . comments [ cs . next ], true
175+ type anchor struct {
176+ node Node
177+ pos int
178+ // marker anchors are emission points (beforeClause, list boundaries)
179+ // rather than printed nodes; they carry no position of their own.
180+ marker bool
85181}
86182
87- // Exhausted reports whether every comment was printed.
88- func (cs * CommentSet ) Exhausted () bool {
89- return cs .next >= len (cs .comments )
183+ // collectAnchors renders the statement once, flat, recording every
184+ // positioned node in the order the printer visits them. Print order is what
185+ // comment emission is defined against, so classifying against it keeps the
186+ // attach-time decision and the print-time emission point identical.
187+ func collectAnchors (n Node , d format.Dialect ) (out []anchor ) {
188+ defer func () {
189+ // A formatter panic here surfaces later, on the real print; anchors
190+ // collected so far still place most comments.
191+ recover ()
192+ }()
193+ tb := NewTrackedBuffer ()
194+ tb .anchors = & out
195+ if ft , ok := n .(nodeFormatter ); ok {
196+ ft .Format (tb , d )
197+ }
198+ return out
90199}
200+
0 commit comments