parser: ParseFile keeps the trivia the grammar never sees

lexer/lexer.go

@@ -23,12 +23,32 @@ import (
 // A NUL byte terminates the input, matching SQLite, whose tokenizer walks a
 // NUL-terminated buffer.
 func Lex(src string) []token.Token {
+	toks, _ := lex(src, false)
+	return toks
+}
+
+// LexFile splits src into tokens like Lex, and also returns the trivia Lex
+// drops: the runs of whitespace and comments, in source order, as SPACE and
+// COMMENT tokens. "Trivia" is Roslyn's name — the C# compiler calls the
+// channel of source text the grammar never sees "syntax trivia", and
+// swift-syntax and rust-analyzer use the same term — adopted here because
+// meyer's consumers (formatters) speak it too.
+//
+// Tokens and trivia together tile the consumed input exactly: every byte
+// before the EOF token's position belongs to exactly one token or trivia
+// span, in order, with no overlap. Both slices come from the single pass
+// Lex already makes; nothing is scanned twice.
+func LexFile(src string) (tokens, trivia []token.Token) {
+	return lex(src, true)
+}
+
+func lex(src string, keepTrivia bool) (toks, trivia []token.Token) {
 	// Across the corpus SQL runs to 3.7 bytes per token including the
 	// separators, but the mean is the wrong statistic to size from: what
 	// costs is the tail that has to grow and copy. A divisor of three
 	// covers 93% of cases in one allocation where four covers 80%, for
 	// about seventeen tokens of slack apiece.
-	toks := make([]token.Token, 0, len(src)/3+8)
+	toks = make([]token.Token, 0, len(src)/3+8)
 	i := 0
 	ambiguous := false
 	for i < len(src) {
@@ -45,6 +65,8 @@ func Lex(src string) []token.Token {
 				ambiguous = true
 			}
 			toks = append(toks, token.Token{Kind: kind, Pos: i, End: i + n})
+		} else if keepTrivia {
+			trivia = append(trivia, token.Token{Kind: kind, Pos: i, End: i + n})
 		}
 		i += n
 	}
@@ -55,7 +77,7 @@ func Lex(src string) []token.Token {
 		resolveWindowKeywords(toks)
 	}
 	toks = append(toks, token.Token{Kind: token.EOF, Pos: i, End: i})
-	return toks
+	return toks, trivia
 }
 
 // resolveWindowKeywords implements analyzeWindowKeyword, analyzeOverKeyword

lexer/lexer_test.go

@@ -351,3 +351,52 @@ func TestLexNeverLoops(t *testing.T) {
 		}
 	}
 }
+
+// TestLexFile pins the trivia channel: the tokens match Lex exactly, the
+// trivia holds only SPACE and COMMENT runs, and both together tile the
+// consumed input with no gap and no overlap.
+func TestLexFile(t *testing.T) {
+	for _, src := range []string{
+		"SELECT 1",
+		"SELECT 1 -- one\n, 2 /* two */;",
+		"-- leading\nSELECT /* mid */ 'a''b' [c d];",
+		"/* unterminated",
+		"SELECT 1;\n\n-- trailing",
+		"",
+	} {
+		toks, trivia := LexFile(src)
+		if !slices.Equal(toks, Lex(src)) {
+			t.Errorf("LexFile(%q) tokens differ from Lex", src)
+		}
+		for _, tr := range trivia {
+			if tr.Kind != token.SPACE && tr.Kind != token.COMMENT {
+				t.Errorf("LexFile(%q): trivia kind %v", src, tr.Kind)
+			}
+		}
+		// Merge the two ordered span lists and require a perfect tiling of
+		// [0, EOF.Pos).
+		at := 0
+		ti, vi := 0, 0
+		body := toks[:len(toks)-1] // drop the zero-width EOF
+		for ti < len(body) || vi < len(trivia) {
+			var next token.Token
+			switch {
+			case ti == len(body):
+				next, vi = trivia[vi], vi+1
+			case vi == len(trivia):
+				next, ti = body[ti], ti+1
+			case body[ti].Pos < trivia[vi].Pos:
+				next, ti = body[ti], ti+1
+			default:
+				next, vi = trivia[vi], vi+1
+			}
+			if next.Pos != at {
+				t.Fatalf("LexFile(%q): tiling broken at %d, next span starts at %d", src, at, next.Pos)
+			}
+			at = next.End
+		}
+		if eof := toks[len(toks)-1]; at != eof.Pos {
+			t.Fatalf("LexFile(%q): tiling ends at %d, EOF at %d", src, at, eof.Pos)
+		}
+	}
+}

parser/parser.go

