//! Tools for importing and exporting to Portable Game Notation.
use crate::{
error::{Err, Result}, positions::{timer::Update, PlyClock, Timers}, prelude::{
ChessPosition,
Color::{Black, White},
PositionStatus, Timer,
}
};
use anyhow::Error;
use fxhash::FxHashMap;
use nom::{
branch::alt,
bytes::complete::{escaped, is_a, is_not, tag},
character::complete::{char, digit1, multispace0, multispace1},
sequence::{delimited, pair, preceded, separated_pair, terminated, tuple},
IResult, Parser,
};
use rayon::iter::{ParallelBridge, ParallelIterator};
use std::{
borrow::{Borrow, Cow},
collections::BTreeMap,
fs::File,
io::{BufRead, BufReader, BufWriter, Lines, Write},
};
fn tag_pair(input: &str) -> IResult<&str, (&str, &str)> {
delimited(
preceded(multispace0, char('[')),
separated_pair(
is_not(" "),
char(' '),
delimited(
char('"'),
escaped(is_not(r#"\""#), '\\', char('"')),
char('"'),
),
),
char(']'),
)(input)
}
fn movetext_token(input: &str) -> Result<(&str, MovetextToken)> {
fn move_number_indication(input: &str) -> IResult<&str, (&str, &str)> {
preceded(multispace0, pair(digit1, is_a(".")))(input)
}
fn san_notation(input: &str) -> IResult<&str, &str> {
preceded(multispace0, terminated(is_not(" {\n"), is_a(" \n")))(input)
}
fn timer_annotation(input: &str) -> IResult<&str, (&str, &str, &str)> {
preceded(
multispace0,
delimited(
tag("{").and(multispace0).and(tag("[%clk")).and(multispace1),
tuple((
terminated(digit1, char(':')), // hours
terminated(digit1, char(':')), // minutes
is_a("0123456789."), // seconds
)),
tag("]").and(multispace0).and(tag("}")),
),
)(input)
}
fn move_annotation(input: &str) -> IResult<&str, &str> {
preceded(multispace0, delimited(is_a("{ "), is_not("}"), is_a("} ")))(input)
}
fn result(input: &str) -> IResult<&str, &str> {
preceded(
multispace0,
alt((tag("*"), tag("1/2-1/2"), tag("0-1"), tag("1-0"))),
)(input)
}
if let Ok((input, (number, turn_indicator))) = move_number_indication(input) {
Ok((input, MovetextToken::MoveIndicator(number, turn_indicator)))
} else if let Ok((input, result)) = result(input) {
Ok((input, MovetextToken::Result(result)))
} else if let Ok((input, san)) = san_notation(input) {
Ok((input, MovetextToken::SanNotation(san)))
} else if let Ok((input, (h, m, s))) = timer_annotation(input) {
Ok((input, MovetextToken::TimerAnnotation(h, m, s)))
} else if let Ok((input, move_annotation)) = move_annotation(input) {
Ok((input, MovetextToken::MoveAnnotation(move_annotation)))
} else {
Err(Err::MovetextParseError)
}
}
#[derive(Debug, Clone, Copy)]
enum MovetextToken<'a> {
MoveIndicator(&'a str, &'a str),
SanNotation(&'a str),
/// Timer annotations are represented as a (hours, minutes, seconds) tuple.
TimerAnnotation(&'a str, &'a str, &'a str),
MoveAnnotation(&'a str),
Result(&'a str),
}
/// Represents whether a PGN string or file is valid.
#[derive(Debug)]
pub enum Validity {
Valid,
Invalid(Error),
}
impl PartialEq for Validity {
fn eq(&self, other: &Self) -> bool {
match self {
Self::Valid => match other {
Self::Valid => true,
Self::Invalid(_) => false,
},
Self::Invalid(_) => match other {
Self::Valid => false,
Self::Invalid(_) => true,
},
}
}
}
/// The parsed contents of a Portable Game Notation.
#[derive(Debug, Clone)]
pub struct Pgn<'a> {
pub tags: FxHashMap<&'a str, Cow<'a, str>>,
pub sans: Vec<Cow<'a, str>>,
pub annotations: BTreeMap<u16, Cow<'a, str>>,
pub result: &'a str,
pub timer_updates: BTreeMap<u16, Update>,
}
impl<'a> TryFrom<&'a str> for Pgn<'a> {
type Error = Err;
fn try_from(value: &'a str) -> std::prelude::v1::Result<Self, Self::Error> {
let mut input = value;
let mut tags = FxHashMap::default();
while let Ok((rem, (title, contents))) = tag_pair(input) {
input = rem;
tags.insert(title, Cow::Borrowed(contents));
}
let mut sans = Vec::new();
let mut annotations = BTreeMap::new();
let mut timer_updates: BTreeMap<u16, Update> = BTreeMap::new();
let mut last_marker: Option<(&str, &str)> = None;
let mut result: Option<&str> = None;
while let Ok((rem, token)) = movetext_token(input) {
input = rem;
match token {
MovetextToken::SanNotation(san) => sans.push(Cow::Borrowed(san)),
MovetextToken::MoveIndicator(num, turn) => last_marker = Some((num, turn)),
MovetextToken::TimerAnnotation(h, m, s) => {
let hours: f32 = h.parse().unwrap();
let minutes: f32 = m.parse().unwrap();
let seconds: f32 = s.parse().unwrap();
let remaining = hours.mul_add(3600.0, minutes * 60.0) + seconds;
let Some(last_num) = last_marker else {
return Err(Err::MissingMoveNumberIndicatorError);
};
let ply_number = PlyClock::ply_number_from_split_pgn_marker(last_num)?;
timer_updates.insert(ply_number, Update::Remaining(remaining));
}
MovetextToken::MoveAnnotation(annotation) => {
let Some(last_num) = last_marker else {
return Err(Err::MissingMoveNumberIndicatorError);
};
let ply_number = PlyClock::ply_number_from_split_pgn_marker(last_num)?;
annotations.insert(ply_number, Cow::Borrowed(annotation));
}
MovetextToken::Result(res) => result = Some(res),
}
}
let Some(result) = result else {
return Err(Err::MissingResultError);
};
Ok(Self {
tags,
sans,
annotations,
result,
timer_updates,
})
}
}
impl TryFrom<&Pgn<'_>> for ChessPosition {
type Error = Err;
fn try_from(value: &Pgn) -> std::prelude::v1::Result<Self, Self::Error> {
let mut position = if let Some(initial_fen) = value.tag("FEN") {
Self::from_fen(initial_fen)?
} else {
Self::new()
};
position.apply_sans(value.sans.iter().map(|san| san as &str))?;
if let Some(result) = value.tag("Result") {
position.cached_result = match result.borrow() {
"1-0" => Some(PositionStatus::ImportedVictory(White)),
"0-1" => Some(PositionStatus::ImportedVictory(Black)),
"1/2-1/2" => Some(PositionStatus::ImportedDraw),
_ => None,
};
}
if let Some(time_control) = value.tag("TimeControl") {
if let Ok(timer) = Timer::from_pgn_tag(time_control) {
let mut white = timer.clone();
let mut black = timer;
for (ply_number, update) in &value.timer_updates {
match PlyClock::turn_from_ply_number(*ply_number) {
White => white.update(*ply_number, *update),
Black => black.update(*ply_number, *update),
}
}
position.timers = Some(Timers { white, black });
}
}
Ok(position)
}
}
impl TryFrom<&ChessPosition> for Pgn<'_> {
type Error = Err;
fn try_from(value: &ChessPosition) -> std::prelude::v1::Result<Self, Self::Error> {
let result: &'static str = value.status().pgn_indicator();
let tags = {
let mut map = FxHashMap::default();
map.insert("Result", Cow::Borrowed(result));
if value.initial_position.spec_fen() {
map.insert("SetUp", Cow::Borrowed("1"));
map.insert("FEN", Cow::Owned(value.initial_position.fen()?.into()));
}
map
};
let compact_sans = value.sans()?;
let sans: Vec<_> = compact_sans
.iter()
.map(|s| Cow::Owned(s.to_string()))
.collect();
let timer_updates = {
let mut map = BTreeMap::new();
if let Some(ref timers) = value.timers {
for (ply, rem) in timers
.white
.history
.iter()
.chain(timers.black.history.iter())
{
map.insert(*ply, Update::Remaining(*rem));
}
}
map
};
Ok(Self {
tags,
sans,
annotations: BTreeMap::new(),
result,
timer_updates,
})
}
}
impl TryFrom<&Pgn<'_>> for String {
type Error = Err;
fn try_from(value: &Pgn) -> Result<Self> {
let mut output = Self::with_capacity(93);
let roster = TagRoster::sort_tags(&value.tags);
roster.push_tags(&mut output);
output.push('\n');
let mut movetext = Self::new();
let mut ply_clock = if let Some(fen) = value.tags.get("FEN") {
PlyClock::from_fen(fen, None)?
} else {
PlyClock::new()
};
for san in &value.sans {
let total = ply_clock.total();
let update = value.timer_updates.get(&total);
let annotation = value.annotations.get(&total);
if ply_clock.infer_turn() == White || update.is_some() || annotation.is_some() {
movetext.push_str(&ply_clock.pgn_marker());
movetext.push(' ');
}
movetext.push_str(san);
movetext.push(' ');
if let Some(update) = update {
movetext.push_str(&Update::write_annotation(update.value()));
movetext.push(' ');
}
if let Some(annotation) = annotation {
movetext.push_str("{ ");
movetext.push_str(annotation);
movetext.push_str(" } ");
}
ply_clock.increment(false);
}
movetext.push_str(value.result);
movetext.push('\n');
textwrap::fill_inplace(&mut movetext, 80);
output.push_str(&movetext);
Ok(output)
}
}
impl Pgn<'_> {
/// Get a tag's value, or return `None`.
#[must_use]
pub fn tag(&self, title: &str) -> Option<&Cow<'_, str>> {
self.tags.get(title)
}
/// Verify that a PGN represents a valid game of chess.
#[must_use]
pub fn verify(&self) -> Validity {
match ChessPosition::try_from(self) {
Ok(_) => Validity::Valid,
Err(e) => Validity::Invalid(e.into()),
}
}
/// Read a `.pgn` file to a `Vec<Pgn>`.
///
/// # Errors
///
/// Returns an error if file cannot be read.
pub fn read_file(path: &str) -> anyhow::Result<Vec<String>> {
let file = File::open(path)?;
let reader = BufReader::new(file);
let mut output = Vec::new();
let mut pgn = String::new();
for line in reader.lines() {
let Ok(line) = line else { break };
if !pgn.is_empty() && line.contains("[Event ") {
output.push(pgn.clone());
pgn.clear();
}
pgn.push_str(&line);
pgn.push('\n');
}
if !pgn.is_empty() {
output.push(pgn);
}
Ok(output)
}
/// Iterate through a file.
///
/// # Errors
///
/// Returns an error if the file cannot be opened or read.
pub fn iter_file(path: &str) -> anyhow::Result<FileIterator> {
FileIterator::new(path)
}
/// Verify that a PGN file is valid.
///
/// # Errors
///
/// Returns an error if the file cannot be opened or read.
pub fn verify_file(path: &str) -> anyhow::Result<Validity> {
let result = Self::iter_file(path)?
.par_bridge()
.try_for_each(|s| -> anyhow::Result<()> {
let pgn = Pgn::try_from(s.as_str())?;
let imported_sans: Vec<_> = pgn.sans.clone();
let position = ChessPosition::try_from(&pgn)?;
let expected_sans: Vec<_> = position
.sans()?
.into_iter()
.map(|cs| cs.to_string())
.collect();
for (i_san, e_san) in std::iter::zip(&imported_sans, &expected_sans) {
if !i_san.contains(e_san) {
return Err(anyhow::Error::msg(format!(
"Imported SANs did not match expected SANs.\n\n\
Imported SANs: {imported_sans:?}\n\n\
Expected SANs: {expected_sans:?}\n\n\
PGN:\n{s}\n"
)));
}
}
Ok(())
});
match result {
Ok(()) => Ok(Validity::Valid),
Err(e) => Ok(Validity::Invalid(e)),
}
}
}
/// Iterate through a `.pgn` file.
pub struct FileIterator {
pgn: String,
iter: Lines<BufReader<File>>,
}
impl FileIterator {
/// Create a `PGNFileIterator` from a path.
///
/// # Errors
///
/// Returns an error if the file at `path` cannot be opened.
pub fn new(path: &str) -> anyhow::Result<Self> {
let file = File::open(path)?;
let buf = BufReader::new(file);
Ok(Self {
pgn: String::new(),
iter: buf.lines(),
})
}
}
impl Iterator for FileIterator {
type Item = String;
fn next(&mut self) -> Option<Self::Item> {
loop {
let line = self.iter.next()?.ok()?;
if !self.pgn.is_empty() && line.contains("[Event ") {
let output = Some(self.pgn.clone());
self.pgn.clear();
return output;
}
self.pgn.push_str(&line);
self.pgn.push('\n');
}
}
}
#[derive(Default)]
struct TagRoster<'a> {
event: Option<Cow<'a, str>>,
site: Option<Cow<'a, str>>,
date: Option<Cow<'a, str>>,
round: Option<Cow<'a, str>>,
white: Option<Cow<'a, str>>,
black: Option<Cow<'a, str>>,
result: Option<Cow<'a, str>>,
other: Vec<(&'a str, Cow<'a, str>)>,
}
impl TagRoster<'_> {
fn push_tags(&self, s: &mut String) {
if let Some(tag) = &self.event {
TagRoster::push_tag(s, "Event", tag.borrow());
} else {
TagRoster::push_tag(s, "Event", "?");
}
if let Some(tag) = &self.site {
TagRoster::push_tag(s, "Site", tag.borrow());
} else {
TagRoster::push_tag(s, "Site", "?");
}
if let Some(tag) = &self.date {
TagRoster::push_tag(s, "Date", tag.borrow());
} else {
TagRoster::push_tag(s, "Date", "????.??.??");
}
if let Some(tag) = &self.round {
TagRoster::push_tag(s, "Round", tag.borrow());
} else {
TagRoster::push_tag(s, "Round", "?");
}
if let Some(tag) = &self.white {
TagRoster::push_tag(s, "White", tag.borrow());
} else {
TagRoster::push_tag(s, "White", "?");
}
if let Some(tag) = &self.black {
TagRoster::push_tag(s, "Black", tag.borrow());
} else {
TagRoster::push_tag(s, "Black", "?");
}
if let Some(tag) = &self.result {
TagRoster::push_tag(s, "Result", tag.borrow());
} else {
TagRoster::push_tag(s, "Result", "?");
}
for (title, value) in &self.other {
TagRoster::push_tag(s, title, value.borrow());
}
}
}
impl From<TagRoster<'_>> for String {
fn from(value: TagRoster) -> Self {
let mut output = Self::with_capacity(93);
value.push_tags(&mut output);
output
}
}
impl<'a> TagRoster<'a> {
fn push_tag(s: &mut String, title: &str, value: &str) {
s.push('[');
s.push_str(title);
s.push_str(" \"");
s.push_str(value);
s.push_str("\"]\n");
}
fn sort_tags(tags: &FxHashMap<&'a str, Cow<'a, str>>) -> Self {
let mut roster = TagRoster::default();
for (title, value) in tags {
match *title {
"Event" => roster.event = Some(value.clone()),
"Site" => roster.site = Some(value.clone()),
"Date" => roster.date = Some(value.clone()),
"Round" => roster.round = Some(value.clone()),
"White" => roster.white = Some(value.clone()),
"Black" => roster.black = Some(value.clone()),
"Result" => roster.result = Some(value.clone()),
_ => roster.other.push((title, value.clone())),
}
}
roster
}
}
/// Writes PGNs to a `.pgn` file using a `BufWriter`.
pub struct FileWriter {
first_iter: bool,
writer: BufWriter<File>,
}
impl FileWriter {
/// Create a `PgnWriter` from a path.
///
/// # Errors
///
/// Returns an error if the file cannot be created.
pub fn new(path: &str) -> anyhow::Result<Self> {
Ok(Self {
first_iter: true,
writer: BufWriter::new(File::create(path)?),
})
}
/// Write a `PGN` to the file.
///
/// # Errors
///
/// Propagates any error from `std::io::BufWriter`.
pub fn write(&mut self, pgn: &str) -> anyhow::Result<()> {
if !self.first_iter {
self.writer.write_all(b"\n\n")?;
}
self.first_iter = false;
self.writer.write_all(pgn.as_bytes())?;
Ok(())
}
}
impl From<BufWriter<File>> for FileWriter {
fn from(value: BufWriter<File>) -> Self {
Self {
first_iter: true,
writer: value,
}
}
}
impl Drop for FileWriter {
fn drop(&mut self) {
self.writer.write_all(b"\n").unwrap();
}
}
#[cfg(test)]
mod tests {
use crate::{pgn::Validity, prelude::*};
use parking_lot::Mutex;
use rayon::prelude::*;
use std::time::{SystemTime, UNIX_EPOCH};
#[test]
#[ignore]
fn pgn_import_perf() {
let pgns = Pgn::iter_file("test_data/mega.pgn").unwrap();
let start_time = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
let writer = Mutex::new(FileWriter::new("test_data/reexported.pgn").unwrap());
pgns.for_each(|pgn| {
let pgn: Pgn<'_> = Pgn::try_from(
&ChessPosition::try_from(&Pgn::try_from(pgn.as_str()).unwrap())
.unwrap_or_else(|e| panic!("Error \"{e}\" on PGN:\n{pgn}")),
)
.unwrap();
writer.lock().write(&String::try_from(&pgn).unwrap()).unwrap();
});
let end_time = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
let elapsed = end_time - start_time;
println!("{elapsed:?} elapsed.");
}
#[test]
#[ignore]
fn assert_san_eq() {
Pgn::iter_file("test_data/100k.pgn")
.unwrap()
.par_bridge()
.for_each(|s| {
let pgn = Pgn::try_from(s.as_str()).unwrap();
let expected_sans: Vec<String> = pgn
.sans
.clone()
.into_iter()
.map(|s| s.replace(['?', '!'], ""))
.collect();
let sans = ChessPosition::try_from(&pgn)
.unwrap_or_else(|e| panic!("Error \"{e}\" on PGN:\n{s}"))
.sans()
.unwrap_or_else(|e| panic!("Error \"{e}\" on PGN:\n{s}"));
assert_eq!(expected_sans, sans, "{s}");
});
}
#[test]
#[ignore]
fn test_verify() {
assert_eq!(
Validity::Valid,
Pgn::verify_file("test_data/mega.pgn").unwrap()
);
}
#[test]
fn test_pgn_from_position() {
const GAME: &str = r#"[Event "Rated Chess960 tournament https://lichess.org/tournament/O5dkHvDT"]
[Site "https://lichess.org/6tARGlXY"]
[Date "2022.08.01"]
[Round "-"]
[White "giganotosaurus"]
[Black "Maracana1950"]
[Result "1-0"]
[UTCDate "2022.08.01"]
[UTCTime "00:24:18"]
[WhiteElo "1925"]
[BlackElo "1685"]
[WhiteRatingDiff "+3"]
[BlackRatingDiff "-5"]
[TimeControl "120+0"]
[Termination "Normal"]
[FEN "rbbnknrq/pppppppp/8/8/8/8/PPPPPPPP/RBBNKNRQ w KQkq - 0 1"]
[SetUp "1"]
[Variant "Chess960"]
1. b3 { [%clk 0:02:00] } 1... e5 { [%clk 0:02:00] } 2. Bb2 { [%clk 0:01:59] } 2... c6 { [%clk 0:01:59] } 3. c4 { [%clk 0:01:58] } 3... Ng6 { [%clk 0:01:57] } 4. g3 { [%clk 0:01:57] } 4... d5 { [%clk 0:01:56] } 5. cxd5 { [%clk 0:01:55] } 5... cxd5 { [%clk 0:01:56] } 6. Qxd5 { [%clk 0:01:53] } 6... Be6 { [%clk 0:01:52] } 7. Qb5+ { [%clk 0:01:45] } 7... Nc6 { [%clk 0:01:51] } 8. Qxb7 { [%clk 0:01:43] } 8... O-O { [%clk 0:01:46] } 9. Qxa8 { [%clk 0:01:41] } 9... Bd5 { [%clk 0:01:40] } 10. Qb7 { [%clk 0:01:35] } 10... Nd4 { [%clk 0:01:32] } 11. Qxd5 { [%clk 0:01:32] } 1-0"#;
let position = ChessPosition::try_from(&Pgn::try_from(GAME).unwrap()).unwrap();
println!("{:?}", position.timers.as_ref().unwrap().white.history);
let pgn = Pgn::try_from(&position).unwrap();
println!("{pgn:?}");
println!("\n\n");
println!("{}", String::try_from(&pgn).unwrap());
}
}