diff --git a/.gitignore b/.gitignore index c8f0442..ad136bf 100644 --- a/.gitignore +++ b/.gitignore @@ -70,3 +70,7 @@ docs/_build/ # Pyenv .python-version + +# Local benchmark data and derived reference dumps +*.parquet +*.pgn diff --git a/Cargo.lock b/Cargo.lock index f253d46..ab5e8f7 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" @@ -1317,16 +1306,16 @@ checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "rust_pgn_reader_python_binding" -version = "4.0.0" +version = "5.0.0" 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..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 @@ -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/benches/parquet_bench.rs b/benches/parquet_bench.rs index 4af6b63..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::prelude::*; -use rayon::ThreadPoolBuilder; 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::{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/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/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 new file mode 100644 index 0000000..5f7e2e0 --- /dev/null +++ b/src/flat.rs @@ -0,0 +1,1237 @@ +//! 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 rayon::ThreadPoolBuilder; +use rayon::prelude::*; +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. +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::*; + + fn default_config() -> ParseConfig { + ParseConfig { + store_comments: false, + store_legal_moves: false, + } + } + + fn full_config() -> ParseConfig { + ParseConfig { + store_comments: true, + store_legal_moves: true, + } + } + + fn parse_one(pgn: &str, config: &ParseConfig) -> FlatOutput { + parse_games_flat(&[pgn], 1, config).unwrap() + } + + fn bits(v: &[f32]) -> Vec { + v.iter().map(|f| f.to_bits()).collect() + } + + /// 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!( + &batch.comments[bm..bm + alloc], + &single.comments[..], + "game {i}: comments" + ); + } + 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!( + &batch.legal_move_from_squares[bs..be], + &single.legal_move_from_squares[ss..se], + "game {i} pos {k}: legal from" + ); + assert_eq!( + &batch.legal_move_to_squares[bs..be], + &single.legal_move_to_squares[ss..se], + "game {i} pos {k}: legal to" + ); + assert_eq!( + &batch.legal_move_promotions[bs..be], + &single.legal_move_promotions[ss..se], + "game {i} pos {k}: legal promo" + ); + } + } + } + + /// 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_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_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_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_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()); + } + // 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] + 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); + } + + #[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 9f257e5..8d74ed8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,86 +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 visitor; +mod tokenizer; -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 { @@ -90,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(); @@ -119,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( @@ -303,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. @@ -332,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 @@ -343,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/tokenizer.rs b/src/tokenizer.rs new file mode 100644 index 0000000..a85fdfc --- /dev/null +++ b/src/tokenizer.rs @@ -0,0 +1,499 @@ +use memchr::memchr; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Outcome { + WhiteWins, + BlackWins, + Draw, + 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); + /// 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); +} + +/// 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: +/// - 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) { + // 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 = 0usize; + + while cursor < len { + let c = bytes[cursor]; + + match c { + // '.', '!' 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 { + // 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';' => { + // Line comment + if in_headers { + in_headers = false; + visitor.end_headers(); + } + // 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 { + 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 variation_depth == 0 { + if in_headers { + in_headers = false; + visitor.end_headers(); + } + visitor.outcome(Outcome::Unknown); + break; // End of game + } + cursor += 1; + } + _ => { + // 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() { + // Stray delimiter without a dedicated arm (']' or '}'). + cursor += 1; + continue; + } + + 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; + } 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 token.iter().all(|b| b.is_ascii_digit()) { + // Move number, skip + } else { + // SAN (or garbage) token — delivered raw, parsed by the visitor. + visitor.san_token(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(); +} + +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'!' + | 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_token(&mut self, token: &[u8]) { + self.events + .push(format!("san {}", 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() + } + + #[test] + fn test_move_numbers_without_space() { + assert_eq!(sans("1.e4 e5 2.Nf3 Nc6 1-0"), ["e4", "e5", "Nf3", "Nc6"]); + } + + #[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"]); + } + + #[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_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), ["0-0", "0-0-0", "0-0+", "0-0-0#"]); + } + + #[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 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", "%junk", "e5"]); + } + + #[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_delivered_raw() { + let pgn = "1. e4 xyzzy9 e5 *"; + assert_eq!(sans(pgn), ["e4", "xyzzy9", "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 deleted file mode 100644 index 523b31c..0000000 --- a/src/visitor.rs +++ /dev/null @@ -1,849 +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::{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 std::collections::HashMap; -use std::ops::ControlFlow; - -/// 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() - } -} - -/// 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<'_> { - type Tags = Vec<(String, String)>; - type Movetext = (); - type Output = bool; - - fn begin_tags(&mut self) -> ControlFlow { - 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; - - // 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(); - ControlFlow::Continue(()) - } - - fn san( - &mut self, - _movetext: &mut Self::Movetext, - san_plus: SanPlus, - ) -> ControlFlow { - if self.valid_moves { - 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)); - } - } - } - ControlFlow::Continue(()) - } - - fn comment( - &mut self, - _movetext: &mut Self::Movetext, - comment: RawComment<'_>, - ) -> ControlFlow { - if let Ok((_, parsed_comments)) = parse_comments(comment.as_bytes()) { - 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); - } - } - } - 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 { - self.current_outcome = Some(match _outcome { - Outcome::Known(known) => match known { - KnownOutcome::Decisive { winner } => format!("{:?}", winner), - KnownOutcome::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 { - // 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 - } -} - -/// Parse a single game directly into Buffers. -pub fn parse_game_to_buffers( - pgn: &str, - 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)), - } -} - -#[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_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_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 - } -}