fix(pg): preserve completion state during paren lookahead (#113)

* fix(pg): preserve completion state during paren lookahead

* fix(pg): add completion-safe backtrack snapshot

Author: rebelice

pg/parser/backtrack.go

@@ -8,10 +8,9 @@ package parser
 // SCOPE: this is a TOKEN-STREAM snapshot, not a complete parser/lexer
 // snapshot. It does NOT cover mid-token-content lexer state (literalbuf,
 // dolqstart, utf16FirstPart, xcdepth, stateBeforeStrStop, warning flags)
-// or completion-mode state (candidates, collecting). Those fields are
-// either reset at token boundaries (lexer internals) or not used during
-// speculative parses (completion mode), so they don't need to be saved
-// here for token-stream rollback to be sound.
+// or completion-mode state (candidates, collecting). Lexer internals are
+// reset at token boundaries. Completion-mode speculative callers that can
+// scan past the cursor should use snapshotTokenStreamAndCompletion.
 //
 // If a future caller needs to roll back from INSIDE a token (e.g., from
 // inside a string literal or dollar-quoted block), this struct is
@@ -42,6 +41,13 @@ type tokenStreamState struct {
 	lexerState         LexerState
 }
 
+type tokenStreamAndCompletionState struct {
+	tokenStream  tokenStreamState
+	collecting   bool
+	collectDepth int
+	candidates   *CandidateSet
+}
+
 // snapshotTokenStream captures the current token-stream position for
 // later restoration via restoreTokenStream. See tokenStreamState for
 // scope and limitations.
@@ -58,16 +64,31 @@ func (p *Parser) snapshotTokenStream() tokenStreamState {
 	}
 }
 
+// snapshotTokenStreamAndCompletion captures token-stream state plus the
+// completion state that advance() can mutate when a speculative walk crosses
+// the cursor. Use this for completion-mode lookahead that can scan arbitrary
+// user input before rolling back.
+func (p *Parser) snapshotTokenStreamAndCompletion() tokenStreamAndCompletionState {
+	s := tokenStreamAndCompletionState{
+		tokenStream:  p.snapshotTokenStream(),
+		collecting:   p.collecting,
+		collectDepth: p.collectDepth,
+	}
+	if p.candidates != nil {
+		s.candidates = p.candidates.snapshot()
+	}
+	return s
+}
+
 // restoreTokenStream rewinds parser + lexer state to a previously
 // captured snapshot. After restore, the next advance() will emit the
 // same token as it would have at the moment snapshotTokenStream() was
 // called.
 //
 // Caller responsibility: do not interleave restore with completion-mode
-// queries or with any operation that mutates lexer state outside the
-// token stream (string literal scanning, etc). The current speculative
-// parse sites in parseFuncArg and parseFuncType only consume keyword
-// tokens and punctuation, so they are safe.
+// queries or with any operation that mutates lexer state outside the token
+// stream. Use restoreTokenStreamAndCompletion for lookahead that may cross
+// the completion cursor.
 func (p *Parser) restoreTokenStream(s tokenStreamState) {
 	p.cur = s.cur
 	p.prev = s.prev
@@ -78,3 +99,14 @@ func (p *Parser) restoreTokenStream(s tokenStreamState) {
 	p.lexer.start = s.lexerStart
 	p.lexer.state = s.lexerState
 }
+
+// restoreTokenStreamAndCompletion rewinds token-stream and completion state
+// captured by snapshotTokenStreamAndCompletion.
+func (p *Parser) restoreTokenStreamAndCompletion(s tokenStreamAndCompletionState) {
+	p.restoreTokenStream(s.tokenStream)
+	p.collecting = s.collecting
+	p.collectDepth = s.collectDepth
+	if p.candidates != nil && s.candidates != nil {
+		p.candidates.restore(s.candidates)
+	}
+}

pg/parser/backtrack_test.go

@@ -75,6 +75,52 @@ func TestSnapshotRestoreIdentity(t *testing.T) {
 	}
 }
 
