ParseFile returns the tree with the input's comments

CLAUDE.md

@@ -27,9 +27,10 @@ the main module; the only module allowed to link pg_query_go (cgo) is
   written by `cmd/generate`.
 - The public API mirrors pg_query_go exactly; consumers must migrate by
   changing one import path. Never rename, add, or remove exported symbols
-  without checking pg_query_go v6.2.2. There is one deliberate addition,
-  `parser.ParseToTree` (PLAN.md § 1): the tree without the protobuf round
-  trip the cgo boundary forces upstream.
+  without checking pg_query_go v6.2.2. There are two deliberate additions,
+  both in the `parser` subpackage (PLAN.md § 1): `ParseToTree`, the tree
+  without the protobuf round trip the cgo boundary forces upstream, and
+  `ParseFile`, that tree plus the comments `base_yylex` drops.
 - Location fields must byte-match the reference; when in doubt, check which
   `@N` the `gram.y` action uses.
 

PLAN.md

@@ -137,9 +137,15 @@ Being conveniences, they are not what the root package builds on. Upstream's
 protobuf is how it crosses the cgo boundary; encoding a tree the Go parser
 already holds only to decode it back costs about three quarters of `Parse`'s
 running time and doubles its allocations, so `Parse` calls
-**`parser.ParseToTree`** — oliphant's one deliberate addition to pg_query_go's
-surface, and the exception to § *The API does not change*. Two consequences,
-both pinned by `parser/tree_test.go`:
+**`parser.ParseToTree`**.
+
+That is the first of oliphant's two deliberate additions to pg_query_go's
+surface — the exception to § *The API does not change*. Both live in the
+`parser` subpackage, so the root package stays a mirror function-for-function,
+and both exist for the same reason: the Go implementation can hand back what
+the cgo boundary forces upstream to throw away and rebuild.
+
+`ParseToTree` has two consequences, both pinned by `parser/tree_test.go`:
 
 - proto3 string fields must be valid UTF-8, so a tree carrying a raw invalid
   byte fails to marshal and upstream reports that as an error from `Parse`.
@@ -153,6 +159,25 @@ both pinned by `parser/tree_test.go`:
   pointers. Callers that only read the tree — all pg_query_go's API admits —
   cannot tell.
 
+**`parser.ParseFile`** is the second: the tree together with the input's
+comments, which the scanner already produces as tokens (libpg_query patch 04)
+and `base_yylex` drops on the way to the grammar. `Scan` exposes them, so
+`Scan` + `Parse` answers the same question — at the cost of lexing twice and
+marshalling the whole token stream to protobuf to recover a handful of
+comments, which measures about twice the time and 1.5× the allocations of
+`ParseFile` on a query file. The consumer is a formatter: comments cannot be
+reprinted where they were written without them. Two notes, pinned by
+`parser/file_test.go`:
+
+- The comments are exactly what `Scan` reports, checked case by case across
+  the corpus, so the oracle-derived scan goldens stand behind them too.
+- `ParseFile` accepts and rejects exactly what `Parse` does, the proto3
+  UTF-8 rejection included. That answer is not a property of the input: an
+  invalid byte only fails when it lands in a *string field*, and a comment is
+  not one — so `SELECT 1 -- \xff` parses for both, which is the case that
+  matters here. `ParseFile` therefore decides on the tree, marshalling it
+  only when `utf8.ValidString(input)` is false.
+
 ### 2. The AST is the generated protobuf types — one dependency, not zero
 
 The siblings' "zero dependencies, ever" rule bends once: pg_query_go's API

internal/lexer/lexer.go

@@ -47,6 +47,14 @@ func (e *Error) Error() string { return e.Message }
 type Scanner struct {
 	input string
 	pos   int
+
+	// comments, when keepComments is set, accumulates every SQL_COMMENT and
+	// C_COMMENT token the scanner produces. The grammar never sees them —
+	// Filter drops them, as base_yylex does — so a consumer that wants them
+	// back (a formatter putting comments where they were written) would
+	// otherwise have to lex the input a second time.
+	keepComments bool
+	comments     []Token
 }
 
 // New returns a Scanner over input. Like the C scanner (which receives a