@@ -144,6 +144,42 @@ func (o Options) ParseString(src string) (stmts []ast.Stmt, err error) {
 	return newParser(src, o).parseScript(), nil
 }
 
+// File is a parsed SQL script together with the source trivia the grammar
+// never sees, so a caller that needs both — a formatter putting comments
+// back where they came from — gets them from the one lexer pass parsing
+// already makes.
+type File struct {
+	Stmts []ast.Stmt
+
+	// Trivia holds the runs of whitespace and comments, in source order,
+	// as SPACE and COMMENT tokens; Token.Text recovers each run verbatim.
+	// "Trivia" is Roslyn's name — the C# compiler calls the channel of
+	// source text that does not affect syntax "syntax trivia", and
+	// swift-syntax and rust-analyzer use the same term. Together with the
+	// parsed tokens, trivia tiles the consumed input exactly.
+	Trivia []token.Token
+}
+
+// ParseFile parses a complete SQL script and keeps its trivia.
+func ParseFile(src string) (*File, error) { return Options{}.ParseFile(src) }
+
+// ParseFile is ParseFile with these options.
+func (o Options) ParseFile(src string) (f *File, err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			b, ok := r.(bail)
+			if !ok {
+				panic(r)
+			}
+			f, err = nil, b.err
+		}
+	}()
+	toks, trivia := lexer.LexFile(src)
+	p := &parser{src: src, toks: toks, opts: o}
+	p.checkIllegal()
+	return &File{Stmts: p.parseScript(), Trivia: trivia}, nil
+}
+
 // ParseStatement parses exactly one statement and rejects trailing input.
 func ParseStatement(src string) (ast.Stmt, error) { return Options{}.ParseStatement(src) }
 

parser/trivia_test.go

@@ -0,0 +1,69 @@
+package parser_test
+
+import (
+	"reflect"
+	"testing"
+
+	"github.com/sqlc-dev/meyer/internal/testfile"
+	"github.com/sqlc-dev/meyer/lexer"
+	"github.com/sqlc-dev/meyer/parser"
+	"github.com/sqlc-dev/meyer/token"
+)
+
+// TestCorpusTrivia checks the LexFile decomposition over the whole corpus:
+// tokens and trivia tile the consumed input exactly — in order, with no
+// gap, no overlap, and no byte belonging to neither — and the trivia
+// channel holds nothing but whitespace and comment runs.
+func TestCorpusTrivia(t *testing.T) {
+	forEachCorpusCase(t, func(t *testing.T, c testfile.Case) {
+		toks, trivia := lexer.LexFile(c.SQL)
+		for _, tr := range trivia {
+			if tr.Kind != token.SPACE && tr.Kind != token.COMMENT {
+				t.Fatalf("%s: trivia kind %v", c.Name, tr.Kind)
+			}
+		}
+		at := 0
+		ti, vi := 0, 0
+		body := toks[:len(toks)-1] // drop the zero-width EOF
+		for ti < len(body) || vi < len(trivia) {
+			var next token.Token
+			switch {
+			case ti == len(body):
+				next, vi = trivia[vi], vi+1
+			case vi == len(trivia):
+				next, ti = body[ti], ti+1
+			case body[ti].Pos < trivia[vi].Pos:
+				next, ti = body[ti], ti+1
+			default:
+				next, vi = trivia[vi], vi+1
+			}
+			if next.Pos != at {
+				t.Fatalf("%s: tiling broken at %d, next span starts at %d", c.Name, at, next.Pos)
+			}
+			at = next.End
+		}
+		if eof := toks[len(toks)-1]; at != eof.Pos {
+			t.Fatalf("%s: tiling ends at %d, EOF at %d", c.Name, at, eof.Pos)
+		}
+	})
+}
+
+// TestCorpusParseFile checks that ParseFile agrees with ParseString on
+// every case: the same statements and the same error, since both run the
+// same grammar over the same token stream.
+func TestCorpusParseFile(t *testing.T) {
+	forEachCorpusCase(t, func(t *testing.T, c testfile.Case) {
+		stmts, err := parser.ParseString(c.SQL)
+		f, ferr := parser.ParseFile(c.SQL)
+		switch {
+		case (err == nil) != (ferr == nil):
+			t.Fatalf("%s: ParseString err=%v, ParseFile err=%v", c.Name, err, ferr)
+		case err != nil:
+			if err.Error() != ferr.Error() {
+				t.Fatalf("%s: error mismatch:\n%v\n%v", c.Name, err, ferr)
+			}
+		case !reflect.DeepEqual(stmts, f.Stmts):
+			t.Fatalf("%s: ParseFile statements differ from ParseString", c.Name)
+		}
+	})
+}