<commit_before>use super::{read_uint, write_uint, ByMode, BySpeed, GameId, Mode, Record, Speed};
use byteorder::{ReadBytesExt as _, WriteBytesExt as _};
use std::cmp::min;
use std::io::{self, Read, Write};
use std::ops::AddAssign;

const MAX_GAMES: usize = 15; // 4 bits

#[derive(Debug, Eq, PartialEq)]
enum Header {
    Group {
        mode: Mode,
        speed: Speed,
        num_games: usize,
    },
    End,
}

impl Record for Header {
    fn read<R: Read>(reader: &mut R) -> io::Result<Header> {
        let n = reader.read_u8()?;
        Ok(Header::Group {
            speed: match n & 7 {
                0 => return Ok(Header::End),
                1 => Speed::Ultrabullet,
                2 => Speed::Bullet,
                3 => Speed::Blitz,
                4 => Speed::Rapid,
                5 => Speed::Classical,
                6 => Speed::Correspondence,
                _ => return Err(io::ErrorKind::InvalidData.into()),
            },
            mode: Mode::from_rated((n >> 3) & 1 == 1),
            num_games: usize::from(n >> 4),
        })
    }

    fn write<W: Write>(&self, writer: &mut W) -> io::Result<()> {
        writer.write_u8(match *self {
            Header::End => 0,
            Header::Group {
                mode,
                speed,
                num_games,
            } => {
                (match speed {
                    Speed::Ultrabullet => 1,
                    Speed::Bullet => 2,
                    Speed::Blitz => 3,
                    Speed::Rapid => 4,
                    Speed::Classical => 5,
                    Speed::Correspondence => 6,
                }) | ((mode.is_rated() as u8) << 3)
                    | ((num_games as u8) << 4)
            }
        })
    }
}

#[derive(Debug, Default)]
struct Stats {
    white: u64,
    draw: u64,
    black: u64,
}

impl AddAssign for Stats {
    fn add_assign(&mut self, rhs: Stats) {
        self.white += rhs.white;
        self.draw += rhs.draw;
        self.black += rhs.black;
    }
}

impl Record for Stats {
    fn read<R: Read>(reader: &mut R) -> io::Result<Stats> {
        Ok(Stats {
            white: read_uint(reader)?,
            draw: read_uint(reader)?,
            black: read_uint(reader)?,
        })
    }

    fn write<W: Write>(&self, writer: &mut W) -> io::Result<()> {
        write_uint(writer, self.white)?;
        write_uint(writer, self.draw)?;
        write_uint(writer, self.black)
    }
}

struct Group {
    stats: Stats,
    games: Vec<GameId>,
}

type Inner = BySpeed<ByMode<Group>>;

struct PersonalRecord {
    inner: Inner,
}

impl Record for PersonalRecord {
    fn read<R: Read>(reader: &mut R) -> io::Result<PersonalRecord> {
        let mut inner = Inner::default();
        loop {
            match Header::read(reader)? {
                Header::Group {
                    speed,
                    mode,
                    num_games,
                } => {
                    let stats = Stats::read(reader)?;
                    let mut games = Vec::with_capacity(num_games);
                    for _ in 0..num_games {
                        games.push(GameId::read(reader)?);
                    }
                    let group = inner.by_speed_mut(speed).by_mode_mut(mode);
                    *group = (stats, games);
                }
                Header::End => break,
            }
        }
        Ok(PersonalRecord { inner })
    }