@@ -61,9 +69,22 @@ func New(input string) *Scanner {
 	return &Scanner{input: input}
 }
 
+// NewKeepingComments returns a Scanner that also records the comment tokens
+// it produces, for Comments to hand back after the scan.
+func NewKeepingComments(input string) *Scanner {
+	s := New(input)
+	s.keepComments = true
+	return s
+}
+
 // Input returns the input string as scanned (truncated at any NUL byte).
 func (s *Scanner) Input() string { return s.input }
 
+// Comments returns the SQL_COMMENT and C_COMMENT tokens scanned so far, in
+// source order, for a Scanner from NewKeepingComments (nil otherwise). They
+// do not overlap, and every one lies within Input.
+func (s *Scanner) Comments() []Token { return s.comments }
+
 // Character classes from scan.l's named patterns.
 
 // scan.l: space
@@ -141,7 +162,16 @@ func (s *Scanner) peek(off int) byte {
 }
 
 func (s *Scanner) token(kind ast.Token, start int) (Token, *Error) {
-	return Token{Kind: kind, Start: int32(start), End: int32(s.pos)}, nil
+	tok := Token{Kind: kind, Start: int32(start), End: int32(s.pos)}
+	if s.keepComments && (kind == ast.Token_SQL_COMMENT || kind == ast.Token_C_COMMENT) {
+		// Recorded where the token is minted, not where Filter drops it.
+		// A successful parse does route every comment through that drop
+		// site, but not every path there goes through Filter.Next — the
+		// UIDENT/UESCAPE resolution pulls tokens off to the side — and this
+		// site cannot be bypassed at all.
+		s.comments = append(s.comments, tok)
+	}
+	return tok, nil
 }
 
 // Next returns the next token, exactly as core_yylex would. At end of input

internal/parse/parser.go

