chess_huffman/
lib.rs

1#![crate_name = "chess_huffman"]
2
3mod codes;
4mod pgn;
5mod psqt;
6mod ranking;
7#[cfg(test)]
8mod tests;
9
10use bitm::BitAccess;
11use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
12use codes::Book;
13use minimum_redundancy::{Decoder, DecodingResult};
14use shakmaty::san::{ParseSanError, SanError};
15use shakmaty::{Chess, Move, PlayError, Position};
16use std::fmt;
17use std::io::Cursor;
18
19/// The result of an encoding operation.
20pub type EncodeResult<T> = Result<T, GameEncodeError>;
21/// The result of a decoding operation.
22pub type DecodeResult<T> = Result<T, GameDecodeError>;
23
24/// Error when encoding a chess game.
25#[derive(Debug)]
26pub struct GameEncodeError {
27    /// The underlying problem that caused the error.
28    pub kind: GameEncodeErrorKind,
29    /// A textual explanation for the error.
30    pub explanation: String,
31}
32
33/// Kind of error when encoding a chess game.
34#[derive(Debug, Copy, Clone, PartialEq, Eq)]
35pub enum GameEncodeErrorKind {
36    /// An I/O error.
37    IoError,
38    /// An error during the Huffman encoding process.
39    HuffmanEncodeError,
40    /// An illegal or ambiguous SAN (= Standard Algebraic Notation, a notation for a chess move).
41    SanError,
42    /// An invalid SAN so it could not be parsed.
43    ParseSanError,
44    /// An illegal move in the sequence of moves to be encoded.
45    IllegalMove,
46}
47
48impl std::error::Error for GameEncodeError {}
49
50impl fmt::Display for GameEncodeError {
51    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
52        write!(f, "Failed to encode chess game: {}", &self.explanation)
53    }
54}
55
56/// Error when decoding an encoded bit vector into a game, because the bit vector is invalid.
57#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58pub struct GameDecodeError {}
59
60impl std::error::Error for GameDecodeError {}
61
62impl fmt::Display for GameDecodeError {
63    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
64        write!(f, "Cannot decode invalid bit vector")
65    }
66}
67
68impl From<std::io::Error> for GameEncodeError {
69    fn from(inner: std::io::Error) -> Self {
70        Self {
71            kind: GameEncodeErrorKind::IoError,
72            explanation: format!("I/O Error: {inner}"),
73        }
74    }
75}
76
77impl From<SanError> for GameEncodeError {
78    fn from(inner: SanError) -> Self {
79        Self {
80            kind: GameEncodeErrorKind::SanError,
81            explanation: format!("Illegal or ambiguous SAN: {inner}"),
82        }
83    }
84}
85
86impl From<ParseSanError> for GameEncodeError {
87    fn from(inner: ParseSanError) -> Self {
88        Self {
89            kind: GameEncodeErrorKind::SanError,
90            explanation: format!("Unable to parse SAN: {inner}"),
91        }
92    }
93}
94
95impl From<PlayError<Chess>> for GameDecodeError {
96    fn from(_: PlayError<Chess>) -> Self {
97        Self {}
98    }
99}
100
101/// Representation of an encoded chess game.
102#[derive(Debug, Eq, Hash, PartialEq, Clone)]
103pub struct EncodedGame {
104    pub inner: Vec<u64>,
105    pub bit_index: usize,
106}
107
108/// Error that may occur when constructing an `EncodedGame` from bytes.
109#[derive(Debug, Copy, Clone, PartialEq, Eq)]
110pub enum EncodedGameConstructionError {
111    EmptyBytes,
112    InvalidBytes,
113}
114
115impl std::error::Error for EncodedGameConstructionError {}
116
117impl fmt::Display for EncodedGameConstructionError {
118    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
119        write!(
120            f,
121            "{}",
122            match self {
123                EncodedGameConstructionError::EmptyBytes =>
124                    "Cannot construct EncodedGame from empty byte slice",
125                EncodedGameConstructionError::InvalidBytes =>
126                    "Invalid bytes for constructing EncodedGame",
127            }
128        )
129    }
130}
131
132impl EncodedGame {
133    /// Convert the encoded chess game to a byte vector. Use `from_bytes` to convert the result back to an `EncodedGame`.
134    #[must_use]
135    pub fn to_bytes(&self) -> Vec<u8> {
136        let byte_count = if self.bit_index % 8 == 0 {
137            self.bit_index / 8
138        } else {
139            self.bit_index / 8 + 1
140        };
141        #[allow(clippy::cast_possible_truncation)]
142        let m = (self.bit_index % 64) as u8;
143        let padding = if m == 0 { 0 } else { 64 - m };
144
145        let mut wrt = Vec::with_capacity(self.inner.len() * 8 + 1);
146        for x in &self.inner {
147            wrt.write_u64::<LittleEndian>(*x).unwrap();
148        }
149        wrt.truncate(byte_count);
150        wrt.push(padding);
151
152        wrt
153    }
154
155    /// Convert a byte vector (that was the output of `to_bytes`) to an `EncodedGame`.
156    ///
157    /// The result being `Ok` does not guarantee that the encoded game is a valid chess game,
158    /// only that the byte vector could be parsed into an `EncodedGame`.
159    #[must_use]
160    pub fn from_bytes(bytes: &[u8]) -> Result<Self, EncodedGameConstructionError> {
161        if bytes.is_empty() {
162            return Err(EncodedGameConstructionError::EmptyBytes);
163        }
164
165        let total_len_minus_one = bytes.len() - 1;
166
167        // `padding` can be at most 63, so occupying 6 bits.
168        // Use only those bits, and ignore the rest.
169        // This allows users of EncodedGame to store arbitrary metadata in the other 2 bits.
170        let padding = (bytes[total_len_minus_one] & 0b0011_1111) as usize;
171
172        if padding != 0 && total_len_minus_one == 0 {
173            return Err(EncodedGameConstructionError::InvalidBytes);
174        }
175
176        let padding_bytes = padding / 8;
177        let bit_index = total_len_minus_one * 8 - (padding % 8);
178        let content_slice = &bytes[0..total_len_minus_one];
179        let full_slice = if padding_bytes != 0 {
180            &[content_slice, &vec![0; padding_bytes]].concat()
181        } else {
182            content_slice
183        };
184        let mut rdr = Cursor::new(full_slice);
185        let mut buffer = vec![];
186        while let Ok(x) = rdr.read_u64::<LittleEndian>() {
187            buffer.push(x);
188        }
189
190        if bit_index > buffer.len() * 64 {
191            return Err(EncodedGameConstructionError::InvalidBytes);
192        }
193
194        Ok(EncodedGame {
195            inner: buffer,
196            bit_index,
197        })
198    }
199
200    fn new() -> Self {
201        Self {
202            inner: vec![0; 256 / 64],
203            bit_index: 0,
204        }
205    }
206}
207
208/// Encodes a chess game into a compressed bit vector.
209///
210/// # Arguments
211///
212/// * `moves` - The sequence of moves that the game consists of.
213///
214/// # Errors
215///
216/// [`GameEncodeError`] if the given sequence of moves is invalid.
217///
218/// # Examples
219///
220/// ```
221/// use shakmaty::{Move, Role, Square};
222/// use chess_huffman::encode_game;
223///
224/// # use chess_huffman::GameEncodeError;
225/// # fn try_main() -> Result<(), GameEncodeError> {
226/// let moves = vec![
227///     Move::Normal {
228///         role: Role::Pawn,
229///         from: Square::D2,
230///         to: Square::D4,
231///         capture: None,
232///         promotion: None,
233///     },
234///     Move::Normal {
235///         role: Role::Pawn,
236///         from: Square::E7,
237///         to: Square::E5,
238///         capture: None,
239///         promotion: None,
240///     }
241/// ];
242///
243/// let encoded = encode_game(&moves)?;
244/// # Ok(())
245/// # }
246/// ```
247pub fn encode_game(moves: &[Move]) -> EncodeResult<EncodedGame> {
248    let mut encoder = MoveByMoveEncoder::new();
249    for &m in moves {
250        encoder.add_move(m)?;
251    }
252    Ok(encoder.result)
253}
254
255/// Encodes a chess game, represented as a PGN (Portable Game Notation), into a compressed bit vector.
256///
257/// # Arguments
258///
259/// * `pgn` - The PGN of the game.
260///
261/// # Errors
262///
263/// [`GameEncodeError`] if the PGN is invalid.
264///
265/// # Examples
266///
267/// ```
268/// # use chess_huffman::encode_pgn;
269/// # use chess_huffman::GameEncodeError;
270/// # fn try_main() -> Result<(), GameEncodeError> {
271/// let encoded = encode_pgn("1. e4 c5 2. Nf3 e6 3. c3 d5 4. e5")?;
272/// let encoded = encode_pgn(b"1. e4 c5 2. Nf3 e6 3. c3 d5 4. e5")?;
273/// let encoded = encode_pgn(String::from("1. e4 c5 2. Nf3 e6 3. c3 d5 4. e5"))?;
274/// # Ok(())
275/// # }
276/// ```
277pub fn encode_pgn<T: AsRef<[u8]>>(pgn: T) -> EncodeResult<EncodedGame> {
278    let mut reader = pgn_reader::Reader::new(Cursor::new(pgn.as_ref()));
279
280    let mut encoder = pgn::Encoder::new();
281    let bits = reader
282        .read_game(&mut encoder)?
283        .unwrap_or_else(|| Ok(EncodedGame::new()))?;
284    Ok(bits)
285}
286
287/// Encodes a chess game, stored in a PGN file, into a compressed bit vector.
288///
289/// # Arguments
290///
291/// * `path` - The path to the PGN file.
292///
293/// # Errors
294///
295/// [`GameEncodeError`] if the file could not be read or if the PGN is invalid,
296/// its `kind` attribute tells more about the error reason.
297pub fn encode_pgn_file<P: AsRef<std::path::Path>>(path: P) -> EncodeResult<EncodedGame> {
298    let file = std::fs::File::open(path)?;
299    let mut reader = pgn_reader::Reader::new(file);
300
301    let mut encoder = pgn::Encoder::new();
302    let bits = reader
303        .read_game(&mut encoder)?
304        .unwrap_or_else(|| Ok(EncodedGame::new()))?;
305    Ok(bits)
306}
307
308/// Decodes a bit vector into a game, returning both a vector of all moves
309/// and all positions. The N'th position in the position vector is the
310/// position after the N'th move in the move vector.
311///
312/// # Arguments
313///
314/// * `bits` - A bit vector of a compressed chess game.
315///
316/// # Errors
317///
318/// [`GameDecodeError`] if the game contains invalid moves.
319///
320/// # Examples
321///
322/// ```
323/// # use chess_huffman::{encode_pgn, decode_game};
324/// # use chess_huffman::GameEncodeError;
325/// # fn try_main() -> Result<(), Box<dyn std::error::Error>> {
326/// let encoded = encode_pgn("1. e4 c5 2. Nf3 e6 3. c3 d5 4. e5")?;
327/// let decoded = decode_game(&encoded)?;
328/// # Ok(())
329/// # }
330pub fn decode_game(encoded: &EncodedGame) -> DecodeResult<(Vec<Move>, Vec<Chess>)> {
331    let mut moves = vec![];
332    let mut positions = vec![];
333
334    let decoder = MoveByMoveDecoder::new(encoded);
335    for d in decoder.into_iter_moves_and_positions() {
336        let (m, pos) = d?;
337        moves.push(m);
338        positions.push(pos);
339    }
340
341    Ok((moves, positions))
342}
343
344/// Iterator to decode a game move by move, rather than all at once.
345/// This allows for more fine-grained processing than [`decode_game`].
346///
347/// # Examples
348///
349/// ```
350/// # use chess_huffman::{encode_pgn, MoveByMoveDecoder};
351/// # use shakmaty::Position;
352///
353/// # fn try_main() -> Result<(), Box<dyn std::error::Error>> {
354/// let encoded = encode_pgn("1. e4 c5 2. Nf3 e6 3. c3 d5 4. exd5")?;
355/// let mut capture_count = 0;
356/// let mut decoder = MoveByMoveDecoder::new(&encoded);
357/// for d in decoder.into_iter_moves_and_positions() {
358///     let (mv, _) = d?;
359///     if mv.is_capture() {
360///         capture_count += 1;
361///     }
362/// }
363/// assert_eq!(capture_count, 1);
364/// # Ok(())
365/// # }
366pub struct MoveByMoveDecoder<'a> {
367    bit_iter: bitm::BitIterator<'a>,
368    huff_decoder: Decoder<'a, u8>,
369    pos: Chess,
370}
371
372impl<'a> MoveByMoveDecoder<'a> {
373    /// Construct a new [`MoveByMoveDecoder`] from an [`EncodedGame`].
374    #[must_use]
375    pub fn new(encoded: &'a EncodedGame) -> Self {
376        let huff_decoder = codes::get_decoder();
377        let bit_iter = encoded.inner.bit_in_range_iter(0..encoded.bit_index);
378        Self {
379            bit_iter,
380            huff_decoder,
381            pos: Chess::default(),
382        }
383    }
384}
385
386impl MoveByMoveDecoder<'_> {
387    /// Returns the next move.
388    pub fn next_move(&mut self) -> Option<DecodeResult<Move>> {
389        match self.huff_decoder.decode_next(&mut self.bit_iter) {
390            DecodingResult::Value(rank) => {
391                let m =
392                    ranking::nth_from_position(*rank as usize, &self.pos).ok_or(GameDecodeError {});
393                match m {
394                    Ok(m) => {
395                        self.pos.play_unchecked(m);
396
397                        Some(Ok(m))
398                    }
399                    Err(e) => Some(Err(e)),
400                }
401            }
402            DecodingResult::Invalid => {
403                Some(Err(GameDecodeError {}))
404                // this shouldn't happen though: according to minimum_redundancy's docs, Invalid can only be returned if the bits per fragment > 1
405            }
406            DecodingResult::Incomplete => {
407                if self.huff_decoder.consumed_fragments() == 0 {
408                    None
409                } else {
410                    Some(Err(GameDecodeError {}))
411                }
412            }
413        }
414    }
415
416    /// Returns the resulting position when the next move is played.
417    pub fn next_position(&mut self) -> Option<DecodeResult<&Chess>> {
418        if let Some(move_result) = self.next_move() {
419            match move_result {
420                Ok(_) => Some(Ok(&self.pos)),
421                Err(e) => Some(Err(e)),
422            }
423        } else {
424            None
425        }
426    }
427
428    /// Returns the next move and the resulting position when the move is played.
429    pub fn next_move_and_position(&mut self) -> Option<DecodeResult<(Move, &Chess)>> {
430        if let Some(move_result) = self.next_move() {
431            match move_result {
432                Ok(m) => Some(Ok((m, &self.pos))),
433                Err(e) => Some(Err(e)),
434            }
435        } else {
436            None
437        }
438    }
439
440    /// Turns the decoder into an iterator over the moves in the chess game.
441    ///
442    /// As soon as an error is encountered, all subsequent iterations yield the same error.
443    pub fn into_iter_moves(self) -> impl Iterator<Item = DecodeResult<Move>> {
444        struct MoveIter<'a> {
445            decoder: MoveByMoveDecoder<'a>,
446            error: bool,
447        }
448        impl Iterator for MoveIter<'_> {
449            type Item = DecodeResult<Move>;
450
451            fn next(&mut self) -> Option<Self::Item> {
452                if self.error {
453                    return Some(Err(GameDecodeError {}));
454                }
455
456                let item = self.decoder.next_move();
457                if let Some(Err(_)) = item {
458                    self.error = true;
459                }
460                item
461            }
462        }
463
464        MoveIter {
465            decoder: self,
466            error: false,
467        }
468    }
469
470    /// Turns the decoder into an iterator over the positions in the chess game.
471    /// The first yielded position is the position after the first move.
472    ///
473    /// As soon as an error is encountered, all subsequent iterations yield the same error.
474    pub fn into_iter_positions(self) -> impl Iterator<Item = DecodeResult<Chess>> {
475        struct PosIter<'a> {
476            decoder: MoveByMoveDecoder<'a>,
477            error: bool,
478        }
479        impl Iterator for PosIter<'_> {
480            type Item = DecodeResult<Chess>;
481
482            fn next(&mut self) -> Option<Self::Item> {
483                if self.error {
484                    return Some(Err(GameDecodeError {}));
485                }
486
487                let item = self.decoder.next_position().map(|r| r.cloned());
488                if let Some(Err(_)) = item {
489                    self.error = true;
490                }
491                item
492            }
493        }
494
495        PosIter {
496            decoder: self,
497            error: false,
498        }
499    }
500
501    /// Turns the decoder into an iterator over the moves and positions in the chess game.
502    /// The yielded position is the position after the move.
503    ///
504    /// As soon as an error is encountered, all subsequent iterations yield the same error.
505    pub fn into_iter_moves_and_positions(
506        self,
507    ) -> impl Iterator<Item = DecodeResult<(Move, Chess)>> {
508        struct MovePosIter<'a> {
509            decoder: MoveByMoveDecoder<'a>,
510            error: bool,
511        }
512        impl Iterator for MovePosIter<'_> {
513            type Item = DecodeResult<(Move, Chess)>;
514
515            fn next(&mut self) -> Option<Self::Item> {
516                if self.error {
517                    return Some(Err(GameDecodeError {}));
518                }
519
520                let item = self
521                    .decoder
522                    .next_move_and_position()
523                    .map(|r| r.map(|(m, p)| (m, p.clone())));
524                if let Some(Err(_)) = item {
525                    self.error = true;
526                }
527                item
528            }
529        }
530
531        MovePosIter {
532            decoder: self,
533            error: false,
534        }
535    }
536}
537
538/// An encoder that lets you add moves one-by-one, rather than a whole game at once.
539///
540/// # Examples
541///
542/// ```
543/// # use chess_huffman::MoveByMoveEncoder;
544/// use shakmaty::Move;
545///
546/// # use chess_huffman::GameEncodeError;
547/// # fn try_main() -> Result<(), GameEncodeError> {
548/// let moves: Vec<Move> = vec![ /* ... */ ];
549/// let mut encoder = MoveByMoveEncoder::new();
550/// for m in moves {
551///    encoder.add_move(m)?;
552/// }
553/// # Ok(())
554/// # }
555/// ```
556pub struct MoveByMoveEncoder<'a> {
557    book: &'a Book,
558    /// The current chess position.
559    pub pos: Chess,
560    /// The resulting encoded game.
561    pub result: EncodedGame,
562}
563
564impl MoveByMoveEncoder<'_> {
565    /// Constructs a new [`MoveByMoveEncoder`].
566    #[must_use]
567    pub fn new() -> Self {
568        let book = &*codes::BOOK_FROM_LICHESS_WEIGHTS;
569        Self {
570            book,
571            pos: Chess::default(),
572            result: EncodedGame::new(),
573        }
574    }
575
576    /// Adds a new move to the encoder. This affects the value of `buffer`.
577    ///
578    /// # Errors
579    ///
580    /// [`GameEncodeError`] if the move is invalid or `pos` is an invalid chess position.
581    ///
582    /// # Examples
583    ///
584    /// ```
585    /// # use chess_huffman::MoveByMoveEncoder;
586    /// use shakmaty::Move;
587    ///
588    /// # use chess_huffman::GameEncodeError;
589    /// # fn try_main() -> Result<(), GameEncodeError> {
590    /// let moves: Vec<Move> = vec![ /* ... */ ];
591    /// let mut encoder = MoveByMoveEncoder::new();
592    /// for m in moves {
593    ///    encoder.add_move(m)?;
594    /// }
595    /// # Ok(())
596    /// # }
597    /// ```
598    pub fn add_move(&mut self, m: Move) -> EncodeResult<()> {
599        match ranking::move_rank(&self.pos, m) {
600            Some(rank) => {
601                if rank > 255 {
602                    return Err(GameEncodeError {
603                        kind: GameEncodeErrorKind::HuffmanEncodeError,
604                        explanation: String::from(
605                            "Too many possible valid moves - is this a valid chess position?",
606                        ),
607                    });
608                }
609                #[allow(clippy::cast_possible_truncation)]
610                self.book.encode(&mut self.result, rank as u8);
611                self.pos.play_unchecked(m);
612            }
613            None => {
614                return Err(GameEncodeError {
615                    kind: GameEncodeErrorKind::IllegalMove,
616                    explanation: format!("Illegal move {m}"),
617                });
618            }
619        }
620
621        Ok(())
622    }
623
624    /// Clears the encoder: restores the game state to the start position and empties `buffer`.
625    pub fn clear(&mut self) {
626        self.pos = Chess::default();
627        self.result = EncodedGame::new();
628    }
629}
630
631impl Default for MoveByMoveEncoder<'_> {
632    fn default() -> Self {
633        Self::new()
634    }
635}
chess_huffman/lib.rs

chess_huffman/
lib.rs