    fn write<W: Write>(&self, writer: &mut W) -> io::Result<()> {
        self.inner.as_ref().try_map(|speed, by_mode| {
            by_mode.as_ref().try_map(|mode, (stats, games)| {
                let num_games = min(games.len(), MAX_GAMES);

                Header::Group {
                    speed,
                    mode,
                    num_games,
                }
                .write(writer)?;

                stats.write(writer)?;

                for game in games.iter().take(num_games) {
                    game.write(writer)?;
                }

                Ok::<_, io::Error>(())
            })
        })?;

        Header::End.write(writer)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;

    #[test]
    fn test_header_roundtrip() {
        let headers = [
            Header::Group {
                mode: Mode::Rated,
                speed: Speed::Correspondence,
                num_games: 15,
            },
            Header::End,<commit_msg>mergable groups<commit_after>use super::{read_uint, write_uint, ByMode, BySpeed, GameId, Mode, Record, Speed};
use byteorder::{ReadBytesExt as _, WriteBytesExt as _};
use std::cmp::min;
use std::io::{self, Read, Write};
use std::ops::AddAssign;

const MAX_GAMES: usize = 15; // 4 bits

#[derive(Debug, Eq, PartialEq)]
enum Header {
    Group {
        mode: Mode,
        speed: Speed,
        num_games: usize,
    },
    End,
}

impl Record for Header {
    fn read<R: Read>(reader: &mut R) -> io::Result<Header> {
        let n = reader.read_u8()?;
        Ok(Header::Group {
            speed: match n & 7 {
                0 => return Ok(Header::End),
                1 => Speed::Ultrabullet,
                2 => Speed::Bullet,
                3 => Speed::Blitz,
                4 => Speed::Rapid,
                5 => Speed::Classical,
                6 => Speed::Correspondence,
                _ => return Err(io::ErrorKind::InvalidData.into()),
            },
            mode: Mode::from_rated((n >> 3) & 1 == 1),
            num_games: usize::from(n >> 4),
        })
    }

    fn write<W: Write>(&self, writer: &mut W) -> io::Result<()> {
        writer.write_u8(match *self {
            Header::End => 0,
            Header::Group {
                mode,
                speed,
                num_games,
            } => {
                (match speed {
                    Speed::Ultrabullet => 1,
                    Speed::Bullet => 2,
                    Speed::Blitz => 3,
                    Speed::Rapid => 4,
                    Speed::Classical => 5,
                    Speed::Correspondence => 6,
                }) | ((mode.is_rated() as u8) << 3)
                    | ((num_games as u8) << 4)
            }
        })
    }
}

#[derive(Debug, Default)]
struct Stats {
    white: u64,
    draw: u64,
    black: u64,
}

impl AddAssign for Stats {
    fn add_assign(&mut self, rhs: Stats) {
        self.white += rhs.white;
        self.draw += rhs.draw;
        self.black += rhs.black;
    }
}

impl Record for Stats {
    fn read<R: Read>(reader: &mut R) -> io::Result<Stats> {
        Ok(Stats {
            white: read_uint(reader)?,
            draw: read_uint(reader)?,
            black: read_uint(reader)?,
        })
    }

    fn write<W: Write>(&self, writer: &mut W) -> io::Result<()> {
        write_uint(writer, self.white)?;
        write_uint(writer, self.draw)?;
        write_uint(writer, self.black)
    }
}

#[derive(Default)]
struct Group {
    stats: Stats,
    games: Vec<GameId>,
}

impl AddAssign for Group {
    fn add_assign(&mut self, rhs: Group) {
        self.stats += rhs.stats;
        self.games.extend(rhs.games);
    }
}

type Inner = BySpeed<ByMode<Group>>;

struct PersonalRecord {
    inner: Inner,
}

impl Record for PersonalRecord {
    fn read<R: Read>(reader: &mut R) -> io::Result<PersonalRecord> {
        let mut inner = Inner::default();
        loop {
            match Header::read(reader)? {
                Header::Group {
                    speed,
                    mode,
                    num_games,
                } => {
                    let stats = Stats::read(reader)?;
                    let mut games = Vec::with_capacity(num_games);
                    for _ in 0..num_games {
                        games.push(GameId::read(reader)?);
                    }
                    let group = inner.by_speed_mut(speed).by_mode_mut(mode);
                    *group = Group { stats, games };
                }
                Header::End => break,
            }
        }
        Ok(PersonalRecord { inner })
    }

    fn write<W: Write>(&self, writer: &mut W) -> io::Result<()> {
        self.inner.as_ref().try_map(|speed, by_mode| {
            by_mode.as_ref().try_map(|mode, group| {
                let num_games = min(group.games.len(), MAX_GAMES);

                Header::Group {
                    speed,
                    mode,
                    num_games,
                }
                .write(writer)?;

                group.stats.write(writer)?;

                for game in group.games.iter().take(num_games) {
                    game.write(writer)?;
                }

                Ok::<_, io::Error>(())
            })
        })?;

        Header::End.write(writer)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;

    #[test]
    fn test_header_roundtrip() {
        let headers = [
            Header::Group {
                mode: Mode::Rated,
                speed: Speed::Correspondence,
                num_games: 15,
            },
            Header::End,
{'fraction_non_alphanumeric': 0.14006190581472475, 'fraction_numerical': 0.008180411231483529, 'mean_word_length': 1.4118901626233005, 'pattern_counts': {'":': 0, '<': 43, '<?xml version=': 0, '>': 91, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4473604', 'n_tokens_mistral': 2807, 'n_tokens_neox': 2640, 'n_words': 575}