From 4a175687725cf1cdd78d1214bdd7358525504db6 Mon Sep 17 00:00:00 2001 From: vladkvit Date: Fri, 10 Jul 2026 22:33:08 -0400 Subject: [PATCH 1/8] replace pgn-reader with custom tokenizer --- Cargo.lock | 13 +-- Cargo.toml | 2 +- src/lib.rs | 1 + src/tokenizer.rs | 207 +++++++++++++++++++++++++++++++++++++++++++++++ src/visitor.rs | 100 ++++++++--------------- 5 files changed, 244 insertions(+), 79 deletions(-) create mode 100644 src/tokenizer.rs diff --git a/Cargo.lock b/Cargo.lock index f253d46..320a5b3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1064,17 +1064,6 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" -[[package]] -name = "pgn-reader" -version = "0.29.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b822a2cfc1c90c109ea35458ad9bff6669198fdd4ea8de1132e702b4d9f64b88" -dependencies = [ - "btoi", - "memchr", - "shakmaty", -] - [[package]] name = "phf" version = "0.12.1" @@ -1322,11 +1311,11 @@ dependencies = [ "arrow", "arrow-array", "criterion", + "memchr", "nom", "num_cpus", "numpy", "parquet", - "pgn-reader", "pyo3", "pyo3-arrow", "rayon", diff --git a/Cargo.toml b/Cargo.toml index 92b6b17..d38f22e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,7 @@ crate-type = ["lib", "cdylib"] [dependencies] pyo3 = { version = "0.29", features = ["multiple-pymethods"] } shakmaty = "0.30" -pgn-reader = "0.29.0" +memchr = "2.7" nom = "8.0" rayon = "1.11" num_cpus = "1.17" diff --git a/src/lib.rs b/src/lib.rs index 9f257e5..d229898 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -8,6 +8,7 @@ use rayon::prelude::*; mod board_serialization; mod comment_parsing; mod python_bindings; +mod tokenizer; mod visitor; use python_bindings::{ChunkData, ParsedGames, ParsedGamesIter, PyChunkView, PyGameView}; diff --git a/src/tokenizer.rs b/src/tokenizer.rs new file mode 100644 index 0000000..94bcd3e --- /dev/null +++ b/src/tokenizer.rs @@ -0,0 +1,207 @@ +use memchr::memchr; +use shakmaty::san::{ParseSanError, SanPlus}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Outcome { + WhiteWins, + BlackWins, + Draw, + Unknown, +} + +pub trait Visitor { + fn begin_game(&mut self); + fn begin_headers(&mut self); + fn header(&mut self, key: &[u8], value: &[u8]); + fn end_headers(&mut self); + fn san(&mut self, san_plus: SanPlus); + fn san_error(&mut self, error: ParseSanError, token: &[u8]); + fn comment(&mut self, comment: &[u8]); + fn outcome(&mut self, outcome: Outcome); + fn end_game(&mut self); +} + +pub fn parse_game(mut bytes: &[u8], visitor: &mut V) { + visitor.begin_game(); + + // Strip UTF-8 BOM + if bytes.starts_with(b"\xEF\xBB\xBF") { + bytes = &bytes[3..]; + } + + let mut in_headers = true; + visitor.begin_headers(); + + let mut cursor = 0; + let len = bytes.len(); + let mut variation_depth = 0; + + while cursor < len { + let c = bytes[cursor]; + + match c { + b' ' | b'\n' | b'\r' | b'\t' => { + cursor += 1; + } + b'[' => { + if !in_headers { + // Start of a new game but we are already in movetext -> ignore trailing games + break; + } + cursor += 1; + + // Parse key + let key_start = cursor; + while cursor < len && bytes[cursor] != b' ' && bytes[cursor] != b'"' && bytes[cursor] != b']' { + cursor += 1; + } + let key = &bytes[key_start..cursor]; + + // Find quote + while cursor < len && bytes[cursor] != b'"' { + cursor += 1; + } + + if cursor < len { + cursor += 1; // skip quote + let val_start = cursor; + // Find closing quote, handling escapes + let mut val_end = cursor; + while cursor < len { + if bytes[cursor] == b'\\' && cursor + 1 < len { + cursor += 2; + } else if bytes[cursor] == b'"' { + val_end = cursor; + cursor += 1; + break; + } else { + cursor += 1; + } + } + + visitor.header(key, &bytes[val_start..val_end]); + } + + // Find closing bracket + while cursor < len && bytes[cursor] != b']' { + cursor += 1; + } + if cursor < len { + cursor += 1; + } + } + b'{' => { + if in_headers { + in_headers = false; + visitor.end_headers(); + } + cursor += 1; + let comment_start = cursor; + if let Some(end_offset) = memchr(b'}', &bytes[cursor..]) { + if variation_depth == 0 { + visitor.comment(&bytes[comment_start..cursor + end_offset]); + } + cursor += end_offset + 1; + } else { + // No closing brace, consume rest of file + if variation_depth == 0 { + visitor.comment(&bytes[comment_start..]); + } + cursor = len; + } + } + b';' | b'%' => { + // Line comment or escape line + if in_headers && c == b'%' { + // Escape lines can appear in headers too + } else if in_headers { + in_headers = false; + visitor.end_headers(); + } + if let Some(end_offset) = memchr(b'\n', &bytes[cursor..]) { + cursor += end_offset + 1; + } else { + cursor = len; + } + } + b'(' => { + if in_headers { + in_headers = false; + visitor.end_headers(); + } + variation_depth += 1; + cursor += 1; + } + b')' => { + if variation_depth > 0 { + variation_depth -= 1; + } + cursor += 1; + } + b'$' => { + if in_headers { + in_headers = false; + visitor.end_headers(); + } + cursor += 1; + while cursor < len && bytes[cursor].is_ascii_digit() { + cursor += 1; + } + } + b'*' => { + if in_headers { + visitor.end_headers(); + } + visitor.outcome(Outcome::Unknown); + break; // End of game + } + _ => { + // Read a token + if in_headers { + in_headers = false; + visitor.end_headers(); + } + let token_start = cursor; + while cursor < len && !is_token_delimiter(bytes[cursor]) { + cursor += 1; + } + let token = &bytes[token_start..cursor]; + + if token.is_empty() { + cursor += 1; + continue; + } + + // Check for move numbers or outcomes + if token == b"1-0" { + visitor.outcome(Outcome::WhiteWins); + break; + } else if token == b"0-1" { + visitor.outcome(Outcome::BlackWins); + break; + } else if token == b"1/2-1/2" { + visitor.outcome(Outcome::Draw); + break; + } else if memchr(b'.', token).is_some() || token.iter().all(|b| b.is_ascii_digit()) { + // Move number, skip + } else if token == b"!" || token == b"?" || token == b"!!" || token == b"??" || token == b"!?" || token == b"?!" { + // Standalone suffix, ignore + } else { + // It's a SAN token + if variation_depth == 0 { + match SanPlus::from_ascii(token) { + Ok(san_plus) => visitor.san(san_plus), + Err(e) => visitor.san_error(e, token), + } + } + } + } + } + } + + visitor.end_game(); +} + +fn is_token_delimiter(c: u8) -> bool { + matches!(c, b' ' | b'\n' | b'\r' | b'\t' | b'[' | b']' | b'{' | b'}' | b'(' | b')' | b';' | b'*') +} diff --git a/src/visitor.rs b/src/visitor.rs index 523b31c..d52e92d 100644 --- a/src/visitor.rs +++ b/src/visitor.rs @@ -8,10 +8,9 @@ use crate::board_serialization::{ get_castling_rights, get_en_passant_file, get_halfmove_clock, get_turn, serialize_board, }; use crate::comment_parsing::{parse_comments, CommentContent, ParsedTag}; -use pgn_reader::{KnownOutcome, Outcome, RawComment, RawTag, SanPlus, Skip, Visitor}; -use shakmaty::{fen::Fen, uci::UciMove, CastlingMode, Chess, Color, Position}; +use crate::tokenizer::{Outcome, Visitor}; +use shakmaty::{fen::Fen, uci::UciMove, CastlingMode, Chess, Color, Position, san::SanPlus}; use std::collections::HashMap; -use std::ops::ControlFlow; /// Configuration for what optional data to store during parsing. #[derive(Clone, Debug)] @@ -318,35 +317,24 @@ impl<'a> GameVisitor<'a> { } impl Visitor for GameVisitor<'_> { - type Tags = Vec<(String, String)>; - type Movetext = (); - type Output = bool; + fn begin_game(&mut self) {} - fn begin_tags(&mut self) -> ControlFlow { + fn begin_headers(&mut self) { self.current_headers.clear(); - ControlFlow::Continue(Vec::with_capacity(10)) - } - - fn tag( - &mut self, - tags: &mut Self::Tags, - key: &[u8], - value: RawTag<'_>, - ) -> ControlFlow { - let key_str = String::from_utf8_lossy(key).into_owned(); - let value_str = String::from_utf8_lossy(value.as_bytes()).into_owned(); - tags.push((key_str, value_str)); - ControlFlow::Continue(()) - } - - fn begin_movetext(&mut self, tags: Self::Tags) -> ControlFlow { - self.current_headers = tags; self.valid_moves = true; self.current_outcome = None; self.current_error = None; self.current_move_count = 0; self.current_position_count = 0; + } + fn header(&mut self, key: &[u8], value: &[u8]) { + let key_str = String::from_utf8_lossy(key).into_owned(); + let value_str = String::from_utf8_lossy(value).into_owned(); + self.current_headers.push((key_str, value_str)); + } + + fn end_headers(&mut self) { // Determine castling mode from Variant header (case-insensitive) let castling_mode = self .current_headers @@ -389,14 +377,9 @@ impl Visitor for GameVisitor<'_> { // Record initial board state self.push_board_state(); - ControlFlow::Continue(()) } - fn san( - &mut self, - _movetext: &mut Self::Movetext, - san_plus: SanPlus, - ) -> ControlFlow { + fn san(&mut self, san_plus: SanPlus) { if self.valid_moves { match san_plus.san.to_move(&self.pos) { Ok(m) => { @@ -424,15 +407,17 @@ impl Visitor for GameVisitor<'_> { } } } - ControlFlow::Continue(()) } - fn comment( - &mut self, - _movetext: &mut Self::Movetext, - comment: RawComment<'_>, - ) -> ControlFlow { - if let Ok((_, parsed_comments)) = parse_comments(comment.as_bytes()) { + fn san_error(&mut self, error: shakmaty::san::ParseSanError, token: &[u8]) { + if self.valid_moves { + let token_str = String::from_utf8_lossy(token); + self.set_error(format!("failed to parse SAN: {} ({})", error, token_str)); + } + } + + fn comment(&mut self, comment: &[u8]) { + if let Ok((_, parsed_comments)) = parse_comments(comment) { let mut move_comments = if self.config.store_comments { Some(String::new()) } else { @@ -489,39 +474,24 @@ impl Visitor for GameVisitor<'_> { } } } - ControlFlow::Continue(()) - } - - fn begin_variation( - &mut self, - _movetext: &mut Self::Movetext, - ) -> ControlFlow { - ControlFlow::Continue(Skip(true)) // Skip variations, stay in mainline } - fn outcome( - &mut self, - _movetext: &mut Self::Movetext, - _outcome: Outcome, - ) -> ControlFlow { + fn outcome(&mut self, _outcome: Outcome) { self.current_outcome = Some(match _outcome { - Outcome::Known(known) => match known { - KnownOutcome::Decisive { winner } => format!("{:?}", winner), - KnownOutcome::Draw => "Draw".to_string(), - }, + Outcome::WhiteWins => "White".to_string(), + Outcome::BlackWins => "Black".to_string(), + Outcome::Draw => "Draw".to_string(), Outcome::Unknown => "Unknown".to_string(), }); self.update_position_status(); - ControlFlow::Continue(()) } - fn end_game(&mut self, _movetext: Self::Movetext) -> Self::Output { + fn end_game(&mut self) { // Handle case where outcome() was not called (e.g., incomplete game) if self.buffers.is_checkmate.len() < self.buffers.headers.len() + 1 { self.update_position_status(); } self.finalize_game(); - self.valid_moves } } @@ -531,17 +501,15 @@ pub fn parse_game_to_buffers( buffers: &mut Buffers, config: &ParseConfig, ) -> Result { - use pgn_reader::Reader; - use std::io::Cursor; - - let mut reader = Reader::new(Cursor::new(pgn)); let mut visitor = GameVisitor::new(buffers, config); - - match reader.read_game(&mut visitor) { - Ok(Some(valid)) => Ok(valid), - Ok(None) => Err("No game found in PGN".to_string()), - Err(err) => Err(format!("Parsing error: {}", err)), + + // Check if the PGN string has any non-whitespace characters to emulate pgn-reader's `Ok(None)` behavior + if pgn.trim().is_empty() { + return Err("No game found in PGN".to_string()); } + + crate::tokenizer::parse_game(pgn.as_bytes(), &mut visitor); + Ok(visitor.valid_moves) } #[cfg(test)] From aa1f753bafde6778445477377f65cb190ed6c4b9 Mon Sep 17 00:00:00 2001 From: vladkvit Date: Fri, 10 Jul 2026 22:53:45 -0400 Subject: [PATCH 2/8] cargo fmt --- benches/parquet_bench.rs | 4 ++-- src/comment_parsing.rs | 2 +- src/tokenizer.rs | 30 ++++++++++++++++++++++-------- src/visitor.rs | 6 +++--- 4 files changed, 28 insertions(+), 14 deletions(-) diff --git a/benches/parquet_bench.rs b/benches/parquet_bench.rs index 4af6b63..fd59e19 100644 --- a/benches/parquet_bench.rs +++ b/benches/parquet_bench.rs @@ -5,13 +5,13 @@ use arrow::array::{Array, StringArray}; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; -use rayon::prelude::*; use rayon::ThreadPoolBuilder; +use rayon::prelude::*; use std::fs::File; use std::path::Path; use std::time::Instant; -use rust_pgn_reader_python_binding::{parse_game_to_buffers, Buffers, ParseConfig}; +use rust_pgn_reader_python_binding::{Buffers, ParseConfig, parse_game_to_buffers}; const FILE_PATH: &str = "2013-07-train-00000-of-00001.parquet"; diff --git a/src/comment_parsing.rs b/src/comment_parsing.rs index bb63c67..956f572 100644 --- a/src/comment_parsing.rs +++ b/src/comment_parsing.rs @@ -1,11 +1,11 @@ use nom::{ + IResult, Parser, branch::alt, bytes::complete::{is_not, tag}, character::complete::{char, digit1}, combinator::{map, map_res, not, opt, peek, recognize}, multi::{many0, many1}, sequence::{delimited, pair, preceded}, - IResult, Parser, }; use std::borrow::Cow; // use nom::character::complete::multispace1; diff --git a/src/tokenizer.rs b/src/tokenizer.rs index 94bcd3e..ce79c0f 100644 --- a/src/tokenizer.rs +++ b/src/tokenizer.rs @@ -49,10 +49,14 @@ pub fn parse_game(mut bytes: &[u8], visitor: &mut V) { break; } cursor += 1; - + // Parse key let key_start = cursor; - while cursor < len && bytes[cursor] != b' ' && bytes[cursor] != b'"' && bytes[cursor] != b']' { + while cursor < len + && bytes[cursor] != b' ' + && bytes[cursor] != b'"' + && bytes[cursor] != b']' + { cursor += 1; } let key = &bytes[key_start..cursor]; @@ -61,7 +65,7 @@ pub fn parse_game(mut bytes: &[u8], visitor: &mut V) { while cursor < len && bytes[cursor] != b'"' { cursor += 1; } - + if cursor < len { cursor += 1; // skip quote let val_start = cursor; @@ -78,7 +82,7 @@ pub fn parse_game(mut bytes: &[u8], visitor: &mut V) { cursor += 1; } } - + visitor.header(key, &bytes[val_start..val_end]); } @@ -166,7 +170,7 @@ pub fn parse_game(mut bytes: &[u8], visitor: &mut V) { cursor += 1; } let token = &bytes[token_start..cursor]; - + if token.is_empty() { cursor += 1; continue; @@ -182,9 +186,16 @@ pub fn parse_game(mut bytes: &[u8], visitor: &mut V) { } else if token == b"1/2-1/2" { visitor.outcome(Outcome::Draw); break; - } else if memchr(b'.', token).is_some() || token.iter().all(|b| b.is_ascii_digit()) { + } else if memchr(b'.', token).is_some() || token.iter().all(|b| b.is_ascii_digit()) + { // Move number, skip - } else if token == b"!" || token == b"?" || token == b"!!" || token == b"??" || token == b"!?" || token == b"?!" { + } else if token == b"!" + || token == b"?" + || token == b"!!" + || token == b"??" + || token == b"!?" + || token == b"?!" + { // Standalone suffix, ignore } else { // It's a SAN token @@ -203,5 +214,8 @@ pub fn parse_game(mut bytes: &[u8], visitor: &mut V) { } fn is_token_delimiter(c: u8) -> bool { - matches!(c, b' ' | b'\n' | b'\r' | b'\t' | b'[' | b']' | b'{' | b'}' | b'(' | b')' | b';' | b'*') + matches!( + c, + b' ' | b'\n' | b'\r' | b'\t' | b'[' | b']' | b'{' | b'}' | b'(' | b')' | b';' | b'*' + ) } diff --git a/src/visitor.rs b/src/visitor.rs index d52e92d..41f8546 100644 --- a/src/visitor.rs +++ b/src/visitor.rs @@ -7,9 +7,9 @@ use crate::board_serialization::{ get_castling_rights, get_en_passant_file, get_halfmove_clock, get_turn, serialize_board, }; -use crate::comment_parsing::{parse_comments, CommentContent, ParsedTag}; +use crate::comment_parsing::{CommentContent, ParsedTag, parse_comments}; use crate::tokenizer::{Outcome, Visitor}; -use shakmaty::{fen::Fen, uci::UciMove, CastlingMode, Chess, Color, Position, san::SanPlus}; +use shakmaty::{CastlingMode, Chess, Color, Position, fen::Fen, san::SanPlus, uci::UciMove}; use std::collections::HashMap; /// Configuration for what optional data to store during parsing. @@ -502,7 +502,7 @@ pub fn parse_game_to_buffers( config: &ParseConfig, ) -> Result { let mut visitor = GameVisitor::new(buffers, config); - + // Check if the PGN string has any non-whitespace characters to emulate pgn-reader's `Ok(None)` behavior if pgn.trim().is_empty() { return Err("No game found in PGN".to_string()); From 0c5479bcfe0a4488ade8352a750adb8722c0cc79 Mon Sep 17 00:00:00 2001 From: vladkvit Date: Fri, 10 Jul 2026 23:10:08 -0400 Subject: [PATCH 3/8] FIx up new tokenizer, add tests --- src/tokenizer.rs | 382 ++++++++++++++++++++++++++++++++++++++++++----- src/visitor.rs | 25 +++- 2 files changed, 368 insertions(+), 39 deletions(-) diff --git a/src/tokenizer.rs b/src/tokenizer.rs index ce79c0f..bb71bb3 100644 --- a/src/tokenizer.rs +++ b/src/tokenizer.rs @@ -1,5 +1,6 @@ use memchr::memchr; -use shakmaty::san::{ParseSanError, SanPlus}; +use shakmaty::CastlingSide; +use shakmaty::san::{ParseSanError, San, SanPlus, Suffix}; #[derive(Debug, Clone, PartialEq, Eq)] pub enum Outcome { @@ -10,7 +11,6 @@ pub enum Outcome { } pub trait Visitor { - fn begin_game(&mut self); fn begin_headers(&mut self); fn header(&mut self, key: &[u8], value: &[u8]); fn end_headers(&mut self); @@ -21,9 +21,18 @@ pub trait Visitor { fn end_game(&mut self); } +/// Parse a single PGN game from `bytes`, dispatching events to `visitor`. +/// +/// Contract: one game per input slice. Content after the first game (a blank +/// line followed by anything, or a new tag section) is ignored, matching the +/// behavior of a single `pgn_reader::Reader::read_game` call. +/// +/// Known intentional divergences from pgn-reader: +/// - Syntactically invalid SAN tokens report `san_error` instead of being +/// silently skipped. +/// - An unterminated `{` comment is delivered as a comment spanning the rest +/// of the input instead of raising a hard error. pub fn parse_game(mut bytes: &[u8], visitor: &mut V) { - visitor.begin_game(); - // Strip UTF-8 BOM if bytes.starts_with(b"\xEF\xBB\xBF") { bytes = &bytes[3..]; @@ -34,14 +43,36 @@ pub fn parse_game(mut bytes: &[u8], visitor: &mut V) { let mut cursor = 0; let len = bytes.len(); - let mut variation_depth = 0; + let mut variation_depth = 0usize; while cursor < len { let c = bytes[cursor]; match c { - b' ' | b'\n' | b'\r' | b'\t' => { + // '.', '!' and '?' are token delimiters skipped like whitespace; + // move-number dots ("1.e4") and attached annotations ("e4!?") + // must not become part of the adjacent token. + b' ' | b'\r' | b'\t' | b'.' | b'!' | b'?' => { + cursor += 1; + } + b'\n' => { cursor += 1; + match bytes.get(cursor).copied() { + // PGN escape: '%' at the start of a line skips the line. + Some(b'%') => { + // Leave the trailing '\n' so line-start logic reapplies. + cursor = memchr(b'\n', &bytes[cursor..]).map_or(len, |o| cursor + o); + } + // Blank line or a new tag section terminates the movetext + // (parity with pgn_reader::Reader::read_game). + Some(b'\n') | Some(b'[') if !in_headers => break, + Some(b'\r') if !in_headers && bytes.get(cursor + 1) == Some(&b'\n') => break, + _ => {} + } + } + // Escape line at the very start of the input. + b'%' if cursor == 0 => { + cursor = memchr(b'\n', bytes).unwrap_or(len); } b'[' => { if !in_headers { @@ -114,19 +145,14 @@ pub fn parse_game(mut bytes: &[u8], visitor: &mut V) { cursor = len; } } - b';' | b'%' => { - // Line comment or escape line - if in_headers && c == b'%' { - // Escape lines can appear in headers too - } else if in_headers { + b';' => { + // Line comment + if in_headers { in_headers = false; visitor.end_headers(); } - if let Some(end_offset) = memchr(b'\n', &bytes[cursor..]) { - cursor += end_offset + 1; - } else { - cursor = len; - } + // Leave the trailing '\n' so line-start logic reapplies. + cursor = memchr(b'\n', &bytes[cursor..]).map_or(len, |o| cursor + o); } b'(' => { if in_headers { @@ -153,11 +179,15 @@ pub fn parse_game(mut bytes: &[u8], visitor: &mut V) { } } b'*' => { - if in_headers { - visitor.end_headers(); + if variation_depth == 0 { + if in_headers { + in_headers = false; + visitor.end_headers(); + } + visitor.outcome(Outcome::Unknown); + break; // End of game } - visitor.outcome(Outcome::Unknown); - break; // End of game + cursor += 1; } _ => { // Read a token @@ -172,11 +202,16 @@ pub fn parse_game(mut bytes: &[u8], visitor: &mut V) { let token = &bytes[token_start..cursor]; if token.is_empty() { + // Stray delimiter without a dedicated arm (']' or '}'). cursor += 1; continue; } - // Check for move numbers or outcomes + if variation_depth > 0 { + continue; // Skip everything inside variations + } + + // Check for outcomes and move numbers if token == b"1-0" { visitor.outcome(Outcome::WhiteWins); break; @@ -186,36 +221,309 @@ pub fn parse_game(mut bytes: &[u8], visitor: &mut V) { } else if token == b"1/2-1/2" { visitor.outcome(Outcome::Draw); break; - } else if memchr(b'.', token).is_some() || token.iter().all(|b| b.is_ascii_digit()) - { + } else if token.iter().all(|b| b.is_ascii_digit()) { // Move number, skip - } else if token == b"!" - || token == b"?" - || token == b"!!" - || token == b"??" - || token == b"!?" - || token == b"?!" - { - // Standalone suffix, ignore + } else if let Some(san_plus) = castling_with_zeros(token) { + visitor.san(san_plus); } else { // It's a SAN token - if variation_depth == 0 { - match SanPlus::from_ascii(token) { - Ok(san_plus) => visitor.san(san_plus), - Err(e) => visitor.san_error(e, token), - } + match SanPlus::from_ascii(token) { + Ok(san_plus) => visitor.san(san_plus), + Err(e) => visitor.san_error(e, token), } } } } } + // A game may consist of headers only (no movetext, no result token). + // The visitor still expects end_headers before end_game. + if in_headers { + visitor.end_headers(); + } visitor.end_game(); } +/// Support castling notated with zeros ("0-0", "0-0-0"), with optional +/// check/checkmate suffix, matching pgn-reader's explicit handling. +fn castling_with_zeros(token: &[u8]) -> Option { + let (body, suffix) = match token.last() { + Some(b'+') => (&token[..token.len() - 1], Some(Suffix::Check)), + Some(b'#') => (&token[..token.len() - 1], Some(Suffix::Checkmate)), + _ => (token, None), + }; + let side = match body { + b"0-0" => CastlingSide::KingSide, + b"0-0-0" => CastlingSide::QueenSide, + _ => return None, + }; + Some(SanPlus { + san: San::Castle(side), + suffix, + }) +} + fn is_token_delimiter(c: u8) -> bool { matches!( c, - b' ' | b'\n' | b'\r' | b'\t' | b'[' | b']' | b'{' | b'}' | b'(' | b')' | b';' | b'*' + b' ' | b'\n' + | b'\r' + | b'\t' + | b'[' + | b']' + | b'{' + | b'}' + | b'(' + | b')' + | b';' + | b'*' + | b'.' + | b'!' + | b'?' + | b'$' ) } + +#[cfg(test)] +mod tests { + use super::*; + + #[derive(Default)] + struct RecordingVisitor { + events: Vec, + } + + impl Visitor for RecordingVisitor { + fn begin_headers(&mut self) { + self.events.push("begin_headers".to_string()); + } + fn header(&mut self, key: &[u8], value: &[u8]) { + self.events.push(format!( + "header {}={}", + String::from_utf8_lossy(key), + String::from_utf8_lossy(value) + )); + } + fn end_headers(&mut self) { + self.events.push("end_headers".to_string()); + } + fn san(&mut self, san_plus: SanPlus) { + self.events.push(format!("san {}", san_plus)); + } + fn san_error(&mut self, _error: ParseSanError, token: &[u8]) { + self.events + .push(format!("san_error {}", String::from_utf8_lossy(token))); + } + fn comment(&mut self, comment: &[u8]) { + self.events + .push(format!("comment {}", String::from_utf8_lossy(comment))); + } + fn outcome(&mut self, outcome: Outcome) { + self.events.push(format!("outcome {:?}", outcome)); + } + fn end_game(&mut self) { + self.events.push("end_game".to_string()); + } + } + + fn events(pgn: &str) -> Vec { + let mut visitor = RecordingVisitor::default(); + parse_game(pgn.as_bytes(), &mut visitor); + visitor.events + } + + fn sans(pgn: &str) -> Vec { + events(pgn) + .iter() + .filter_map(|e| e.strip_prefix("san ").map(str::to_owned)) + .collect() + } + + fn san_errors(pgn: &str) -> Vec { + events(pgn) + .iter() + .filter_map(|e| e.strip_prefix("san_error ").map(str::to_owned)) + .collect() + } + + #[test] + fn test_move_numbers_without_space() { + assert_eq!(sans("1.e4 e5 2.Nf3 Nc6 1-0"), ["e4", "e5", "Nf3", "Nc6"]); + assert!(san_errors("1.e4 e5 2.Nf3 Nc6 1-0").is_empty()); + } + + #[test] + fn test_black_move_number_continuation() { + assert_eq!( + sans("1.e4 1...e5 2. Nf3 2... Nc6 *"), + ["e4", "e5", "Nf3", "Nc6"] + ); + } + + #[test] + fn test_attached_annotations() { + let pgn = "1. e4!? e5?! 2. Nf3! Nc6?? 3. Bb5$14 a6 *"; + assert_eq!(sans(pgn), ["e4", "e5", "Nf3", "Nc6", "Bb5", "a6"]); + assert!(san_errors(pgn).is_empty()); + } + + #[test] + fn test_standalone_nags_and_annotations() { + assert_eq!( + sans("1. e4 $1 e5 $139 2. Nf3 !? Nc6 ?? *"), + ["e4", "e5", "Nf3", "Nc6"] + ); + } + + #[test] + fn test_castling_with_zeros() { + let pgn = "1. 0-0 0-0-0 2. 0-0+ 0-0-0# 1-0"; + assert_eq!(sans(pgn), ["O-O", "O-O-O", "O-O+", "O-O-O#"]); + assert!(san_errors(pgn).is_empty()); + } + + #[test] + fn test_check_suffix_kept() { + assert_eq!(sans("1. e4+ e5# *"), ["e4+", "e5#"]); + } + + #[test] + fn test_bom_stripped() { + let pgn = "\u{feff}[Event \"x\"]\n\n1. e4 *"; + let evs = events(pgn); + assert!(evs.contains(&"header Event=x".to_string())); + assert_eq!(sans(pgn), ["e4"]); + } + + #[test] + fn test_crlf() { + let pgn = "[Event \"x\"]\r\n[Site \"y\"]\r\n\r\n1. e4 e5 1-0\r\n"; + let evs = events(pgn); + assert!(evs.contains(&"header Event=x".to_string())); + assert!(evs.contains(&"header Site=y".to_string())); + assert_eq!(sans(pgn), ["e4", "e5"]); + assert!(evs.contains(&"outcome WhiteWins".to_string())); + } + + #[test] + fn test_tag_value_escapes() { + // Escaped quotes/backslashes delimit correctly; value passed raw (undecoded). + let pgn = "[Event \"a \\\"quoted\\\" name\"]\n[Site \"back\\\\slash\"]\n\n*"; + let evs = events(pgn); + assert!(evs.contains(&"header Event=a \\\"quoted\\\" name".to_string())); + assert!(evs.contains(&"header Site=back\\\\slash".to_string())); + } + + #[test] + fn test_semicolon_line_comment() { + assert_eq!(sans("1. e4 ; rest ignored e5 Nf3\nc5 *"), ["e4", "c5"]); + } + + #[test] + fn test_escape_line_at_input_start() { + let pgn = "% ignored line\n[Event \"x\"]\n\n1. e4 *"; + let evs = events(pgn); + assert!(evs.contains(&"header Event=x".to_string())); + assert_eq!(sans(pgn), ["e4"]); + } + + #[test] + fn test_escape_line_mid_game() { + assert_eq!(sans("1. e4\n% skipped e5\nc5 *"), ["e4", "c5"]); + } + + #[test] + fn test_percent_mid_line_is_not_escape() { + // '%' not at line start is not an escape; the junk token errors + // (intentionally stricter than pgn-reader's silent skip) and + // parsing continues to the end of the game. + let pgn = "1. e4 %junk e5 *"; + assert_eq!(sans(pgn), ["e4", "e5"]); + assert_eq!(san_errors(pgn), ["%junk"]); + } + + #[test] + fn test_variations_skipped() { + let pgn = "1. e4 (1. d4 {inner} d5 (1... Nf6 2. c4) 2. c4) e5 2. Nf3 1-0"; + assert_eq!(sans(pgn), ["e4", "e5", "Nf3"]); + let evs = events(pgn); + assert!(!evs.iter().any(|e| e.starts_with("comment"))); + assert!(evs.contains(&"outcome WhiteWins".to_string())); + } + + #[test] + fn test_result_token_inside_variation_ignored() { + let pgn = "1. e4 (1. d4 1-0) e5 (1... c5 *) 2. Nf3 *"; + assert_eq!(sans(pgn), ["e4", "e5", "Nf3"]); + let evs = events(pgn); + assert_eq!(evs.iter().filter(|e| e.starts_with("outcome")).count(), 1); + assert!(evs.contains(&"outcome Unknown".to_string())); + } + + #[test] + fn test_trailing_game_ignored_after_result() { + let pgn = "[Event \"a\"]\n\n1. e4 1-0\n\n[Event \"b\"]\n\n1. d4 0-1"; + let evs = events(pgn); + assert!(evs.contains(&"header Event=a".to_string())); + assert!(!evs.contains(&"header Event=b".to_string())); + assert_eq!(sans(pgn), ["e4"]); + assert!(evs.contains(&"outcome WhiteWins".to_string())); + assert!(!evs.contains(&"outcome BlackWins".to_string())); + } + + #[test] + fn test_blank_line_terminates_movetext() { + // Without a result token, a blank line still ends the game + // (parity with pgn-reader). + assert_eq!(sans("1. e4 e5\n\n2. Nf3 Nc6"), ["e4", "e5"]); + assert_eq!(sans("1. e4 e5\r\n\r\n2. Nf3"), ["e4", "e5"]); + } + + #[test] + fn test_headers_only_game_flushes_end_headers() { + let evs = events("[Event \"x\"]\n[Site \"y\"]\n"); + assert_eq!( + evs, + [ + "begin_headers", + "header Event=x", + "header Site=y", + "end_headers", + "end_game" + ] + ); + } + + #[test] + fn test_empty_input() { + assert_eq!(events(""), ["begin_headers", "end_headers", "end_game"]); + } + + #[test] + fn test_unterminated_comment_is_lenient() { + // Intentional divergence: pgn-reader raises a hard error here. + let evs = events("1. e4 {never closed"); + assert!(evs.contains(&"comment never closed".to_string())); + assert_eq!(sans("1. e4 {never closed"), ["e4"]); + } + + #[test] + fn test_garbage_token_reports_san_error() { + let pgn = "1. e4 xyzzy9 e5 *"; + assert_eq!(san_errors(pgn), ["xyzzy9"]); + assert_eq!(sans(pgn), ["e4", "e5"]); + } + + #[test] + fn test_stray_closing_delimiters_no_hang() { + assert_eq!(sans("1. e4 } ] e5 *"), ["e4", "e5"]); + } + + #[test] + fn test_comment_between_headers_and_moves() { + let pgn = "[Event \"x\"]\n\n{pre-game comment} 1. e4 *"; + let evs = events(pgn); + assert!(evs.contains(&"comment pre-game comment".to_string())); + assert_eq!(sans(pgn), ["e4"]); + } +} diff --git a/src/visitor.rs b/src/visitor.rs index 41f8546..95c007a 100644 --- a/src/visitor.rs +++ b/src/visitor.rs @@ -317,10 +317,9 @@ impl<'a> GameVisitor<'a> { } impl Visitor for GameVisitor<'_> { - fn begin_game(&mut self) {} - fn begin_headers(&mut self) { self.current_headers.clear(); + self.current_headers.reserve(10); self.valid_moves = true; self.current_outcome = None; self.current_error = None; @@ -694,6 +693,28 @@ mod tests { assert_eq!(buffers.legal_move_from_squares.len(), 40); } + #[test] + fn test_headers_only_game() { + // A game with headers but no movetext at all: the initial position + // (from the FEN header) must still be recorded. + let pgn = r#"[Event "Test"] +[FEN "r1bqkbnr/pppp1ppp/2n5/4p3/4P3/5N2/PPPP1PPP/RNBQKB1R w KQkq - 2 3"] +"#; + + let config = default_config(); + let mut buffers = Buffers::with_capacity(1, 70, &config); + let result = parse_game_to_buffers(pgn, &mut buffers, &config); + + assert!(result.is_ok()); + assert!(result.unwrap()); // valid game (no moves, no errors) + assert_eq!(buffers.num_games(), 1); + assert_eq!(buffers.move_counts[0], 0); + assert_eq!(buffers.position_counts[0], 1); // initial position only + assert_eq!(buffers.headers[0].get("Event"), Some(&"Test".to_string())); + assert_eq!(buffers.parse_errors[0], None); // FEN header was applied + assert_eq!(buffers.outcome[0], None); + } + #[test] fn test_parse_game_without_headers() { let pgn = "1. Nf3 d5 2. e4 c5 3. exd5 e5 4. dxe6 0-1"; From ba15febd1c278fbdd8bbe85f1c50452571f3ad60 Mon Sep 17 00:00:00 2001 From: vladkvit Date: Sat, 11 Jul 2026 08:41:58 -0400 Subject: [PATCH 4/8] Make tokenizer purely lexical, move SAN parsing into visitor --- src/tokenizer.rs | 84 ++++++++++--------------------- src/visitor.rs | 126 ++++++++++++++++++++++++++++++++++++----------- 2 files changed, 123 insertions(+), 87 deletions(-) diff --git a/src/tokenizer.rs b/src/tokenizer.rs index bb71bb3..a85fdfc 100644 --- a/src/tokenizer.rs +++ b/src/tokenizer.rs @@ -1,6 +1,4 @@ use memchr::memchr; -use shakmaty::CastlingSide; -use shakmaty::san::{ParseSanError, San, SanPlus, Suffix}; #[derive(Debug, Clone, PartialEq, Eq)] pub enum Outcome { @@ -10,12 +8,20 @@ pub enum Outcome { Unknown, } +/// Purely lexical PGN event consumer. +/// +/// The tokenizer performs no chess-semantic work: SAN tokens are delivered +/// as raw bytes via `san_token` and it is the visitor's job to parse them +/// (e.g. with `shakmaty::san::SanPlus::from_ascii`). This keeps a pure +/// counting pass free of SAN-parsing cost. pub trait Visitor { fn begin_headers(&mut self); fn header(&mut self, key: &[u8], value: &[u8]); fn end_headers(&mut self); - fn san(&mut self, san_plus: SanPlus); - fn san_error(&mut self, error: ParseSanError, token: &[u8]); + /// A mainline movetext token that is not a move number, NAG, result or + /// comment. Usually a SAN move, but syntactically invalid tokens are + /// delivered raw as well. + fn san_token(&mut self, token: &[u8]); fn comment(&mut self, comment: &[u8]); fn outcome(&mut self, outcome: Outcome); fn end_game(&mut self); @@ -28,8 +34,9 @@ pub trait Visitor { /// behavior of a single `pgn_reader::Reader::read_game` call. /// /// Known intentional divergences from pgn-reader: -/// - Syntactically invalid SAN tokens report `san_error` instead of being -/// silently skipped. +/// - SAN tokens are not parsed here; they are delivered raw via `san_token` +/// (including syntactically invalid ones, which pgn-reader silently skips — +/// the visitor decides how to handle them). /// - An unterminated `{` comment is delivered as a comment spanning the rest /// of the input instead of raising a hard error. pub fn parse_game(mut bytes: &[u8], visitor: &mut V) { @@ -223,14 +230,9 @@ pub fn parse_game(mut bytes: &[u8], visitor: &mut V) { break; } else if token.iter().all(|b| b.is_ascii_digit()) { // Move number, skip - } else if let Some(san_plus) = castling_with_zeros(token) { - visitor.san(san_plus); } else { - // It's a SAN token - match SanPlus::from_ascii(token) { - Ok(san_plus) => visitor.san(san_plus), - Err(e) => visitor.san_error(e, token), - } + // SAN (or garbage) token — delivered raw, parsed by the visitor. + visitor.san_token(token); } } } @@ -244,25 +246,6 @@ pub fn parse_game(mut bytes: &[u8], visitor: &mut V) { visitor.end_game(); } -/// Support castling notated with zeros ("0-0", "0-0-0"), with optional -/// check/checkmate suffix, matching pgn-reader's explicit handling. -fn castling_with_zeros(token: &[u8]) -> Option { - let (body, suffix) = match token.last() { - Some(b'+') => (&token[..token.len() - 1], Some(Suffix::Check)), - Some(b'#') => (&token[..token.len() - 1], Some(Suffix::Checkmate)), - _ => (token, None), - }; - let side = match body { - b"0-0" => CastlingSide::KingSide, - b"0-0-0" => CastlingSide::QueenSide, - _ => return None, - }; - Some(SanPlus { - san: San::Castle(side), - suffix, - }) -} - fn is_token_delimiter(c: u8) -> bool { matches!( c, @@ -307,12 +290,9 @@ mod tests { fn end_headers(&mut self) { self.events.push("end_headers".to_string()); } - fn san(&mut self, san_plus: SanPlus) { - self.events.push(format!("san {}", san_plus)); - } - fn san_error(&mut self, _error: ParseSanError, token: &[u8]) { + fn san_token(&mut self, token: &[u8]) { self.events - .push(format!("san_error {}", String::from_utf8_lossy(token))); + .push(format!("san {}", String::from_utf8_lossy(token))); } fn comment(&mut self, comment: &[u8]) { self.events @@ -339,17 +319,9 @@ mod tests { .collect() } - fn san_errors(pgn: &str) -> Vec { - events(pgn) - .iter() - .filter_map(|e| e.strip_prefix("san_error ").map(str::to_owned)) - .collect() - } - #[test] fn test_move_numbers_without_space() { assert_eq!(sans("1.e4 e5 2.Nf3 Nc6 1-0"), ["e4", "e5", "Nf3", "Nc6"]); - assert!(san_errors("1.e4 e5 2.Nf3 Nc6 1-0").is_empty()); } #[test] @@ -364,7 +336,6 @@ mod tests { fn test_attached_annotations() { let pgn = "1. e4!? e5?! 2. Nf3! Nc6?? 3. Bb5$14 a6 *"; assert_eq!(sans(pgn), ["e4", "e5", "Nf3", "Nc6", "Bb5", "a6"]); - assert!(san_errors(pgn).is_empty()); } #[test] @@ -376,10 +347,11 @@ mod tests { } #[test] - fn test_castling_with_zeros() { + fn test_castling_with_zeros_delivered_raw() { + // Zeros-castling normalization is semantic and lives in the visitor + // (see visitor.rs); the tokenizer delivers the tokens verbatim. let pgn = "1. 0-0 0-0-0 2. 0-0+ 0-0-0# 1-0"; - assert_eq!(sans(pgn), ["O-O", "O-O-O", "O-O+", "O-O-O#"]); - assert!(san_errors(pgn).is_empty()); + assert_eq!(sans(pgn), ["0-0", "0-0-0", "0-0+", "0-0-0#"]); } #[test] @@ -434,12 +406,11 @@ mod tests { #[test] fn test_percent_mid_line_is_not_escape() { - // '%' not at line start is not an escape; the junk token errors - // (intentionally stricter than pgn-reader's silent skip) and - // parsing continues to the end of the game. + // '%' not at line start is not an escape; the junk token is delivered + // raw like any other token (the visitor decides how to handle it) and + // tokenizing continues to the end of the game. let pgn = "1. e4 %junk e5 *"; - assert_eq!(sans(pgn), ["e4", "e5"]); - assert_eq!(san_errors(pgn), ["%junk"]); + assert_eq!(sans(pgn), ["e4", "%junk", "e5"]); } #[test] @@ -508,10 +479,9 @@ mod tests { } #[test] - fn test_garbage_token_reports_san_error() { + fn test_garbage_token_delivered_raw() { let pgn = "1. e4 xyzzy9 e5 *"; - assert_eq!(san_errors(pgn), ["xyzzy9"]); - assert_eq!(sans(pgn), ["e4", "e5"]); + assert_eq!(sans(pgn), ["e4", "xyzzy9", "e5"]); } #[test] diff --git a/src/visitor.rs b/src/visitor.rs index 95c007a..e70fbc4 100644 --- a/src/visitor.rs +++ b/src/visitor.rs @@ -9,7 +9,12 @@ use crate::board_serialization::{ }; use crate::comment_parsing::{CommentContent, ParsedTag, parse_comments}; use crate::tokenizer::{Outcome, Visitor}; -use shakmaty::{CastlingMode, Chess, Color, Position, fen::Fen, san::SanPlus, uci::UciMove}; +use shakmaty::{ + CastlingMode, CastlingSide, Chess, Color, Position, + fen::Fen, + san::{San, SanPlus, Suffix}, + uci::UciMove, +}; use std::collections::HashMap; /// Configuration for what optional data to store during parsing. @@ -187,6 +192,25 @@ impl Buffers { } } +/// Support castling notated with zeros ("0-0", "0-0-0"), with optional +/// check/checkmate suffix, matching pgn-reader's explicit handling. +fn castling_with_zeros(token: &[u8]) -> Option { + let (body, suffix) = match token.last() { + Some(b'+') => (&token[..token.len() - 1], Some(Suffix::Check)), + Some(b'#') => (&token[..token.len() - 1], Some(Suffix::Checkmate)), + _ => (token, None), + }; + let side = match body { + b"0-0" => CastlingSide::KingSide, + b"0-0-0" => CastlingSide::QueenSide, + _ => return None, + }; + Some(SanPlus { + san: San::Castle(side), + suffix, + }) +} + /// Visitor that writes directly to shared Buffers. /// /// This visitor does not allocate any per-game Vec structures. @@ -378,40 +402,48 @@ impl Visitor for GameVisitor<'_> { self.push_board_state(); } - fn san(&mut self, san_plus: SanPlus) { - if self.valid_moves { - match san_plus.san.to_move(&self.pos) { - Ok(m) => { - self.pos.play_unchecked(m); + fn san_token(&mut self, token: &[u8]) { + // Early abort: once the game is invalid, skip SAN parsing entirely. + if !self.valid_moves { + return; + } - // Record board state after move - self.push_board_state(); + let san_plus = match castling_with_zeros(token) { + Some(san_plus) => san_plus, + None => match SanPlus::from_ascii(token) { + Ok(san_plus) => san_plus, + Err(error) => { + let token_str = String::from_utf8_lossy(token); + self.set_error(format!("failed to parse SAN: {} ({})", error, token_str)); + return; + } + }, + }; - let uci_move_obj = UciMove::from_standard(m); - match uci_move_obj { - UciMove::Normal { - from, - to, - promotion, - } => { - self.push_move(from as u8, to as u8, promotion.map(|p| p as u8)); - } - _ => { - self.set_error(format!("unexpected UCI move type: {:?}", uci_move_obj)); - } + match san_plus.san.to_move(&self.pos) { + Ok(m) => { + self.pos.play_unchecked(m); + + // Record board state after move + self.push_board_state(); + + let uci_move_obj = UciMove::from_standard(m); + match uci_move_obj { + UciMove::Normal { + from, + to, + promotion, + } => { + self.push_move(from as u8, to as u8, promotion.map(|p| p as u8)); + } + _ => { + self.set_error(format!("unexpected UCI move type: {:?}", uci_move_obj)); } - } - Err(err) => { - self.set_error(format!("illegal move: {} {}", err, san_plus)); } } - } - } - - fn san_error(&mut self, error: shakmaty::san::ParseSanError, token: &[u8]) { - if self.valid_moves { - let token_str = String::from_utf8_lossy(token); - self.set_error(format!("failed to parse SAN: {} ({})", error, token_str)); + Err(err) => { + self.set_error(format!("illegal move: {} {}", err, san_plus)); + } } } @@ -608,6 +640,40 @@ mod tests { assert_eq!(pos_offsets, vec![0, 5, 12, 16]); } + #[test] + fn test_castling_with_zeros_normalized() { + // Zeros notation is handled in the visitor (moved from the tokenizer). + let pgn = "1. e4 e5 2. Nf3 Nf6 3. Bc4 Bc5 4. 0-0 0-0 1-0"; + + let config = default_config(); + let mut buffers = Buffers::with_capacity(1, 70, &config); + let result = parse_game_to_buffers(pgn, &mut buffers, &config); + + assert!(result.is_ok()); + assert!(result.unwrap(), "zeros castling should parse as O-O"); + assert_eq!(buffers.total_moves(), 8); + assert_eq!(buffers.parse_errors[0], None); + } + + #[test] + fn test_invalid_san_token_sets_parse_error() { + let pgn = "1. e4 xyzzy9 e5 1-0"; + + let config = default_config(); + let mut buffers = Buffers::with_capacity(1, 70, &config); + let result = parse_game_to_buffers(pgn, &mut buffers, &config); + + assert!(result.is_ok()); + assert!(!result.unwrap(), "garbage SAN token should invalidate game"); + let err = buffers.parse_errors[0].as_deref().unwrap(); + assert!( + err.starts_with("failed to parse SAN:") && err.contains("xyzzy9"), + "unexpected error message: {err}" + ); + // Moves stop being recorded after the first error. + assert_eq!(buffers.total_moves(), 1); + } + #[test] fn test_outcome_without_headers() { // PGN without Result header - outcome comes from movetext From 48c74f9a5c443040f2171d1f4ef26ce710c5d2ad Mon Sep 17 00:00:00 2001 From: vladkvit Date: Sat, 11 Jul 2026 08:55:43 -0400 Subject: [PATCH 5/8] Add reference dump/compare scripts to test before/after --- .gitignore | 5 + src/compare_reference.py | 40 ++++++ src/dump_reference.py | 38 ++++++ src/reference_lib.py | 287 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 370 insertions(+) create mode 100644 src/compare_reference.py create mode 100644 src/dump_reference.py create mode 100644 src/reference_lib.py diff --git a/.gitignore b/.gitignore index c8f0442..01327b8 100644 --- a/.gitignore +++ b/.gitignore @@ -70,3 +70,8 @@ docs/_build/ # Pyenv .python-version + +# Local benchmark data and derived reference dumps +*.parquet +*.pgn +reference_*.json diff --git a/src/compare_reference.py b/src/compare_reference.py new file mode 100644 index 0000000..dcac99e --- /dev/null +++ b/src/compare_reference.py @@ -0,0 +1,40 @@ +"""Compare the CURRENT parser build's outputs against a saved reference. + +Usage: python src/compare_reference.py [reference_json] + +Exits nonzero and prints mismatches if the outputs differ. +""" + +import json +import sys +import time + +from reference_lib import build_reference, diff_dicts + +DEFAULT_REFERENCE = "reference_2013-07.json" + + +def main(): + ref_path = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_REFERENCE + + with open(ref_path, encoding="utf-8") as f: + ref = json.load(f) + + t0 = time.time() + new = build_reference(ref["parquet"], subset_n=ref["subset_n"]) + print(f"Recomputed summary in {time.time() - t0:.1f}s") + + mismatches = diff_dicts(ref, new) + if mismatches: + print(f"MISMATCH: {len(mismatches)} difference(s):") + for m in mismatches[:50]: + print(f" {m}") + if len(mismatches) > 50: + print(f" ... and {len(mismatches) - 50} more") + sys.exit(1) + + print("OK: new output matches reference") + + +if __name__ == "__main__": + main() diff --git a/src/dump_reference.py b/src/dump_reference.py new file mode 100644 index 0000000..ba37432 --- /dev/null +++ b/src/dump_reference.py @@ -0,0 +1,38 @@ +"""Dump reference outputs of the CURRENT parser build to a JSON file. + +Run this against the last-known-good build (e.g. before a rearchitecture), +then verify the new build with compare_reference.py. + +Usage: python src/dump_reference.py [parquet_path] [output_json] +""" + +import json +import sys +import time + +from reference_lib import build_reference + +DEFAULT_PARQUET = "2013-07-train-00000-of-00001.parquet" +DEFAULT_OUTPUT = "reference_2013-07.json" + + +def main(): + parquet = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_PARQUET + output = sys.argv[2] if len(sys.argv) > 2 else DEFAULT_OUTPUT + + t0 = time.time() + ref = build_reference(parquet) + with open(output, "w", encoding="utf-8") as f: + json.dump(ref, f, indent=1, ensure_ascii=False) + + main_part = ref["main"] + print(f"Wrote {output} in {time.time() - t0:.1f}s") + print( + f" games={main_part['num_games']} invalid={main_part['num_invalid']} " + f"valid_moves={main_part['valid_moves_total']} " + f"valid_positions={main_part['valid_positions_total']}" + ) + + +if __name__ == "__main__": + main() diff --git a/src/reference_lib.py b/src/reference_lib.py new file mode 100644 index 0000000..f5a7912 --- /dev/null +++ b/src/reference_lib.py @@ -0,0 +1,287 @@ +"""Shared logic for dumping/comparing reference outputs of the PGN parser. + +Used by dump_reference.py (run against the OLD build to create a reference +file) and compare_reference.py (run against the NEW build to verify parity). + +Works with both the chunked API (<= 4.x: result.chunks / PyChunkView) and the +flat API (>= 5.0: global arrays directly on ParsedGames), so the exact same +digest computation runs against both builds. + +Comparison semantics for invalid games (parse_errors is not None): +- 5.0 allocates per-game array space from pass-1 token counts and zero-fills + the unwritten tail, so whole-array digests would differ by design. +- Therefore digests cover valid games only, and invalid games are recorded + individually with digests of their actually-parsed prefix slices. +""" + +import hashlib +import json + +import numpy as np + + +def _sha(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +def digest_array(arr) -> str: + a = np.ascontiguousarray(arr) + return _sha(a.tobytes()) + + +def digest_json(obj) -> str: + return _sha(json.dumps(obj, sort_keys=True, ensure_ascii=False).encode("utf-8")) + + +class Adapted: + """Uniform view over old (chunked) and new (flat) ParsedGames results.""" + + def __init__(self, result): + self.legal_moves_stored = False + if hasattr(result, "chunks"): + self._from_chunked(result) + else: + self._from_flat(result) + self.num_games = len(self.valid) + # Positions per game are always parsed moves + 1 (initial position is + # recorded even for games that error out immediately). + self.parsed_position_counts = self.parsed_move_counts + 1 + + def _from_chunked(self, result): + chunks = result.chunks + + def cat(name, axis=0): + return np.concatenate([np.asarray(getattr(c, name)) for c in chunks], axis=axis) + + def cat_list(name): + out = [] + for c in chunks: + out.extend(getattr(c, name)) + return out + + for name in ( + "boards", + "castling", + "en_passant", + "halfmove_clock", + "turn", + "from_squares", + "to_squares", + "promotions", + "clocks", + "evals", + "is_checkmate", + "is_stalemate", + "is_insufficient", + "legal_move_count", + "valid", + ): + setattr(self, name, cat(name)) + + self.headers = cat_list("headers") + self.outcome = cat_list("outcome") + self.parse_errors = cat_list("parse_errors") + self.comments = cat_list("comments") + + self.move_counts = np.concatenate( + [np.diff(np.asarray(c.move_offsets, dtype=np.int64)) for c in chunks] + ) + self.position_counts = np.concatenate( + [np.diff(np.asarray(c.position_offsets, dtype=np.int64)) for c in chunks] + ) + # Old builds only record actually-parsed moves: allocated == parsed. + self.parsed_move_counts = self.move_counts + + if any(len(np.asarray(c.legal_move_offsets)) > 0 for c in chunks): + self.legal_moves_stored = True + self.legal_move_from_squares = cat("legal_move_from_squares") + self.legal_move_to_squares = cat("legal_move_to_squares") + self.legal_move_promotions = cat("legal_move_promotions") + self.legal_move_counts = np.concatenate( + [np.diff(np.asarray(c.legal_move_offsets, dtype=np.int64)) for c in chunks] + ) + + def _from_flat(self, result): + for name in ( + "boards", + "castling", + "en_passant", + "halfmove_clock", + "turn", + "from_squares", + "to_squares", + "promotions", + "clocks", + "evals", + "is_checkmate", + "is_stalemate", + "is_insufficient", + "legal_move_count", + "valid", + ): + setattr(self, name, np.asarray(getattr(result, name))) + + self.headers = result.headers + self.outcome = result.outcome + self.parse_errors = result.parse_errors + self.comments = result.comments + + self.move_counts = np.diff(np.asarray(result.move_offsets, dtype=np.int64)) + self.position_counts = np.diff(np.asarray(result.position_offsets, dtype=np.int64)) + self.parsed_move_counts = np.asarray(result.parsed_move_counts, dtype=np.int64) + + legal_offsets = np.asarray(result.legal_move_offsets) + if len(legal_offsets) > 0: + self.legal_moves_stored = True + self.legal_move_from_squares = np.asarray(result.legal_move_from_squares) + self.legal_move_to_squares = np.asarray(result.legal_move_to_squares) + self.legal_move_promotions = np.asarray(result.legal_move_promotions) + self.legal_move_counts = np.diff(legal_offsets.astype(np.int64)) + + +def _masked_digest(arr, mask): + if mask.all(): + return digest_array(arr) + return digest_array(arr[mask]) + + +def summarize(result, include_optional=False) -> dict: + """Compute a comparable summary dict for a ParsedGames result.""" + a = Adapted(result) + n = a.num_games + valid = np.asarray(a.valid, dtype=bool) + + pos_off = np.concatenate([[0], np.cumsum(a.position_counts)]).astype(np.int64) + move_off = np.concatenate([[0], np.cumsum(a.move_counts)]).astype(np.int64) + + game_of_pos = np.repeat(np.arange(n), a.position_counts) + game_of_move = np.repeat(np.arange(n), a.move_counts) + pos_valid = valid[game_of_pos] + move_valid = valid[game_of_move] + + digests = { + # Per-position arrays (valid games only) + "boards": _masked_digest(a.boards, pos_valid), + "castling": _masked_digest(a.castling, pos_valid), + "en_passant": _masked_digest(a.en_passant, pos_valid), + "halfmove_clock": _masked_digest(a.halfmove_clock, pos_valid), + "turn": _masked_digest(a.turn, pos_valid), + # Per-move arrays (valid games only) + "from_squares": _masked_digest(a.from_squares, move_valid), + "to_squares": _masked_digest(a.to_squares, move_valid), + "promotions": _masked_digest(a.promotions, move_valid), + "clocks": _masked_digest(a.clocks, move_valid), + "evals": _masked_digest(a.evals, move_valid), + # Per-game arrays (all games; identical for old/new by design) + "is_checkmate": digest_array(a.is_checkmate), + "is_stalemate": digest_array(a.is_stalemate), + "is_insufficient": digest_array(a.is_insufficient), + "legal_move_count": digest_array(a.legal_move_count), + "valid": digest_array(valid), + "parsed_move_counts_valid": digest_array(a.parsed_move_counts[valid]), + # Per-game metadata (all games) + "headers": digest_json(a.headers), + "outcome": digest_json(a.outcome), + "parse_errors": digest_json(a.parse_errors), + } + + if include_optional: + move_valid_list = move_valid.tolist() + digests["comments"] = digest_json( + [c for c, ok in zip(a.comments, move_valid_list) if ok] + ) + if a.legal_moves_stored: + entry_valid = np.repeat(pos_valid, a.legal_move_counts) + digests["legal_move_counts"] = _masked_digest(a.legal_move_counts, pos_valid) + digests["legal_move_from_squares"] = _masked_digest( + a.legal_move_from_squares, entry_valid + ) + digests["legal_move_to_squares"] = _masked_digest( + a.legal_move_to_squares, entry_valid + ) + digests["legal_move_promotions"] = _masked_digest( + a.legal_move_promotions, entry_valid + ) + + invalid_games = [] + for idx in np.flatnonzero(~valid): + idx = int(idx) + pm = int(a.parsed_move_counts[idx]) + pp = pm + 1 + ps, ms = int(pos_off[idx]), int(move_off[idx]) + invalid_games.append( + { + "index": idx, + "parse_error": a.parse_errors[idx], + "outcome": a.outcome[idx], + "parsed_moves": pm, + "digests": { + "boards": digest_array(a.boards[ps : ps + pp]), + "castling": digest_array(a.castling[ps : ps + pp]), + "en_passant": digest_array(a.en_passant[ps : ps + pp]), + "halfmove_clock": digest_array(a.halfmove_clock[ps : ps + pp]), + "turn": digest_array(a.turn[ps : ps + pp]), + "from_squares": digest_array(a.from_squares[ms : ms + pm]), + "to_squares": digest_array(a.to_squares[ms : ms + pm]), + "promotions": digest_array(a.promotions[ms : ms + pm]), + "clocks": digest_array(a.clocks[ms : ms + pm]), + "evals": digest_array(a.evals[ms : ms + pm]), + }, + } + ) + + return { + "num_games": int(n), + "num_invalid": int((~valid).sum()), + "valid_moves_total": int(a.parsed_move_counts[valid].sum()), + "valid_positions_total": int(a.parsed_position_counts[valid].sum()), + "digests": digests, + "invalid_games": invalid_games, + } + + +def build_reference(parquet_path: str, subset_n: int = 2000) -> dict: + import pyarrow.parquet as pq + import rust_pgn_reader_python_binding as rp + + col = pq.ParquetFile(parquet_path).read(columns=["movetext"]).column("movetext") + + result = rp.parse_games(col) + main = summarize(result) + del result + + subset_col = col.slice(0, subset_n) + subset_result = rp.parse_games( + subset_col, store_comments=True, store_legal_moves=True + ) + subset = summarize(subset_result, include_optional=True) + del subset_result + + return { + "parquet": parquet_path, + "subset_n": subset_n, + "main": main, + "subset": subset, + } + + +def diff_dicts(ref: dict, new: dict, path: str = "") -> list: + """Recursively diff two summary dicts, returning mismatch descriptions.""" + mismatches = [] + if isinstance(ref, dict) and isinstance(new, dict): + for key in sorted(set(ref) | set(new)): + if key not in ref: + mismatches.append(f"{path}.{key}: missing in reference") + elif key not in new: + mismatches.append(f"{path}.{key}: missing in new output") + else: + mismatches.extend(diff_dicts(ref[key], new[key], f"{path}.{key}")) + elif isinstance(ref, list) and isinstance(new, list): + if len(ref) != len(new): + mismatches.append(f"{path}: length {len(ref)} != {len(new)}") + else: + for i, (r, n) in enumerate(zip(ref, new)): + mismatches.extend(diff_dicts(r, n, f"{path}[{i}]")) + elif ref != new: + mismatches.append(f"{path}: {ref!r} != {new!r}") + return mismatches From 37316a12dc7d3c1031e39f5dde97d405675c6151 Mon Sep 17 00:00:00 2001 From: vladkvit Date: Sat, 11 Jul 2026 09:08:23 -0400 Subject: [PATCH 6/8] Add two-pass flat-output parser (parse_games_flat) Pass 1 counts mainline SAN tokens per game (tokenizer-only, no SAN parsing) to compute exact output sizes and per-game CSR offsets. Pass 2 parses --- src/flat.rs | 988 +++++++++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 2 + src/visitor.rs | 2 +- 3 files changed, 991 insertions(+), 1 deletion(-) create mode 100644 src/flat.rs diff --git a/src/flat.rs b/src/flat.rs new file mode 100644 index 0000000..642b85e --- /dev/null +++ b/src/flat.rs @@ -0,0 +1,988 @@ +//! Two-pass parallel parsing into single exact-size flat output arrays. +//! +//! Pass 1 (parallel, per game): tokenizer-only scan counting mainline SAN +//! tokens per game. Prefix sums give exact total array sizes and per-game +//! write offsets. No SAN parsing, no board state. +//! +//! Pass 2 (parallel, work-stealing over tasks of contiguous games): full +//! parse; every game writes into its precomputed disjoint `&mut` sub-slice +//! of the final output arrays (partitioned up front with `split_at_mut`, +//! no unsafe, no atomics). Output order == input order by construction. +//! +//! Invalid games (illegal/unparsable move mid-game): pass-1 counts remain +//! authoritative for the CSR offsets; the unwritten tail of the game's +//! range is filled with sentinels (promotions/en_passant = -1, clocks/evals +//! = NaN, everything else zero) and the actually-parsed move count is +//! recorded in `parsed_move_counts`. + +use crate::board_serialization::{ + get_castling_rights, get_en_passant_file, get_halfmove_clock, get_turn, serialize_board, +}; +use crate::comment_parsing::{CommentContent, ParsedTag, parse_comments}; +use crate::tokenizer::{self, Outcome, Visitor}; +use crate::visitor::{ParseConfig, castling_with_zeros}; +use rayon::ThreadPoolBuilder; +use rayon::prelude::*; +use shakmaty::{CastlingMode, Chess, Color, Position, fen::Fen, san::SanPlus, uci::UciMove}; +use std::collections::HashMap; + +/// Number of games per pass-2 task. Small enough that rayon work stealing +/// evens out length imbalance between tasks (no straggler tail), large +/// enough that per-task overhead is negligible. +const GAMES_PER_TASK: usize = 256; + +/// Final flat output of a parse run. All arrays are exact-size, +/// allocated once, and indexed by the CSR offset arrays. +#[derive(Default)] +pub struct FlatOutput { + // Per-position arrays (total_positions entries) + pub boards: Vec, // 64 bytes per position + pub castling: Vec, // 4 bools per position [K,Q,k,q] + pub en_passant: Vec, // -1 or file 0-7 + pub halfmove_clock: Vec, // per position + pub turn: Vec, // true = white + + // Per-move arrays (total_moves entries) + pub from_squares: Vec, + pub to_squares: Vec, + pub promotions: Vec, // -1 for no promotion + pub clocks: Vec, // NaN for missing + pub evals: Vec, // NaN for missing + + // CSR offsets (num_games + 1 entries), from pass-1 counts + pub move_offsets: Vec, + pub position_offsets: Vec, + + // Per-game arrays (num_games entries) + pub parsed_move_counts: Vec, // == offsets diff for valid games + pub is_checkmate: Vec, + pub is_stalemate: Vec, + pub is_insufficient: Vec, // 2 bools per game [white, black] + pub legal_move_count: Vec, + pub valid: Vec, + pub headers: Vec>, + pub outcome: Vec>, + pub parse_errors: Vec>, + + // Optional: per-move comments (only when store_comments) + pub comments: Vec>, + + // Optional: legal moves per position (only when store_legal_moves) + pub legal_move_from_squares: Vec, + pub legal_move_to_squares: Vec, + pub legal_move_promotions: Vec, + pub legal_move_offsets: Vec, // total_positions + 1 entries + + pub num_games: usize, + pub total_moves: usize, + pub total_positions: usize, +} + +// --------------------------------------------------------------------------- +// Pass 1: counting +// --------------------------------------------------------------------------- + +struct CountingVisitor { + moves: u32, +} + +impl Visitor for CountingVisitor { + fn begin_headers(&mut self) {} + fn header(&mut self, _key: &[u8], _value: &[u8]) {} + fn end_headers(&mut self) {} + fn san_token(&mut self, _token: &[u8]) { + self.moves += 1; + } + fn comment(&mut self, _comment: &[u8]) {} + fn outcome(&mut self, _outcome: Outcome) {} + fn end_game(&mut self) {} +} + +/// Count mainline SAN tokens in one game (upper bound on parsed moves; +/// exact for games without parse errors). +fn count_san_tokens(pgn: &str) -> u32 { + let mut counter = CountingVisitor { moves: 0 }; + tokenizer::parse_game(pgn.as_bytes(), &mut counter); + counter.moves +} + +// --------------------------------------------------------------------------- +// Pass 2: slice-writing visitor +// --------------------------------------------------------------------------- + +/// Disjoint mutable sub-slices of the output arrays for one task +/// (a contiguous run of games). +struct TaskOutput<'a> { + /// Pass-1 allocated move counts for this task's games. + move_counts_alloc: &'a [u32], + + boards: &'a mut [u8], + castling: &'a mut [bool], + en_passant: &'a mut [i8], + halfmove_clock: &'a mut [u8], + turn: &'a mut [bool], + + from_squares: &'a mut [u8], + to_squares: &'a mut [u8], + promotions: &'a mut [i8], + clocks: &'a mut [f32], + evals: &'a mut [f32], + comments: &'a mut [Option], // empty when store_comments=false + + parsed_move_counts: &'a mut [u32], + is_checkmate: &'a mut [bool], + is_stalemate: &'a mut [bool], + is_insufficient: &'a mut [bool], + legal_move_count: &'a mut [u16], + valid: &'a mut [bool], + headers: &'a mut [HashMap], + outcome: &'a mut [Option], + parse_errors: &'a mut [Option], +} + +struct Task<'a> { + out: TaskOutput<'a>, + games: &'a [&'a str], +} + +/// Task-local legal-move accumulation (sizes unknowable in pass 1; +/// concatenated after pass 2). `counts` has one entry per allocated +/// position in the task, including zero entries for unwritten tails. +#[derive(Default)] +struct TaskLegal { + counts: Vec, + from: Vec, + to: Vec, + promotions: Vec, +} + +/// Visitor writing one task's games into its output slices. +struct FlatVisitor<'a> { + out: TaskOutput<'a>, + store_comments: bool, + store_legal_moves: bool, + pos: Chess, + valid_moves: bool, + current_headers: Vec<(String, String)>, + current_outcome: Option, + current_error: Option, + status_written: bool, + /// Index of the current game within the task. + game_idx: usize, + /// Task-local start of the current game's allocated ranges. + move_base: usize, + pos_base: usize, + /// Task-local write cursors. + move_cursor: usize, + pos_cursor: usize, + legal: TaskLegal, +} + +impl<'a> FlatVisitor<'a> { + fn new(out: TaskOutput<'a>, config: &ParseConfig) -> Self { + let legal = if config.store_legal_moves { + TaskLegal { + counts: Vec::with_capacity(out.en_passant.len()), + from: Vec::with_capacity(out.en_passant.len() * 30), + to: Vec::with_capacity(out.en_passant.len() * 30), + promotions: Vec::with_capacity(out.en_passant.len() * 30), + } + } else { + TaskLegal::default() + }; + FlatVisitor { + out, + store_comments: config.store_comments, + store_legal_moves: config.store_legal_moves, + pos: Chess::default(), + valid_moves: true, + current_headers: Vec::with_capacity(10), + current_outcome: None, + current_error: None, + status_written: false, + game_idx: 0, + move_base: 0, + pos_base: 0, + move_cursor: 0, + pos_cursor: 0, + legal, + } + } + + fn finish(self) -> TaskLegal { + debug_assert_eq!(self.game_idx, self.out.move_counts_alloc.len()); + debug_assert_eq!(self.move_cursor, self.out.from_squares.len()); + debug_assert_eq!(self.pos_cursor, self.out.en_passant.len()); + self.legal + } + + fn push_board_state(&mut self) { + let p = self.pos_cursor; + self.out.boards[p * 64..(p + 1) * 64].copy_from_slice(&serialize_board(&self.pos)); + self.out.castling[p * 4..(p + 1) * 4].copy_from_slice(&get_castling_rights(&self.pos)); + self.out.en_passant[p] = get_en_passant_file(&self.pos); + self.out.halfmove_clock[p] = get_halfmove_clock(&self.pos); + self.out.turn[p] = get_turn(&self.pos); + self.pos_cursor += 1; + + if self.store_legal_moves { + self.push_legal_moves(); + } + } + + fn push_legal_moves(&mut self) { + let legal_moves = self.pos.legal_moves(); + let mut count: u32 = 0; + for m in legal_moves { + let uci_move_obj = UciMove::from_standard(m); + if let UciMove::Normal { + from, + to, + promotion, + } = uci_move_obj + { + self.legal.from.push(from as u8); + self.legal.to.push(to as u8); + self.legal + .promotions + .push(promotion.map(|p| p as i8).unwrap_or(-1)); + count += 1; + } + } + self.legal.counts.push(count); + } + + fn push_move(&mut self, from: u8, to: u8, promotion: Option) { + let m = self.move_cursor; + self.out.from_squares[m] = from; + self.out.to_squares[m] = to; + self.out.promotions[m] = promotion.map(|p| p as i8).unwrap_or(-1); + self.out.clocks[m] = f32::NAN; + self.out.evals[m] = f32::NAN; + // comments[m] is already None from allocation + self.move_cursor += 1; + } + + fn update_position_status(&mut self) { + let gi = self.game_idx; + self.out.is_checkmate[gi] = self.pos.is_checkmate(); + self.out.is_stalemate[gi] = self.pos.is_stalemate(); + self.out.is_insufficient[gi * 2] = self.pos.has_insufficient_material(Color::White); + self.out.is_insufficient[gi * 2 + 1] = self.pos.has_insufficient_material(Color::Black); + self.out.legal_move_count[gi] = self.pos.legal_moves().len() as u16; + self.status_written = true; + } + + fn set_error(&mut self, msg: String) { + self.valid_moves = false; + self.current_error = Some(msg); + } +} + +impl Visitor for FlatVisitor<'_> { + fn begin_headers(&mut self) { + self.current_headers.clear(); + self.valid_moves = true; + self.current_outcome = None; + self.current_error = None; + self.status_written = false; + } + + fn header(&mut self, key: &[u8], value: &[u8]) { + let key_str = String::from_utf8_lossy(key).into_owned(); + let value_str = String::from_utf8_lossy(value).into_owned(); + self.current_headers.push((key_str, value_str)); + } + + fn end_headers(&mut self) { + // Determine castling mode from Variant header (case-insensitive) + let castling_mode = self + .current_headers + .iter() + .find(|(k, _)| k.eq_ignore_ascii_case("Variant")) + .and_then(|(_, v)| { + let v_lower = v.to_lowercase(); + if v_lower == "chess960" { + Some(CastlingMode::Chess960) + } else { + None + } + }) + .unwrap_or(CastlingMode::Standard); + + // Try to parse FEN from headers, fall back to default position + let fen_header = self + .current_headers + .iter() + .find(|(k, _)| k.eq_ignore_ascii_case("FEN")) + .map(|(_, v)| v.as_str()); + + if let Some(fen_str) = fen_header { + match fen_str.parse::() { + Ok(fen) => match fen.into_position(castling_mode) { + Ok(pos) => self.pos = pos, + Err(e) => { + self.set_error(format!("invalid FEN position: {}", e)); + self.pos = Chess::default(); + } + }, + Err(e) => { + self.set_error(format!("failed to parse FEN: {}", e)); + self.pos = Chess::default(); + } + } + } else { + self.pos = Chess::default(); + } + + // Record initial board state + self.push_board_state(); + } + + fn san_token(&mut self, token: &[u8]) { + // Early abort: once the game is invalid, skip SAN parsing entirely. + if !self.valid_moves { + return; + } + + let san_plus = match castling_with_zeros(token) { + Some(san_plus) => san_plus, + None => match SanPlus::from_ascii(token) { + Ok(san_plus) => san_plus, + Err(error) => { + let token_str = String::from_utf8_lossy(token); + self.set_error(format!("failed to parse SAN: {} ({})", error, token_str)); + return; + } + }, + }; + + match san_plus.san.to_move(&self.pos) { + Ok(m) => { + self.pos.play_unchecked(m); + + // Record board state after move + self.push_board_state(); + + let uci_move_obj = UciMove::from_standard(m); + match uci_move_obj { + UciMove::Normal { + from, + to, + promotion, + } => { + self.push_move(from as u8, to as u8, promotion.map(|p| p as u8)); + } + _ => { + self.set_error(format!("unexpected UCI move type: {:?}", uci_move_obj)); + } + } + } + Err(err) => { + self.set_error(format!("illegal move: {} {}", err, san_plus)); + } + } + } + + fn comment(&mut self, comment: &[u8]) { + // Comments annotate the last-written move of the CURRENT game; a + // comment before the game's first move has no slot to update. + // (The old Buffers code could touch the previous game's last slot + // in that case — a chunk-boundary-dependent quirk, deliberately + // fixed here.) + if self.move_cursor == self.move_base { + return; + } + let last_move = self.move_cursor - 1; + + if let Ok((_, parsed_comments)) = parse_comments(comment) { + let mut move_comments = if self.store_comments { + Some(String::new()) + } else { + None + }; + + for content in parsed_comments { + match content { + CommentContent::Tag(tag_content) => match tag_content { + ParsedTag::Eval(eval_value) => { + self.out.evals[last_move] = eval_value as f32; + } + ParsedTag::ClkTime { + hours, + minutes, + seconds, + } => { + self.out.clocks[last_move] = + hours as f32 * 3600.0 + minutes as f32 * 60.0 + seconds as f32; + } + ParsedTag::Mate(mate_value) => { + // Mate scores stored as text in comments (matching old API behavior) + if let Some(ref mut comments) = move_comments { + if !comments.is_empty() && !comments.ends_with(' ') { + comments.push(' '); + } + comments.push_str(&format!("[Mate {}]", mate_value)); + } + } + }, + CommentContent::Text(text) => { + if let Some(ref mut comments) = move_comments + && !text.trim().is_empty() + { + if !comments.is_empty() { + comments.push(' '); + } + comments.push_str(&text); + } + } + } + } + + if let Some(comment_text) = move_comments { + self.out.comments[last_move] = Some(comment_text); + } + } + } + + fn outcome(&mut self, outcome: Outcome) { + self.current_outcome = Some(match outcome { + Outcome::WhiteWins => "White".to_string(), + Outcome::BlackWins => "Black".to_string(), + Outcome::Draw => "Draw".to_string(), + Outcome::Unknown => "Unknown".to_string(), + }); + self.update_position_status(); + } + + fn end_game(&mut self) { + if !self.status_written { + self.update_position_status(); + } + + let gi = self.game_idx; + let parsed_moves = (self.move_cursor - self.move_base) as u32; + self.out.parsed_move_counts[gi] = parsed_moves; + self.out.valid[gi] = self.valid_moves; + self.out.outcome[gi] = self.current_outcome.take(); + self.out.parse_errors[gi] = self.current_error.take(); + self.out.headers[gi] = self.current_headers.drain(..).collect(); + + // Fill the unwritten tail of the allocated ranges with sentinels + // (only non-empty for games that errored mid-parse). + let alloc_moves = self.out.move_counts_alloc[gi] as usize; + let alloc_positions = alloc_moves + 1; + for m in self.move_cursor..self.move_base + alloc_moves { + self.out.promotions[m] = -1; + self.out.clocks[m] = f32::NAN; + self.out.evals[m] = f32::NAN; + } + for p in self.pos_cursor..self.pos_base + alloc_positions { + self.out.en_passant[p] = -1; + } + if self.store_legal_moves { + for _ in self.pos_cursor..self.pos_base + alloc_positions { + self.legal.counts.push(0); + } + } + + // Advance to the next game's allocated ranges. + self.move_base += alloc_moves; + self.pos_base += alloc_positions; + self.move_cursor = self.move_base; + self.pos_cursor = self.pos_base; + self.game_idx += 1; + } +} + +// --------------------------------------------------------------------------- +// Driver +// --------------------------------------------------------------------------- + +fn prefix_sum_u32(counts: impl Iterator) -> Vec { + let mut offsets = Vec::with_capacity(counts.size_hint().0 + 1); + let mut acc: u32 = 0; + offsets.push(0); + for count in counts { + acc += count; + offsets.push(acc); + } + offsets +} + +/// Split `head_len` elements off the front of a mutable rest-slice. +macro_rules! split_mut { + ($rest:expr, $len:expr) => {{ + let (head, tail) = std::mem::take(&mut $rest).split_at_mut($len); + $rest = tail; + head + }}; +} + +/// Split `head_len` elements off the front of a shared rest-slice. +macro_rules! split_ref { + ($rest:expr, $len:expr) => {{ + let (head, tail) = $rest.split_at($len); + $rest = tail; + head + }}; +} + +/// Parse all games in `slices` (one game per string; empty/whitespace-only +/// strings are skipped) into exact-size flat arrays using two passes. +pub fn parse_games_flat( + slices: &[&str], + num_threads: usize, + config: &ParseConfig, +) -> Result { + // Skip empty inputs (parity with the previous per-game "No game found + // in PGN" behavior, where such inputs produced no game entry). + let games: Vec<&str> = slices + .iter() + .copied() + .filter(|s| !s.trim().is_empty()) + .collect(); + let n = games.len(); + + let pool = ThreadPoolBuilder::new() + .num_threads(num_threads) + .build() + .map_err(|e| format!("Failed to build thread pool: {}", e))?; + + // ---- Pass 1: count mainline SAN tokens per game ---- + let move_counts: Vec = + pool.install(|| games.par_iter().map(|g| count_san_tokens(g)).collect()); + + let move_offsets = prefix_sum_u32(move_counts.iter().copied()); + let position_offsets = prefix_sum_u32(move_counts.iter().map(|&c| c + 1)); + let total_moves = move_offsets[n] as usize; + let total_positions = position_offsets[n] as usize; + + // ---- Allocate final arrays once, at exact size ---- + // vec![0; n] uses alloc_zeroed: no memset for large allocations; the + // first-touch page faults happen on the worker threads during pass 2. + let mut out = FlatOutput { + boards: vec![0; total_positions * 64], + castling: vec![false; total_positions * 4], + en_passant: vec![0; total_positions], + halfmove_clock: vec![0; total_positions], + turn: vec![false; total_positions], + from_squares: vec![0; total_moves], + to_squares: vec![0; total_moves], + promotions: vec![0; total_moves], + clocks: vec![0.0; total_moves], + evals: vec![0.0; total_moves], + move_offsets, + position_offsets, + parsed_move_counts: vec![0; n], + is_checkmate: vec![false; n], + is_stalemate: vec![false; n], + is_insufficient: vec![false; n * 2], + legal_move_count: vec![0; n], + valid: vec![false; n], + headers: (0..n).map(|_| HashMap::new()).collect(), + outcome: vec![None; n], + parse_errors: vec![None; n], + comments: if config.store_comments { + vec![None; total_moves] + } else { + Vec::new() + }, + legal_move_from_squares: Vec::new(), + legal_move_to_squares: Vec::new(), + legal_move_promotions: Vec::new(), + legal_move_offsets: Vec::new(), + num_games: n, + total_moves, + total_positions, + }; + + // ---- Partition the output arrays into per-task disjoint slices ---- + let mut tasks: Vec = Vec::with_capacity(n.div_ceil(GAMES_PER_TASK.max(1))); + { + let mut boards_rest: &mut [u8] = &mut out.boards; + let mut castling_rest: &mut [bool] = &mut out.castling; + let mut en_passant_rest: &mut [i8] = &mut out.en_passant; + let mut halfmove_rest: &mut [u8] = &mut out.halfmove_clock; + let mut turn_rest: &mut [bool] = &mut out.turn; + let mut from_rest: &mut [u8] = &mut out.from_squares; + let mut to_rest: &mut [u8] = &mut out.to_squares; + let mut promo_rest: &mut [i8] = &mut out.promotions; + let mut clocks_rest: &mut [f32] = &mut out.clocks; + let mut evals_rest: &mut [f32] = &mut out.evals; + let mut comments_rest: &mut [Option] = &mut out.comments; + let mut parsed_rest: &mut [u32] = &mut out.parsed_move_counts; + let mut checkmate_rest: &mut [bool] = &mut out.is_checkmate; + let mut stalemate_rest: &mut [bool] = &mut out.is_stalemate; + let mut insufficient_rest: &mut [bool] = &mut out.is_insufficient; + let mut legal_count_rest: &mut [u16] = &mut out.legal_move_count; + let mut valid_rest: &mut [bool] = &mut out.valid; + let mut headers_rest: &mut [HashMap] = &mut out.headers; + let mut outcome_rest: &mut [Option] = &mut out.outcome; + let mut errors_rest: &mut [Option] = &mut out.parse_errors; + let mut counts_rest: &[u32] = &move_counts; + let mut games_rest: &[&str] = &games; + + let mut g = 0; + while g < n { + let g_end = (g + GAMES_PER_TASK).min(n); + let task_games = g_end - g; + let task_moves = (out.move_offsets[g_end] - out.move_offsets[g]) as usize; + let task_positions = task_moves + task_games; + + let task = Task { + out: TaskOutput { + move_counts_alloc: split_ref!(counts_rest, task_games), + boards: split_mut!(boards_rest, task_positions * 64), + castling: split_mut!(castling_rest, task_positions * 4), + en_passant: split_mut!(en_passant_rest, task_positions), + halfmove_clock: split_mut!(halfmove_rest, task_positions), + turn: split_mut!(turn_rest, task_positions), + from_squares: split_mut!(from_rest, task_moves), + to_squares: split_mut!(to_rest, task_moves), + promotions: split_mut!(promo_rest, task_moves), + clocks: split_mut!(clocks_rest, task_moves), + evals: split_mut!(evals_rest, task_moves), + comments: if config.store_comments { + split_mut!(comments_rest, task_moves) + } else { + &mut [] + }, + parsed_move_counts: split_mut!(parsed_rest, task_games), + is_checkmate: split_mut!(checkmate_rest, task_games), + is_stalemate: split_mut!(stalemate_rest, task_games), + is_insufficient: split_mut!(insufficient_rest, task_games * 2), + legal_move_count: split_mut!(legal_count_rest, task_games), + valid: split_mut!(valid_rest, task_games), + headers: split_mut!(headers_rest, task_games), + outcome: split_mut!(outcome_rest, task_games), + parse_errors: split_mut!(errors_rest, task_games), + }, + games: split_ref!(games_rest, task_games), + }; + tasks.push(task); + g = g_end; + } + + // ---- Pass 2: full parse, work-stealing over tasks ---- + let legal_parts: Vec = pool.install(|| { + tasks + .into_par_iter() + .map(|task| { + let mut visitor = FlatVisitor::new(task.out, config); + for &pgn in task.games { + tokenizer::parse_game(pgn.as_bytes(), &mut visitor); + } + visitor.finish() + }) + .collect() + }); + + // ---- Concatenate task-local legal-move arrays (optional path) ---- + if config.store_legal_moves { + let total_legal: usize = legal_parts.iter().map(|p| p.from.len()).sum(); + let total_counts: usize = legal_parts.iter().map(|p| p.counts.len()).sum(); + debug_assert_eq!(total_counts, total_positions); + + out.legal_move_from_squares = Vec::with_capacity(total_legal); + out.legal_move_to_squares = Vec::with_capacity(total_legal); + out.legal_move_promotions = Vec::with_capacity(total_legal); + let mut legal_counts: Vec = Vec::with_capacity(total_counts); + for part in &legal_parts { + out.legal_move_from_squares.extend_from_slice(&part.from); + out.legal_move_to_squares.extend_from_slice(&part.to); + out.legal_move_promotions + .extend_from_slice(&part.promotions); + legal_counts.extend_from_slice(&part.counts); + } + out.legal_move_offsets = prefix_sum_u32(legal_counts.into_iter()); + } + } + + Ok(out) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::visitor::{Buffers, parse_game_to_buffers}; + + /// Differential check: the flat two-pass path must produce the same + /// per-game data as the Buffers path for every game. For invalid games + /// the flat arrays are compared over the actually-parsed prefix and the + /// allocated tail is checked for sentinels. + fn assert_flat_matches_buffers(pgns: &[&str], config: &ParseConfig, num_threads: usize) { + let flat = parse_games_flat(pgns, num_threads, config).unwrap(); + + let mut buffers = Buffers::with_capacity(pgns.len().max(1), 70, config); + for &p in pgns { + let _ = parse_game_to_buffers(p, &mut buffers, config); + } + + assert_eq!(flat.num_games, buffers.num_games(), "game count"); + let n = flat.num_games; + assert_eq!(flat.move_offsets.len(), n + 1); + assert_eq!(flat.position_offsets.len(), n + 1); + + let b_move_off = buffers.compute_move_offsets(); + let b_pos_off = buffers.compute_position_offsets(); + let b_legal_off = if config.store_legal_moves { + buffers.compute_legal_move_offsets() + } else { + Vec::new() + }; + + for i in 0..n { + let parsed = flat.parsed_move_counts[i] as usize; + assert_eq!( + parsed as u32, buffers.move_counts[i], + "game {i}: parsed moves" + ); + assert_eq!(flat.valid[i], buffers.valid[i], "game {i}: valid"); + assert_eq!(flat.headers[i], buffers.headers[i], "game {i}: headers"); + assert_eq!(flat.outcome[i], buffers.outcome[i], "game {i}: outcome"); + assert_eq!( + flat.parse_errors[i], buffers.parse_errors[i], + "game {i}: parse_errors" + ); + assert_eq!(flat.is_checkmate[i], buffers.is_checkmate[i], "game {i}"); + assert_eq!(flat.is_stalemate[i], buffers.is_stalemate[i], "game {i}"); + assert_eq!( + flat.is_insufficient[i * 2..i * 2 + 2], + buffers.is_insufficient[i * 2..i * 2 + 2], + "game {i}: is_insufficient" + ); + assert_eq!( + flat.legal_move_count[i], buffers.legal_move_count[i], + "game {i}: legal_move_count" + ); + + // --- Moves: parsed prefix identical, tail sentinel-filled --- + let fm = flat.move_offsets[i] as usize; + let bm = b_move_off[i] as usize; + let alloc = (flat.move_offsets[i + 1] - flat.move_offsets[i]) as usize; + assert!(parsed <= alloc, "game {i}: parsed > allocated"); + + assert_eq!( + &flat.from_squares[fm..fm + parsed], + &buffers.from_squares[bm..bm + parsed], + "game {i}: from_squares" + ); + assert_eq!( + &flat.to_squares[fm..fm + parsed], + &buffers.to_squares[bm..bm + parsed], + "game {i}: to_squares" + ); + assert_eq!( + &flat.promotions[fm..fm + parsed], + &buffers.promotions[bm..bm + parsed], + "game {i}: promotions" + ); + for k in 0..parsed { + assert_eq!( + flat.clocks[fm + k].to_bits(), + buffers.clocks[bm + k].to_bits(), + "game {i} move {k}: clocks" + ); + assert_eq!( + flat.evals[fm + k].to_bits(), + buffers.evals[bm + k].to_bits(), + "game {i} move {k}: evals" + ); + } + if config.store_comments { + assert_eq!( + &flat.comments[fm..fm + parsed], + &buffers.comments[bm..bm + parsed], + "game {i}: comments" + ); + } + for k in parsed..alloc { + assert_eq!(flat.from_squares[fm + k], 0, "game {i}: tail from"); + assert_eq!(flat.to_squares[fm + k], 0, "game {i}: tail to"); + assert_eq!(flat.promotions[fm + k], -1, "game {i}: tail promo"); + assert!(flat.clocks[fm + k].is_nan(), "game {i}: tail clock"); + assert!(flat.evals[fm + k].is_nan(), "game {i}: tail eval"); + if config.store_comments { + assert_eq!(flat.comments[fm + k], None, "game {i}: tail comment"); + } + } + + // --- Positions: parsed prefix identical, tail sentinel-filled --- + let parsed_pos = parsed + 1; + assert_eq!( + parsed_pos as u32, buffers.position_counts[i], + "game {i}: position count" + ); + let fp = flat.position_offsets[i] as usize; + let bp = b_pos_off[i] as usize; + let alloc_pos = (flat.position_offsets[i + 1] - flat.position_offsets[i]) as usize; + assert_eq!(alloc_pos, alloc + 1); + + assert_eq!( + &flat.boards[fp * 64..(fp + parsed_pos) * 64], + &buffers.boards[bp * 64..(bp + parsed_pos) * 64], + "game {i}: boards" + ); + assert_eq!( + &flat.castling[fp * 4..(fp + parsed_pos) * 4], + &buffers.castling[bp * 4..(bp + parsed_pos) * 4], + "game {i}: castling" + ); + assert_eq!( + &flat.en_passant[fp..fp + parsed_pos], + &buffers.en_passant[bp..bp + parsed_pos], + "game {i}: en_passant" + ); + assert_eq!( + &flat.halfmove_clock[fp..fp + parsed_pos], + &buffers.halfmove_clock[bp..bp + parsed_pos], + "game {i}: halfmove_clock" + ); + assert_eq!( + &flat.turn[fp..fp + parsed_pos], + &buffers.turn[bp..bp + parsed_pos], + "game {i}: turn" + ); + for k in parsed_pos..alloc_pos { + assert_eq!(flat.en_passant[fp + k], -1, "game {i}: tail en_passant"); + assert_eq!( + &flat.boards[(fp + k) * 64..(fp + k + 1) * 64], + &[0u8; 64], + "game {i}: tail board" + ); + } + + // --- Legal moves per position (optional) --- + if config.store_legal_moves { + for k in 0..parsed_pos { + let fs = flat.legal_move_offsets[fp + k] as usize; + let fe = flat.legal_move_offsets[fp + k + 1] as usize; + let bs = b_legal_off[bp + k] as usize; + let be = b_legal_off[bp + k + 1] as usize; + assert_eq!(fe - fs, be - bs, "game {i} pos {k}: legal count"); + assert_eq!( + &flat.legal_move_from_squares[fs..fe], + &buffers.legal_move_from_squares[bs..be], + "game {i} pos {k}: legal from" + ); + assert_eq!( + &flat.legal_move_to_squares[fs..fe], + &buffers.legal_move_to_squares[bs..be], + "game {i} pos {k}: legal to" + ); + assert_eq!( + &flat.legal_move_promotions[fs..fe], + &buffers.legal_move_promotions[bs..be], + "game {i} pos {k}: legal promo" + ); + } + for k in parsed_pos..alloc_pos { + assert_eq!( + flat.legal_move_offsets[fp + k], + flat.legal_move_offsets[fp + k + 1], + "game {i}: tail position must have zero legal moves" + ); + } + } + } + } + + fn default_config() -> ParseConfig { + ParseConfig { + store_comments: false, + store_legal_moves: false, + } + } + + fn full_config() -> ParseConfig { + ParseConfig { + store_comments: true, + store_legal_moves: true, + } + } + + /// Corpus exercising every tricky path: comments with eval/clk/mate + /// tags, empty inputs, FEN starts, Chess960, invalid SAN, illegal + /// moves, invalid FEN, variations, zeros castling, trailing games, + /// escape lines, unterminated comments, headers-only, no headers. + const TRICKY: &[&str] = &[ + "[Event \"A\"]\n[White \"x\"]\n[Black \"y\"]\n\n1. e4 { [%eval 0.17] [%clk 0:03:00] } 1... e5 { [%eval 0.19] [%clk 0:02:58] } 2. Nf3 { [%eval -0.2] } Nc6 { some text } 1-0", + "1. e4 e5 2. Nf3 Nc6 1-0", + "", + " \n\t ", + "[Event \"HeadersOnly\"]\n[Site \"nowhere\"]\n", + "[FEN \"r1bqkbnr/pppp1ppp/2n5/4p3/4P3/5N2/PPPP1PPP/RNBQKB1R w KQkq - 2 3\"]\n\n3. Bb5 a6 4. Ba4 Nf6 1-0", + "[Variant \"chess960\"]\n[FEN \"brkrqnnb/pppppppp/8/8/8/8/PPPPPPPP/BRKRQNNB w KQkq - 0 1\"]\n\n1. g3 d5 2. d4 g6 1-0", + "1. e4 xyzzy9 e5 Nf3 1-0", + "1. e4 e4 2. Nf3 1-0", + "[FEN \"invalid fen string\"]\n\n1. e4 e5 1-0", + "1. e4 (1. d4 {inner} d5 (1... Nf6 2. c4) 2. c4) e5 2. Nf3 1-0", + "1. e4 e5 2. Nf3 Nf6 3. Bc4 Bc5 4. 0-0 0-0 { [%clk 0:01:00] } 1/2-1/2", + "[Event \"a\"]\n\n1. e4 1-0\n\n[Event \"b\"]\n\n1. d4 0-1", + "% escape line\n[Event \"esc\"]\n\n1. e4 c5 { [%eval 0.3] } *", + "1. e4 {unterminated comment", + "1. f3 e5 2. g4 Qh4# 0-1", + "[Event \"Mate tag\"]\n\n1. e4 e5 { [%eval #12] } *", + ]; + + #[test] + fn test_flat_matches_buffers_tricky() { + assert_flat_matches_buffers(TRICKY, &default_config(), 2); + } + + #[test] + fn test_flat_matches_buffers_comments_and_legal_moves() { + assert_flat_matches_buffers(TRICKY, &full_config(), 2); + } + + #[test] + fn test_flat_matches_buffers_single_threaded() { + assert_flat_matches_buffers(TRICKY, &default_config(), 1); + } + + #[test] + fn test_flat_matches_buffers_many_games_multi_task() { + // More games than GAMES_PER_TASK so pass 2 spans multiple tasks, + // with invalid games sprinkled across task boundaries. + let mut pgns: Vec<&str> = Vec::new(); + for i in 0..(3 * GAMES_PER_TASK + 17) { + pgns.push(TRICKY[i % TRICKY.len()]); + } + assert_flat_matches_buffers(&pgns, &default_config(), 4); + assert_flat_matches_buffers(&pgns, &full_config(), 4); + } + + #[test] + fn test_pass1_count_is_upper_bound() { + let config = default_config(); + // Invalid SAN mid-game: 4 tokens counted, 1 move parsed. + let flat = parse_games_flat(&["1. e4 xyzzy9 e5 Nf3 1-0"], 1, &config).unwrap(); + assert_eq!(flat.num_games, 1); + assert_eq!(flat.move_offsets[1] - flat.move_offsets[0], 4); + assert_eq!(flat.parsed_move_counts[0], 1); + assert!(!flat.valid[0]); + + // Valid game: count == parsed. + let flat = parse_games_flat(&["1. e4 e5 2. Nf3 Nc6 1-0"], 1, &config).unwrap(); + assert_eq!(flat.move_offsets[1] - flat.move_offsets[0], 4); + assert_eq!(flat.parsed_move_counts[0], 4); + assert!(flat.valid[0]); + } + + #[test] + fn test_empty_input() { + let flat = parse_games_flat(&[], 1, &default_config()).unwrap(); + assert_eq!(flat.num_games, 0); + assert_eq!(flat.move_offsets, vec![0]); + assert_eq!(flat.position_offsets, vec![0]); + assert!(flat.boards.is_empty()); + } + + #[test] + fn test_empty_strings_skipped() { + let flat = parse_games_flat(&["", "1. e4 e5 1-0", " \n "], 1, &default_config()).unwrap(); + assert_eq!(flat.num_games, 1); + assert_eq!(flat.parsed_move_counts[0], 2); + } +} diff --git a/src/lib.rs b/src/lib.rs index d229898..5af15f0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,10 +7,12 @@ use rayon::prelude::*; mod board_serialization; mod comment_parsing; +mod flat; mod python_bindings; mod tokenizer; mod visitor; +pub use flat::{FlatOutput, parse_games_flat}; use python_bindings::{ChunkData, ParsedGames, ParsedGamesIter, PyChunkView, PyGameView}; pub use visitor::{Buffers, ParseConfig, parse_game_to_buffers}; diff --git a/src/visitor.rs b/src/visitor.rs index e70fbc4..5605a73 100644 --- a/src/visitor.rs +++ b/src/visitor.rs @@ -194,7 +194,7 @@ impl Buffers { /// Support castling notated with zeros ("0-0", "0-0-0"), with optional /// check/checkmate suffix, matching pgn-reader's explicit handling. -fn castling_with_zeros(token: &[u8]) -> Option { +pub(crate) fn castling_with_zeros(token: &[u8]) -> Option { let (body, suffix) = match token.last() { Some(b'+') => (&token[..token.len() - 1], Some(Suffix::Check)), Some(b'#') => (&token[..token.len() - 1], Some(Suffix::Checkmate)), From 5692e92002e5ad6fa8401233053b675b214729f6 Mon Sep 17 00:00:00 2001 From: vladkvit Date: Sat, 11 Jul 2026 09:23:54 -0400 Subject: [PATCH 7/8] Rearchitect output to single exact-size flat arrays Two-pass parsing replaces the per-chunk Buffers architecture: pass 1 counts SAN tokens per game to size and offset every output array exactly; pass 2 parses in parallel over small game tasks with rayon work stealing, each game writing into its precomputed disjoint slice. No merge step, no reallocs, no per-thread straggler tail. The GIL is released during both passes. Breaking Python API changes: - ParsedGames now exposes flat global arrays directly (boards, from_squares, move_offsets, ..., parsed_move_counts); .chunks, num_chunks, PyChunkView and the chunk_multiplier parameter are gone. - Invalid games keep their pre-counted array range with sentinel- filled tails; parsed_move_counts records actual moves. PyGameView still exposes only the parsed prefix, matching 4.x semantics. - Buffers/GameVisitor (src/visitor.rs) deleted; core entry point is parse_games_flat, shared by the pyo3 layer and the native bench. Validation: byte-identical outputs vs 4.0.0 on the lichess 2013-07 corpus (293k games, compare_reference.py), cargo test + test.py green. --- Cargo.lock | 2 +- Cargo.toml | 2 +- benches/parquet_bench.rs | 69 +-- rust_pgn_reader_python_binding.pyi | 193 +++--- src/bench_data_access.py | 42 +- src/example_parse_games.py | 30 +- src/flat.rs | 639 +++++++++++++------- src/lib.rs | 317 +++------- src/python_bindings.rs | 510 ++++++---------- src/test.py | 116 ++-- src/visitor.rs | 904 ----------------------------- 11 files changed, 953 insertions(+), 1871 deletions(-) delete mode 100644 src/visitor.rs diff --git a/Cargo.lock b/Cargo.lock index 320a5b3..ab5e8f7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1306,7 +1306,7 @@ checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "rust_pgn_reader_python_binding" -version = "4.0.0" +version = "5.0.0" dependencies = [ "arrow", "arrow-array", diff --git a/Cargo.toml b/Cargo.toml index d38f22e..e02ed97 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "rust_pgn_reader_python_binding" -version = "4.0.0" +version = "5.0.0" edition = "2024" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html diff --git a/benches/parquet_bench.rs b/benches/parquet_bench.rs index fd59e19..c465a27 100644 --- a/benches/parquet_bench.rs +++ b/benches/parquet_bench.rs @@ -5,21 +5,14 @@ use arrow::array::{Array, StringArray}; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; -use rayon::ThreadPoolBuilder; -use rayon::prelude::*; use std::fs::File; use std::path::Path; use std::time::Instant; -use rust_pgn_reader_python_binding::{Buffers, ParseConfig, parse_game_to_buffers}; +use rust_pgn_reader_python_binding::{ParseConfig, parse_games_flat}; const FILE_PATH: &str = "2013-07-train-00000-of-00001.parquet"; -/// Chunk multiplier for explicit chunking. -/// 1 = exactly num_threads chunks (minimal overhead) -/// Higher values provide better load balancing at cost of more buffers. -const CHUNK_MULTIPLIER: usize = 1; - /// Read parquet file and return the raw Arrow StringArrays. /// This preserves Arrow's memory layout for zero-copy string access. fn read_parquet_to_string_arrays(file_path: &str) -> Vec { @@ -49,7 +42,7 @@ fn read_parquet_to_string_arrays(file_path: &str) -> Vec { } /// Extract &str slices from Arrow StringArrays (zero-copy). -fn extract_str_slices<'a>(arrays: &'a [StringArray]) -> Vec<&'a str> { +fn extract_str_slices(arrays: &[StringArray]) -> Vec<&str> { let total_len: usize = arrays.iter().map(|a| a.len()).sum(); let mut slices = Vec::with_capacity(total_len); @@ -67,9 +60,7 @@ fn extract_str_slices<'a>(arrays: &'a [StringArray]) -> Vec<&'a str> { /// /// 1. Read parquet to Arrow arrays /// 2. Extract &str slices from StringArray -/// 3. Parse in parallel with explicit chunking (par_chunks) -> fixed number of Buffers -/// -/// No merge step - the chunked architecture keeps per-thread buffers as-is. +/// 3. Two-pass parallel parse into exact-size flat arrays pub fn bench_parse_api() { let config = ParseConfig { store_comments: false, @@ -83,64 +74,24 @@ pub fn bench_parse_api() { let pgn_slices = extract_str_slices(&arrays); println!("Read {} games from parquet.", pgn_slices.len()); - // Step 3: Build thread pool and compute capacity estimates let num_threads = num_cpus::get(); - let n_games = pgn_slices.len(); - let moves_per_game = 70; - - // Calculate chunk size for explicit chunking - let num_chunks = num_threads * CHUNK_MULTIPLIER; - let chunk_size = (n_games + num_chunks - 1) / num_chunks; - let chunk_size = chunk_size.max(1); - let games_per_chunk = chunk_size; - - println!( - "Using {} threads, {} chunks, {} games/chunk", - num_threads, num_chunks, games_per_chunk - ); + println!("Using {} threads (two-pass flat output)", num_threads); - let thread_pool = ThreadPoolBuilder::new() - .num_threads(num_threads) - .build() - .expect("Failed to build Rayon thread pool"); - - // Step 4: Parse in parallel using par_chunks + // Step 3: Parse with the two-pass flat driver let start = Instant::now(); - let chunk_results: Vec = thread_pool.install(|| { - pgn_slices - .par_chunks(chunk_size) - .map(|chunk| { - let mut buffers = Buffers::with_capacity(games_per_chunk, moves_per_game, &config); - for &pgn in chunk { - let _ = parse_game_to_buffers(pgn, &mut buffers, &config); - } - buffers - }) - .collect() - }); + let flat = parse_games_flat(&pgn_slices, num_threads, &config).expect("parse failed"); let duration_parallel = start.elapsed(); println!("Parallel parsing time: {:?}", duration_parallel); println!( - "Created {} Buffers chunks (no merge needed)", - chunk_results.len() - ); - - // Compute totals from chunks - let total_games: usize = chunk_results.iter().map(|b| b.num_games()).sum(); - let total_positions: usize = chunk_results.iter().map(|b| b.total_positions()).sum(); - - let duration_total = start.elapsed(); - println!("Total time (parsing, no merge): {:?}", duration_total); - println!( - "Parsed {} games, {} total positions.", - total_games, total_positions + "Parsed {} games, {} total positions, {} total moves.", + flat.num_games, flat.total_positions, flat.total_moves ); // Measure cleanup time let drop_start = Instant::now(); - drop(chunk_results); + drop(flat); let drop_duration = drop_start.elapsed(); let total_duration = start.elapsed(); @@ -149,6 +100,6 @@ pub fn bench_parse_api() { } fn main() { - println!("=== Parse API (Buffers) ===\n"); + println!("=== Parse API (flat two-pass) ===\n"); bench_parse_api(); } diff --git a/rust_pgn_reader_python_binding.pyi b/rust_pgn_reader_python_binding.pyi index f3daeb5..bb6c947 100644 --- a/rust_pgn_reader_python_binding.pyi +++ b/rust_pgn_reader_python_binding.pyi @@ -162,119 +162,174 @@ class ParsedGamesIter: def __iter__(self) -> "ParsedGamesIter": ... def __next__(self) -> PyGameView: ... -class PyChunkView: - """View into a single chunk's raw numpy arrays. +class ParsedGames: + """Flat container for parsed chess games, optimized for ML training. + + All data lives in single flat NumPy arrays indexed by the CSR offset + arrays ``move_offsets`` / ``position_offsets`` (length num_games + 1). - Access via ``parsed_games.chunks[i]``. Each chunk corresponds to one - parsing thread's output. Use this for advanced access patterns like - manual concatenation or custom batching. + Invalid games (``parse_errors[i] is not None``): their array ranges are + sized from a pre-parse token count; only the first + ``parsed_move_counts[i]`` moves (and that many + 1 positions) contain + data, the remainder of the range is sentinel-filled + (promotions/en_passant = -1, clocks/evals = NaN, zeros elsewhere). + + Board layout: + Boards use square indexing: a1=0, b1=1, ..., h8=63 + Piece encoding: 0=empty, 1-6=white PNBRQK, 7-12=black pnbrqk """ + # === Computed properties === + @property - def num_games(self) -> int: ... + def num_games(self) -> int: + """Number of games in the result.""" + ... + @property - def num_moves(self) -> int: ... + def num_moves(self) -> int: + """Total number of move slots across all games (== move_offsets[-1]).""" + ... + @property - def num_positions(self) -> int: ... + def num_positions(self) -> int: + """Total number of board position slots recorded.""" + ... + + # === Global flat arrays === + @property def boards(self) -> NDArray[np.uint8]: - """Board positions, shape (N_positions, 8, 8), dtype uint8.""" + """Board positions, shape (num_positions, 8, 8), dtype uint8.""" ... + @property def castling(self) -> NDArray[np.bool_]: - """Castling rights [K,Q,k,q], shape (N_positions, 4), dtype bool.""" + """Castling rights [K,Q,k,q], shape (num_positions, 4), dtype bool.""" ... + @property - def en_passant(self) -> NDArray[np.int8]: ... - @property - def halfmove_clock(self) -> NDArray[np.uint8]: ... - @property - def turn(self) -> NDArray[np.bool_]: ... - @property - def from_squares(self) -> NDArray[np.uint8]: ... - @property - def to_squares(self) -> NDArray[np.uint8]: ... - @property - def promotions(self) -> NDArray[np.int8]: ... - @property - def clocks(self) -> NDArray[np.float32]: ... - @property - def evals(self) -> NDArray[np.float32]: ... - @property - def move_offsets(self) -> NDArray[np.uint32]: ... - @property - def position_offsets(self) -> NDArray[np.uint32]: ... - @property - def is_checkmate(self) -> NDArray[np.bool_]: ... + def en_passant(self) -> NDArray[np.int8]: + """En passant file (-1 if none), shape (num_positions,).""" + ... + @property - def is_stalemate(self) -> NDArray[np.bool_]: ... + def halfmove_clock(self) -> NDArray[np.uint8]: + """Halfmove clock, shape (num_positions,).""" + ... + @property - def is_insufficient(self) -> NDArray[np.bool_]: ... + def turn(self) -> NDArray[np.bool_]: + """Side to move (True=white), shape (num_positions,).""" + ... + @property - def legal_move_count(self) -> NDArray[np.uint16]: ... + def from_squares(self) -> NDArray[np.uint8]: + """From squares, shape (num_moves,).""" + ... + @property - def valid(self) -> NDArray[np.bool_]: ... + def to_squares(self) -> NDArray[np.uint8]: + """To squares, shape (num_moves,).""" + ... + @property - def headers(self) -> List[Dict[str, str]]: ... + def promotions(self) -> NDArray[np.int8]: + """Promotions (-1=none, 2=N, 3=B, 4=R, 5=Q), shape (num_moves,).""" + ... + @property - def outcome(self) -> List[Optional[str]]: ... + def clocks(self) -> NDArray[np.float32]: + """Clock times in seconds (NaN if missing), shape (num_moves,).""" + ... + @property - def parse_errors(self) -> List[Optional[str]]: ... + def evals(self) -> NDArray[np.float32]: + """Engine evals (NaN if missing), shape (num_moves,).""" + ... + @property - def comments(self) -> List[Optional[str]]: ... + def move_offsets(self) -> NDArray[np.uint32]: + """CSR offsets into the move arrays, shape (num_games + 1,).""" + ... + @property - def legal_move_from_squares(self) -> NDArray[np.uint8]: ... + def position_offsets(self) -> NDArray[np.uint32]: + """CSR offsets into the position arrays, shape (num_games + 1,).""" + ... + @property - def legal_move_to_squares(self) -> NDArray[np.uint8]: ... + def parsed_move_counts(self) -> NDArray[np.uint32]: + """Actually-parsed moves per game, shape (num_games,). + + Equals ``np.diff(move_offsets)`` for valid games; smaller for + invalid games (whose allocated tail is sentinel-filled).""" + ... + @property - def legal_move_promotions(self) -> NDArray[np.int8]: ... + def is_checkmate(self) -> NDArray[np.bool_]: + """Final position is checkmate, shape (num_games,).""" + ... + @property - def legal_move_offsets(self) -> NDArray[np.uint32]: ... - def __repr__(self) -> str: ... + def is_stalemate(self) -> NDArray[np.bool_]: + """Final position is stalemate, shape (num_games,).""" + ... -class ParsedGames: - """Chunked container for parsed chess games, optimized for ML training. + @property + def is_insufficient(self) -> NDArray[np.bool_]: + """Insufficient material [white, black], shape (num_games, 2).""" + ... - Internally stores data in multiple chunks (one per parsing thread) to - avoid the cost of merging. Per-game access is O(log(num_chunks)) via - binary search on precomputed boundaries. + @property + def legal_move_count(self) -> NDArray[np.uint16]: + """Legal move count in final position, shape (num_games,).""" + ... - Board layout: - Boards use square indexing: a1=0, b1=1, ..., h8=63 - Piece encoding: 0=empty, 1-6=white PNBRQK, 7-12=black pnbrqk - """ + @property + def valid(self) -> NDArray[np.bool_]: + """Whether each game parsed successfully, shape (num_games,).""" + ... - # === Computed properties === + @property + def headers(self) -> List[Dict[str, str]]: + """Raw PGN headers per game.""" + ... @property - def num_games(self) -> int: - """Number of games in the result.""" + def outcome(self) -> List[Optional[str]]: + """Outcome per game: 'White', 'Black', 'Draw', 'Unknown', or None.""" ... @property - def num_moves(self) -> int: - """Total number of moves across all games.""" + def parse_errors(self) -> List[Optional[str]]: + """Parse error message per game (None if valid).""" ... @property - def num_positions(self) -> int: - """Total number of board positions recorded.""" + def comments(self) -> List[Optional[str]]: + """Raw text comments per move (only populated when store_comments=True).""" ... @property - def num_chunks(self) -> int: - """Number of internal chunks.""" + def legal_move_from_squares(self) -> NDArray[np.uint8]: + """Legal-move from squares (only when store_legal_moves=True).""" ... - # === Escape hatch: raw chunk access === + @property + def legal_move_to_squares(self) -> NDArray[np.uint8]: + """Legal-move to squares (only when store_legal_moves=True).""" + ... @property - def chunks(self) -> List[PyChunkView]: - """Access raw per-chunk data. + def legal_move_promotions(self) -> NDArray[np.int8]: + """Legal-move promotions (only when store_legal_moves=True).""" + ... - Each chunk corresponds to one parsing thread's output. Use this - for advanced access patterns like manual concatenation. - """ + @property + def legal_move_offsets(self) -> NDArray[np.uint32]: + """CSR offsets into the legal-move arrays per position, + shape (num_positions + 1,) (only when store_legal_moves=True).""" ... # === Sequence protocol === @@ -348,7 +403,6 @@ def parse_game( def parse_games( pgn_chunked_array: pyarrow.ChunkedArray, num_threads: Optional[int] = None, - chunk_multiplier: Optional[int] = None, store_comments: bool = False, store_legal_moves: bool = False, ) -> ParsedGames: @@ -360,7 +414,6 @@ def parse_games( Args: pgn_chunked_array: PyArrow ChunkedArray containing PGN strings num_threads: Number of threads for parallel parsing (default: all CPUs) - chunk_multiplier: Multiplier for number of chunks (default: 1) store_comments: Whether to store raw text comments (default: False) store_legal_moves: Whether to store legal moves at each position (default: False) diff --git a/src/bench_data_access.py b/src/bench_data_access.py index 5c1e46b..c73a737 100644 --- a/src/bench_data_access.py +++ b/src/bench_data_access.py @@ -31,16 +31,14 @@ def main(): f" {result.num_games:,} games, {result.num_moves:,} moves, {result.num_positions:,} positions" ) print(f" {result.num_games / elapsed:,.0f} games/sec") - print(f" {result.num_chunks} chunks") - # Data access: chunk-level array access + # Data access: global flat arrays start = time.perf_counter() - for chunk in result.chunks: - _ = chunk.boards.sum() - _ = chunk.from_squares.sum() - _ = chunk.to_squares.sum() + _ = result.boards.sum() + _ = result.from_squares.sum() + _ = result.to_squares.sum() elapsed = time.perf_counter() - start - print(f"\nChunk array access: {elapsed:.3f}s") + print(f"\nFlat array access: {elapsed:.3f}s") # Data access: per-game views n_access = min(1000, result.num_games) @@ -60,22 +58,20 @@ def main(): print(f"Position-to-game (1000 lookups): {elapsed * 1000:.3f}ms") # Memory usage - total_bytes = 0 - for chunk in result.chunks: - total_bytes += ( - chunk.boards.nbytes - + chunk.castling.nbytes - + chunk.en_passant.nbytes - + chunk.halfmove_clock.nbytes - + chunk.turn.nbytes - + chunk.from_squares.nbytes - + chunk.to_squares.nbytes - + chunk.promotions.nbytes - + chunk.clocks.nbytes - + chunk.evals.nbytes - + chunk.move_offsets.nbytes - + chunk.position_offsets.nbytes - ) + total_bytes = ( + result.boards.nbytes + + result.castling.nbytes + + result.en_passant.nbytes + + result.halfmove_clock.nbytes + + result.turn.nbytes + + result.from_squares.nbytes + + result.to_squares.nbytes + + result.promotions.nbytes + + result.clocks.nbytes + + result.evals.nbytes + + result.move_offsets.nbytes + + result.position_offsets.nbytes + ) print( f"\nMemory: {total_bytes / 1024 / 1024:.1f} MB ({total_bytes / result.num_positions:.0f} bytes/position)" ) diff --git a/src/example_parse_games.py b/src/example_parse_games.py index 7981306..f34b7f6 100644 --- a/src/example_parse_games.py +++ b/src/example_parse_games.py @@ -66,20 +66,17 @@ def main(): print(f"Number of games: {result.num_games}") print(f"Total moves: {result.num_moves}") print(f"Total positions: {result.num_positions}") - print(f"Number of chunks: {result.num_chunks}") - # === Chunk Details (escape hatch for raw array access) === - print(f"\n--- Chunk Details ---") - for i, chunk in enumerate(result.chunks): - print( - f" Chunk {i}: {chunk.num_games} games, " - f"{chunk.num_moves} moves, " - f"{chunk.num_positions} positions" - ) - print(f" boards: {chunk.boards.shape} ({chunk.boards.dtype})") - print( - f" from_squares: {chunk.from_squares.shape} ({chunk.from_squares.dtype})" - ) + # === Global Flat Arrays (raw array access) === + print(f"\n--- Flat Arrays ---") + print(f" boards: {result.boards.shape} ({result.boards.dtype})") + print( + f" from_squares: {result.from_squares.shape} ({result.from_squares.dtype})" + ) + print( + f" move_offsets: {result.move_offsets.shape} ({result.move_offsets.dtype})" + ) + print(f" position_offsets: {result.position_offsets.shape}") # === Iterate Over Games === print(f"\n--- Game Details ---") @@ -105,10 +102,9 @@ def main(): # === Direct Array Access for ML === print(f"\n--- ML-Ready Data Access ---") - # Access boards via chunks (no single merged array) - # To concatenate all boards: np.concatenate([c.boards for c in result.chunks]) - chunk0_boards = result.chunks[0].boards - print(f"Chunk 0 boards: {chunk0_boards.shape}") + # All boards live in one flat global array + all_boards = result.boards + print(f"All boards: {all_boards.shape}") # Get initial position of first game game0 = result[0] diff --git a/src/flat.rs b/src/flat.rs index 642b85e..5f7e2e0 100644 --- a/src/flat.rs +++ b/src/flat.rs @@ -20,12 +20,42 @@ use crate::board_serialization::{ }; use crate::comment_parsing::{CommentContent, ParsedTag, parse_comments}; use crate::tokenizer::{self, Outcome, Visitor}; -use crate::visitor::{ParseConfig, castling_with_zeros}; use rayon::ThreadPoolBuilder; use rayon::prelude::*; -use shakmaty::{CastlingMode, Chess, Color, Position, fen::Fen, san::SanPlus, uci::UciMove}; +use shakmaty::{ + CastlingMode, CastlingSide, Chess, Color, Position, + fen::Fen, + san::{San, SanPlus, Suffix}, + uci::UciMove, +}; use std::collections::HashMap; +/// Configuration for what optional data to store during parsing. +#[derive(Clone, Debug)] +pub struct ParseConfig { + pub store_comments: bool, + pub store_legal_moves: bool, +} + +/// Support castling notated with zeros ("0-0", "0-0-0"), with optional +/// check/checkmate suffix, matching pgn-reader's explicit handling. +fn castling_with_zeros(token: &[u8]) -> Option { + let (body, suffix) = match token.last() { + Some(b'+') => (&token[..token.len() - 1], Some(Suffix::Check)), + Some(b'#') => (&token[..token.len() - 1], Some(Suffix::Checkmate)), + _ => (token, None), + }; + let side = match body { + b"0-0" => CastlingSide::KingSide, + b"0-0-0" => CastlingSide::QueenSide, + _ => return None, + }; + Some(SanPlus { + san: San::Castle(side), + suffix, + }) +} + /// Number of games per pass-2 task. Small enough that rayon work stealing /// evens out length imbalance between tasks (no straggler tail), large /// enough that per-task overhead is negligible. @@ -705,200 +735,176 @@ pub fn parse_games_flat( #[cfg(test)] mod tests { use super::*; - use crate::visitor::{Buffers, parse_game_to_buffers}; - - /// Differential check: the flat two-pass path must produce the same - /// per-game data as the Buffers path for every game. For invalid games - /// the flat arrays are compared over the actually-parsed prefix and the - /// allocated tail is checked for sentinels. - fn assert_flat_matches_buffers(pgns: &[&str], config: &ParseConfig, num_threads: usize) { - let flat = parse_games_flat(pgns, num_threads, config).unwrap(); - - let mut buffers = Buffers::with_capacity(pgns.len().max(1), 70, config); - for &p in pgns { - let _ = parse_game_to_buffers(p, &mut buffers, config); + + fn default_config() -> ParseConfig { + ParseConfig { + store_comments: false, + store_legal_moves: false, } + } - assert_eq!(flat.num_games, buffers.num_games(), "game count"); - let n = flat.num_games; - assert_eq!(flat.move_offsets.len(), n + 1); - assert_eq!(flat.position_offsets.len(), n + 1); + fn full_config() -> ParseConfig { + ParseConfig { + store_comments: true, + store_legal_moves: true, + } + } - let b_move_off = buffers.compute_move_offsets(); - let b_pos_off = buffers.compute_position_offsets(); - let b_legal_off = if config.store_legal_moves { - buffers.compute_legal_move_offsets() - } else { - Vec::new() - }; + fn parse_one(pgn: &str, config: &ParseConfig) -> FlatOutput { + parse_games_flat(&[pgn], 1, config).unwrap() + } - for i in 0..n { - let parsed = flat.parsed_move_counts[i] as usize; - assert_eq!( - parsed as u32, buffers.move_counts[i], - "game {i}: parsed moves" - ); - assert_eq!(flat.valid[i], buffers.valid[i], "game {i}: valid"); - assert_eq!(flat.headers[i], buffers.headers[i], "game {i}: headers"); - assert_eq!(flat.outcome[i], buffers.outcome[i], "game {i}: outcome"); - assert_eq!( - flat.parse_errors[i], buffers.parse_errors[i], - "game {i}: parse_errors" - ); - assert_eq!(flat.is_checkmate[i], buffers.is_checkmate[i], "game {i}"); - assert_eq!(flat.is_stalemate[i], buffers.is_stalemate[i], "game {i}"); - assert_eq!( - flat.is_insufficient[i * 2..i * 2 + 2], - buffers.is_insufficient[i * 2..i * 2 + 2], - "game {i}: is_insufficient" - ); - assert_eq!( - flat.legal_move_count[i], buffers.legal_move_count[i], - "game {i}: legal_move_count" - ); + fn bits(v: &[f32]) -> Vec { + v.iter().map(|f| f.to_bits()).collect() + } - // --- Moves: parsed prefix identical, tail sentinel-filled --- - let fm = flat.move_offsets[i] as usize; - let bm = b_move_off[i] as usize; - let alloc = (flat.move_offsets[i + 1] - flat.move_offsets[i]) as usize; - assert!(parsed <= alloc, "game {i}: parsed > allocated"); + /// Full structural equality between two runs (f32 compared bitwise). + fn assert_outputs_equal(a: &FlatOutput, b: &FlatOutput) { + assert_eq!(a.num_games, b.num_games); + assert_eq!(a.total_moves, b.total_moves); + assert_eq!(a.total_positions, b.total_positions); + assert_eq!(a.boards, b.boards); + assert_eq!(a.castling, b.castling); + assert_eq!(a.en_passant, b.en_passant); + assert_eq!(a.halfmove_clock, b.halfmove_clock); + assert_eq!(a.turn, b.turn); + assert_eq!(a.from_squares, b.from_squares); + assert_eq!(a.to_squares, b.to_squares); + assert_eq!(a.promotions, b.promotions); + assert_eq!(bits(&a.clocks), bits(&b.clocks)); + assert_eq!(bits(&a.evals), bits(&b.evals)); + assert_eq!(a.move_offsets, b.move_offsets); + assert_eq!(a.position_offsets, b.position_offsets); + assert_eq!(a.parsed_move_counts, b.parsed_move_counts); + assert_eq!(a.is_checkmate, b.is_checkmate); + assert_eq!(a.is_stalemate, b.is_stalemate); + assert_eq!(a.is_insufficient, b.is_insufficient); + assert_eq!(a.legal_move_count, b.legal_move_count); + assert_eq!(a.valid, b.valid); + assert_eq!(a.headers, b.headers); + assert_eq!(a.outcome, b.outcome); + assert_eq!(a.parse_errors, b.parse_errors); + assert_eq!(a.comments, b.comments); + assert_eq!(a.legal_move_from_squares, b.legal_move_from_squares); + assert_eq!(a.legal_move_to_squares, b.legal_move_to_squares); + assert_eq!(a.legal_move_promotions, b.legal_move_promotions); + assert_eq!(a.legal_move_offsets, b.legal_move_offsets); + } + /// Assert game `i` of `batch` has identical content to game 0 of + /// `single` (a one-game parse of the same PGN). + fn assert_game_matches_single(batch: &FlatOutput, i: usize, single: &FlatOutput) { + assert_eq!(single.num_games, 1); + let parsed = batch.parsed_move_counts[i] as usize; + assert_eq!(parsed, single.parsed_move_counts[0] as usize, "game {i}"); + + let alloc = (batch.move_offsets[i + 1] - batch.move_offsets[i]) as usize; + assert_eq!(alloc, single.move_offsets[1] as usize, "game {i}: alloc"); + + let bm = batch.move_offsets[i] as usize; + let bp = batch.position_offsets[i] as usize; + let alloc_pos = alloc + 1; + + assert_eq!( + &batch.boards[bp * 64..(bp + alloc_pos) * 64], + &single.boards[..], + "game {i}: boards" + ); + assert_eq!( + &batch.castling[bp * 4..(bp + alloc_pos) * 4], + &single.castling[..], + "game {i}: castling" + ); + assert_eq!( + &batch.en_passant[bp..bp + alloc_pos], + &single.en_passant[..], + "game {i}: en_passant" + ); + assert_eq!( + &batch.halfmove_clock[bp..bp + alloc_pos], + &single.halfmove_clock[..], + "game {i}: halfmove_clock" + ); + assert_eq!( + &batch.turn[bp..bp + alloc_pos], + &single.turn[..], + "game {i}: turn" + ); + assert_eq!( + &batch.from_squares[bm..bm + alloc], + &single.from_squares[..], + "game {i}: from_squares" + ); + assert_eq!( + &batch.to_squares[bm..bm + alloc], + &single.to_squares[..], + "game {i}: to_squares" + ); + assert_eq!( + &batch.promotions[bm..bm + alloc], + &single.promotions[..], + "game {i}: promotions" + ); + assert_eq!( + bits(&batch.clocks[bm..bm + alloc]), + bits(&single.clocks), + "game {i}: clocks" + ); + assert_eq!( + bits(&batch.evals[bm..bm + alloc]), + bits(&single.evals), + "game {i}: evals" + ); + + assert_eq!(batch.valid[i], single.valid[0], "game {i}: valid"); + assert_eq!(batch.headers[i], single.headers[0], "game {i}: headers"); + assert_eq!(batch.outcome[i], single.outcome[0], "game {i}: outcome"); + assert_eq!( + batch.parse_errors[i], single.parse_errors[0], + "game {i}: parse_errors" + ); + assert_eq!(batch.is_checkmate[i], single.is_checkmate[0], "game {i}"); + assert_eq!(batch.is_stalemate[i], single.is_stalemate[0], "game {i}"); + assert_eq!( + batch.is_insufficient[i * 2..i * 2 + 2], + single.is_insufficient[..], + "game {i}: is_insufficient" + ); + assert_eq!( + batch.legal_move_count[i], single.legal_move_count[0], + "game {i}: legal_move_count" + ); + + if !batch.comments.is_empty() { assert_eq!( - &flat.from_squares[fm..fm + parsed], - &buffers.from_squares[bm..bm + parsed], - "game {i}: from_squares" + &batch.comments[bm..bm + alloc], + &single.comments[..], + "game {i}: comments" ); - assert_eq!( - &flat.to_squares[fm..fm + parsed], - &buffers.to_squares[bm..bm + parsed], - "game {i}: to_squares" - ); - assert_eq!( - &flat.promotions[fm..fm + parsed], - &buffers.promotions[bm..bm + parsed], - "game {i}: promotions" - ); - for k in 0..parsed { + } + if !batch.legal_move_offsets.is_empty() { + for k in 0..alloc_pos { + let bs = batch.legal_move_offsets[bp + k] as usize; + let be = batch.legal_move_offsets[bp + k + 1] as usize; + let ss = single.legal_move_offsets[k] as usize; + let se = single.legal_move_offsets[k + 1] as usize; + assert_eq!(be - bs, se - ss, "game {i} pos {k}: legal count"); assert_eq!( - flat.clocks[fm + k].to_bits(), - buffers.clocks[bm + k].to_bits(), - "game {i} move {k}: clocks" + &batch.legal_move_from_squares[bs..be], + &single.legal_move_from_squares[ss..se], + "game {i} pos {k}: legal from" ); assert_eq!( - flat.evals[fm + k].to_bits(), - buffers.evals[bm + k].to_bits(), - "game {i} move {k}: evals" + &batch.legal_move_to_squares[bs..be], + &single.legal_move_to_squares[ss..se], + "game {i} pos {k}: legal to" ); - } - if config.store_comments { assert_eq!( - &flat.comments[fm..fm + parsed], - &buffers.comments[bm..bm + parsed], - "game {i}: comments" + &batch.legal_move_promotions[bs..be], + &single.legal_move_promotions[ss..se], + "game {i} pos {k}: legal promo" ); } - for k in parsed..alloc { - assert_eq!(flat.from_squares[fm + k], 0, "game {i}: tail from"); - assert_eq!(flat.to_squares[fm + k], 0, "game {i}: tail to"); - assert_eq!(flat.promotions[fm + k], -1, "game {i}: tail promo"); - assert!(flat.clocks[fm + k].is_nan(), "game {i}: tail clock"); - assert!(flat.evals[fm + k].is_nan(), "game {i}: tail eval"); - if config.store_comments { - assert_eq!(flat.comments[fm + k], None, "game {i}: tail comment"); - } - } - - // --- Positions: parsed prefix identical, tail sentinel-filled --- - let parsed_pos = parsed + 1; - assert_eq!( - parsed_pos as u32, buffers.position_counts[i], - "game {i}: position count" - ); - let fp = flat.position_offsets[i] as usize; - let bp = b_pos_off[i] as usize; - let alloc_pos = (flat.position_offsets[i + 1] - flat.position_offsets[i]) as usize; - assert_eq!(alloc_pos, alloc + 1); - - assert_eq!( - &flat.boards[fp * 64..(fp + parsed_pos) * 64], - &buffers.boards[bp * 64..(bp + parsed_pos) * 64], - "game {i}: boards" - ); - assert_eq!( - &flat.castling[fp * 4..(fp + parsed_pos) * 4], - &buffers.castling[bp * 4..(bp + parsed_pos) * 4], - "game {i}: castling" - ); - assert_eq!( - &flat.en_passant[fp..fp + parsed_pos], - &buffers.en_passant[bp..bp + parsed_pos], - "game {i}: en_passant" - ); - assert_eq!( - &flat.halfmove_clock[fp..fp + parsed_pos], - &buffers.halfmove_clock[bp..bp + parsed_pos], - "game {i}: halfmove_clock" - ); - assert_eq!( - &flat.turn[fp..fp + parsed_pos], - &buffers.turn[bp..bp + parsed_pos], - "game {i}: turn" - ); - for k in parsed_pos..alloc_pos { - assert_eq!(flat.en_passant[fp + k], -1, "game {i}: tail en_passant"); - assert_eq!( - &flat.boards[(fp + k) * 64..(fp + k + 1) * 64], - &[0u8; 64], - "game {i}: tail board" - ); - } - - // --- Legal moves per position (optional) --- - if config.store_legal_moves { - for k in 0..parsed_pos { - let fs = flat.legal_move_offsets[fp + k] as usize; - let fe = flat.legal_move_offsets[fp + k + 1] as usize; - let bs = b_legal_off[bp + k] as usize; - let be = b_legal_off[bp + k + 1] as usize; - assert_eq!(fe - fs, be - bs, "game {i} pos {k}: legal count"); - assert_eq!( - &flat.legal_move_from_squares[fs..fe], - &buffers.legal_move_from_squares[bs..be], - "game {i} pos {k}: legal from" - ); - assert_eq!( - &flat.legal_move_to_squares[fs..fe], - &buffers.legal_move_to_squares[bs..be], - "game {i} pos {k}: legal to" - ); - assert_eq!( - &flat.legal_move_promotions[fs..fe], - &buffers.legal_move_promotions[bs..be], - "game {i} pos {k}: legal promo" - ); - } - for k in parsed_pos..alloc_pos { - assert_eq!( - flat.legal_move_offsets[fp + k], - flat.legal_move_offsets[fp + k + 1], - "game {i}: tail position must have zero legal moves" - ); - } - } - } - } - - fn default_config() -> ParseConfig { - ParseConfig { - store_comments: false, - store_legal_moves: false, - } - } - - fn full_config() -> ParseConfig { - ParseConfig { - store_comments: true, - store_legal_moves: true, } } @@ -927,30 +933,264 @@ mod tests { ]; #[test] - fn test_flat_matches_buffers_tricky() { - assert_flat_matches_buffers(TRICKY, &default_config(), 2); + fn test_batch_matches_individual_parses() { + for config in [default_config(), full_config()] { + let batch = parse_games_flat(TRICKY, 2, &config).unwrap(); + let non_empty: Vec<&str> = TRICKY + .iter() + .copied() + .filter(|s| !s.trim().is_empty()) + .collect(); + assert_eq!(batch.num_games, non_empty.len()); + for (i, pgn) in non_empty.iter().enumerate() { + let single = parse_one(pgn, &config); + assert_game_matches_single(&batch, i, &single); + } + } + } + + #[test] + fn test_thread_count_invariance_multi_task() { + // More games than GAMES_PER_TASK so pass 2 spans multiple tasks, + // with invalid games sprinkled across task boundaries. Output + // must be identical regardless of thread/task scheduling. + let mut pgns: Vec<&str> = Vec::new(); + for i in 0..(3 * GAMES_PER_TASK + 17) { + pgns.push(TRICKY[i % TRICKY.len()]); + } + for config in [default_config(), full_config()] { + let a = parse_games_flat(&pgns, 1, &config).unwrap(); + let b = parse_games_flat(&pgns, 4, &config).unwrap(); + assert_outputs_equal(&a, &b); + } + } + + #[test] + fn test_parse_simple_game() { + let pgn = "[Event \"Test\"]\n[White \"Player1\"]\n[Black \"Player2\"]\n[Result \"1-0\"]\n\n1. e4 e5 2. Nf3 Nc6 1-0"; + let flat = parse_one(pgn, &default_config()); + + assert_eq!(flat.num_games, 1); + assert!(flat.valid[0]); + assert_eq!(flat.parsed_move_counts[0], 4); + assert_eq!(flat.move_offsets, vec![0, 4]); + assert_eq!(flat.position_offsets, vec![0, 5]); + assert_eq!(flat.total_moves, 4); + assert_eq!(flat.total_positions, 5); + assert_eq!(flat.outcome[0], Some("White".to_string())); + assert_eq!(flat.headers[0].get("Event"), Some(&"Test".to_string())); } #[test] - fn test_flat_matches_buffers_comments_and_legal_moves() { - assert_flat_matches_buffers(TRICKY, &full_config(), 2); + fn test_parse_game_with_annotations() { + let pgn = "[Event \"Test\"]\n[Result \"1-0\"]\n\n1. e4 { [%eval 0.17] [%clk 0:03:00] } 1... e5 { [%eval 0.19] [%clk 0:02:58] } 1-0"; + let flat = parse_one(pgn, &default_config()); + + assert_eq!(flat.total_moves, 2); + assert!((flat.evals[0] - 0.17).abs() < 0.01); + assert!((flat.evals[1] - 0.19).abs() < 0.01); + assert!((flat.clocks[0] - 180.0).abs() < 0.01); + assert!((flat.clocks[1] - 178.0).abs() < 0.01); } #[test] - fn test_flat_matches_buffers_single_threaded() { - assert_flat_matches_buffers(TRICKY, &default_config(), 1); + fn test_multiple_games() { + let pgns = &[ + "[Event \"Game1\"]\n[Result \"1-0\"]\n\n1. e4 e5 1-0", + "[Event \"Game2\"]\n[Result \"0-1\"]\n\n1. d4 d5 2. c4 0-1", + ]; + let flat = parse_games_flat(pgns, 1, &default_config()).unwrap(); + + assert_eq!(flat.num_games, 2); + assert_eq!(flat.total_moves, 5); + assert_eq!(flat.move_offsets, vec![0, 2, 5]); + assert_eq!(flat.position_offsets, vec![0, 3, 7]); + assert_eq!(flat.outcome[0], Some("White".to_string())); + assert_eq!(flat.outcome[1], Some("Black".to_string())); } #[test] - fn test_flat_matches_buffers_many_games_multi_task() { - // More games than GAMES_PER_TASK so pass 2 spans multiple tasks, - // with invalid games sprinkled across task boundaries. - let mut pgns: Vec<&str> = Vec::new(); - for i in 0..(3 * GAMES_PER_TASK + 17) { - pgns.push(TRICKY[i % TRICKY.len()]); + fn test_outcome_without_headers() { + let flat = parse_one("1. e4 e5 2. Nf3 Nc6 0-1", &default_config()); + assert_eq!(flat.outcome[0], Some("Black".to_string())); + } + + #[test] + fn test_outcome_draw() { + let flat = parse_one("1. e4 e5 1/2-1/2", &default_config()); + assert_eq!(flat.outcome[0], Some("Draw".to_string())); + } + + #[test] + fn test_comments_disabled() { + let flat = parse_one("1. e4 { a comment } e5 1-0", &default_config()); + assert!(flat.comments.is_empty()); + } + + #[test] + fn test_comments_enabled() { + let config = ParseConfig { + store_comments: true, + store_legal_moves: false, + }; + let flat = parse_one("1. e4 { a comment } 1... e5 { [%eval 0.19] } 1-0", &config); + + assert_eq!(flat.comments.len(), 2); + // Raw text from PGN includes surrounding spaces from the parser + assert_eq!(flat.comments[0], Some(" a comment ".to_string())); + // The second comment only has an eval tag, so text portion is empty + assert_eq!(flat.comments[1], Some("".to_string())); + } + + #[test] + fn test_legal_moves_disabled() { + let flat = parse_one("1. e4 e5 1-0", &default_config()); + assert!(flat.legal_move_from_squares.is_empty()); + assert!(flat.legal_move_offsets.is_empty()); + } + + #[test] + fn test_legal_moves_enabled() { + let config = ParseConfig { + store_comments: false, + store_legal_moves: true, + }; + let flat = parse_one("1. e4 1-0", &config); + + // 2 positions: initial + after e4, offsets len = positions + 1 + assert_eq!(flat.legal_move_offsets.len(), 3); + // Initial position has 20 legal moves; after e4, black has 20 + assert_eq!(flat.legal_move_offsets, vec![0, 20, 40]); + assert_eq!(flat.legal_move_from_squares.len(), 40); + } + + #[test] + fn test_headers_only_game() { + // A game with headers but no movetext at all: the initial position + // (from the FEN header) must still be recorded. + let pgn = "[Event \"Test\"]\n[FEN \"r1bqkbnr/pppp1ppp/2n5/4p3/4P3/5N2/PPPP1PPP/RNBQKB1R w KQkq - 2 3\"]\n"; + let flat = parse_one(pgn, &default_config()); + + assert_eq!(flat.num_games, 1); + assert!(flat.valid[0]); + assert_eq!(flat.parsed_move_counts[0], 0); + assert_eq!(flat.position_offsets, vec![0, 1]); + assert_eq!(flat.headers[0].get("Event"), Some(&"Test".to_string())); + assert_eq!(flat.parse_errors[0], None); + assert_eq!(flat.outcome[0], None); + } + + #[test] + fn test_parse_game_without_headers() { + let flat = parse_one( + "1. Nf3 d5 2. e4 c5 3. exd5 e5 4. dxe6 0-1", + &default_config(), + ); + assert!(flat.valid[0]); + assert_eq!(flat.total_moves, 7); + assert_eq!(flat.outcome[0], Some("Black".to_string())); + } + + #[test] + fn test_parse_game_with_standard_fen() { + let pgn = "[FEN \"r1bqkbnr/pppp1ppp/2n5/4p3/4P3/5N2/PPPP1PPP/RNBQKB1R w KQkq - 2 3\"]\n\n3. Bb5 a6 4. Ba4 Nf6 1-0"; + let flat = parse_one(pgn, &default_config()); + assert!(flat.valid[0]); + assert_eq!(flat.total_moves, 4); + } + + #[test] + fn test_parse_chess960_game() { + let pgn = "[Variant \"chess960\"]\n[FEN \"brkrqnnb/pppppppp/8/8/8/8/PPPPPPPP/BRKRQNNB w KQkq - 0 1\"]\n\n1. g3 d5 2. d4 g6 3. b3 Nf6 1-0"; + let flat = parse_one(pgn, &default_config()); + assert!(flat.valid[0], "Chess960 moves should be valid with FEN"); + assert_eq!(flat.total_moves, 6); + } + + #[test] + fn test_parse_chess960_variant_case_insensitive() { + let pgn = "[Variant \"Chess960\"]\n[FEN \"brkrqnnb/pppppppp/8/8/8/8/PPPPPPPP/BRKRQNNB w KQkq - 0 1\"]\n\n1. g3 d5 1-0"; + let flat = parse_one(pgn, &default_config()); + assert!(flat.valid[0], "Should handle Chess960 case variations"); + } + + #[test] + fn test_parse_invalid_fen_falls_back() { + let pgn = "[FEN \"invalid fen string\"]\n\n1. e4 e5 1-0"; + let flat = parse_one(pgn, &default_config()); + assert!(!flat.valid[0], "invalid FEN should mark game invalid"); + assert!( + flat.parse_errors[0] + .as_deref() + .unwrap() + .starts_with("failed to parse FEN:") + ); + // No moves parsed, but initial (default) position recorded. + assert_eq!(flat.parsed_move_counts[0], 0); + } + + #[test] + fn test_fen_header_case_insensitive() { + let pgn = "[fen \"r1bqkbnr/pppp1ppp/2n5/4p3/4P3/5N2/PPPP1PPP/RNBQKB1R w KQkq - 2 3\"]\n\n3. Bb5 1-0"; + let flat = parse_one(pgn, &default_config()); + assert!(flat.valid[0], "Should handle lowercase 'fen' header"); + } + + #[test] + fn test_parse_game_with_custom_fen_no_variant() { + let pgn = "[Event \"Test Game\"]\n[FEN \"r1bqkbnr/pppp1ppp/2n5/1B2p3/4P3/5N2/PPPP1PPP/RNBQK2R b KQkq - 3 3\"]\n\n3... a6 4. Ba4 Nf6 5. O-O Be7 1-0"; + let flat = parse_one(pgn, &default_config()); + assert!(flat.valid[0]); + assert_eq!(flat.total_moves, 5); // a6, Ba4, Nf6, O-O, Be7 + } + + #[test] + fn test_castling_with_zeros_normalized() { + let pgn = "1. e4 e5 2. Nf3 Nf6 3. Bc4 Bc5 4. 0-0 0-0 1-0"; + let flat = parse_one(pgn, &default_config()); + assert!(flat.valid[0], "zeros castling should parse as O-O"); + assert_eq!(flat.total_moves, 8); + assert_eq!(flat.parse_errors[0], None); + } + + #[test] + fn test_invalid_san_sets_parse_error_and_fills_tail() { + let flat = parse_one("1. e4 xyzzy9 e5 1-0", &default_config()); + + assert!(!flat.valid[0]); + let err = flat.parse_errors[0].as_deref().unwrap(); + assert!( + err.starts_with("failed to parse SAN:") && err.contains("xyzzy9"), + "unexpected error message: {err}" + ); + // Pass 1 counted 3 tokens; only 1 move parsed. + assert_eq!(flat.move_offsets, vec![0, 3]); + assert_eq!(flat.parsed_move_counts[0], 1); + // Tail move slots are sentinel-filled. + for m in 1..3 { + assert_eq!(flat.from_squares[m], 0); + assert_eq!(flat.promotions[m], -1); + assert!(flat.clocks[m].is_nan()); + assert!(flat.evals[m].is_nan()); } - assert_flat_matches_buffers(&pgns, &default_config(), 4); - assert_flat_matches_buffers(&pgns, &full_config(), 4); + // Tail position slots: zero board, en_passant sentinel. + for p in 2..4 { + assert_eq!(&flat.boards[p * 64..(p + 1) * 64], &[0u8; 64]); + assert_eq!(flat.en_passant[p], -1); + } + } + + #[test] + fn test_illegal_move_sets_parse_error() { + let flat = parse_one("1. e4 e4 2. Nf3 1-0", &default_config()); + assert!(!flat.valid[0]); + assert!( + flat.parse_errors[0] + .as_deref() + .unwrap() + .starts_with("illegal move:") + ); + assert_eq!(flat.parsed_move_counts[0], 1); } #[test] @@ -985,4 +1225,13 @@ mod tests { assert_eq!(flat.num_games, 1); assert_eq!(flat.parsed_move_counts[0], 2); } + + #[test] + fn test_checkmate_status() { + let flat = parse_one("1. f3 e5 2. g4 Qh4# 0-1", &default_config()); + assert!(flat.valid[0]); + assert!(flat.is_checkmate[0]); + assert!(!flat.is_stalemate[0]); + assert_eq!(flat.legal_move_count[0], 0); + } } diff --git a/src/lib.rs b/src/lib.rs index 5af15f0..8d74ed8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,89 +2,45 @@ use arrow_array::{Array, LargeStringArray, StringArray}; use numpy::{PyArray1, PyArrayMethods}; use pyo3::prelude::*; use pyo3_arrow::PyChunkedArray; -use rayon::ThreadPoolBuilder; -use rayon::prelude::*; mod board_serialization; mod comment_parsing; mod flat; mod python_bindings; mod tokenizer; -mod visitor; -pub use flat::{FlatOutput, parse_games_flat}; -use python_bindings::{ChunkData, ParsedGames, ParsedGamesIter, PyChunkView, PyGameView}; -pub use visitor::{Buffers, ParseConfig, parse_game_to_buffers}; +pub use flat::{FlatOutput, ParseConfig, parse_games_flat}; +use python_bindings::{ParsedGames, ParsedGamesIter, PyGameView}; /// Shared parallel parsing logic for a slice of PGN strings. /// /// Both `parse_games` and `parse_games_from_strings` delegate here after -/// extracting their `&str` slices from different input types. +/// extracting their `&str` slices from different input types. The GIL is +/// released for the duration of both parsing passes. fn parse_str_slices( py: Python<'_>, slices: &[&str], num_threads: usize, - chunk_multiplier: usize, config: &ParseConfig, ) -> PyResult { - let n_games = slices.len(); - if n_games == 0 { - let empty_chunk = buffers_to_chunk_data(py, Buffers::default())?; - return build_parsed_games(py, vec![empty_chunk]); - } - - let thread_pool = ThreadPoolBuilder::new() - .num_threads(num_threads) - .build() - .map_err(|e| { - PyErr::new::(format!( - "Failed to build thread pool: {}", - e - )) - })?; - - // num_chunks = num_threads * chunk_multiplier (e.g., 16 threads * 1 = 16 chunks) - let num_chunks = num_threads * chunk_multiplier; - let chunk_size = ((n_games + num_chunks - 1) / num_chunks).max(1); - let moves_per_game = 70; - - let chunk_results: Vec = thread_pool.install(|| { - slices - .par_chunks(chunk_size) - .map(|chunk| { - let mut buffers = Buffers::with_capacity(chunk_size, moves_per_game, config); - for &pgn in chunk { - let _ = parse_game_to_buffers(pgn, &mut buffers, config); - } - buffers - }) - .collect() - }); - - let chunk_data_vec: Vec = chunk_results - .into_iter() - .map(|buf| buffers_to_chunk_data(py, buf)) - .collect::>>()?; - - build_parsed_games(py, chunk_data_vec) + let flat = py + .detach(|| parse_games_flat(slices, num_threads, config)) + .map_err(PyErr::new::)?; + build_parsed_games(py, flat) } -/// Parse games from Arrow chunked array into a chunked ParsedGames container. -/// -/// This implementation uses explicit chunking with a fixed number of chunks -/// (num_chunks = num_threads * chunk_multiplier) to avoid the allocation storm -/// caused by Rayon's dynamic work-stealing with fold_with. +/// Parse games from an Arrow chunked array into a flat ParsedGames container. /// -/// Each chunk gets exactly one Buffers instance. Instead of merging all -/// chunks into a single buffer (which was memory-bandwidth-bound), we keep -/// the per-thread buffers and provide virtual indexing across them. +/// Two-pass implementation: a parallel counting pass computes exact output +/// sizes and per-game offsets, then a parallel parsing pass writes every +/// game into its precomputed disjoint range of single flat output arrays +/// (rayon work stealing over small tasks; output order == input order). #[pyfunction] -#[pyo3(signature = (pgn_chunked_array, num_threads=None, chunk_multiplier=None, store_comments=false, store_legal_moves=false))] +#[pyo3(signature = (pgn_chunked_array, num_threads=None, store_comments=false, store_legal_moves=false))] fn parse_games( py: Python<'_>, pgn_chunked_array: PyChunkedArray, num_threads: Option, - chunk_multiplier: Option, store_comments: bool, store_legal_moves: bool, ) -> PyResult { @@ -93,10 +49,6 @@ fn parse_games( store_legal_moves, }; let num_threads = num_threads.unwrap_or_else(num_cpus::get); - // Default multiplier of 1 means exactly num_threads chunks (one per thread). - // Higher values (e.g., 4) create more chunks for better load balancing - // at the cost of slightly more complex indexing. - let chunk_multiplier = chunk_multiplier.unwrap_or(1); // Extract PGN strings from Arrow chunks let num_elements: usize = pgn_chunked_array.chunks().iter().map(|c| c.len()).sum(); @@ -122,182 +74,101 @@ fn parse_games( } } - parse_str_slices(py, &pgn_str_slices, num_threads, chunk_multiplier, &config) + parse_str_slices(py, &pgn_str_slices, num_threads, &config) } -/// Convert a single Buffers into a ChunkData with NumPy arrays. -fn buffers_to_chunk_data(py: Python<'_>, buffers: Buffers) -> PyResult { - let n_games = buffers.num_games(); - let total_positions = buffers.total_positions(); - let total_moves = buffers.total_moves(); +/// Convert a FlatOutput into a ParsedGames with NumPy arrays (zero-copy). +fn build_parsed_games(py: Python<'_>, flat: FlatOutput) -> PyResult { + let num_games = flat.num_games; + let num_moves = flat.total_moves; + let num_positions = flat.total_positions; + let num_legal_moves = flat.legal_move_from_squares.len(); - // Compute all CSR offsets BEFORE any from_vec calls consume buffer fields - let move_offsets_vec = buffers.compute_move_offsets(); - let position_offsets_vec = buffers.compute_position_offsets(); - let has_legal_moves = !buffers.legal_move_counts.is_empty(); - let legal_move_offsets_vec = if has_legal_moves { - buffers.compute_legal_move_offsets() - } else { - Vec::new() - }; - let total_legal_moves_count = buffers.total_legal_moves(); + // Rust-side offset copies for fast PyGameView construction. + let move_offsets_rs = flat.move_offsets.clone(); + let position_offsets_rs = flat.position_offsets.clone(); + let parsed_move_counts_rs = flat.parsed_move_counts.clone(); // Boards: reshape from flat to (N_positions, 8, 8) - let boards_array = PyArray1::from_vec(py, buffers.boards); + let boards_array = PyArray1::from_vec(py, flat.boards); let boards_reshaped = boards_array - .reshape([total_positions, 8, 8]) + .reshape([num_positions, 8, 8]) .map_err(|e| PyErr::new::(e.to_string()))?; // Castling: reshape from flat to (N_positions, 4) - let castling_array = PyArray1::from_vec(py, buffers.castling); + let castling_array = PyArray1::from_vec(py, flat.castling); let castling_reshaped = castling_array - .reshape([total_positions, 4]) + .reshape([num_positions, 4]) .map_err(|e| PyErr::new::(e.to_string()))?; - // 1D arrays - let en_passant_array = PyArray1::from_vec(py, buffers.en_passant); - let halfmove_clock_array = PyArray1::from_vec(py, buffers.halfmove_clock); - let turn_array = PyArray1::from_vec(py, buffers.turn); - - let from_squares_array = PyArray1::from_vec(py, buffers.from_squares); - let to_squares_array = PyArray1::from_vec(py, buffers.to_squares); - let promotions_array = PyArray1::from_vec(py, buffers.promotions); - let clocks_array = PyArray1::from_vec(py, buffers.clocks); - let evals_array = PyArray1::from_vec(py, buffers.evals); - - let move_offsets_array = PyArray1::from_vec(py, move_offsets_vec); - let position_offsets_array = PyArray1::from_vec(py, position_offsets_vec); - - let is_checkmate_array = PyArray1::from_vec(py, buffers.is_checkmate); - let is_stalemate_array = PyArray1::from_vec(py, buffers.is_stalemate); - - let is_insufficient_array = PyArray1::from_vec(py, buffers.is_insufficient); + // Insufficient material: reshape from flat to (N_games, 2) + let is_insufficient_array = PyArray1::from_vec(py, flat.is_insufficient); let is_insufficient_reshaped = is_insufficient_array - .reshape([n_games, 2]) + .reshape([num_games, 2]) .map_err(|e| PyErr::new::(e.to_string()))?; - let legal_move_count_array = PyArray1::from_vec(py, buffers.legal_move_count); - let valid_array = PyArray1::from_vec(py, buffers.valid); - - let legal_move_from_squares_array = PyArray1::from_vec(py, buffers.legal_move_from_squares); - let legal_move_to_squares_array = PyArray1::from_vec(py, buffers.legal_move_to_squares); - let legal_move_promotions_array = PyArray1::from_vec(py, buffers.legal_move_promotions); - let legal_move_offsets_array = PyArray1::from_vec(py, legal_move_offsets_vec); - - Ok(ChunkData { + Ok(ParsedGames { boards: boards_reshaped.unbind().into_any(), castling: castling_reshaped.unbind().into_any(), - en_passant: en_passant_array.unbind().into_any(), - halfmove_clock: halfmove_clock_array.unbind().into_any(), - turn: turn_array.unbind().into_any(), - from_squares: from_squares_array.unbind().into_any(), - to_squares: to_squares_array.unbind().into_any(), - promotions: promotions_array.unbind().into_any(), - clocks: clocks_array.unbind().into_any(), - evals: evals_array.unbind().into_any(), - move_offsets: move_offsets_array.unbind().into_any(), - position_offsets: position_offsets_array.unbind().into_any(), - is_checkmate: is_checkmate_array.unbind().into_any(), - is_stalemate: is_stalemate_array.unbind().into_any(), + en_passant: PyArray1::from_vec(py, flat.en_passant).unbind().into_any(), + halfmove_clock: PyArray1::from_vec(py, flat.halfmove_clock) + .unbind() + .into_any(), + turn: PyArray1::from_vec(py, flat.turn).unbind().into_any(), + from_squares: PyArray1::from_vec(py, flat.from_squares) + .unbind() + .into_any(), + to_squares: PyArray1::from_vec(py, flat.to_squares).unbind().into_any(), + promotions: PyArray1::from_vec(py, flat.promotions).unbind().into_any(), + clocks: PyArray1::from_vec(py, flat.clocks).unbind().into_any(), + evals: PyArray1::from_vec(py, flat.evals).unbind().into_any(), + move_offsets: PyArray1::from_vec(py, flat.move_offsets) + .unbind() + .into_any(), + position_offsets: PyArray1::from_vec(py, flat.position_offsets) + .unbind() + .into_any(), + parsed_move_counts: PyArray1::from_vec(py, flat.parsed_move_counts) + .unbind() + .into_any(), + is_checkmate: PyArray1::from_vec(py, flat.is_checkmate) + .unbind() + .into_any(), + is_stalemate: PyArray1::from_vec(py, flat.is_stalemate) + .unbind() + .into_any(), is_insufficient: is_insufficient_reshaped.unbind().into_any(), - legal_move_count: legal_move_count_array.unbind().into_any(), - valid: valid_array.unbind().into_any(), - headers: buffers.headers, - outcome: buffers.outcome, - parse_errors: buffers.parse_errors, - comments: buffers.comments, - legal_move_from_squares: legal_move_from_squares_array.unbind().into_any(), - legal_move_to_squares: legal_move_to_squares_array.unbind().into_any(), - legal_move_promotions: legal_move_promotions_array.unbind().into_any(), - legal_move_offsets: legal_move_offsets_array.unbind().into_any(), - num_games: n_games, - num_moves: total_moves, - num_positions: total_positions, - num_legal_moves: total_legal_moves_count, - }) -} - -/// Build a ParsedGames from a Vec of ChunkData. -/// -/// Computes prefix-sum boundary arrays and global CSR offsets for -/// position_to_game / move_to_game. -fn build_parsed_games(py: Python<'_>, chunks: Vec) -> PyResult { - let mut total_games: usize = 0; - let mut total_moves: usize = 0; - let mut total_positions: usize = 0; - - // Build boundary arrays (prefix sums) - let mut game_boundaries = Vec::with_capacity(chunks.len() + 1); - let mut move_boundaries = Vec::with_capacity(chunks.len() + 1); - let mut position_boundaries = Vec::with_capacity(chunks.len() + 1); - - game_boundaries.push(0); - move_boundaries.push(0); - position_boundaries.push(0); - - for chunk in &chunks { - total_games += chunk.num_games; - total_moves += chunk.num_moves; - total_positions += chunk.num_positions; - game_boundaries.push(total_games); - move_boundaries.push(total_moves); - position_boundaries.push(total_positions); - } - - // Build global CSR offsets for position_to_game / move_to_game. - // These are the per-chunk local offsets shifted by the chunk's base offset, - // concatenated into a single array. Length = total_games + 1. - let mut global_move_offsets_vec: Vec = Vec::with_capacity(total_games + 1); - let mut global_position_offsets_vec: Vec = Vec::with_capacity(total_games + 1); - - for (chunk_idx, chunk) in chunks.iter().enumerate() { - let move_base = move_boundaries[chunk_idx] as u32; - let pos_base = position_boundaries[chunk_idx] as u32; - - // Read the chunk's local offsets - let local_move_offsets = chunk.move_offsets.bind(py); - let local_move_offsets: &Bound<'_, PyArray1> = local_move_offsets.cast()?; - let local_move_ro = local_move_offsets.readonly(); - let local_move_slice = local_move_ro.as_slice()?; - - let local_pos_offsets = chunk.position_offsets.bind(py); - let local_pos_offsets: &Bound<'_, PyArray1> = local_pos_offsets.cast()?; - let local_pos_ro = local_pos_offsets.readonly(); - let local_pos_slice = local_pos_ro.as_slice()?; - - // Append all but the last offset (which is the total for this chunk) - // The last chunk's final offset will be added after the loop. - for &offset in &local_move_slice[..local_move_slice.len() - 1] { - global_move_offsets_vec.push(move_base + offset); - } - for &offset in &local_pos_slice[..local_pos_slice.len() - 1] { - global_position_offsets_vec.push(pos_base + offset); - } - } - - // Final sentinel value - global_move_offsets_vec.push(total_moves as u32); - global_position_offsets_vec.push(total_positions as u32); - - let global_move_offsets = PyArray1::from_vec(py, global_move_offsets_vec); - let global_position_offsets = PyArray1::from_vec(py, global_position_offsets_vec); - - Ok(ParsedGames { - chunks, - game_boundaries, - move_boundaries, - position_boundaries, - total_games, - total_moves, - total_positions, - global_move_offsets: global_move_offsets.unbind().into_any(), - global_position_offsets: global_position_offsets.unbind().into_any(), + legal_move_count: PyArray1::from_vec(py, flat.legal_move_count) + .unbind() + .into_any(), + valid: PyArray1::from_vec(py, flat.valid).unbind().into_any(), + headers: flat.headers, + outcome: flat.outcome, + parse_errors: flat.parse_errors, + comments: flat.comments, + legal_move_from_squares: PyArray1::from_vec(py, flat.legal_move_from_squares) + .unbind() + .into_any(), + legal_move_to_squares: PyArray1::from_vec(py, flat.legal_move_to_squares) + .unbind() + .into_any(), + legal_move_promotions: PyArray1::from_vec(py, flat.legal_move_promotions) + .unbind() + .into_any(), + legal_move_offsets: PyArray1::from_vec(py, flat.legal_move_offsets) + .unbind() + .into_any(), + move_offsets_rs, + position_offsets_rs, + parsed_move_counts_rs, + num_games, + num_moves, + num_positions, + num_legal_moves, }) } /// Parse a single PGN game string. -/// -/// Convenience wrapper that creates a single-element Arrow array internally. #[pyfunction] #[pyo3(signature = (pgn, store_comments=false, store_legal_moves=false))] fn parse_game( @@ -306,15 +177,16 @@ fn parse_game( store_comments: bool, store_legal_moves: bool, ) -> PyResult { + if pgn.trim().is_empty() { + return Err(pyo3::exceptions::PyValueError::new_err( + "No game found in PGN", + )); + } let config = ParseConfig { store_comments, store_legal_moves, }; - let mut buffers = Buffers::with_capacity(1, 70, &config); - parse_game_to_buffers(pgn, &mut buffers, &config) - .map_err(pyo3::exceptions::PyValueError::new_err)?; - let chunk = buffers_to_chunk_data(py, buffers)?; - build_parsed_games(py, vec![chunk]) + parse_str_slices(py, &[pgn], 1, &config) } /// Parse multiple PGN game strings in parallel. @@ -335,7 +207,7 @@ fn parse_games_from_strings( }; let num_threads = num_threads.unwrap_or_else(num_cpus::get); let str_slices: Vec<&str> = pgns.iter().map(|s| s.as_str()).collect(); - parse_str_slices(py, &str_slices, num_threads, 1, &config) + parse_str_slices(py, &str_slices, num_threads, &config) } /// Parser for chess PGN notation @@ -346,7 +218,6 @@ fn rust_pgn_reader_python_binding(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(parse_games_from_strings, m)?)?; m.add_class::()?; m.add_class::()?; - m.add_class::()?; m.add_class::()?; Ok(()) } diff --git a/src/python_bindings.rs b/src/python_bindings.rs index 740f2ac..b9f1f16 100644 --- a/src/python_bindings.rs +++ b/src/python_bindings.rs @@ -3,51 +3,6 @@ use pyo3::prelude::*; use pyo3::types::{IntoPyDict, PyList, PySlice}; use std::collections::HashMap; -/// Internal per-chunk data. Not exposed to Python directly. -/// Each chunk corresponds to one thread's output during parallel parsing. -pub struct ChunkData { - // Per-position numpy arrays - pub boards: Py, // (N_positions, 8, 8) u8 - pub castling: Py, // (N_positions, 4) bool - pub en_passant: Py, // (N_positions,) i8 - pub halfmove_clock: Py, // (N_positions,) u8 - pub turn: Py, // (N_positions,) bool - - // Per-move numpy arrays - pub from_squares: Py, // (N_moves,) u8 - pub to_squares: Py, // (N_moves,) u8 - pub promotions: Py, // (N_moves,) i8 - pub clocks: Py, // (N_moves,) f32 - pub evals: Py, // (N_moves,) f32 - - // Per-game arrays (CSR offsets are local to this chunk) - pub move_offsets: Py, // (N_games + 1,) u32 - pub position_offsets: Py, // (N_games + 1,) u32 - pub is_checkmate: Py, // (N_games,) bool - pub is_stalemate: Py, // (N_games,) bool - pub is_insufficient: Py, // (N_games, 2) bool - pub legal_move_count: Py, // (N_games,) u16 - pub valid: Py, // (N_games,) bool - pub headers: Vec>, - pub outcome: Vec>, // Per-game: "White", "Black", "Draw", "Unknown", or None - pub parse_errors: Vec>, // Per-game: None if valid, Some(msg) if not - - // Optional: raw text comments (per-move), only populated when store_comments=true - pub comments: Vec>, - - // Optional: legal moves at each position (CSR arrays + offsets) - pub legal_move_from_squares: Py, // (N_legal_moves,) u8 - pub legal_move_to_squares: Py, // (N_legal_moves,) u8 - pub legal_move_promotions: Py, // (N_legal_moves,) i8 - pub legal_move_offsets: Py, // (N_positions + 1,) u32 - - // Metadata - pub num_games: usize, - pub num_moves: usize, - pub num_positions: usize, - pub num_legal_moves: usize, -} - // --------------------------------------------------------------------------- // Helper: searchsorted-based index mapping (used by position_to_game / move_to_game) // --------------------------------------------------------------------------- @@ -114,108 +69,175 @@ fn format_uci(from_sq: u8, to_sq: u8, promo: i8) -> String { // ParsedGames // --------------------------------------------------------------------------- -/// Chunked container for parsed chess games, optimized for ML training. +/// Flat container for parsed chess games, optimized for ML training. /// -/// Stores parsed PGN data across multiple internal chunks (one per parsing -/// thread). Supports integer indexing, slicing, and iteration to access -/// individual games as ``PyGameView`` objects. +/// All data lives in single flat NumPy arrays indexed by the CSR offset +/// arrays ``move_offsets`` / ``position_offsets`` (length num_games + 1). +/// Supports integer indexing, slicing, and iteration to access individual +/// games as ``PyGameView`` objects. +/// +/// Invalid games (``parse_errors[i] is not None``): their array ranges are +/// sized from a pre-parse token count; only the first ``parsed_move_counts[i]`` +/// moves (and that many + 1 positions) contain data, the remainder of the +/// range is sentinel-filled (promotions/en_passant = -1, clocks/evals = NaN, +/// zeros elsewhere). /// /// Properties: /// num_games: Total number of games. -/// num_moves: Total half-moves across all games. -/// num_positions: Total board positions recorded. -/// num_chunks: Number of internal chunks. -/// chunks: List of ``PyChunkView`` for direct array access. +/// num_moves: Total move slots across all games (== move_offsets[-1]). +/// num_positions: Total position slots recorded. +/// boards, castling, ...: Flat global arrays (see PyGameView docs for +/// encodings). /// /// Methods: /// position_to_game(indices): Map position indices to game indices. /// move_to_game(indices): Map move indices to game indices. #[pyclass] pub struct ParsedGames { - pub chunks: Vec, - - // Global prefix sums for O(1) chunk lookup. - // game_boundaries[i] = total games in chunks 0..i - // So game_boundaries = [0, chunk0.num_games, chunk0+chunk1, ..., total_games] - pub game_boundaries: Vec, - #[allow(dead_code)] - pub move_boundaries: Vec, - #[allow(dead_code)] - pub position_boundaries: Vec, - - pub total_games: usize, - pub total_moves: usize, - pub total_positions: usize, - - // Global offsets for position_to_game / move_to_game (precomputed) - pub global_move_offsets: Py, - pub global_position_offsets: Py, + // Per-position NumPy arrays + pub boards: Py, // (N_positions, 8, 8) u8 + pub castling: Py, // (N_positions, 4) bool + pub en_passant: Py, // (N_positions,) i8 + pub halfmove_clock: Py, // (N_positions,) u8 + pub turn: Py, // (N_positions,) bool + + // Per-move NumPy arrays + pub from_squares: Py, // (N_moves,) u8 + pub to_squares: Py, // (N_moves,) u8 + pub promotions: Py, // (N_moves,) i8 + pub clocks: Py, // (N_moves,) f32 + pub evals: Py, // (N_moves,) f32 + + // Global CSR offsets and per-game NumPy arrays + pub move_offsets: Py, // (N_games + 1,) u32 + pub position_offsets: Py, // (N_games + 1,) u32 + pub parsed_move_counts: Py, // (N_games,) u32 + pub is_checkmate: Py, // (N_games,) bool + pub is_stalemate: Py, // (N_games,) bool + pub is_insufficient: Py, // (N_games, 2) bool + pub legal_move_count: Py, // (N_games,) u16 + pub valid: Py, // (N_games,) bool + + // Per-game Rust-side data + pub headers: Vec>, + pub outcome: Vec>, // "White", "Black", "Draw", "Unknown", or None + pub parse_errors: Vec>, // None if valid, Some(msg) if not + + // Optional: raw text comments (per-move), only populated when store_comments=true + pub comments: Vec>, + + // Optional: legal moves at each position (CSR arrays + offsets) + pub legal_move_from_squares: Py, // (N_legal_moves,) u8 + pub legal_move_to_squares: Py, // (N_legal_moves,) u8 + pub legal_move_promotions: Py, // (N_legal_moves,) i8 + pub legal_move_offsets: Py, // (N_positions + 1,) u32 + + // Rust-side offset copies for fast game-view construction + pub move_offsets_rs: Vec, + pub position_offsets_rs: Vec, + pub parsed_move_counts_rs: Vec, + + pub num_games: usize, + pub num_moves: usize, + pub num_positions: usize, + pub num_legal_moves: usize, } -impl ParsedGames { - /// Locate which chunk a global game index belongs to. - /// Returns (chunk_index, local_game_index). - fn locate_game(&self, global_idx: usize) -> (usize, usize) { - // Binary search: find the last boundary <= global_idx - // game_boundaries = [0, n0, n0+n1, ...], length = num_chunks + 1 - let chunk_idx = match self.game_boundaries.binary_search(&global_idx) { - Ok(i) => { - // Exact match. If it's the last boundary, back up one. - if i >= self.chunks.len() { - self.chunks.len() - 1 - } else { - i +/// Generate a `#[pymethods]` block with NumPy-array getters for ParsedGames. +macro_rules! parsed_games_array_getters { + ($($name:ident),+ $(,)?) => { + #[pymethods] + impl ParsedGames { + $( + #[getter] + fn $name(&self, py: Python<'_>) -> Py { + self.$name.clone_ref(py) } - } - Err(i) => i - 1, // insertion point - 1 - }; - let local_idx = global_idx - self.game_boundaries[chunk_idx]; - (chunk_idx, local_idx) - } + )+ + } + }; } +/// Generate a `#[pymethods]` block with Vec-cloning getters for ParsedGames. +macro_rules! parsed_games_vec_getters { + ($($name:ident -> $ret:ty),+ $(,)?) => { + #[pymethods] + impl ParsedGames { + $( + #[getter] + fn $name(&self) -> $ret { + self.$name.clone() + } + )+ + } + }; +} + +parsed_games_array_getters!( + boards, + castling, + en_passant, + halfmove_clock, + turn, + from_squares, + to_squares, + promotions, + clocks, + evals, + move_offsets, + position_offsets, + parsed_move_counts, + is_checkmate, + is_stalemate, + is_insufficient, + legal_move_count, + valid, + legal_move_from_squares, + legal_move_to_squares, + legal_move_promotions, + legal_move_offsets, +); + +parsed_games_vec_getters!( + headers -> Vec>, + outcome -> Vec>, + parse_errors -> Vec>, + comments -> Vec>, +); + #[pymethods] impl ParsedGames { /// Number of games in the result. #[getter] fn num_games(&self) -> usize { - self.total_games + self.num_games } - /// Total number of moves across all games. + /// Total number of move slots across all games. #[getter] fn num_moves(&self) -> usize { - self.total_moves + self.num_moves } - /// Total number of board positions recorded. + /// Total number of board position slots recorded. #[getter] fn num_positions(&self) -> usize { - self.total_positions - } - - /// Number of internal chunks. - #[getter] - fn num_chunks(&self) -> usize { - self.chunks.len() + self.num_positions } fn __len__(&self) -> usize { - self.total_games + self.num_games } fn __repr__(&self) -> String { format!( - "", - self.total_games, - self.total_moves, - self.total_positions, - self.chunks.len() + "", + self.num_games, self.num_moves, self.num_positions ) } fn __getitem__(slf: Py, py: Python<'_>, idx: &Bound<'_, PyAny>) -> PyResult> { - let n_games = slf.borrow(py).total_games; + let n_games = slf.borrow(py).num_games; // Handle integer index if let Ok(mut i) = idx.extract::() { @@ -256,7 +278,7 @@ impl ParsedGames { } fn __iter__(slf: Py, py: Python<'_>) -> PyResult { - let total = slf.borrow(py).total_games; + let total = slf.borrow(py).num_games; Ok(ParsedGamesIter { data: slf, index: 0, @@ -264,26 +286,6 @@ impl ParsedGames { }) } - /// Escape hatch: access raw per-chunk data. - /// - /// Returns a list of chunk view objects, each exposing numpy arrays - /// for that chunk's data. Use this for advanced/custom access patterns. - #[getter] - fn chunks(slf: Py, py: Python<'_>) -> PyResult> { - let n_chunks = slf.borrow(py).chunks.len(); - let mut views: Vec> = Vec::with_capacity(n_chunks); - for i in 0..n_chunks { - views.push(Py::new( - py, - PyChunkView { - parent: slf.clone_ref(py), - chunk_idx: i, - }, - )?); - } - Ok(PyList::new(py, views)?.into_any().unbind()) - } - /// Map position indices to game indices. /// /// Given an array of indices into the global position space, returns @@ -301,7 +303,7 @@ impl ParsedGames { py: Python<'py>, position_indices: &Bound<'py, PyAny>, ) -> PyResult>> { - index_to_game(py, &self.global_position_offsets, position_indices) + index_to_game(py, &self.position_offsets, position_indices) } /// Map move indices to game indices. @@ -320,115 +322,7 @@ impl ParsedGames { py: Python<'py>, move_indices: &Bound<'py, PyAny>, ) -> PyResult>> { - index_to_game(py, &self.global_move_offsets, move_indices) - } -} - -// --------------------------------------------------------------------------- -// PyChunkView — macros + impl -// --------------------------------------------------------------------------- - -/// Lightweight view into a single parsing chunk's raw numpy arrays. -/// -/// Access via ``parsed_games.chunks[i]``. Each chunk corresponds to one -/// parsing thread's output. All numpy array properties are zero-copy -/// references into the parent data. -/// -/// Use this for advanced access patterns like manual concatenation, -/// custom batching, or direct array-level ML pipelines. -#[pyclass] -pub struct PyChunkView { - parent: Py, - chunk_idx: usize, -} - -/// Generate a `#[pymethods]` block with numpy-array getters for PyChunkView. -macro_rules! chunk_array_getters { - ($($name:ident),+ $(,)?) => { - #[pymethods] - impl PyChunkView { - $( - #[getter] - fn $name(&self, py: Python<'_>) -> Py { - self.parent.borrow(py).chunks[self.chunk_idx].$name.clone_ref(py) - } - )+ - } - }; -} - -/// Generate a `#[pymethods]` block with Vec-cloning getters for PyChunkView. -macro_rules! chunk_vec_getters { - ($($name:ident -> $ret:ty),+ $(,)?) => { - #[pymethods] - impl PyChunkView { - $( - #[getter] - fn $name(&self, py: Python<'_>) -> $ret { - self.parent.borrow(py).chunks[self.chunk_idx].$name.clone() - } - )+ - } - }; -} - -/// Generate a `#[pymethods]` block with scalar getters for PyChunkView. -macro_rules! chunk_scalar_getters { - ($($name:ident),+ $(,)?) => { - #[pymethods] - impl PyChunkView { - $( - #[getter] - fn $name(&self, py: Python<'_>) -> usize { - self.parent.borrow(py).chunks[self.chunk_idx].$name - } - )+ - } - }; -} - -chunk_scalar_getters!(num_games, num_moves, num_positions); - -chunk_array_getters!( - boards, - castling, - en_passant, - halfmove_clock, - turn, - from_squares, - to_squares, - promotions, - clocks, - evals, - move_offsets, - position_offsets, - is_checkmate, - is_stalemate, - is_insufficient, - legal_move_count, - valid, - legal_move_from_squares, - legal_move_to_squares, - legal_move_promotions, - legal_move_offsets, -); - -chunk_vec_getters!( - headers -> Vec>, - outcome -> Vec>, - parse_errors -> Vec>, - comments -> Vec>, -); - -#[pymethods] -impl PyChunkView { - fn __repr__(&self, py: Python<'_>) -> String { - let borrowed = self.parent.borrow(py); - let chunk = &borrowed.chunks[self.chunk_idx]; - format!( - "", - self.chunk_idx, chunk.num_games, chunk.num_moves, chunk.num_positions - ) + index_to_game(py, &self.move_offsets, move_indices) } } @@ -468,7 +362,8 @@ impl ParsedGamesIter { /// /// Provides access to board positions, moves, metadata, and annotations /// for one game. All array properties return numpy array slices (views, -/// not copies) into the parent chunk's data. +/// not copies) into the parent data. For invalid games the view covers +/// only the actually-parsed prefix of the game's range. /// /// Board encoding: /// Piece values: 0=empty, 1=P, 2=N, 3=B, 4=R, 5=Q, 6=K (white), @@ -486,60 +381,40 @@ impl ParsedGamesIter { #[pyclass] pub struct PyGameView { data: Py, - /// Index of the chunk this game lives in. - chunk_idx: usize, - /// Local game index within the chunk. - local_idx: usize, - /// Move range within the chunk's move arrays. + /// Global game index. + game_idx: usize, + /// Move range within the global move arrays (parsed moves only). move_start: usize, move_end: usize, - /// Position range within the chunk's position arrays. + /// Position range within the global position arrays (parsed only). pos_start: usize, pos_end: usize, } impl PyGameView { - /// Create a new game view for global game index `global_idx`. - pub fn new(py: Python<'_>, data: Py, global_idx: usize) -> PyResult { + /// Create a new game view for game index `game_idx`. + pub fn new(py: Python<'_>, data: Py, game_idx: usize) -> PyResult { let borrowed = data.borrow(py); - let (chunk_idx, local_idx) = borrowed.locate_game(global_idx); - let chunk = &borrowed.chunks[chunk_idx]; - - // Read this chunk's local CSR offsets - let move_offsets = chunk.move_offsets.bind(py); - let move_offsets: &Bound<'_, PyArray1> = move_offsets.cast()?; - let pos_offsets = chunk.position_offsets.bind(py); - let pos_offsets: &Bound<'_, PyArray1> = pos_offsets.cast()?; - - let move_offsets_ro = move_offsets.readonly(); - let move_offsets_slice = move_offsets_ro.as_slice()?; - let pos_offsets_ro = pos_offsets.readonly(); - let pos_offsets_slice = pos_offsets_ro.as_slice()?; + if game_idx >= borrowed.num_games { + return Err(pyo3::exceptions::PyIndexError::new_err( + "Invalid game index", + )); + } - let move_start = move_offsets_slice - .get(local_idx) - .ok_or_else(|| pyo3::exceptions::PyIndexError::new_err("Invalid game index"))?; - let move_end = move_offsets_slice - .get(local_idx + 1) - .ok_or_else(|| pyo3::exceptions::PyIndexError::new_err("Invalid game index"))?; - let pos_start = pos_offsets_slice - .get(local_idx) - .ok_or_else(|| pyo3::exceptions::PyIndexError::new_err("Invalid game index"))?; - let pos_end = pos_offsets_slice - .get(local_idx + 1) - .ok_or_else(|| pyo3::exceptions::PyIndexError::new_err("Invalid game index"))?; + let move_start = borrowed.move_offsets_rs[game_idx] as usize; + let pos_start = borrowed.position_offsets_rs[game_idx] as usize; + let parsed = borrowed.parsed_move_counts_rs[game_idx] as usize; drop(borrowed); Ok(Self { data, - chunk_idx, - local_idx, - move_start: *move_start as usize, - move_end: *move_end as usize, - pos_start: *pos_start as usize, - pos_end: *pos_end as usize, + game_idx, + move_start, + move_end: move_start + parsed, + pos_start, + pos_end: pos_start + parsed + 1, }) } } @@ -553,7 +428,7 @@ macro_rules! game_view_pos_getters { #[getter] fn $name<'py>(&self, py: Python<'py>) -> PyResult> { let borrowed = self.data.borrow(py); - let arr = borrowed.chunks[self.chunk_idx].$name.bind(py); + let arr = borrowed.$name.bind(py); let slice_obj = PySlice::new(py, self.pos_start as isize, self.pos_end as isize, 1); let slice = arr.call_method1("__getitem__", (slice_obj,))?; Ok(slice.unbind()) @@ -572,7 +447,7 @@ macro_rules! game_view_move_getters { #[getter] fn $name<'py>(&self, py: Python<'py>) -> PyResult> { let borrowed = self.data.borrow(py); - let arr = borrowed.chunks[self.chunk_idx].$name.bind(py); + let arr = borrowed.$name.bind(py); let slice_obj = PySlice::new(py, self.move_start as isize, self.move_end as isize, 1); let slice = arr.call_method1("__getitem__", (slice_obj,))?; Ok(slice.unbind()) @@ -602,7 +477,7 @@ impl PyGameView { #[getter] fn initial_board<'py>(&self, py: Python<'py>) -> PyResult> { let borrowed = self.data.borrow(py); - let boards = borrowed.chunks[self.chunk_idx].boards.bind(py); + let boards = borrowed.boards.bind(py); let slice = boards.call_method1("__getitem__", (self.pos_start,))?; Ok(slice.unbind()) } @@ -611,7 +486,7 @@ impl PyGameView { #[getter] fn final_board<'py>(&self, py: Python<'py>) -> PyResult> { let borrowed = self.data.borrow(py); - let boards = borrowed.chunks[self.chunk_idx].boards.bind(py); + let boards = borrowed.boards.bind(py); let slice = boards.call_method1("__getitem__", (self.pos_end - 1,))?; Ok(slice.unbind()) } @@ -622,19 +497,19 @@ impl PyGameView { #[getter] fn headers(&self, py: Python<'_>) -> PyResult> { let borrowed = self.data.borrow(py); - Ok(borrowed.chunks[self.chunk_idx].headers[self.local_idx].clone()) + Ok(borrowed.headers[self.game_idx].clone()) } /// Final position is checkmate. #[getter] fn is_checkmate(&self, py: Python<'_>) -> PyResult { let borrowed = self.data.borrow(py); - let arr = borrowed.chunks[self.chunk_idx].is_checkmate.bind(py); + let arr = borrowed.is_checkmate.bind(py); let arr: &Bound<'_, PyArray1> = arr.cast()?; let readonly = arr.readonly(); let slice = readonly.as_slice()?; slice - .get(self.local_idx) + .get(self.game_idx) .copied() .ok_or_else(|| pyo3::exceptions::PyIndexError::new_err("Invalid game index")) } @@ -643,12 +518,12 @@ impl PyGameView { #[getter] fn is_stalemate(&self, py: Python<'_>) -> PyResult { let borrowed = self.data.borrow(py); - let arr = borrowed.chunks[self.chunk_idx].is_stalemate.bind(py); + let arr = borrowed.is_stalemate.bind(py); let arr: &Bound<'_, PyArray1> = arr.cast()?; let readonly = arr.readonly(); let slice = readonly.as_slice()?; slice - .get(self.local_idx) + .get(self.game_idx) .copied() .ok_or_else(|| pyo3::exceptions::PyIndexError::new_err("Invalid game index")) } @@ -657,11 +532,11 @@ impl PyGameView { #[getter] fn is_insufficient(&self, py: Python<'_>) -> PyResult<(bool, bool)> { let borrowed = self.data.borrow(py); - let arr = borrowed.chunks[self.chunk_idx].is_insufficient.bind(py); + let arr = borrowed.is_insufficient.bind(py); let arr: &Bound<'_, PyArray2> = arr.cast()?; let readonly = arr.readonly(); let slice = readonly.as_slice()?; - let base = self.local_idx * 2; + let base = self.game_idx * 2; let white = slice .get(base) .copied() @@ -677,12 +552,12 @@ impl PyGameView { #[getter] fn legal_move_count(&self, py: Python<'_>) -> PyResult { let borrowed = self.data.borrow(py); - let arr = borrowed.chunks[self.chunk_idx].legal_move_count.bind(py); + let arr = borrowed.legal_move_count.bind(py); let arr: &Bound<'_, PyArray1> = arr.cast()?; let readonly = arr.readonly(); let slice = readonly.as_slice()?; slice - .get(self.local_idx) + .get(self.game_idx) .copied() .ok_or_else(|| pyo3::exceptions::PyIndexError::new_err("Invalid game index")) } @@ -691,12 +566,12 @@ impl PyGameView { #[getter] fn is_valid(&self, py: Python<'_>) -> PyResult { let borrowed = self.data.borrow(py); - let arr = borrowed.chunks[self.chunk_idx].valid.bind(py); + let arr = borrowed.valid.bind(py); let arr: &Bound<'_, PyArray1> = arr.cast()?; let readonly = arr.readonly(); let slice = readonly.as_slice()?; slice - .get(self.local_idx) + .get(self.game_idx) .copied() .ok_or_else(|| pyo3::exceptions::PyIndexError::new_err("Invalid game index")) } @@ -705,31 +580,30 @@ impl PyGameView { #[getter] fn is_game_over(&self, py: Python<'_>) -> PyResult { let borrowed = self.data.borrow(py); - let chunk = &borrowed.chunks[self.chunk_idx]; - let cm_arr = chunk.is_checkmate.bind(py); + let cm_arr = borrowed.is_checkmate.bind(py); let cm_arr: &Bound<'_, PyArray1> = cm_arr.cast()?; let checkmate = cm_arr .readonly() .as_slice()? - .get(self.local_idx) + .get(self.game_idx) .copied() .unwrap_or(false); - let sm_arr = chunk.is_stalemate.bind(py); + let sm_arr = borrowed.is_stalemate.bind(py); let sm_arr: &Bound<'_, PyArray1> = sm_arr.cast()?; let stalemate = sm_arr .readonly() .as_slice()? - .get(self.local_idx) + .get(self.game_idx) .copied() .unwrap_or(false); - let ins_arr = chunk.is_insufficient.bind(py); + let ins_arr = borrowed.is_insufficient.bind(py); let ins_arr: &Bound<'_, PyArray2> = ins_arr.cast()?; let ins_slice = ins_arr.readonly(); let ins_slice = ins_slice.as_slice()?; - let base = self.local_idx * 2; + let base = self.game_idx * 2; let insuf_white = ins_slice.get(base).copied().unwrap_or(false); let insuf_black = ins_slice.get(base + 1).copied().unwrap_or(false); @@ -740,14 +614,14 @@ impl PyGameView { #[getter] fn outcome(&self, py: Python<'_>) -> PyResult> { let borrowed = self.data.borrow(py); - Ok(borrowed.chunks[self.chunk_idx].outcome[self.local_idx].clone()) + Ok(borrowed.outcome[self.game_idx].clone()) } /// Parse error message if the game failed to parse, or None if valid. #[getter] fn parse_error(&self, py: Python<'_>) -> PyResult> { let borrowed = self.data.borrow(py); - Ok(borrowed.chunks[self.chunk_idx].parse_errors[self.local_idx].clone()) + Ok(borrowed.parse_errors[self.game_idx].clone()) } /// Raw text comments per move (only populated when store_comments=true). @@ -755,11 +629,10 @@ impl PyGameView { #[getter] fn comments(&self, py: Python<'_>) -> PyResult>> { let borrowed = self.data.borrow(py); - let chunk_comments = &borrowed.chunks[self.chunk_idx].comments; - if chunk_comments.is_empty() { + if borrowed.comments.is_empty() { return Ok(Vec::new()); } - Ok(chunk_comments[self.move_start..self.move_end].to_vec()) + Ok(borrowed.comments[self.move_start..self.move_end].to_vec()) } /// Legal moves at each position in this game. @@ -768,28 +641,27 @@ impl PyGameView { #[getter] fn legal_moves(&self, py: Python<'_>) -> PyResult>> { let borrowed = self.data.borrow(py); - let chunk = &borrowed.chunks[self.chunk_idx]; - if chunk.num_legal_moves == 0 { + if borrowed.num_legal_moves == 0 { return Ok(Vec::new()); } - let offsets_arr = chunk.legal_move_offsets.bind(py); + let offsets_arr = borrowed.legal_move_offsets.bind(py); let offsets_arr: &Bound<'_, PyArray1> = offsets_arr.cast()?; let offsets_ro = offsets_arr.readonly(); let offsets_slice = offsets_ro.as_slice()?; - let from_arr = chunk.legal_move_from_squares.bind(py); + let from_arr = borrowed.legal_move_from_squares.bind(py); let from_arr: &Bound<'_, PyArray1> = from_arr.cast()?; let from_ro = from_arr.readonly(); let from_slice = from_ro.as_slice()?; - let to_arr = chunk.legal_move_to_squares.bind(py); + let to_arr = borrowed.legal_move_to_squares.bind(py); let to_arr: &Bound<'_, PyArray1> = to_arr.cast()?; let to_ro = to_arr.readonly(); let to_slice = to_ro.as_slice()?; - let promo_arr = chunk.legal_move_promotions.bind(py); + let promo_arr = borrowed.legal_move_promotions.bind(py); let promo_arr: &Bound<'_, PyArray1> = promo_arr.cast()?; let promo_ro = promo_arr.readonly(); let promo_slice = promo_ro.as_slice()?; @@ -820,12 +692,11 @@ impl PyGameView { } let borrowed = self.data.borrow(py); - let chunk = &borrowed.chunks[self.chunk_idx]; - let from_arr = chunk.from_squares.bind(py); + let from_arr = borrowed.from_squares.bind(py); let from_arr: &Bound<'_, PyArray1> = from_arr.cast()?; - let to_arr = chunk.to_squares.bind(py); + let to_arr = borrowed.to_squares.bind(py); let to_arr: &Bound<'_, PyArray1> = to_arr.cast()?; - let promo_arr = chunk.promotions.bind(py); + let promo_arr = borrowed.promotions.bind(py); let promo_arr: &Bound<'_, PyArray1> = promo_arr.cast()?; let abs_idx = self.move_start + move_idx; @@ -839,13 +710,12 @@ impl PyGameView { /// Get all moves as UCI strings (e.g. ["e2e4", "e7e5", "g1f3"]). fn moves_uci(&self, py: Python<'_>) -> PyResult> { let borrowed = self.data.borrow(py); - let chunk = &borrowed.chunks[self.chunk_idx]; - let from_arr = chunk.from_squares.bind(py); + let from_arr = borrowed.from_squares.bind(py); let from_arr: &Bound<'_, PyArray1> = from_arr.cast()?; - let to_arr = chunk.to_squares.bind(py); + let to_arr = borrowed.to_squares.bind(py); let to_arr: &Bound<'_, PyArray1> = to_arr.cast()?; - let promo_arr = chunk.promotions.bind(py); + let promo_arr = borrowed.promotions.bind(py); let promo_arr: &Bound<'_, PyArray1> = promo_arr.cast()?; let from_slice = from_arr.readonly(); diff --git a/src/test.py b/src/test.py index 7090ef9..6f719a7 100644 --- a/src/test.py +++ b/src/test.py @@ -15,7 +15,6 @@ def test_basic_structure(self): "1. d4 d5 2. c4 e6 0-1", ] chunked = pa.chunked_array([pa.array(pgns)]) - # Use 1 thread to get a single chunk for predictable array shapes result = rust_pgn_reader_python_binding.parse_games(chunked, num_threads=1) # Check game count @@ -33,16 +32,14 @@ def test_basic_structure(self): self.assertEqual(len(game1), 4) # Game 2: 4 half-moves self.assertEqual(game1.num_positions, 5) - # With 1 thread, all games are in a single chunk - self.assertEqual(result.num_chunks, 1) - chunk = result.chunks[0] + # Global flat arrays total_moves = 9 total_positions = 9 + 2 - self.assertEqual(chunk.boards.shape, (total_positions, 8, 8)) - self.assertEqual(chunk.castling.shape, (total_positions, 4)) - self.assertEqual(chunk.en_passant.shape, (total_positions,)) - self.assertEqual(chunk.from_squares.shape, (total_moves,)) - self.assertEqual(chunk.valid.shape, (2,)) + self.assertEqual(result.boards.shape, (total_positions, 8, 8)) + self.assertEqual(result.castling.shape, (total_positions, 4)) + self.assertEqual(result.en_passant.shape, (total_positions,)) + self.assertEqual(result.from_squares.shape, (total_moves,)) + self.assertEqual(result.valid.shape, (2,)) def test_initial_board_encoding(self): """Test initial board state encoding.""" @@ -764,8 +761,8 @@ def test_parse_games_from_strings_with_status(self): self.assertEqual(moves1[0], "e2e4") def test_multithreaded_correctness(self): - """Test that multithreaded parsing produces correct results across chunk boundaries.""" - # Generate enough games to force multiple chunks with 4 threads + """Test that multithreaded parsing produces identical results.""" + # Generate enough games to spread across multiple threads pgns = [] for i in range(40): if i % 3 == 0: @@ -791,9 +788,6 @@ def test_multithreaded_correctness(self): self.assertEqual(result_1t.num_positions, result_4t.num_positions) self.assertEqual(result_4t.num_games, 40) - # Multiple chunks were actually created - self.assertGreater(result_4t.num_chunks, 1) - # Every game matches: moves, outcome, validity for i in range(40): g1 = result_1t[i] @@ -802,18 +796,26 @@ def test_multithreaded_correctness(self): self.assertEqual(g1.outcome, g4.outcome, f"Game {i} outcome differs") self.assertEqual(g1.is_valid, g4.is_valid, f"Game {i} validity differs") - # position_to_game works across chunk boundaries + # position_to_game agrees between thread counts all_pos = np.arange(result_4t.num_positions) game_ids = result_4t.position_to_game(all_pos) game_ids_1t = result_1t.position_to_game(all_pos) np.testing.assert_array_equal(game_ids, game_ids_1t) - # move_to_game works across chunk boundaries + # move_to_game agrees between thread counts all_moves = np.arange(result_4t.num_moves) move_game_ids = result_4t.move_to_game(all_moves) move_game_ids_1t = result_1t.move_to_game(all_moves) np.testing.assert_array_equal(move_game_ids, move_game_ids_1t) + # Global flat arrays are identical regardless of thread count + np.testing.assert_array_equal(result_1t.boards, result_4t.boards) + np.testing.assert_array_equal(result_1t.from_squares, result_4t.from_squares) + np.testing.assert_array_equal(result_1t.move_offsets, result_4t.move_offsets) + np.testing.assert_array_equal( + result_1t.position_offsets, result_4t.position_offsets + ) + def test_empty_inputs(self): """Test empty inputs return valid empty results.""" # parse_games_from_strings with empty list @@ -829,52 +831,54 @@ def test_empty_inputs(self): self.assertEqual(result2.num_games, 0) self.assertEqual(len(result2), 0) - def test_chunk_view_getters(self): - """Test PyChunkView exposes correct array shapes and values.""" + def test_flat_array_getters(self): + """Test ParsedGames exposes correct global flat arrays.""" pgns = ["1. e4 e5 1-0", "1. d4 d5 2. c4 0-1"] chunked = pa.chunked_array([pa.array(pgns)]) result = rust_pgn_reader_python_binding.parse_games(chunked, num_threads=1) - self.assertEqual(result.num_chunks, 1) - chunk = result.chunks[0] - # Scalar getters - self.assertEqual(chunk.num_games, 2) - self.assertEqual(chunk.num_moves, 5) - self.assertEqual(chunk.num_positions, 7) + self.assertEqual(result.num_games, 2) + self.assertEqual(result.num_moves, 5) + self.assertEqual(result.num_positions, 7) # Array shapes - self.assertEqual(chunk.boards.shape, (7, 8, 8)) - self.assertEqual(chunk.castling.shape, (7, 4)) - self.assertEqual(chunk.en_passant.shape, (7,)) - self.assertEqual(chunk.halfmove_clock.shape, (7,)) - self.assertEqual(chunk.turn.shape, (7,)) - self.assertEqual(chunk.from_squares.shape, (5,)) - self.assertEqual(chunk.to_squares.shape, (5,)) - self.assertEqual(chunk.promotions.shape, (5,)) - self.assertEqual(chunk.clocks.shape, (5,)) - self.assertEqual(chunk.evals.shape, (5,)) - self.assertEqual(chunk.is_checkmate.shape, (2,)) - self.assertEqual(chunk.is_stalemate.shape, (2,)) - self.assertEqual(chunk.is_insufficient.shape, (2, 2)) - self.assertEqual(chunk.legal_move_count.shape, (2,)) - self.assertEqual(chunk.valid.shape, (2,)) - - # CSR offset shapes - self.assertEqual(chunk.move_offsets.shape, (3,)) # 2 games + 1 - self.assertEqual(chunk.position_offsets.shape, (3,)) - - # Vec getters - self.assertEqual(len(chunk.headers), 2) - self.assertEqual(len(chunk.outcome), 2) - self.assertEqual(chunk.outcome[0], "White") - self.assertEqual(chunk.outcome[1], "Black") - - # Chunk values match per-game views + self.assertEqual(result.boards.shape, (7, 8, 8)) + self.assertEqual(result.castling.shape, (7, 4)) + self.assertEqual(result.en_passant.shape, (7,)) + self.assertEqual(result.halfmove_clock.shape, (7,)) + self.assertEqual(result.turn.shape, (7,)) + self.assertEqual(result.from_squares.shape, (5,)) + self.assertEqual(result.to_squares.shape, (5,)) + self.assertEqual(result.promotions.shape, (5,)) + self.assertEqual(result.clocks.shape, (5,)) + self.assertEqual(result.evals.shape, (5,)) + self.assertEqual(result.is_checkmate.shape, (2,)) + self.assertEqual(result.is_stalemate.shape, (2,)) + self.assertEqual(result.is_insufficient.shape, (2, 2)) + self.assertEqual(result.legal_move_count.shape, (2,)) + self.assertEqual(result.valid.shape, (2,)) + + # CSR offset shapes and values + self.assertEqual(result.move_offsets.shape, (3,)) # 2 games + 1 + self.assertEqual(result.position_offsets.shape, (3,)) + np.testing.assert_array_equal(result.move_offsets, [0, 2, 5]) + np.testing.assert_array_equal(result.position_offsets, [0, 3, 7]) + self.assertEqual(result.parsed_move_counts.shape, (2,)) + np.testing.assert_array_equal(result.parsed_move_counts, [2, 3]) + + # List getters + self.assertEqual(len(result.headers), 2) + self.assertEqual(len(result.outcome), 2) + self.assertEqual(result.outcome[0], "White") + self.assertEqual(result.outcome[1], "Black") + self.assertEqual(result.parse_errors, [None, None]) + + # Flat array values match per-game views game0 = result[0] game1 = result[1] - np.testing.assert_array_equal(game0.boards, chunk.boards[:3]) - np.testing.assert_array_equal(game1.boards, chunk.boards[3:]) + np.testing.assert_array_equal(game0.boards, result.boards[:3]) + np.testing.assert_array_equal(game1.boards, result.boards[3:]) def test_legal_moves_from_python(self): """Test legal_moves property returns correct moves for known positions.""" @@ -927,7 +931,7 @@ def test_parse_error_invalid_fen(self): self.assertIn("FEN", result[0].parse_error) def test_repr(self): - """Test __repr__ on ParsedGames, PyGameView, PyChunkView.""" + """Test __repr__ on ParsedGames and PyGameView.""" pgns = ["1. e4 e5 1-0"] result = rust_pgn_reader_python_binding.parse_games_from_strings(pgns) @@ -940,10 +944,6 @@ def test_repr(self): game_repr = repr(result[0]) self.assertIn("PyGameView", game_repr) - # PyChunkView repr - chunk_repr = repr(result.chunks[0]) - self.assertIn("PyChunkView", chunk_repr) - if __name__ == "__main__": unittest.main() diff --git a/src/visitor.rs b/src/visitor.rs deleted file mode 100644 index 5605a73..0000000 --- a/src/visitor.rs +++ /dev/null @@ -1,904 +0,0 @@ -//! SoA (Struct-of-Arrays) visitor for PGN parsing. -//! -//! This module provides a memory-efficient parsing approach that writes -//! directly to shared buffers instead of allocating per-game Vec structures. -//! Used by `parse_games` for optimal performance. - -use crate::board_serialization::{ - get_castling_rights, get_en_passant_file, get_halfmove_clock, get_turn, serialize_board, -}; -use crate::comment_parsing::{CommentContent, ParsedTag, parse_comments}; -use crate::tokenizer::{Outcome, Visitor}; -use shakmaty::{ - CastlingMode, CastlingSide, Chess, Color, Position, - fen::Fen, - san::{San, SanPlus, Suffix}, - uci::UciMove, -}; -use std::collections::HashMap; - -/// Configuration for what optional data to store during parsing. -#[derive(Clone, Debug)] -pub struct ParseConfig { - pub store_comments: bool, - pub store_legal_moves: bool, -} - -/// Compute CSR-style prefix-sum offsets from a slice of counts. -/// Returns a Vec of length `counts.len() + 1`, starting with 0. -fn prefix_sum(counts: &[u32]) -> Vec { - let mut offsets = Vec::with_capacity(counts.len() + 1); - let mut acc: u32 = 0; - offsets.push(0); - for &count in counts { - acc += count; - offsets.push(acc); - } - offsets -} - -/// Accumulated buffers for multiple parsed games. -/// -/// This struct holds all data in a struct-of-arrays layout, optimized for: -/// - Efficient thread-local accumulation during parallel parsing -/// - Fast merging of thread-local buffers via `extend_from_slice` -/// - Direct conversion to NumPy arrays without intermediate allocations -#[derive(Default, Clone)] -pub struct Buffers { - // Board state arrays (one entry per position) - pub boards: Vec, // Flattened: 64 bytes per position - pub castling: Vec, // Flattened: 4 bools per position [K,Q,k,q] - pub en_passant: Vec, // Per position: -1 or file 0-7 - pub halfmove_clock: Vec, // Per position - pub turn: Vec, // Per position: true=white - - // Move arrays (one entry per move) - pub from_squares: Vec, - pub to_squares: Vec, - pub promotions: Vec, // -1 for no promotion - pub clocks: Vec, // NaN for missing - pub evals: Vec, // NaN for missing - - // Per-game data - pub move_counts: Vec, // Number of moves per game - pub position_counts: Vec, // Number of positions per game - pub is_checkmate: Vec, - pub is_stalemate: Vec, - pub is_insufficient: Vec, // Flattened: 2 bools per game [white, black] - pub legal_move_count: Vec, - pub valid: Vec, - pub headers: Vec>, - pub outcome: Vec>, // "White", "Black", "Draw", "Unknown", or None - pub parse_errors: Vec>, // Per-game: None if valid, Some(msg) if not - - // Optional: raw text comments (per-move), only populated when store_comments=true - pub comments: Vec>, - - // Optional: legal moves at each position, only populated when store_legal_moves=true - // Stored as flat arrays with CSR-style offsets - pub legal_move_from_squares: Vec, - pub legal_move_to_squares: Vec, - pub legal_move_promotions: Vec, - pub legal_move_counts: Vec, // Number of legal moves per position -} - -impl Buffers { - /// Create a new Buffers with pre-allocated capacity. - /// - /// # Arguments - /// * `estimated_games` - Expected number of games - /// * `moves_per_game` - Expected average moves per game (default: 70) - /// * `config` - Configuration for optional features - pub fn with_capacity( - estimated_games: usize, - moves_per_game: usize, - config: &ParseConfig, - ) -> Self { - let estimated_moves = estimated_games * moves_per_game; - let estimated_positions = estimated_moves + estimated_games; // +1 initial position per game - - Buffers { - // Board state arrays - boards: Vec::with_capacity(estimated_positions * 64), - castling: Vec::with_capacity(estimated_positions * 4), - en_passant: Vec::with_capacity(estimated_positions), - halfmove_clock: Vec::with_capacity(estimated_positions), - turn: Vec::with_capacity(estimated_positions), - - // Move arrays - from_squares: Vec::with_capacity(estimated_moves), - to_squares: Vec::with_capacity(estimated_moves), - promotions: Vec::with_capacity(estimated_moves), - clocks: Vec::with_capacity(estimated_moves), - evals: Vec::with_capacity(estimated_moves), - - // Per-game data - move_counts: Vec::with_capacity(estimated_games), - position_counts: Vec::with_capacity(estimated_games), - is_checkmate: Vec::with_capacity(estimated_games), - is_stalemate: Vec::with_capacity(estimated_games), - is_insufficient: Vec::with_capacity(estimated_games * 2), - legal_move_count: Vec::with_capacity(estimated_games), - valid: Vec::with_capacity(estimated_games), - headers: Vec::with_capacity(estimated_games), - outcome: Vec::with_capacity(estimated_games), - parse_errors: Vec::with_capacity(estimated_games), - - // Optional comments - comments: if config.store_comments { - Vec::with_capacity(estimated_moves) - } else { - Vec::new() - }, - - // Optional legal moves - legal_move_from_squares: if config.store_legal_moves { - Vec::with_capacity(estimated_positions * 30) - } else { - Vec::new() - }, - legal_move_to_squares: if config.store_legal_moves { - Vec::with_capacity(estimated_positions * 30) - } else { - Vec::new() - }, - legal_move_promotions: if config.store_legal_moves { - Vec::with_capacity(estimated_positions * 30) - } else { - Vec::new() - }, - legal_move_counts: if config.store_legal_moves { - Vec::with_capacity(estimated_positions) - } else { - Vec::new() - }, - } - } - - /// Number of games in this buffer. - pub fn num_games(&self) -> usize { - self.headers.len() - } - - /// Total number of moves across all games. - #[allow(dead_code)] - pub fn total_moves(&self) -> usize { - self.from_squares.len() - } - - /// Total number of positions across all games. - pub fn total_positions(&self) -> usize { - self.boards.len() / 64 - } - - /// Compute CSR-style offsets from move counts. - pub fn compute_move_offsets(&self) -> Vec { - prefix_sum(&self.move_counts) - } - - /// Compute CSR-style offsets from position counts. - pub fn compute_position_offsets(&self) -> Vec { - prefix_sum(&self.position_counts) - } - - /// Compute CSR-style offsets from legal move counts (per position). - pub fn compute_legal_move_offsets(&self) -> Vec { - prefix_sum(&self.legal_move_counts) - } - - /// Total number of legal moves stored across all positions. - pub fn total_legal_moves(&self) -> usize { - self.legal_move_from_squares.len() - } -} - -/// Support castling notated with zeros ("0-0", "0-0-0"), with optional -/// check/checkmate suffix, matching pgn-reader's explicit handling. -pub(crate) fn castling_with_zeros(token: &[u8]) -> Option { - let (body, suffix) = match token.last() { - Some(b'+') => (&token[..token.len() - 1], Some(Suffix::Check)), - Some(b'#') => (&token[..token.len() - 1], Some(Suffix::Checkmate)), - _ => (token, None), - }; - let side = match body { - b"0-0" => CastlingSide::KingSide, - b"0-0-0" => CastlingSide::QueenSide, - _ => return None, - }; - Some(SanPlus { - san: San::Castle(side), - suffix, - }) -} - -/// Visitor that writes directly to shared Buffers. -/// -/// This visitor does not allocate any per-game Vec structures. -/// All data is appended directly to the shared Buffers. -pub struct GameVisitor<'a> { - buffers: &'a mut Buffers, - config: ParseConfig, - pos: Chess, - valid_moves: bool, - current_headers: Vec<(String, String)>, - current_outcome: Option, - current_error: Option, - // Track counts for current game - current_move_count: u32, - current_position_count: u32, -} - -impl<'a> GameVisitor<'a> { - pub fn new(buffers: &'a mut Buffers, config: &ParseConfig) -> Self { - GameVisitor { - buffers, - config: config.clone(), - pos: Chess::default(), - valid_moves: true, - current_headers: Vec::with_capacity(10), - current_outcome: None, - current_error: None, - current_move_count: 0, - current_position_count: 0, - } - } - - /// Record current board state to buffers. - fn push_board_state(&mut self) { - self.buffers - .boards - .extend_from_slice(&serialize_board(&self.pos)); - let castling = get_castling_rights(&self.pos); - self.buffers.castling.extend_from_slice(&castling); - self.buffers.en_passant.push(get_en_passant_file(&self.pos)); - self.buffers - .halfmove_clock - .push(get_halfmove_clock(&self.pos)); - self.buffers.turn.push(get_turn(&self.pos)); - self.current_position_count += 1; - - // Store legal moves if enabled - if self.config.store_legal_moves { - self.push_legal_moves(); - } - } - - /// Record legal moves at current position to buffers. - fn push_legal_moves(&mut self) { - let legal_moves = self.pos.legal_moves(); - let mut count: u32 = 0; - for m in legal_moves { - let uci_move_obj = UciMove::from_standard(m); - if let UciMove::Normal { - from, - to, - promotion, - } = uci_move_obj - { - self.buffers.legal_move_from_squares.push(from as u8); - self.buffers.legal_move_to_squares.push(to as u8); - self.buffers - .legal_move_promotions - .push(promotion.map(|p| p as i8).unwrap_or(-1)); - count += 1; - } - } - self.buffers.legal_move_counts.push(count); - } - - /// Record move data to buffers. - fn push_move(&mut self, from: u8, to: u8, promotion: Option) { - self.buffers.from_squares.push(from); - self.buffers.to_squares.push(to); - self.buffers - .promotions - .push(promotion.map(|p| p as i8).unwrap_or(-1)); - // Push placeholders for clock and eval (will be overwritten by comment()) - self.buffers.clocks.push(f32::NAN); - self.buffers.evals.push(f32::NAN); - // Push comment placeholder if enabled (will be overwritten by comment()) - if self.config.store_comments { - self.buffers.comments.push(None); - } - self.current_move_count += 1; - } - - /// Record final position status. - fn update_position_status(&mut self) { - self.buffers.is_checkmate.push(self.pos.is_checkmate()); - self.buffers.is_stalemate.push(self.pos.is_stalemate()); - self.buffers - .is_insufficient - .push(self.pos.has_insufficient_material(Color::White)); - self.buffers - .is_insufficient - .push(self.pos.has_insufficient_material(Color::Black)); - self.buffers - .legal_move_count - .push(self.pos.legal_moves().len() as u16); - } - - /// Record a parse error for the current game. - fn set_error(&mut self, msg: String) { - self.valid_moves = false; - self.current_error = Some(msg); - } - - /// Finalize current game - record per-game data. - fn finalize_game(&mut self) { - self.buffers.move_counts.push(self.current_move_count); - self.buffers - .position_counts - .push(self.current_position_count); - self.buffers.valid.push(self.valid_moves); - self.buffers.outcome.push(self.current_outcome.take()); - self.buffers.parse_errors.push(self.current_error.take()); - - // Convert headers to HashMap - let header_map: HashMap = self.current_headers.drain(..).collect(); - self.buffers.headers.push(header_map); - } -} - -impl Visitor for GameVisitor<'_> { - fn begin_headers(&mut self) { - self.current_headers.clear(); - self.current_headers.reserve(10); - self.valid_moves = true; - self.current_outcome = None; - self.current_error = None; - self.current_move_count = 0; - self.current_position_count = 0; - } - - fn header(&mut self, key: &[u8], value: &[u8]) { - let key_str = String::from_utf8_lossy(key).into_owned(); - let value_str = String::from_utf8_lossy(value).into_owned(); - self.current_headers.push((key_str, value_str)); - } - - fn end_headers(&mut self) { - // Determine castling mode from Variant header (case-insensitive) - let castling_mode = self - .current_headers - .iter() - .find(|(k, _)| k.eq_ignore_ascii_case("Variant")) - .and_then(|(_, v)| { - let v_lower = v.to_lowercase(); - if v_lower == "chess960" { - Some(CastlingMode::Chess960) - } else { - None - } - }) - .unwrap_or(CastlingMode::Standard); - - // Try to parse FEN from headers, fall back to default position - let fen_header = self - .current_headers - .iter() - .find(|(k, _)| k.eq_ignore_ascii_case("FEN")) - .map(|(_, v)| v.as_str()); - - if let Some(fen_str) = fen_header { - match fen_str.parse::() { - Ok(fen) => match fen.into_position(castling_mode) { - Ok(pos) => self.pos = pos, - Err(e) => { - self.set_error(format!("invalid FEN position: {}", e)); - self.pos = Chess::default(); - } - }, - Err(e) => { - self.set_error(format!("failed to parse FEN: {}", e)); - self.pos = Chess::default(); - } - } - } else { - self.pos = Chess::default(); - } - - // Record initial board state - self.push_board_state(); - } - - fn san_token(&mut self, token: &[u8]) { - // Early abort: once the game is invalid, skip SAN parsing entirely. - if !self.valid_moves { - return; - } - - let san_plus = match castling_with_zeros(token) { - Some(san_plus) => san_plus, - None => match SanPlus::from_ascii(token) { - Ok(san_plus) => san_plus, - Err(error) => { - let token_str = String::from_utf8_lossy(token); - self.set_error(format!("failed to parse SAN: {} ({})", error, token_str)); - return; - } - }, - }; - - match san_plus.san.to_move(&self.pos) { - Ok(m) => { - self.pos.play_unchecked(m); - - // Record board state after move - self.push_board_state(); - - let uci_move_obj = UciMove::from_standard(m); - match uci_move_obj { - UciMove::Normal { - from, - to, - promotion, - } => { - self.push_move(from as u8, to as u8, promotion.map(|p| p as u8)); - } - _ => { - self.set_error(format!("unexpected UCI move type: {:?}", uci_move_obj)); - } - } - } - Err(err) => { - self.set_error(format!("illegal move: {} {}", err, san_plus)); - } - } - } - - fn comment(&mut self, comment: &[u8]) { - if let Ok((_, parsed_comments)) = parse_comments(comment) { - let mut move_comments = if self.config.store_comments { - Some(String::new()) - } else { - None - }; - - for content in parsed_comments { - match content { - CommentContent::Tag(tag_content) => match tag_content { - ParsedTag::Eval(eval_value) => { - // Update the last eval entry - if let Some(last_eval) = self.buffers.evals.last_mut() { - *last_eval = eval_value as f32; - } - } - ParsedTag::ClkTime { - hours, - minutes, - seconds, - } => { - // Convert to seconds and update the last clock entry - if let Some(last_clk) = self.buffers.clocks.last_mut() { - *last_clk = - hours as f32 * 3600.0 + minutes as f32 * 60.0 + seconds as f32; - } - } - ParsedTag::Mate(mate_value) => { - // Mate scores stored as text in comments (matching old API behavior) - if let Some(ref mut comments) = move_comments { - if !comments.is_empty() && !comments.ends_with(' ') { - comments.push(' '); - } - comments.push_str(&format!("[Mate {}]", mate_value)); - } - } - }, - CommentContent::Text(text) => { - if let Some(ref mut comments) = move_comments { - if !text.trim().is_empty() { - if !comments.is_empty() { - comments.push(' '); - } - comments.push_str(&text); - } - } - } - } - } - - // Update the last comment entry if comments are enabled - if let Some(comment_text) = move_comments { - if let Some(last_comment) = self.buffers.comments.last_mut() { - *last_comment = Some(comment_text); - } - } - } - } - - fn outcome(&mut self, _outcome: Outcome) { - self.current_outcome = Some(match _outcome { - Outcome::WhiteWins => "White".to_string(), - Outcome::BlackWins => "Black".to_string(), - Outcome::Draw => "Draw".to_string(), - Outcome::Unknown => "Unknown".to_string(), - }); - self.update_position_status(); - } - - fn end_game(&mut self) { - // Handle case where outcome() was not called (e.g., incomplete game) - if self.buffers.is_checkmate.len() < self.buffers.headers.len() + 1 { - self.update_position_status(); - } - self.finalize_game(); - } -} - -/// Parse a single game directly into Buffers. -pub fn parse_game_to_buffers( - pgn: &str, - buffers: &mut Buffers, - config: &ParseConfig, -) -> Result { - let mut visitor = GameVisitor::new(buffers, config); - - // Check if the PGN string has any non-whitespace characters to emulate pgn-reader's `Ok(None)` behavior - if pgn.trim().is_empty() { - return Err("No game found in PGN".to_string()); - } - - crate::tokenizer::parse_game(pgn.as_bytes(), &mut visitor); - Ok(visitor.valid_moves) -} - -#[cfg(test)] -mod tests { - use super::*; - - fn default_config() -> ParseConfig { - ParseConfig { - store_comments: false, - store_legal_moves: false, - } - } - - #[test] - fn test_parse_simple_game() { - let pgn = r#"[Event "Test"] -[White "Player1"] -[Black "Player2"] -[Result "1-0"] - -1. e4 e5 2. Nf3 Nc6 1-0"#; - - let config = default_config(); - let mut buffers = Buffers::with_capacity(1, 70, &config); - let result = parse_game_to_buffers(pgn, &mut buffers, &config); - - assert!(result.is_ok()); - assert!(result.unwrap()); // valid game - assert_eq!(buffers.num_games(), 1); - assert_eq!(buffers.move_counts[0], 4); // 4 moves - assert_eq!(buffers.position_counts[0], 5); // 5 positions (initial + 4 moves) - assert_eq!(buffers.total_moves(), 4); - assert_eq!(buffers.total_positions(), 5); - assert_eq!(buffers.outcome[0], Some("White".to_string())); - } - - #[test] - fn test_parse_game_with_annotations() { - let pgn = r#"[Event "Test"] -[Result "1-0"] - -1. e4 { [%eval 0.17] [%clk 0:03:00] } 1... e5 { [%eval 0.19] [%clk 0:02:58] } 1-0"#; - - let config = default_config(); - let mut buffers = Buffers::with_capacity(1, 70, &config); - let result = parse_game_to_buffers(pgn, &mut buffers, &config); - - assert!(result.is_ok()); - assert_eq!(buffers.total_moves(), 2); - - // Check that evals were parsed - assert!(!buffers.evals[0].is_nan()); - assert!((buffers.evals[0] - 0.17).abs() < 0.01); - assert!(!buffers.evals[1].is_nan()); - assert!((buffers.evals[1] - 0.19).abs() < 0.01); - - // Check that clocks were parsed (3 minutes = 180 seconds) - assert!(!buffers.clocks[0].is_nan()); - assert!((buffers.clocks[0] - 180.0).abs() < 0.01); - } - - #[test] - fn test_multiple_games_in_one_buffer() { - let pgn1 = r#"[Event "Game1"] -[Result "1-0"] - -1. e4 e5 1-0"#; - - let pgn2 = r#"[Event "Game2"] -[Result "0-1"] - -1. d4 d5 2. c4 0-1"#; - - let config = default_config(); - let mut buffers = Buffers::with_capacity(2, 70, &config); - - parse_game_to_buffers(pgn1, &mut buffers, &config).unwrap(); - parse_game_to_buffers(pgn2, &mut buffers, &config).unwrap(); - - assert_eq!(buffers.num_games(), 2); - assert_eq!(buffers.total_moves(), 5); // 2 + 3 moves - assert_eq!(buffers.move_counts, vec![2, 3]); - assert_eq!(buffers.outcome[0], Some("White".to_string())); - assert_eq!(buffers.outcome[1], Some("Black".to_string())); - } - - #[test] - fn test_compute_offsets() { - let mut buffers = Buffers::default(); - buffers.move_counts = vec![4, 6, 3]; - buffers.position_counts = vec![5, 7, 4]; - - let move_offsets = buffers.compute_move_offsets(); - assert_eq!(move_offsets, vec![0, 4, 10, 13]); - - let pos_offsets = buffers.compute_position_offsets(); - assert_eq!(pos_offsets, vec![0, 5, 12, 16]); - } - - #[test] - fn test_castling_with_zeros_normalized() { - // Zeros notation is handled in the visitor (moved from the tokenizer). - let pgn = "1. e4 e5 2. Nf3 Nf6 3. Bc4 Bc5 4. 0-0 0-0 1-0"; - - let config = default_config(); - let mut buffers = Buffers::with_capacity(1, 70, &config); - let result = parse_game_to_buffers(pgn, &mut buffers, &config); - - assert!(result.is_ok()); - assert!(result.unwrap(), "zeros castling should parse as O-O"); - assert_eq!(buffers.total_moves(), 8); - assert_eq!(buffers.parse_errors[0], None); - } - - #[test] - fn test_invalid_san_token_sets_parse_error() { - let pgn = "1. e4 xyzzy9 e5 1-0"; - - let config = default_config(); - let mut buffers = Buffers::with_capacity(1, 70, &config); - let result = parse_game_to_buffers(pgn, &mut buffers, &config); - - assert!(result.is_ok()); - assert!(!result.unwrap(), "garbage SAN token should invalidate game"); - let err = buffers.parse_errors[0].as_deref().unwrap(); - assert!( - err.starts_with("failed to parse SAN:") && err.contains("xyzzy9"), - "unexpected error message: {err}" - ); - // Moves stop being recorded after the first error. - assert_eq!(buffers.total_moves(), 1); - } - - #[test] - fn test_outcome_without_headers() { - // PGN without Result header - outcome comes from movetext - let pgn = "1. e4 e5 2. Nf3 Nc6 0-1"; - - let config = default_config(); - let mut buffers = Buffers::with_capacity(1, 70, &config); - parse_game_to_buffers(pgn, &mut buffers, &config).unwrap(); - - assert_eq!(buffers.outcome[0], Some("Black".to_string())); - } - - #[test] - fn test_outcome_draw() { - let pgn = "1. e4 e5 1/2-1/2"; - - let config = default_config(); - let mut buffers = Buffers::with_capacity(1, 70, &config); - parse_game_to_buffers(pgn, &mut buffers, &config).unwrap(); - - assert_eq!(buffers.outcome[0], Some("Draw".to_string())); - } - - #[test] - fn test_comments_disabled() { - let pgn = r#"1. e4 { a comment } e5 1-0"#; - - let config = default_config(); - let mut buffers = Buffers::with_capacity(1, 70, &config); - parse_game_to_buffers(pgn, &mut buffers, &config).unwrap(); - - assert!(buffers.comments.is_empty()); - } - - #[test] - fn test_comments_enabled() { - let pgn = r#"1. e4 { a comment } 1... e5 { [%eval 0.19] } 1-0"#; - - let config = ParseConfig { - store_comments: true, - store_legal_moves: false, - }; - let mut buffers = Buffers::with_capacity(1, 70, &config); - parse_game_to_buffers(pgn, &mut buffers, &config).unwrap(); - - assert_eq!(buffers.comments.len(), 2); - // Raw text from PGN includes surrounding spaces from the parser - assert_eq!(buffers.comments[0], Some(" a comment ".to_string())); - // The second comment only has an eval tag, so text portion is empty - assert_eq!(buffers.comments[1], Some("".to_string())); - } - - #[test] - fn test_legal_moves_disabled() { - let pgn = "1. e4 e5 1-0"; - - let config = default_config(); - let mut buffers = Buffers::with_capacity(1, 70, &config); - parse_game_to_buffers(pgn, &mut buffers, &config).unwrap(); - - assert!(buffers.legal_move_from_squares.is_empty()); - assert!(buffers.legal_move_counts.is_empty()); - } - - #[test] - fn test_legal_moves_enabled() { - let pgn = "1. e4 1-0"; - - let config = ParseConfig { - store_comments: false, - store_legal_moves: true, - }; - let mut buffers = Buffers::with_capacity(1, 70, &config); - parse_game_to_buffers(pgn, &mut buffers, &config).unwrap(); - - // 2 positions: initial + after e4 - assert_eq!(buffers.legal_move_counts.len(), 2); - // Initial position has 20 legal moves - assert_eq!(buffers.legal_move_counts[0], 20); - // After e4, black has 20 legal moves - assert_eq!(buffers.legal_move_counts[1], 20); - // Total legal moves stored - assert_eq!(buffers.legal_move_from_squares.len(), 40); - } - - #[test] - fn test_headers_only_game() { - // A game with headers but no movetext at all: the initial position - // (from the FEN header) must still be recorded. - let pgn = r#"[Event "Test"] -[FEN "r1bqkbnr/pppp1ppp/2n5/4p3/4P3/5N2/PPPP1PPP/RNBQKB1R w KQkq - 2 3"] -"#; - - let config = default_config(); - let mut buffers = Buffers::with_capacity(1, 70, &config); - let result = parse_game_to_buffers(pgn, &mut buffers, &config); - - assert!(result.is_ok()); - assert!(result.unwrap()); // valid game (no moves, no errors) - assert_eq!(buffers.num_games(), 1); - assert_eq!(buffers.move_counts[0], 0); - assert_eq!(buffers.position_counts[0], 1); // initial position only - assert_eq!(buffers.headers[0].get("Event"), Some(&"Test".to_string())); - assert_eq!(buffers.parse_errors[0], None); // FEN header was applied - assert_eq!(buffers.outcome[0], None); - } - - #[test] - fn test_parse_game_without_headers() { - let pgn = "1. Nf3 d5 2. e4 c5 3. exd5 e5 4. dxe6 0-1"; - - let config = default_config(); - let mut buffers = Buffers::with_capacity(1, 70, &config); - let result = parse_game_to_buffers(pgn, &mut buffers, &config); - - assert!(result.is_ok()); - assert!(result.unwrap()); // valid game - assert_eq!(buffers.num_games(), 1); - assert_eq!(buffers.total_moves(), 7); - assert_eq!(buffers.outcome[0], Some("Black".to_string())); - } - - #[test] - fn test_parse_game_with_standard_fen() { - // A game starting from a mid-game position - let pgn = r#"[FEN "r1bqkbnr/pppp1ppp/2n5/4p3/4P3/5N2/PPPP1PPP/RNBQKB1R w KQkq - 2 3"] - -3. Bb5 a6 4. Ba4 Nf6 1-0"#; - - let config = default_config(); - let mut buffers = Buffers::with_capacity(1, 70, &config); - let result = parse_game_to_buffers(pgn, &mut buffers, &config); - - assert!(result.is_ok()); - assert!(result.unwrap()); // valid game - assert_eq!(buffers.total_moves(), 4); - } - - #[test] - fn test_parse_chess960_game() { - // Chess960 game with custom starting position - let pgn = r#"[Variant "chess960"] -[FEN "brkrqnnb/pppppppp/8/8/8/8/PPPPPPPP/BRKRQNNB w KQkq - 0 1"] - -1. g3 d5 2. d4 g6 3. b3 Nf6 1-0"#; - - let config = default_config(); - let mut buffers = Buffers::with_capacity(1, 70, &config); - let result = parse_game_to_buffers(pgn, &mut buffers, &config); - - assert!(result.is_ok()); - assert!( - result.unwrap(), - "Chess960 moves should be valid with proper FEN" - ); - assert_eq!(buffers.total_moves(), 6); - } - - #[test] - fn test_parse_chess960_variant_case_insensitive() { - // Test that variant detection is case-insensitive - let pgn = r#"[Variant "Chess960"] -[FEN "brkrqnnb/pppppppp/8/8/8/8/PPPPPPPP/BRKRQNNB w KQkq - 0 1"] - -1. g3 d5 1-0"#; - - let config = default_config(); - let mut buffers = Buffers::with_capacity(1, 70, &config); - let result = parse_game_to_buffers(pgn, &mut buffers, &config); - - assert!(result.is_ok()); - assert!(result.unwrap(), "Should handle Chess960 case variations"); - } - - #[test] - fn test_parse_invalid_fen_falls_back() { - // Invalid FEN should fall back to default and mark invalid - let pgn = r#"[FEN "invalid fen string"] - -1. e4 e5 1-0"#; - - let config = default_config(); - let mut buffers = Buffers::with_capacity(1, 70, &config); - let result = parse_game_to_buffers(pgn, &mut buffers, &config); - - assert!(result.is_ok()); - assert!( - !result.unwrap(), - "Should mark as invalid when FEN parsing fails" - ); - } - - #[test] - fn test_fen_header_case_insensitive() { - // FEN header key should be case-insensitive - let pgn = r#"[fen "r1bqkbnr/pppp1ppp/2n5/4p3/4P3/5N2/PPPP1PPP/RNBQKB1R w KQkq - 2 3"] - -3. Bb5 1-0"#; - - let config = default_config(); - let mut buffers = Buffers::with_capacity(1, 70, &config); - let result = parse_game_to_buffers(pgn, &mut buffers, &config); - - assert!(result.is_ok()); - assert!(result.unwrap(), "Should handle lowercase 'fen' header"); - } - - #[test] - fn test_parse_game_with_custom_fen_no_variant() { - // Standard chess from a mid-game position (no Variant header) - // Position after 1.e4 e5 2.Nf3 Nc6 3.Bb5 (Ruy Lopez) - let pgn = r#"[Event "Test Game"] -[FEN "r1bqkbnr/pppp1ppp/2n5/1B2p3/4P3/5N2/PPPP1PPP/RNBQK2R b KQkq - 3 3"] - -3... a6 4. Ba4 Nf6 5. O-O Be7 1-0"#; - - let config = default_config(); - let mut buffers = Buffers::with_capacity(1, 70, &config); - let result = parse_game_to_buffers(pgn, &mut buffers, &config); - - assert!(result.is_ok()); - assert!( - result.unwrap(), - "Standard game with custom FEN should be valid" - ); - assert_eq!(buffers.total_moves(), 5); // a6, Ba4, Nf6, O-O, Be7 - } -} From 3ba51135bc21fd15a688f3176d14834e36439dbf Mon Sep 17 00:00:00 2001 From: vladkvit Date: Sat, 11 Jul 2026 12:23:49 -0400 Subject: [PATCH 8/8] Remove comparison code for verification --- .gitignore | 1 - src/compare_reference.py | 40 ------ src/dump_reference.py | 38 ------ src/reference_lib.py | 287 --------------------------------------- 4 files changed, 366 deletions(-) delete mode 100644 src/compare_reference.py delete mode 100644 src/dump_reference.py delete mode 100644 src/reference_lib.py diff --git a/.gitignore b/.gitignore index 01327b8..ad136bf 100644 --- a/.gitignore +++ b/.gitignore @@ -74,4 +74,3 @@ docs/_build/ # Local benchmark data and derived reference dumps *.parquet *.pgn -reference_*.json diff --git a/src/compare_reference.py b/src/compare_reference.py deleted file mode 100644 index dcac99e..0000000 --- a/src/compare_reference.py +++ /dev/null @@ -1,40 +0,0 @@ -"""Compare the CURRENT parser build's outputs against a saved reference. - -Usage: python src/compare_reference.py [reference_json] - -Exits nonzero and prints mismatches if the outputs differ. -""" - -import json -import sys -import time - -from reference_lib import build_reference, diff_dicts - -DEFAULT_REFERENCE = "reference_2013-07.json" - - -def main(): - ref_path = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_REFERENCE - - with open(ref_path, encoding="utf-8") as f: - ref = json.load(f) - - t0 = time.time() - new = build_reference(ref["parquet"], subset_n=ref["subset_n"]) - print(f"Recomputed summary in {time.time() - t0:.1f}s") - - mismatches = diff_dicts(ref, new) - if mismatches: - print(f"MISMATCH: {len(mismatches)} difference(s):") - for m in mismatches[:50]: - print(f" {m}") - if len(mismatches) > 50: - print(f" ... and {len(mismatches) - 50} more") - sys.exit(1) - - print("OK: new output matches reference") - - -if __name__ == "__main__": - main() diff --git a/src/dump_reference.py b/src/dump_reference.py deleted file mode 100644 index ba37432..0000000 --- a/src/dump_reference.py +++ /dev/null @@ -1,38 +0,0 @@ -"""Dump reference outputs of the CURRENT parser build to a JSON file. - -Run this against the last-known-good build (e.g. before a rearchitecture), -then verify the new build with compare_reference.py. - -Usage: python src/dump_reference.py [parquet_path] [output_json] -""" - -import json -import sys -import time - -from reference_lib import build_reference - -DEFAULT_PARQUET = "2013-07-train-00000-of-00001.parquet" -DEFAULT_OUTPUT = "reference_2013-07.json" - - -def main(): - parquet = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_PARQUET - output = sys.argv[2] if len(sys.argv) > 2 else DEFAULT_OUTPUT - - t0 = time.time() - ref = build_reference(parquet) - with open(output, "w", encoding="utf-8") as f: - json.dump(ref, f, indent=1, ensure_ascii=False) - - main_part = ref["main"] - print(f"Wrote {output} in {time.time() - t0:.1f}s") - print( - f" games={main_part['num_games']} invalid={main_part['num_invalid']} " - f"valid_moves={main_part['valid_moves_total']} " - f"valid_positions={main_part['valid_positions_total']}" - ) - - -if __name__ == "__main__": - main() diff --git a/src/reference_lib.py b/src/reference_lib.py deleted file mode 100644 index f5a7912..0000000 --- a/src/reference_lib.py +++ /dev/null @@ -1,287 +0,0 @@ -"""Shared logic for dumping/comparing reference outputs of the PGN parser. - -Used by dump_reference.py (run against the OLD build to create a reference -file) and compare_reference.py (run against the NEW build to verify parity). - -Works with both the chunked API (<= 4.x: result.chunks / PyChunkView) and the -flat API (>= 5.0: global arrays directly on ParsedGames), so the exact same -digest computation runs against both builds. - -Comparison semantics for invalid games (parse_errors is not None): -- 5.0 allocates per-game array space from pass-1 token counts and zero-fills - the unwritten tail, so whole-array digests would differ by design. -- Therefore digests cover valid games only, and invalid games are recorded - individually with digests of their actually-parsed prefix slices. -""" - -import hashlib -import json - -import numpy as np - - -def _sha(data: bytes) -> str: - return hashlib.sha256(data).hexdigest() - - -def digest_array(arr) -> str: - a = np.ascontiguousarray(arr) - return _sha(a.tobytes()) - - -def digest_json(obj) -> str: - return _sha(json.dumps(obj, sort_keys=True, ensure_ascii=False).encode("utf-8")) - - -class Adapted: - """Uniform view over old (chunked) and new (flat) ParsedGames results.""" - - def __init__(self, result): - self.legal_moves_stored = False - if hasattr(result, "chunks"): - self._from_chunked(result) - else: - self._from_flat(result) - self.num_games = len(self.valid) - # Positions per game are always parsed moves + 1 (initial position is - # recorded even for games that error out immediately). - self.parsed_position_counts = self.parsed_move_counts + 1 - - def _from_chunked(self, result): - chunks = result.chunks - - def cat(name, axis=0): - return np.concatenate([np.asarray(getattr(c, name)) for c in chunks], axis=axis) - - def cat_list(name): - out = [] - for c in chunks: - out.extend(getattr(c, name)) - return out - - for name in ( - "boards", - "castling", - "en_passant", - "halfmove_clock", - "turn", - "from_squares", - "to_squares", - "promotions", - "clocks", - "evals", - "is_checkmate", - "is_stalemate", - "is_insufficient", - "legal_move_count", - "valid", - ): - setattr(self, name, cat(name)) - - self.headers = cat_list("headers") - self.outcome = cat_list("outcome") - self.parse_errors = cat_list("parse_errors") - self.comments = cat_list("comments") - - self.move_counts = np.concatenate( - [np.diff(np.asarray(c.move_offsets, dtype=np.int64)) for c in chunks] - ) - self.position_counts = np.concatenate( - [np.diff(np.asarray(c.position_offsets, dtype=np.int64)) for c in chunks] - ) - # Old builds only record actually-parsed moves: allocated == parsed. - self.parsed_move_counts = self.move_counts - - if any(len(np.asarray(c.legal_move_offsets)) > 0 for c in chunks): - self.legal_moves_stored = True - self.legal_move_from_squares = cat("legal_move_from_squares") - self.legal_move_to_squares = cat("legal_move_to_squares") - self.legal_move_promotions = cat("legal_move_promotions") - self.legal_move_counts = np.concatenate( - [np.diff(np.asarray(c.legal_move_offsets, dtype=np.int64)) for c in chunks] - ) - - def _from_flat(self, result): - for name in ( - "boards", - "castling", - "en_passant", - "halfmove_clock", - "turn", - "from_squares", - "to_squares", - "promotions", - "clocks", - "evals", - "is_checkmate", - "is_stalemate", - "is_insufficient", - "legal_move_count", - "valid", - ): - setattr(self, name, np.asarray(getattr(result, name))) - - self.headers = result.headers - self.outcome = result.outcome - self.parse_errors = result.parse_errors - self.comments = result.comments - - self.move_counts = np.diff(np.asarray(result.move_offsets, dtype=np.int64)) - self.position_counts = np.diff(np.asarray(result.position_offsets, dtype=np.int64)) - self.parsed_move_counts = np.asarray(result.parsed_move_counts, dtype=np.int64) - - legal_offsets = np.asarray(result.legal_move_offsets) - if len(legal_offsets) > 0: - self.legal_moves_stored = True - self.legal_move_from_squares = np.asarray(result.legal_move_from_squares) - self.legal_move_to_squares = np.asarray(result.legal_move_to_squares) - self.legal_move_promotions = np.asarray(result.legal_move_promotions) - self.legal_move_counts = np.diff(legal_offsets.astype(np.int64)) - - -def _masked_digest(arr, mask): - if mask.all(): - return digest_array(arr) - return digest_array(arr[mask]) - - -def summarize(result, include_optional=False) -> dict: - """Compute a comparable summary dict for a ParsedGames result.""" - a = Adapted(result) - n = a.num_games - valid = np.asarray(a.valid, dtype=bool) - - pos_off = np.concatenate([[0], np.cumsum(a.position_counts)]).astype(np.int64) - move_off = np.concatenate([[0], np.cumsum(a.move_counts)]).astype(np.int64) - - game_of_pos = np.repeat(np.arange(n), a.position_counts) - game_of_move = np.repeat(np.arange(n), a.move_counts) - pos_valid = valid[game_of_pos] - move_valid = valid[game_of_move] - - digests = { - # Per-position arrays (valid games only) - "boards": _masked_digest(a.boards, pos_valid), - "castling": _masked_digest(a.castling, pos_valid), - "en_passant": _masked_digest(a.en_passant, pos_valid), - "halfmove_clock": _masked_digest(a.halfmove_clock, pos_valid), - "turn": _masked_digest(a.turn, pos_valid), - # Per-move arrays (valid games only) - "from_squares": _masked_digest(a.from_squares, move_valid), - "to_squares": _masked_digest(a.to_squares, move_valid), - "promotions": _masked_digest(a.promotions, move_valid), - "clocks": _masked_digest(a.clocks, move_valid), - "evals": _masked_digest(a.evals, move_valid), - # Per-game arrays (all games; identical for old/new by design) - "is_checkmate": digest_array(a.is_checkmate), - "is_stalemate": digest_array(a.is_stalemate), - "is_insufficient": digest_array(a.is_insufficient), - "legal_move_count": digest_array(a.legal_move_count), - "valid": digest_array(valid), - "parsed_move_counts_valid": digest_array(a.parsed_move_counts[valid]), - # Per-game metadata (all games) - "headers": digest_json(a.headers), - "outcome": digest_json(a.outcome), - "parse_errors": digest_json(a.parse_errors), - } - - if include_optional: - move_valid_list = move_valid.tolist() - digests["comments"] = digest_json( - [c for c, ok in zip(a.comments, move_valid_list) if ok] - ) - if a.legal_moves_stored: - entry_valid = np.repeat(pos_valid, a.legal_move_counts) - digests["legal_move_counts"] = _masked_digest(a.legal_move_counts, pos_valid) - digests["legal_move_from_squares"] = _masked_digest( - a.legal_move_from_squares, entry_valid - ) - digests["legal_move_to_squares"] = _masked_digest( - a.legal_move_to_squares, entry_valid - ) - digests["legal_move_promotions"] = _masked_digest( - a.legal_move_promotions, entry_valid - ) - - invalid_games = [] - for idx in np.flatnonzero(~valid): - idx = int(idx) - pm = int(a.parsed_move_counts[idx]) - pp = pm + 1 - ps, ms = int(pos_off[idx]), int(move_off[idx]) - invalid_games.append( - { - "index": idx, - "parse_error": a.parse_errors[idx], - "outcome": a.outcome[idx], - "parsed_moves": pm, - "digests": { - "boards": digest_array(a.boards[ps : ps + pp]), - "castling": digest_array(a.castling[ps : ps + pp]), - "en_passant": digest_array(a.en_passant[ps : ps + pp]), - "halfmove_clock": digest_array(a.halfmove_clock[ps : ps + pp]), - "turn": digest_array(a.turn[ps : ps + pp]), - "from_squares": digest_array(a.from_squares[ms : ms + pm]), - "to_squares": digest_array(a.to_squares[ms : ms + pm]), - "promotions": digest_array(a.promotions[ms : ms + pm]), - "clocks": digest_array(a.clocks[ms : ms + pm]), - "evals": digest_array(a.evals[ms : ms + pm]), - }, - } - ) - - return { - "num_games": int(n), - "num_invalid": int((~valid).sum()), - "valid_moves_total": int(a.parsed_move_counts[valid].sum()), - "valid_positions_total": int(a.parsed_position_counts[valid].sum()), - "digests": digests, - "invalid_games": invalid_games, - } - - -def build_reference(parquet_path: str, subset_n: int = 2000) -> dict: - import pyarrow.parquet as pq - import rust_pgn_reader_python_binding as rp - - col = pq.ParquetFile(parquet_path).read(columns=["movetext"]).column("movetext") - - result = rp.parse_games(col) - main = summarize(result) - del result - - subset_col = col.slice(0, subset_n) - subset_result = rp.parse_games( - subset_col, store_comments=True, store_legal_moves=True - ) - subset = summarize(subset_result, include_optional=True) - del subset_result - - return { - "parquet": parquet_path, - "subset_n": subset_n, - "main": main, - "subset": subset, - } - - -def diff_dicts(ref: dict, new: dict, path: str = "") -> list: - """Recursively diff two summary dicts, returning mismatch descriptions.""" - mismatches = [] - if isinstance(ref, dict) and isinstance(new, dict): - for key in sorted(set(ref) | set(new)): - if key not in ref: - mismatches.append(f"{path}.{key}: missing in reference") - elif key not in new: - mismatches.append(f"{path}.{key}: missing in new output") - else: - mismatches.extend(diff_dicts(ref[key], new[key], f"{path}.{key}")) - elif isinstance(ref, list) and isinstance(new, list): - if len(ref) != len(new): - mismatches.append(f"{path}: length {len(ref)} != {len(new)}") - else: - for i, (r, n) in enumerate(zip(ref, new)): - mismatches.extend(diff_dicts(r, n, f"{path}[{i}]")) - elif ref != new: - mismatches.append(f"{path}: {ref!r} != {new!r}") - return mismatches