fix(splitter): improve EOF handling (#556)

Author: juleswritescode

crates/pgt_lexer/src/lexed.rs

@@ -23,6 +23,7 @@ pub struct LexDiagnostic {
 }
 
 /// Result of lexing a string, providing access to tokens and diagnostics
+#[derive(Debug)]
 pub struct Lexed<'a> {
  pub(crate) text: &'a str,
  pub(crate) kind: Vec<SyntaxKind>,

crates/pgt_statement_splitter/src/lib.rs

@@ -19,7 +19,7 @@ pub fn split(sql: &str) -> SplitResult {
 
  let mut splitter = Splitter::new(&lexed);
 
- source(&mut splitter);
+ let _ = source(&mut splitter);
 
  let split_result = splitter.finish();
 
@@ -587,4 +587,10 @@ VALUES
  TextRange::new(0.into(), 6.into()),
  )]);
  }
+
+ #[test]
+ fn does_not_panic_on_incomplete_statements() {
+ // does not panic
+ let _ = Tester::from("select case ");
+ }
 }

crates/pgt_statement_splitter/src/splitter.rs

@@ -8,6 +8,8 @@ pub use common::source;
 use pgt_lexer::{Lexed, SyntaxKind};
 use pgt_text_size::TextRange;
 
+use crate::splitter::common::{ReachedEOFException, SplitterResult};
+
 pub struct SplitResult {
  pub ranges: Vec<TextRange>,
  pub errors: Vec<SplitError>,
@@ -29,6 +31,7 @@ pub struct SplitError {
  pub token: usize,
 }
 
+#[derive(Debug)]
 pub struct Splitter<'a> {
  lexed: &'a Lexed<'a>,
  current_pos: usize,
@@ -102,12 +105,12 @@ impl<'a> Splitter<'a> {
  self.lexed.kind(self.current_pos)
  }
 
- fn eat(&mut self, kind: SyntaxKind) -> bool {
+ fn eat(&mut self, kind: SyntaxKind) -> Result<bool, ReachedEOFException> {
  if self.current() == kind {
- self.advance();
- true
+ self.advance()?;
+ Ok(true)
  } else {
- false
+ Ok(false)
  }
  }
 
@@ -118,13 +121,17 @@ impl<'a> Splitter<'a> {
  /// Advances the parser to the next relevant token and returns it.
  ///
  /// NOTE: This will skip trivia tokens.
- fn advance(&mut self) -> SyntaxKind {
+ fn advance(&mut self) -> Result<SyntaxKind, ReachedEOFException> {
+ if self.current() == SyntaxKind::EOF {
+ return Err(ReachedEOFException);
+ }
+
  let pos = (self.current_pos + 1..self.lexed.len())
  .find(|&idx| !self.is_trivia(idx))
- .expect("lexed should have non-trivia eof token");
+ .unwrap();
 
  self.current_pos = pos;
- self.lexed.kind(pos)
+ Ok(self.lexed.kind(pos))
  }
 
  fn look_ahead(&self, ignore_trivia: bool) -> SyntaxKind {
@@ -164,9 +171,9 @@ impl<'a> Splitter<'a> {
 
  /// Will advance if the `kind` matches the current token.
  /// Otherwise, will add a diagnostic to the internal `errors`.
- fn expect(&mut self, kind: SyntaxKind) {
+ fn expect(&mut self, kind: SyntaxKind) -> SplitterResult {
  if self.current() == kind {
- self.advance();
+ self.advance()?;
  } else {
  let token = if self.current() == SyntaxKind::EOF {
  self.current_pos - 1
@@ -178,6 +185,8 @@ impl<'a> Splitter<'a> {
  msg: format!("Expected {:#?}", kind),
  token,
  });
- }
+ };
+
+ Ok(())
  }
 }