diff --git a/parser/src/earley/parser.rs b/parser/src/earley/parser.rs index 35838058..87032d61 100644 --- a/parser/src/earley/parser.rs +++ b/parser/src/earley/parser.rs @@ -1043,11 +1043,10 @@ impl ParserState { let mut applied_idx = state.byte_to_token_idx.len(); let tok_env = state.tok_env.clone(); let trie = tok_env.tok_trie(); - let eos = trie.eos_token(); let mut recog = ParserRecognizer { state }; for (tidx, &tok) in tokens.iter().enumerate() { let state = &mut recog.state; - if tok == eos { + if trie.eos_tokens().contains(&tok) { if applied_idx == state.bytes.len() && state.is_accepting_inner() { return tidx + 1; } else { diff --git a/parser/tests/test_raw_parser.rs b/parser/tests/test_raw_parser.rs index 41e82669..8a6efc19 100644 --- a/parser/tests/test_raw_parser.rs +++ b/parser/tests/test_raw_parser.rs @@ -398,3 +398,42 @@ fn test_multi_eos_mask_when_stopped() { ); assert!(matcher.is_stopped()); } + +#[test] +fn test_multi_eos_validate_tokens() { + let base = ApproximateTokEnv::single_byte(); + let base_trie = base.tok_trie(); + let primary_eos = base_trie.eos_token(); + let extra_eos = primary_eos - 1; + let multi_trie = base_trie.clone().with_eos_tokens(&[primary_eos, extra_eos]); + let tok_env: TokEnv = Arc::new(ApproximateTokEnv::new(multi_trie)); + + let factory = ParserFactory::new( + &tok_env, + InferenceCapabilities::default(), + &SlicedBiasComputer::general_slices(), + ) + .unwrap(); + + let grm = TopLevelGrammar::from_lark(r#"start: /[a-z]*/"#.to_string()); + let mut parser = factory.create_parser(grm).unwrap(); + parser.start_without_prompt(); + let mut matcher = Matcher::new(Ok(parser)); + + matcher.consume_token(b'a' as u32).unwrap(); + assert!( + !matcher.is_stopped(), + "parser should still be running while lexeme bytes are pending" + ); + + assert_eq!( + matcher.validate_tokens(&[primary_eos]).unwrap(), + 1, + "primary EOS must be accepted by validate_tokens" + ); + assert_eq!( + matcher.validate_tokens(&[extra_eos]).unwrap(), + 1, + "extra EOS must also be accepted by validate_tokens" + ); +}