Skip to main content

chess/formats/pgn/
stream.rs

1use std::{
2    io::{self, BufRead, BufReader, Read},
3    iter::Peekable,
4};
5
6use encoding_rs::WINDOWS_1252;
7use winnow::Parser as _;
8
9use super::{Game, Result, parse};
10
11pub fn games<R: Read>(reader: R) -> impl Iterator<Item = io::Result<Result<Game>>> {
12    GameShaped::new(reader).map(|chunk| {
13        chunk.map(|chunk| {
14            parse::game
15                .parse(chunk.text.as_str())
16                .map_err(|error| parse::Error::from(&chunk.text, chunk.line, error).into())
17        })
18    })
19}
20
21struct GameShaped<R: Read> {
22    lines: Peekable<Lines<R>>,
23}
24
25impl<R: Read> GameShaped<R> {
26    fn new(reader: R) -> Self {
27        Self { lines: Lines::new(reader).peekable() }
28    }
29}
30
31impl<R: Read> Iterator for GameShaped<R> {
32    type Item = io::Result<Chunk>;
33
34    fn next(&mut self) -> Option<Self::Item> {
35        game_shaped(&mut self.lines).transpose()
36    }
37}
38
39struct Lines<R: Read> {
40    reader: BufReader<R>,
41    bytes: Vec<u8>,
42    line: usize,
43}
44
45impl<R: Read> Lines<R> {
46    fn new(reader: R) -> Self {
47        Self { reader: BufReader::new(reader), bytes: Vec::new(), line: 0 }
48    }
49}
50
51impl<R: Read> Iterator for Lines<R> {
52    type Item = io::Result<SourceLine>;
53
54    fn next(&mut self) -> Option<Self::Item> {
55        self.bytes.clear();
56        match self.reader.read_until(b'\n', &mut self.bytes) {
57            Ok(0) => None,
58            Ok(_) => {
59                self.line += 1;
60                Some(Ok(SourceLine { number: self.line, text: decode_line(&self.bytes) }))
61            }
62            Err(error) => Some(Err(error)),
63        }
64    }
65}
66
67struct SourceLine {
68    number: usize,
69    text: String,
70}
71
72struct Chunk {
73    line: usize,
74    text: String,
75}
76
77fn decode_line(bytes: &[u8]) -> String {
78    let bytes = bytes.strip_suffix(b"\n").unwrap_or(bytes);
79    let bytes = bytes.strip_suffix(b"\r").unwrap_or(bytes);
80
81    match std::str::from_utf8(bytes) {
82        Ok(text) => text.to_string(),
83        Err(_) => {
84            let (text, _, _) = WINDOWS_1252.decode(bytes);
85            text.into_owned()
86        }
87    }
88}
89
90fn game_shaped<I>(lines: &mut Peekable<I>) -> io::Result<Option<Chunk>>
91where
92    I: Iterator<Item = io::Result<SourceLine>>,
93{
94    let Some(line) = lines.next().transpose()? else {
95        return Ok(None);
96    };
97    let first_line = line.number;
98    let mut buffer = Buffer::new(line.text);
99
100    loop {
101        let next = match lines.peek() {
102            Some(Ok(next)) => Some(next.text.as_str()),
103            Some(Err(_)) | None => None,
104        };
105        if buffer.is_complete(next) {
106            return Ok(Some(Chunk { line: first_line, text: buffer.take() }));
107        }
108
109        let Some(line) = lines.next().transpose()? else {
110            return Ok(Some(Chunk { line: first_line, text: buffer.take() }));
111        };
112        buffer.push(line.text);
113    }
114}
115
116struct Buffer {
117    text: String,
118    in_comment: bool,
119    movetext: bool,
120}
121
122impl Buffer {
123    fn new(line: String) -> Self {
124        let mut buffer = Self { text: String::new(), in_comment: false, movetext: false };
125        buffer.push(line);
126        buffer
127    }
128
129    fn push(&mut self, line: String) {
130        self.text.push_str(&line);
131        self.text.push('\n');
132
133        let line = if self.in_comment { line.as_str() } else { parse::strip_tags(&line) };
134        for c in line.chars() {
135            if self.in_comment {
136                if c == '}' {
137                    self.in_comment = false;
138                }
139            } else {
140                match c {
141                    '{' => {
142                        self.movetext = true;
143                        self.in_comment = true;
144                    }
145                    ';' => {
146                        self.movetext = true;
147                        break;
148                    }
149                    c if c.is_whitespace() => {}
150                    _ => self.movetext = true,
151                }
152            }
153        }
154    }
155
156    fn is_game_shaped(&self) -> bool {
157        self.movetext && !self.in_comment
158    }
159
160    fn is_complete(&self, next: Option<&str>) -> bool {
161        self.is_game_shaped() && next.is_none_or(parse::tag_start_ok)
162    }
163
164    fn take(self) -> String {
165        self.text
166    }
167}