+func TestSnapshotRestoreCompletionState(t *testing.T) {
+	sql := "SELECT * FROM (SELECT  FROM t1) a"
+	cursor := len("SELECT * FROM (SELECT ")
+	cs := newCandidateSet()
+	p := &Parser{
+		lexer:      NewLexer(sql),
+		source:     sql,
+		completing: true,
+		cursorOff:  cursor,
+		candidates: cs,
+	}
+	p.advance()
+
+	snap := p.snapshotTokenStreamAndCompletion()
+	for p.cur.Type != lex_EOF && !p.collecting {
+		p.advance()
+	}
+	if !p.collecting {
+		t.Fatal("expected speculative walk to cross cursor")
+	}
+	p.addRuleCandidate("leaked_rule")
+	p.addTokenCandidate(SELECT)
+	p.addCTEPosition(123)
+	p.addSelectAliasPosition(456)
+
+	p.restoreTokenStreamAndCompletion(snap)
+	if p.collecting {
+		t.Fatal("expected collecting to be restored to false")
+	}
+	if cs.HasRule("leaked_rule") {
+		t.Fatal("expected rule candidates to be restored")
+	}
+	if cs.HasToken(SELECT) {
+		t.Fatal("expected token candidates to be restored")
+	}
+	if len(cs.CTEPositions) != 0 {
+		t.Fatalf("expected CTE positions to be restored, got %v", cs.CTEPositions)
+	}
+	if len(cs.SelectAliasPositions) != 0 {
+		t.Fatalf("expected select alias positions to be restored, got %v", cs.SelectAliasPositions)
+	}
+	if p.cur.Type != SELECT {
+		t.Fatalf("expected token stream to be restored to SELECT, got %d", p.cur.Type)
+	}
+}
+
 // walkTokens parses sql and returns the full token sequence (excluding EOF).
 func walkTokens(t *testing.T, sql string) []Token {
 	t.Helper()

pg/parser/complete.go

@@ -75,10 +75,10 @@ type RuleCandidate struct {
 // CandidateSet holds the token and rule candidates collected during a
 // completion-mode parse.
 type CandidateSet struct {
-	Tokens []int            // token type candidates
-	Rules  []RuleCandidate  // grammar rule candidates
-	seen   map[int]bool     // dedup tokens
-	seenR  map[string]bool  // dedup rules
+	Tokens []int           // token type candidates
+	Rules  []RuleCandidate // grammar rule candidates
+	seen   map[int]bool    // dedup tokens
+	seenR  map[string]bool // dedup rules
 
 	// CTEPositions holds the byte offsets of WITH clause starts encountered
 	// before the cursor. Bytebase uses these to re-parse CTE definitions
@@ -246,13 +246,17 @@ func (p *Parser) addKeywordsByCategory(categories ...KeywordCategory) {
 // snapshot returns a copy of the current candidate set state.
 func (cs *CandidateSet) snapshot() *CandidateSet {
 	s := &CandidateSet{
-		Tokens: make([]int, len(cs.Tokens)),
-		Rules:  make([]RuleCandidate, len(cs.Rules)),
-		seen:   make(map[int]bool, len(cs.seen)),
-		seenR:  make(map[string]bool, len(cs.seenR)),
+		Tokens:               make([]int, len(cs.Tokens)),
+		Rules:                make([]RuleCandidate, len(cs.Rules)),
+		seen:                 make(map[int]bool, len(cs.seen)),
+		seenR:                make(map[string]bool, len(cs.seenR)),
+		CTEPositions:         make([]int, len(cs.CTEPositions)),
+		SelectAliasPositions: make([]int, len(cs.SelectAliasPositions)),
 	}
 	copy(s.Tokens, cs.Tokens)
 	copy(s.Rules, cs.Rules)
+	copy(s.CTEPositions, cs.CTEPositions)
+	copy(s.SelectAliasPositions, cs.SelectAliasPositions)
 	for k, v := range cs.seen {
 		s.seen[k] = v
 	}
@@ -262,6 +266,23 @@ func (cs *CandidateSet) snapshot() *CandidateSet {
 	return s
 }
 
+// restore replaces cs with snapshot's contents while preserving cs's identity.
+func (cs *CandidateSet) restore(snapshot *CandidateSet) {
+	cs.Tokens = append(cs.Tokens[:0], snapshot.Tokens...)
+	cs.Rules = append(cs.Rules[:0], snapshot.Rules...)
+	cs.CTEPositions = append(cs.CTEPositions[:0], snapshot.CTEPositions...)
+	cs.SelectAliasPositions = append(cs.SelectAliasPositions[:0], snapshot.SelectAliasPositions...)
+
+	cs.seen = make(map[int]bool, len(snapshot.seen))
+	for k, v := range snapshot.seen {
+		cs.seen[k] = v
+	}
+	cs.seenR = make(map[string]bool, len(snapshot.seenR))
+	for k, v := range snapshot.seenR {
+		cs.seenR[k] = v
+	}
+}
+
 // diff returns candidates in cs that are not in before.
 func (cs *CandidateSet) diff(before *CandidateSet) *CandidateSet {
 	d := newCandidateSet()

pg/parser/complete_test.go

@@ -72,6 +72,21 @@ func TestCollectAfterSelect(t *testing.T) {
 	}
 }
 
+func TestCollectNestedFromSubqueryAfterSelect(t *testing.T) {
+	prefix := "SELECT * FROM (SELECT * FROM (SELECT "
+	sql := prefix + " FROM t1) a) b"
+	candidates := Collect(sql, len(prefix))
+	if candidates == nil {
+		t.Fatal("expected non-nil candidates")
+	}
+	if !candidates.HasRule("columnref") {
+		t.Error("expected columnref rule candidate in nested SELECT target list")
+	}
+	if !candidates.HasRule("func_name") {
+		t.Error("expected func_name rule candidate in nested SELECT target list")
+	}
+}
+
 func TestCollectAfterFrom(t *testing.T) {
 	candidates := Collect("SELECT 1 FROM ", 14)
 	if candidates == nil {

pg/parser/select.go

@@ -1268,8 +1268,8 @@ func (p *Parser) parenBeginsSubquery() bool {
 	if p.cur.Type != '(' {
 		return false
 	}
-	snap := p.snapshotTokenStream()
-	defer p.restoreTokenStream(snap)
+	snap := p.snapshotTokenStreamAndCompletion()
+	defer p.restoreTokenStreamAndCompletion(snap)
 	return p.consumeMatchedParenIsSubquery()
 }