@@ -58,8 +58,26 @@ func Parse(input string) (res *ast.ParseResult, err *lexer.Error) {
 // ParseTracked is Parse plus the present-but-empty string set (see
 // parser.emptyStrings) for the fingerprint walk.
 func ParseTracked(input string) (res *ast.ParseResult, emptyStrings map[any]bool, err *lexer.Error) {
-	s := lexer.New(input)
-	p := &parser{src: s.Input(), filter: lexer.NewFilter(s), toks: make([]lexer.Token, 0, tokenCap(input))}
+	return parseFrom(lexer.New(input))
+}
+
+// ParseWithComments is Parse plus the comment tokens the scanner produced
+// and the grammar never saw, from the one pass parsing already makes.
+// Comments are returned in source order; a failed parse returns none, since
+// the scan stops where the grammar does.
+func ParseWithComments(input string) (res *ast.ParseResult, comments []lexer.Token, err *lexer.Error) {
+	s := lexer.NewKeepingComments(input)
+	res, _, err = parseFrom(s)
+	if err != nil {
+		return nil, nil, err
+	}
+	return res, s.Comments(), nil
+}
+
+// parseFrom runs parse_toplevel over s, turning the bail a production panics
+// with back into an error.
+func parseFrom(s *lexer.Scanner) (res *ast.ParseResult, emptyStrings map[any]bool, err *lexer.Error) {
+	p := &parser{src: s.Input(), filter: lexer.NewFilter(s), toks: make([]lexer.Token, 0, tokenCap(s.Input()))}
 	defer func() {
 		if r := recover(); r != nil {
 			if b, ok := r.(bail); ok {

parser/file_test.go

@@ -0,0 +1,234 @@
+package parser_test
+
+import (
+	"errors"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"google.golang.org/protobuf/proto"
+
+	pg_query "github.com/sqlc-dev/oliphant"
+	"github.com/sqlc-dev/oliphant/ast"
+	"github.com/sqlc-dev/oliphant/internal/testfile"
+	"github.com/sqlc-dev/oliphant/parser"
+)
+
+// scanComments is the comment tokens Scan reports, which is what ParseFile
+// must reproduce: Scan runs the core scanner over the whole input with no
+// grammar attached, and its token stream is pinned against the oracle by the
+// scan corpus.
+func scanComments(t *testing.T, input string) []*ast.ScanToken {
+	t.Helper()
+	res, err := pg_query.Scan(input)
+	if err != nil {
+		return nil
+	}
+	var out []*ast.ScanToken
+	for _, tok := range res.Tokens {
+		if tok.Token == ast.Token_SQL_COMMENT || tok.Token == ast.Token_C_COMMENT {
+			out = append(out, tok)
+		}
+	}
+	return out
+}
+
+// TestParseFileMatchesScan is the contract: on every corpus input that
+// parses, the comments ParseFile keeps from the parse's own scan are exactly
+// the ones a separate Scan reports, and the tree is the one ParseToTree
+// returns. Nothing is dropped by the grammar's filter, nothing is invented.
+func TestParseFileMatchesScan(t *testing.T) {
+	var parsed, withComments int
+	for _, path := range treeCorpusFiles(t) {
+		cases, err := testfile.Read(path)
+		if err != nil {
+			t.Fatal(err)
+		}
+		rel, _ := filepath.Rel("testdata", path)
+		for _, c := range cases {
+			file, fileErr := parser.ParseFile(c.Input)
+			tree, treeErr := parser.ParseToTree(c.Input)
+			if (fileErr == nil) != (treeErr == nil) {
+				t.Errorf("%s case %s: ParseFile err=%v, ParseToTree err=%v\ninput:\n%s",
+					rel, c.Name, fileErr, treeErr, c.Input)
+				continue
+			}
+			if fileErr != nil {
+				if fileErr.Error() != treeErr.Error() {
+					t.Errorf("%s case %s: ParseFile err %q, ParseToTree err %q",
+						rel, c.Name, fileErr, treeErr)
+				}
+				continue
+			}
+			parsed++
+			if !proto.Equal(file.ParseResult, tree) {
+				t.Errorf("%s case %s: ParseFile tree differs from ParseToTree\ninput:\n%s",
+					rel, c.Name, c.Input)
+			}
+			want := scanComments(t, c.Input)
+			if len(want) > 0 {
+				withComments++
+			}
+			if len(file.Comments) != len(want) {
+				t.Errorf("%s case %s: got %d comments, Scan reports %d\ninput:\n%s",
+					rel, c.Name, len(file.Comments), len(want), c.Input)
+				continue
+			}
+			for i := range want {
+				if !proto.Equal(file.Comments[i], want[i]) {
+					t.Errorf("%s case %s: comment %d = %v, Scan reports %v\ninput:\n%s",
+						rel, c.Name, i, file.Comments[i], want[i], c.Input)
+				}
+			}
+			// Comments are ordered, disjoint and inside the input.
+			prev := int32(-1)
+			for i, c2 := range file.Comments {
+				if c2.Start < prev || c2.End < c2.Start || int(c2.End) > len(c.Input) {
+					t.Errorf("%s case %s: comment %d spans [%d,%d), previous ended %d, input is %d bytes",
+						rel, c.Name, i, c2.Start, c2.End, prev, len(c.Input))
+				}
+				prev = c2.End
+			}
+		}
+	}
+	t.Logf("compared %d parsed cases, %d of them carrying comments", parsed, withComments)
+}
+
+// TestParseFilePlacement covers the positions a query file puts comments in:
+// above the first statement, trailing a statement on its terminator's line,
+// inside a statement, between two statements, and after the last one. The
+// last is the one a parse could plausibly miss, since nothing follows it but
+// end of input.
+func TestParseFilePlacement(t *testing.T) {
+	const input = `-- name: GetAuthor :one
+SELECT id, -- the primary key
+       name
+FROM authors
+WHERE id = $1; -- the lookup
+
+/* a block comment
+   over two lines */
+SELECT 2;
+-- trailing`
+
+	file, err := parser.ParseFile(input)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(file.Stmts) != 2 {
+		t.Fatalf("got %d statements, want 2", len(file.Stmts))
+	}
+	var got []string
+	for _, c := range file.Comments {
+		got = append(got, input[c.Start:c.End])
+	}
+	want := []string{
+		"-- name: GetAuthor :one",
+		"-- the primary key",
+		"-- the lookup",
+		"/* a block comment\n   over two lines */",
+		"-- trailing",
+	}
+	if len(got) != len(want) {
+		t.Fatalf("got %d comments %q, want %d %q", len(got), got, len(want), want)
+	}
+	for i := range want {
+		if got[i] != want[i] {
+			t.Errorf("comment %d = %q, want %q", i, got[i], want[i])
+		}
+	}
+}
+
+// TestParseFileNested checks a nested block comment arrives as the single
+// token PostgreSQL scans it as, rather than ending at the inner "*/".
+func TestParseFileNested(t *testing.T) {
+	const input = "SELECT /* outer /* inner */ still outer */ 1"
+	file, err := parser.ParseFile(input)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(file.Comments) != 1 {
+		t.Fatalf("got %d comments, want 1", len(file.Comments))
+	}
+	if got, want := input[file.Comments[0].Start:file.Comments[0].End], "/* outer /* inner */ still outer */"; got != want {
+		t.Errorf("comment = %q, want %q", got, want)
+	}
+}
+
+// TestParseFileError pins that a failed parse yields no half-scanned comment
+// list: the scan stops where the grammar does, so what it collected covers
+// only part of the input.
+func TestParseFileError(t *testing.T) {
+	file, err := parser.ParseFile("-- leading\nSELECT FROM WHERE; -- trailing")
+	if err == nil {
+		t.Fatalf("ParseFile = %v, want a syntax error", file)
+	}
+	if file != nil {
+		t.Errorf("ParseFile returned %v alongside %v", file, err)
+	}
+	var perr *parser.Error
+	if !errors.As(err, &perr) {
+		t.Fatalf("error is %T, want *parser.Error", err)
+	}
+	if perr.Cursorpos == 0 {
+		t.Errorf("error carries no cursor position: %+v", perr)
+	}
+}
+
+// TestParseFileInvalidUTF8 pins that ParseFile answers exactly as Parse
+// does. The answer is not a property of the input: an invalid byte only
+// fails when it lands in a string field, and a comment is not one — which
+// is the case that matters, since comments are what ParseFile is for.
+func TestParseFileInvalidUTF8(t *testing.T) {
+	for _, tt := range []struct {
+		input    string
+		wantErr  bool
+		comments int
+	}{
+		{"SELECT '\xff'", true, 0},               // in a string field
+		{"SELECT * FROM \"tab\xffle\"", true, 0}, // in an identifier
+		{"SELECT 1 -- \xff", false, 1},           // only in a line comment
+		{"SELECT 1 /* \xff */", false, 1},        // only in a block comment
+	} {
+		file, err := parser.ParseFile(tt.input)
+		_, parseErr := pg_query.Parse(tt.input)
+
+		if (err == nil) != (parseErr == nil) {
+			t.Errorf("ParseFile(%q) err=%v, Parse err=%v — the two must agree",
+				tt.input, err, parseErr)
+			continue
+		}
+		if err != nil && parseErr != nil && err.Error() != parseErr.Error() {
+			t.Errorf("ParseFile(%q) err %q, Parse err %q", tt.input, err, parseErr)
+		}
+		if tt.wantErr {
+			if err == nil {
+				t.Errorf("ParseFile(%q) = %v, want the protobuf UTF-8 error", tt.input, file)
+			} else if !strings.Contains(err.Error(), "invalid UTF-8") {
+				t.Errorf("ParseFile(%q) err = %q, want the protobuf UTF-8 error", tt.input, err)
+			}
+			if file != nil {
+				t.Errorf("ParseFile(%q) returned a result alongside %v", tt.input, err)
+			}
+			continue
+		}
+		if err != nil {
+			t.Errorf("ParseFile(%q) = %v, want no error", tt.input, err)
+			continue
+		}
+		if len(file.Comments) != tt.comments {
+			t.Errorf("ParseFile(%q) kept %d comments, want %d", tt.input, len(file.Comments), tt.comments)
+		}
+	}
+}
+
+// TestParseFileNoComments checks the common case allocates nothing.
+func TestParseFileNoComments(t *testing.T) {
+	file, err := parser.ParseFile("SELECT 1")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if file.Comments != nil {
+		t.Errorf("Comments = %v, want nil", file.Comments)
+	}
+}