<commit_before>use std::{io::Cursor, path::Path};

use rocksdb::{
    merge_operator::MergeFn, BlockBasedIndexType, BlockBasedOptions, ColumnFamily,
    ColumnFamilyDescriptor, IteratorMode, MergeOperands, Options, ReadOptions, SliceTransform,
    WriteBatch, DB,
};

use crate::model::{
    GameId, Key, KeyPrefix, LichessEntry, LichessGame, MastersEntry, MastersGame, Month,
    PlayerEntry, PlayerStatus, UserId, Year,
};

#[derive(Debug)]
pub struct Database {
    pub inner: DB,
}

fn column_family(
    name: &str,
    merge: Option<&str>,
    merge_fn: impl MergeFn + Clone,
    prefix: Option<usize>,
    block_size: usize,
    bloom_filter: i32,
) -> ColumnFamilyDescriptor {
    let mut opts = Options::default();
    if let Some(merge) = merge {
        opts.set_merge_operator_associative(merge, merge_fn);
    }
    opts.set_prefix_extractor(match prefix {
        Some(prefix) => SliceTransform::create_fixed_prefix(prefix),
        None => SliceTransform::create_noop(),
    });
    let mut block_opts = BlockBasedOptions::default();
    block_opts.set_index_type(BlockBasedIndexType::HashSearch);
    block_opts.set_block_size(block_size);
    if bloom_filter > 0 {
        block_opts.set_bloom_filter(bloom_filter, true);
    }
    opts.set_block_based_table_factory(&block_opts);
    ColumnFamilyDescriptor::new(name, opts)
}

impl Database {
    pub fn open<P: AsRef<Path>>(path: P) -> Result<Database, rocksdb::Error> {
        let mut db_opts = Options::default();
        db_opts.create_if_missing(true);
        db_opts.create_missing_column_families(true);

        let inner = DB::open_cf_descriptors(
            &db_opts,
            path,
            vec![
                // Masters database
                column_family(
                    "masters",
                    Some("masters_merge"),
                    masters_merge,
                    Some(KeyPrefix::SIZE),
                    8 * 1024,
                    3,
                ),
                column_family("masters_game", None, void_merge, None, 4 * 1024, 0),
                // Lichess database
                column_family(
                    "lichess",
                    Some("lichess_merge"),
                    lichess_merge,
                    Some(KeyPrefix::SIZE),
                    16 * 1024,
                    2,
                ),
                column_family(
                    "lichess_game",
                    Some("lichess_game_merge"),
                    lichess_game_merge,
                    None,
                    4 * 1024,
                    0,
                ),
                // Player database (also shares lichess_game)
                column_family(
                    "player",
                    Some("player_merge"),
                    player_merge,
                    Some(KeyPrefix::SIZE),
                    16 * 1024,
                    2,
                ),
                column_family("player_status", None, void_merge, None, 4 * 1024, 0),
            ],
        )?;

        Ok(Database { inner })
    }

    pub fn masters(&self) -> MastersDatabase<'_> {
        MastersDatabase {
            inner: &self.inner,
            cf_masters: self.inner.cf_handle("masters").expect("cf masters"),
            cf_masters_game: self
                .inner
                .cf_handle("masters_game")
                .expect("cf masters_game"),
        }
    }

    pub fn lichess(&self) -> LichessDatabase<'_> {
        LichessDatabase {
            inner: &self.inner,
            cf_lichess: self.inner.cf_handle("lichess").expect("cf lichess"),
            cf_lichess_game: self
                .inner
                .cf_handle("lichess_game")
                .expect("cf lichess_game"),

            cf_player: self.inner.cf_handle("player").expect("cf player"),
            cf_player_status: self
                .inner
                .cf_handle("player_status")
                .expect("cf player_status"),
        }
    }
}

pub struct MastersDatabase<'a> {
    inner: &'a DB,
    cf_masters: &'a ColumnFamily,
    cf_masters_game: &'a ColumnFamily,
}

impl MastersDatabase<'_> {
    pub fn has_game(&self, id: GameId) -> Result<bool, rocksdb::Error> {
        self.inner
            .get_cf(self.cf_masters_game, id.to_bytes())
            .map(|maybe_entry| maybe_entry.is_some())
    }

    pub fn game(&self, id: GameId) -> Result<Option<MastersGame>, rocksdb::Error> {
        Ok(self
            .inner
            .get_cf(self.cf_masters_game, id.to_bytes())?
            .map(|buf| serde_json::from_slice(&buf).expect("deserialize masters game")))
    }

    pub fn has(&self, key: Key) -> Result<bool, rocksdb::Error> {
        self.inner
            .get_cf(self.cf_masters, key.into_bytes())
            .map(|maybe_entry| maybe_entry.is_some())
    }

    pub fn read(
        &self,
        key: KeyPrefix,
        since: Year,
        until: Year,
    ) -> Result<MastersEntry, rocksdb::Error> {
        let mut opt = ReadOptions::default();
        opt.set_prefix_same_as_start(true);
        opt.set_iterate_lower_bound(key.with_year(since).into_bytes());
        opt.set_iterate_upper_bound(key.with_year(until.add_years_saturating(1)).into_bytes());

        let iterator = self
            .inner
            .iterator_cf_opt(self.cf_masters, opt, IteratorMode::Start);

        let mut entry = MastersEntry::default();
        for (_key, value) in iterator {
            let mut cursor = Cursor::new(value);
            entry
                .extend_from_reader(&mut cursor)
                .expect("deserialize masters entry");
        }

        Ok(entry)
    }

    pub fn batch(&self) -> MastersBatch<'_> {
        MastersBatch {
            db: self,
            batch: WriteBatch::default(),
        }
    }
}

pub struct MastersBatch<'a> {
    db: &'a MastersDatabase<'a>,
    batch: WriteBatch,
}

impl MastersBatch<'_> {
    pub fn merge(&mut self, key: Key, entry: MastersEntry) {
        let mut cursor = Cursor::new(Vec::with_capacity(MastersEntry::SIZE_HINT));
        entry.write(&mut cursor).expect("serialize masters entry");
        self.batch
            .merge_cf(self.db.cf_masters, key.into_bytes(), cursor.into_inner());
    }

    pub fn put_game(&mut self, id: GameId, game: &MastersGame) {
        self.batch.put_cf(
            self.db.cf_masters_game,
            id.to_bytes(),
            serde_json::to_vec(game).expect("serialize masters game"),
        );
    }

    pub fn commit(self) -> Result<(), rocksdb::Error> {
        self.db.inner.write(self.batch)
    }
}

pub struct LichessDatabase<'a> {
    inner: &'a DB,
    cf_lichess: &'a ColumnFamily,
    cf_lichess_game: &'a ColumnFamily,

    cf_player: &'a ColumnFamily,
    cf_player_status: &'a ColumnFamily,
}

impl LichessDatabase<'_> {
    pub fn game(&self, id: GameId) -> Result<Option<LichessGame>, rocksdb::Error> {
        Ok(self
            .inner
            .get_cf(self.cf_lichess_game, id.to_bytes())?
            .map(|buf| {
                let mut cursor = Cursor::new(buf);
                LichessGame::read(&mut cursor).expect("deserialize game info")
            }))
    }

    pub fn read_lichess(
        &self,
        key: &KeyPrefix,
        since: Month,
        until: Month,
    ) -> Result<LichessEntry, rocksdb::Error> {
        let mut opt = ReadOptions::default();
        opt.set_prefix_same_as_start(true);
        opt.set_iterate_lower_bound(key.with_month(since).into_bytes());
        opt.set_iterate_upper_bound(key.with_month(until.add_months_saturating(1)).into_bytes());

        let iterator = self
            .inner
            .iterator_cf_opt(self.cf_lichess, opt, IteratorMode::Start);

        let mut entry = LichessEntry::default();
        for (_key, value) in iterator {
            let mut cursor = Cursor::new(value);
            entry
                .extend_from_reader(&mut cursor)
                .expect("deserialize lichess entry");
        }

        Ok(entry)
    }

    pub fn read_player(
        &self,
        key: &KeyPrefix,
        since: Month,
        until: Month,
    ) -> Result<PlayerEntry, rocksdb::Error> {
        let mut opt = ReadOptions::default();
        opt.set_prefix_same_as_start(true);
        opt.set_iterate_lower_bound(key.with_month(since).into_bytes());
        opt.set_iterate_upper_bound(key.with_month(until.add_months_saturating(1)).into_bytes());

        let iterator = self
            .inner
            .iterator_cf_opt(self.cf_player, opt, IteratorMode::Start);

        let mut entry = PlayerEntry::default();
        for (_key, value) in iterator {
            let mut cursor = Cursor::new(value);
            entry
                .extend_from_reader(&mut cursor)
                .expect("deserialize player entry");
        }

        Ok(entry)
    }

    pub fn player_status(&self, id: &UserId) -> Result<Option<PlayerStatus>, rocksdb::Error> {
        Ok(self
            .inner
            .get_cf(self.cf_player_status, id.as_lowercase_str())?
            .map(|buf| {
                let mut cursor = Cursor::new(buf);
                PlayerStatus::read(&mut cursor).expect("deserialize status")
            }))
    }

    pub fn put_player_status(
        &self,
        id: &UserId,
        status: &PlayerStatus,
    ) -> Result<(), rocksdb::Error> {
        let mut cursor = Cursor::new(Vec::with_capacity(PlayerStatus::SIZE_HINT));
        status.write(&mut cursor).expect("serialize status");
        self.inner.put_cf(
            self.cf_player_status,
            id.as_lowercase_str(),
            cursor.into_inner(),
        )
    }

    pub fn batch(&self) -> LichessBatch<'_> {
        LichessBatch {
            inner: self,
            batch: WriteBatch::default(),
        }
    }
}

pub struct LichessBatch<'a> {
    inner: &'a LichessDatabase<'a>,
    batch: WriteBatch,
}

impl LichessBatch<'_> {
    pub fn merge_lichess(&mut self, key: Key, entry: LichessEntry) {
        let mut cursor = Cursor::new(Vec::with_capacity(LichessEntry::SIZE_HINT));
        entry.write(&mut cursor).expect("serialize lichess entry");
        self.batch
            .merge_cf(self.inner.cf_lichess, key.into_bytes(), cursor.into_inner());
    }

    pub fn merge_game(&mut self, id: GameId, info: LichessGame) {
        let mut cursor = Cursor::new(Vec::with_capacity(LichessGame::SIZE_HINT));
        info.write(&mut cursor).expect("serialize game info");
        self.batch.merge_cf(
            self.inner.cf_lichess_game,
            id.to_bytes(),
            cursor.into_inner(),
        );
    }

    pub fn merge_player(&mut self, key: Key, entry: PlayerEntry) {
        let mut cursor = Cursor::new(Vec::with_capacity(PlayerEntry::SIZE_HINT));
        entry.write(&mut cursor).expect("serialize player entry");
        self.batch
            .merge_cf(self.inner.cf_player, key.into_bytes(), cursor.into_inner());
    }

    pub fn commit(self) -> Result<(), rocksdb::Error> {
        self.inner.inner.write(self.batch)
    }
}

fn lichess_merge(
    _key: &[u8],
    existing: Option<&[u8]>,
    operands: &mut MergeOperands,
) -> Option<Vec<u8>> {
    let mut entry = LichessEntry::default();
    let mut size_hint = 0;
    for op in existing.into_iter().chain(operands.into_iter()) {
        let mut cursor = Cursor::new(op);
        entry
            .extend_from_reader(&mut cursor)
            .expect("deserialize for lichess merge");
        size_hint += op.len();
    }
    let mut cursor = Cursor::new(Vec::with_capacity(size_hint));
    entry.write(&mut cursor).expect("write lichess entry");
    Some(cursor.into_inner())
}

fn lichess_game_merge(
    _key: &[u8],
    existing: Option<&[u8]>,
    operands: &mut MergeOperands,
) -> Option<Vec<u8>> {
    // Take latest game info, but merge index status.
    let mut info: Option<LichessGame> = None;
    let mut size_hint = 0;
    for op in existing.into_iter().chain(operands.into_iter()) {
        let mut cursor = Cursor::new(op);
        let mut new_info = LichessGame::read(&mut cursor).expect("read for lichess game merge");
        if let Some(old_info) = info {
            new_info.indexed_player.white |= old_info.indexed_player.white;
            new_info.indexed_player.black |= old_info.indexed_player.black;
            new_info.indexed_lichess |= old_info.indexed_lichess;
        }
        info = Some(new_info);
        size_hint = op.len();
    }
    info.map(|info| {
        let mut cursor = Cursor::new(Vec::with_capacity(size_hint));
        info.write(&mut cursor).expect("write lichess game");
        cursor.into_inner()
    })
}

fn player_merge(
    _key: &[u8],
    existing: Option<&[u8]>,
    operands: &mut MergeOperands,
) -> Option<Vec<u8>> {
    let mut entry = PlayerEntry::default();
    let mut size_hint = 0;
    for op in existing.into_iter().chain(operands.into_iter()) {
        let mut cursor = Cursor::new(op);
        entry
            .extend_from_reader(&mut cursor)
            .expect("deserialize for player merge");
        size_hint += op.len();
    }
    let mut cursor = Cursor::new(Vec::with_capacity(size_hint));
    entry.write(&mut cursor).expect("write player entry");
    Some(cursor.into_inner())
}

fn masters_merge(
    _key: &[u8],
    existing: Option<&[u8]>,
    operands: &mut MergeOperands,
) -> Option<Vec<u8>> {
    let mut entry = MastersEntry::default();
    let mut size_hint = 0;
    for op in existing.into_iter().chain(operands.into_iter()) {
        let mut cursor = Cursor::new(op);
        entry
            .extend_from_reader(&mut cursor)
            .expect("deserialize for masters merge");
        size_hint += op.len();
    }
    let mut cursor = Cursor::new(Vec::with_capacity(size_hint));
    entry.write(&mut cursor).expect("write masters entry");
    Some(cursor.into_inner())
}

fn void_merge(
    _key: &[u8],
    _existing: Option<&[u8]>,
    _operands: &mut MergeOperands,
) -> Option<Vec<u8>> {
    unreachable!("void merge")
}
<commit_msg>increase block cache size from 8 to 512 mb<commit_after>use std::{io::Cursor, path::Path};

use rocksdb::{
    merge_operator::MergeFn, BlockBasedIndexType, BlockBasedOptions, Cache, ColumnFamily,
    ColumnFamilyDescriptor, IteratorMode, MergeOperands, Options, ReadOptions, SliceTransform,
    WriteBatch, DB,
};

use crate::model::{
    GameId, Key, KeyPrefix, LichessEntry, LichessGame, MastersEntry, MastersGame, Month,
    PlayerEntry, PlayerStatus, UserId, Year,
};

#[derive(Debug)]
pub struct Database {
    pub inner: DB,
}

fn column_family(
    name: &str,
    merge: Option<&str>,
    merge_fn: impl MergeFn + Clone,
    prefix: Option<usize>,
    block_size: usize,
    bloom_filter: i32,
    cache: &Cache,
) -> ColumnFamilyDescriptor {
    let mut opts = Options::default();
    if let Some(merge) = merge {
        opts.set_merge_operator_associative(merge, merge_fn);
    }
    opts.set_prefix_extractor(match prefix {
        Some(prefix) => SliceTransform::create_fixed_prefix(prefix),
        None => SliceTransform::create_noop(),
    });
    let mut block_opts = BlockBasedOptions::default();
    block_opts.set_block_cache(cache);
    block_opts.set_index_type(BlockBasedIndexType::HashSearch);
    block_opts.set_block_size(block_size);
    if bloom_filter > 0 {
        block_opts.set_bloom_filter(bloom_filter, true);
    }
    opts.set_block_based_table_factory(&block_opts);
    ColumnFamilyDescriptor::new(name, opts)
}

impl Database {
    pub fn open<P: AsRef<Path>>(path: P) -> Result<Database, rocksdb::Error> {
        let mut db_opts = Options::default();
        db_opts.create_if_missing(true);
        db_opts.create_missing_column_families(true);

        let cache = Cache::new_lru_cache(512 * 1024 * 1024)?;

        let inner = DB::open_cf_descriptors(
            &db_opts,
            path,
            vec![
                // Masters database
                column_family(
                    "masters",
                    Some("masters_merge"),
                    masters_merge,
                    Some(KeyPrefix::SIZE),
                    8 * 1024,
                    3,
                    &cache,
                ),
                column_family("masters_game", None, void_merge, None, 4 * 1024, 0, &cache),
                // Lichess database
                column_family(
                    "lichess",
                    Some("lichess_merge"),
                    lichess_merge,
                    Some(KeyPrefix::SIZE),
                    16 * 1024,
                    2,
                    &cache,
                ),
                column_family(
                    "lichess_game",
                    Some("lichess_game_merge"),
                    lichess_game_merge,
                    None,
                    4 * 1024,
                    0,
                    &cache,
                ),
                // Player database (also shares lichess_game)
                column_family(
                    "player",
                    Some("player_merge"),
                    player_merge,
                    Some(KeyPrefix::SIZE),
                    16 * 1024,
                    2,
                    &cache,
                ),
                column_family("player_status", None, void_merge, None, 4 * 1024, 0, &cache),
            ],
        )?;

        Ok(Database { inner })
    }

    pub fn masters(&self) -> MastersDatabase<'_> {
        MastersDatabase {
            inner: &self.inner,
            cf_masters: self.inner.cf_handle("masters").expect("cf masters"),
            cf_masters_game: self
                .inner
                .cf_handle("masters_game")
                .expect("cf masters_game"),
        }
    }

    pub fn lichess(&self) -> LichessDatabase<'_> {
        LichessDatabase {
            inner: &self.inner,
            cf_lichess: self.inner.cf_handle("lichess").expect("cf lichess"),
            cf_lichess_game: self
                .inner
                .cf_handle("lichess_game")
                .expect("cf lichess_game"),

            cf_player: self.inner.cf_handle("player").expect("cf player"),
            cf_player_status: self
                .inner
                .cf_handle("player_status")
                .expect("cf player_status"),
        }
    }
}

pub struct MastersDatabase<'a> {
    inner: &'a DB,
    cf_masters: &'a ColumnFamily,
    cf_masters_game: &'a ColumnFamily,
}

impl MastersDatabase<'_> {
    pub fn has_game(&self, id: GameId) -> Result<bool, rocksdb::Error> {
        self.inner
            .get_cf(self.cf_masters_game, id.to_bytes())
            .map(|maybe_entry| maybe_entry.is_some())
    }

    pub fn game(&self, id: GameId) -> Result<Option<MastersGame>, rocksdb::Error> {
        Ok(self
            .inner
            .get_cf(self.cf_masters_game, id.to_bytes())?
            .map(|buf| serde_json::from_slice(&buf).expect("deserialize masters game")))
    }

    pub fn has(&self, key: Key) -> Result<bool, rocksdb::Error> {
        self.inner
            .get_cf(self.cf_masters, key.into_bytes())
            .map(|maybe_entry| maybe_entry.is_some())
    }

    pub fn read(
        &self,
        key: KeyPrefix,
        since: Year,
        until: Year,
    ) -> Result<MastersEntry, rocksdb::Error> {
        let mut opt = ReadOptions::default();
        opt.set_prefix_same_as_start(true);
        opt.set_iterate_lower_bound(key.with_year(since).into_bytes());
        opt.set_iterate_upper_bound(key.with_year(until.add_years_saturating(1)).into_bytes());

        let iterator = self
            .inner
            .iterator_cf_opt(self.cf_masters, opt, IteratorMode::Start);

        let mut entry = MastersEntry::default();
        for (_key, value) in iterator {
            let mut cursor = Cursor::new(value);
            entry
                .extend_from_reader(&mut cursor)
                .expect("deserialize masters entry");
        }

        Ok(entry)
    }

    pub fn batch(&self) -> MastersBatch<'_> {
        MastersBatch {
            db: self,
            batch: WriteBatch::default(),
        }
    }
}

pub struct MastersBatch<'a> {
    db: &'a MastersDatabase<'a>,
    batch: WriteBatch,
}

impl MastersBatch<'_> {
    pub fn merge(&mut self, key: Key, entry: MastersEntry) {
        let mut cursor = Cursor::new(Vec::with_capacity(MastersEntry::SIZE_HINT));
        entry.write(&mut cursor).expect("serialize masters entry");
        self.batch
            .merge_cf(self.db.cf_masters, key.into_bytes(), cursor.into_inner());
    }

    pub fn put_game(&mut self, id: GameId, game: &MastersGame) {
        self.batch.put_cf(
            self.db.cf_masters_game,
            id.to_bytes(),
            serde_json::to_vec(game).expect("serialize masters game"),
        );
    }

    pub fn commit(self) -> Result<(), rocksdb::Error> {
        self.db.inner.write(self.batch)
    }
}

pub struct LichessDatabase<'a> {
    inner: &'a DB,
    cf_lichess: &'a ColumnFamily,
    cf_lichess_game: &'a ColumnFamily,

    cf_player: &'a ColumnFamily,
    cf_player_status: &'a ColumnFamily,
}

impl LichessDatabase<'_> {
    pub fn game(&self, id: GameId) -> Result<Option<LichessGame>, rocksdb::Error> {
        Ok(self
            .inner
            .get_cf(self.cf_lichess_game, id.to_bytes())?
            .map(|buf| {
                let mut cursor = Cursor::new(buf);
                LichessGame::read(&mut cursor).expect("deserialize game info")
            }))
    }

    pub fn read_lichess(
        &self,
        key: &KeyPrefix,
        since: Month,
        until: Month,
    ) -> Result<LichessEntry, rocksdb::Error> {
        let mut opt = ReadOptions::default();
        opt.set_prefix_same_as_start(true);
        opt.set_iterate_lower_bound(key.with_month(since).into_bytes());
        opt.set_iterate_upper_bound(key.with_month(until.add_months_saturating(1)).into_bytes());

        let iterator = self
            .inner
            .iterator_cf_opt(self.cf_lichess, opt, IteratorMode::Start);

        let mut entry = LichessEntry::default();
        for (_key, value) in iterator {
            let mut cursor = Cursor::new(value);
            entry
                .extend_from_reader(&mut cursor)
                .expect("deserialize lichess entry");
        }

        Ok(entry)
    }

    pub fn read_player(
        &self,
        key: &KeyPrefix,
        since: Month,
        until: Month,
    ) -> Result<PlayerEntry, rocksdb::Error> {
        let mut opt = ReadOptions::default();
        opt.set_prefix_same_as_start(true);
        opt.set_iterate_lower_bound(key.with_month(since).into_bytes());
        opt.set_iterate_upper_bound(key.with_month(until.add_months_saturating(1)).into_bytes());

        let iterator = self
            .inner
            .iterator_cf_opt(self.cf_player, opt, IteratorMode::Start);

        let mut entry = PlayerEntry::default();
        for (_key, value) in iterator {
            let mut cursor = Cursor::new(value);
            entry
                .extend_from_reader(&mut cursor)
                .expect("deserialize player entry");
        }

        Ok(entry)
    }

    pub fn player_status(&self, id: &UserId) -> Result<Option<PlayerStatus>, rocksdb::Error> {
        Ok(self
            .inner
            .get_cf(self.cf_player_status, id.as_lowercase_str())?
            .map(|buf| {
                let mut cursor = Cursor::new(buf);
                PlayerStatus::read(&mut cursor).expect("deserialize status")
            }))
    }

    pub fn put_player_status(
        &self,
        id: &UserId,
        status: &PlayerStatus,
    ) -> Result<(), rocksdb::Error> {
        let mut cursor = Cursor::new(Vec::with_capacity(PlayerStatus::SIZE_HINT));
        status.write(&mut cursor).expect("serialize status");
        self.inner.put_cf(
            self.cf_player_status,
            id.as_lowercase_str(),
            cursor.into_inner(),
        )
    }

    pub fn batch(&self) -> LichessBatch<'_> {
        LichessBatch {
            inner: self,
            batch: WriteBatch::default(),
        }
    }
}

pub struct LichessBatch<'a> {
    inner: &'a LichessDatabase<'a>,
    batch: WriteBatch,
}

impl LichessBatch<'_> {
    pub fn merge_lichess(&mut self, key: Key, entry: LichessEntry) {
        let mut cursor = Cursor::new(Vec::with_capacity(LichessEntry::SIZE_HINT));
        entry.write(&mut cursor).expect("serialize lichess entry");
        self.batch
            .merge_cf(self.inner.cf_lichess, key.into_bytes(), cursor.into_inner());
    }

    pub fn merge_game(&mut self, id: GameId, info: LichessGame) {
        let mut cursor = Cursor::new(Vec::with_capacity(LichessGame::SIZE_HINT));
        info.write(&mut cursor).expect("serialize game info");
        self.batch.merge_cf(
            self.inner.cf_lichess_game,
            id.to_bytes(),
            cursor.into_inner(),
        );
    }

    pub fn merge_player(&mut self, key: Key, entry: PlayerEntry) {
        let mut cursor = Cursor::new(Vec::with_capacity(PlayerEntry::SIZE_HINT));
        entry.write(&mut cursor).expect("serialize player entry");
        self.batch
            .merge_cf(self.inner.cf_player, key.into_bytes(), cursor.into_inner());
    }

    pub fn commit(self) -> Result<(), rocksdb::Error> {
        self.inner.inner.write(self.batch)
    }
}

fn lichess_merge(
    _key: &[u8],
    existing: Option<&[u8]>,
    operands: &mut MergeOperands,
) -> Option<Vec<u8>> {
    let mut entry = LichessEntry::default();
    let mut size_hint = 0;
    for op in existing.into_iter().chain(operands.into_iter()) {
        let mut cursor = Cursor::new(op);
        entry
            .extend_from_reader(&mut cursor)
            .expect("deserialize for lichess merge");
        size_hint += op.len();
    }
    let mut cursor = Cursor::new(Vec::with_capacity(size_hint));
    entry.write(&mut cursor).expect("write lichess entry");
    Some(cursor.into_inner())
}

fn lichess_game_merge(
    _key: &[u8],
    existing: Option<&[u8]>,
    operands: &mut MergeOperands,
) -> Option<Vec<u8>> {
    // Take latest game info, but merge index status.
    let mut info: Option<LichessGame> = None;
    let mut size_hint = 0;
    for op in existing.into_iter().chain(operands.into_iter()) {
        let mut cursor = Cursor::new(op);
        let mut new_info = LichessGame::read(&mut cursor).expect("read for lichess game merge");
        if let Some(old_info) = info {
            new_info.indexed_player.white |= old_info.indexed_player.white;
            new_info.indexed_player.black |= old_info.indexed_player.black;
            new_info.indexed_lichess |= old_info.indexed_lichess;
        }
        info = Some(new_info);
        size_hint = op.len();
    }
    info.map(|info| {
        let mut cursor = Cursor::new(Vec::with_capacity(size_hint));
        info.write(&mut cursor).expect("write lichess game");
        cursor.into_inner()
    })
}

fn player_merge(
    _key: &[u8],
    existing: Option<&[u8]>,
    operands: &mut MergeOperands,
) -> Option<Vec<u8>> {
    let mut entry = PlayerEntry::default();
    let mut size_hint = 0;
    for op in existing.into_iter().chain(operands.into_iter()) {
        let mut cursor = Cursor::new(op);
        entry
            .extend_from_reader(&mut cursor)
            .expect("deserialize for player merge");
        size_hint += op.len();
    }
    let mut cursor = Cursor::new(Vec::with_capacity(size_hint));
    entry.write(&mut cursor).expect("write player entry");
    Some(cursor.into_inner())
}

fn masters_merge(
    _key: &[u8],
    existing: Option<&[u8]>,
    operands: &mut MergeOperands,
) -> Option<Vec<u8>> {
    let mut entry = MastersEntry::default();
    let mut size_hint = 0;
    for op in existing.into_iter().chain(operands.into_iter()) {
        let mut cursor = Cursor::new(op);
        entry
            .extend_from_reader(&mut cursor)
            .expect("deserialize for masters merge");
        size_hint += op.len();
    }
    let mut cursor = Cursor::new(Vec::with_capacity(size_hint));
    entry.write(&mut cursor).expect("write masters entry");
    Some(cursor.into_inner())
}

fn void_merge(
    _key: &[u8],
    _existing: Option<&[u8]>,
    _operands: &mut MergeOperands,
) -> Option<Vec<u8>> {
    unreachable!("void merge")
}
<|endoftext|>


<commit_before>use context::PortMidi;
use ffi;
use std::os::raw::c_int;
use std::ptr;
use types::*;
use ffi::MaybeError;
use device::DeviceInfo;
use std::marker::Send;

/// Represents the input port of a PortMidi device.
pub struct InputPort<'a> {
    stream: *const ffi::PortMidiStream,
    buffer_size: usize,
    _context: &'a PortMidi, // Used for lifetime pinning
    device: DeviceInfo,
}
impl<'a> InputPort<'a> {
    /// Construct a new `InputPort` for the given device and buffer size.
    ///
    /// If the `device` is not an input device an `Error::NotAnInputDevice` is returned.
    pub fn new(context: &'a PortMidi, device: DeviceInfo, buffer_size: usize) -> Result<InputPort> {
        if device.is_output() {
            return Err(Error::NotAnInputDevice);
        }
        let raw_stream: *const ffi::PortMidiStream = ptr::null();
        try!(Result::from(unsafe {
            ffi::Pm_OpenInput(&raw_stream as *const *const _,
                              device.id(),
                              ptr::null(), // *inputDriverInfo, not needed for normal operation
                              buffer_size as c_int,
                              ptr::null(), // PmTimeProcPtr, a procedure that returns time in ms
                              ptr::null()) // time_info, a pointer passed to the time procedure
        }));

        Ok(InputPort {
            stream: raw_stream,
            buffer_size: buffer_size,
            _context: context,
            device,
        })
    }

    /// Returns a `Vec<MidiEvent>` with at most `cnt` elements.
    /// If there was no Midi event available, `None` is returned.
    /// If PortMidi fails to read from the device an `Error::PortMidi(_)` is returned.
    pub fn read_n(&self, cnt: usize) -> Result<Option<Vec<MidiEvent>>> {
        let read_cnt = if cnt > self.buffer_size {
            self.buffer_size as c_int
        } else {
            cnt as c_int
        };
        let mut event_buffer = vec![ffi::PmEvent::default(); self.buffer_size];
        let res = unsafe { ffi::Pm_Read(self.stream, event_buffer.as_mut_ptr(), read_cnt) };
        match ffi::PmError::try_from(res) {
            Ok(event_cnt) => {
                let events = (0..event_cnt as usize)
                                 .map(|i| MidiEvent::from(event_buffer[i].clone()))
                                 .collect::<Vec<MidiEvent>>();
                Ok(Some(events))
            }
            Err(ffi::PmError::PmNoError) => Ok(None),
            Err(err) => return Err(Error::PortMidi(err)),
        }
    }

    /// Reads a single `MidiEvent` if one is avaible.
    ///
    /// A `Result` of `None` means no event was available.
    pub fn read(&mut self) -> Result<Option<MidiEvent>> {
        match self.read_n(1) {
            Ok(Some(mut vec)) => Ok(vec.pop()),
            Ok(_) => Ok(None),
            Err(e) => Err(e),
        }
    }

    /// Polls for available Midi events.
    /// Returns `true` if there are events available, otherwise `false` is returned.
    /// If the polling fails an `Error::PortMidi(_)` is returned.
    pub fn poll(&self) -> Result<bool> {
        let pm_error = unsafe { ffi::Pm_Poll(self.stream) };
        match pm_error {
            ffi::PmError::PmNoError => Ok(false),
            ffi::PmError::PmGotData => Ok(true),
            err @ _ => Err(Error::PortMidi(err)),
        }
    }

    /// Returns the `DeviceInfo` of the Midi device that owns this port.
    pub fn device(&self) -> DeviceInfo {
        return self.device.clone();
    }
}
impl<'a> Drop for InputPort<'a> {
    fn drop(&mut self) {
        if let Err(err) = Result::from(unsafe { ffi::Pm_Close(self.stream) }) {
            println!("{}", err);
        }
    }
}
unsafe impl<'a> Send for InputPort<'a> {}

/// Represents the output port of a PortMidi device.
pub struct OutputPort<'a> {
    stream: *const ffi::PortMidiStream,
    _context: &'a PortMidi, // Used for lifetime pinning
    device: DeviceInfo,
}
impl<'a> OutputPort<'a> {
    /// Construct a new `OutputPort` for the given device and buffer size.
    ///
    /// If the `device` is not an output device an `Error::NotAnOutputDevice` is returned.
    pub fn new(context: &'a PortMidi, device: DeviceInfo, buffer_size: usize) -> Result<OutputPort> {
        if device.is_input() {
            return Err(Error::NotAnOutputDevice);
        }
        let raw_stream: *const ffi::PortMidiStream = ptr::null();
        try!(Result::from(unsafe {
            ffi::Pm_OpenOutput(&raw_stream as *const *const _,
                               device.id(),
                               ptr::null(), // *inputDriverInfo, not needed for normal operation
                               buffer_size as c_int,
                               ptr::null(), // PmTimeProcPtr, a procedure that returns time in ms,
                               ptr::null(), // time_info, a pointer passed to the time procedure
                               0) //latency
        }));

        Ok(OutputPort {
            stream: raw_stream,
            _context: context,
            device,
        })
    }

    /// Write a single `MidiEvent`.
    /// Returns an `Error::PortMidi(_)` if something went wrong.
    pub fn write_event<T: Into<MidiEvent>>(&mut self, midi_event: T) -> Result<()> {
        self.write_events(vec![midi_event])
    }

    /// Write a buffer of midi events to the output port.
    /// Returns an `Error::PortMidi(_)` if something went wrong.
    pub fn write_events<T: Into<MidiEvent>>(&mut self, midi_events: Vec<T>) -> Result<()> {
        let events: Vec<ffi::PmEvent> = midi_events.into_iter()
                                                   .map(|event| event.into().into())
                                                   .collect();
        Result::from(unsafe { ffi::Pm_Write(self.stream, events.as_ptr(), events.len() as c_int) })
    }

    /// Write a single `MidiMessage`.
    /// Returns an `Error::PortMidi(_)` if something went wrong.
    pub fn write_message<T: Into<MidiMessage>>(&mut self, midi_message: T) -> Result<()> {
        Result::from(unsafe { ffi::Pm_WriteShort(self.stream, 0, midi_message.into().into()) })
    }

    /// Returns the `DeviceInfo` of the Midi device that owns this port.
    pub fn device(&self) -> DeviceInfo {
        return self.device.clone();
    }

    // Write arbitrarily long EOX-terminated data
    pub fn write_sysex(&self, timestamp: ffi::PmTimestamp, msg: &[u8]) -> Result<()> {
        // Sysex writes MUST be EOX-terminated
        if Some(&ffi::MIDI_EOX) != msg.last() {
            Err(Error::Invalid)
        } else {
            Result::from(unsafe { ffi::Pm_WriteSysEx(self.stream, timestamp, msg.as_ptr())})
        }
    }
}
impl<'a> Drop for OutputPort<'a> {
    fn drop(&mut self) {
        if let Err(err) = Result::from(unsafe { ffi::Pm_Close(self.stream) }) {
            println!("{}", err);
        }
    }
}
unsafe impl<'a> Send for OutputPort<'a> {}
<commit_msg>Don't construct a vector when reading a single event.<commit_after>use context::PortMidi;
use device::DeviceInfo;
use ffi;
use ffi::MaybeError;
use std::marker::Send;
use std::os::raw::c_int;
use std::ptr;
use types::*;

/// Represents the input port of a PortMidi device.
pub struct InputPort<'a> {
    stream: *const ffi::PortMidiStream,
    buffer_size: usize,
    _context: &'a PortMidi, // Used for lifetime pinning
    device: DeviceInfo,
}
impl<'a> InputPort<'a> {
    /// Construct a new `InputPort` for the given device and buffer size.
    ///
    /// If the `device` is not an input device an `Error::NotAnInputDevice` is returned.
    pub fn new(context: &'a PortMidi, device: DeviceInfo, buffer_size: usize) -> Result<InputPort> {
        if device.is_output() {
            return Err(Error::NotAnInputDevice);
        }
        let raw_stream: *const ffi::PortMidiStream = ptr::null();
        try!(Result::from(unsafe {
            ffi::Pm_OpenInput(
                &raw_stream as *const *const _,
                device.id(),
                ptr::null(), // *inputDriverInfo, not needed for normal operation
                buffer_size as c_int,
                ptr::null(), // PmTimeProcPtr, a procedure that returns time in ms
                ptr::null(),
            ) // time_info, a pointer passed to the time procedure
        }));

        Ok(InputPort {
            stream: raw_stream,
            buffer_size: buffer_size,
            _context: context,
            device,
        })
    }

    /// Returns a `Vec<MidiEvent>` with at most `cnt` elements.
    /// If there was no Midi event available, `None` is returned.
    /// If PortMidi fails to read from the device an `Error::PortMidi(_)` is returned.
    pub fn read_n(&self, cnt: usize) -> Result<Option<Vec<MidiEvent>>> {
        let read_cnt = if cnt > self.buffer_size {
            self.buffer_size as c_int
        } else {
            cnt as c_int
        };
        let mut event_buffer = vec![ffi::PmEvent::default(); self.buffer_size];
        let res = unsafe { ffi::Pm_Read(self.stream, event_buffer.as_mut_ptr(), read_cnt) };
        match ffi::PmError::try_from(res) {
            Ok(event_cnt) => {
                let events = (0..event_cnt as usize)
                    .map(|i| MidiEvent::from(event_buffer[i].clone()))
                    .collect::<Vec<MidiEvent>>();
                Ok(Some(events))
            }
            Err(ffi::PmError::PmNoError) => Ok(None),
            Err(err) => return Err(Error::PortMidi(err)),
        }
    }

    /// Reads a single `MidiEvent` if one is avaible.
    ///
    /// A `Result` of `None` means no event was available.
    pub fn read(&mut self) -> Result<Option<MidiEvent>> {
        let mut event = ffi::PmEvent::default();
        let read_cnt = 1;
        let res = unsafe { ffi::Pm_Read(self.stream, &mut event, read_cnt) };
        match ffi::PmError::try_from(res) {
            Ok(0) => Ok(None),
            Ok(_) => Ok(Some(MidiEvent::from(event))),
            Err(ffi::PmError::PmNoError) => Ok(None),
            Err(err) => Err(Error::PortMidi(err)),
        }
    }

    /// Polls for available Midi events.
    /// Returns `true` if there are events available, otherwise `false` is returned.
    /// If the polling fails an `Error::PortMidi(_)` is returned.
    pub fn poll(&self) -> Result<bool> {
        let pm_error = unsafe { ffi::Pm_Poll(self.stream) };
        match pm_error {
            ffi::PmError::PmNoError => Ok(false),
            ffi::PmError::PmGotData => Ok(true),
            err @ _ => Err(Error::PortMidi(err)),
        }
    }

    /// Returns the `DeviceInfo` of the Midi device that owns this port.
    pub fn device(&self) -> DeviceInfo {
        return self.device.clone();
    }
}
impl<'a> Drop for InputPort<'a> {
    fn drop(&mut self) {
        if let Err(err) = Result::from(unsafe { ffi::Pm_Close(self.stream) }) {
            println!("{}", err);
        }
    }
}
unsafe impl<'a> Send for InputPort<'a> {}

/// Represents the output port of a PortMidi device.
pub struct OutputPort<'a> {
    stream: *const ffi::PortMidiStream,
    _context: &'a PortMidi, // Used for lifetime pinning
    device: DeviceInfo,
}
impl<'a> OutputPort<'a> {
    /// Construct a new `OutputPort` for the given device and buffer size.
    ///
    /// If the `device` is not an output device an `Error::NotAnOutputDevice` is returned.
    pub fn new(
        context: &'a PortMidi,
        device: DeviceInfo,
        buffer_size: usize,
    ) -> Result<OutputPort> {
        if device.is_input() {
            return Err(Error::NotAnOutputDevice);
        }
        let raw_stream: *const ffi::PortMidiStream = ptr::null();
        try!(Result::from(unsafe {
            ffi::Pm_OpenOutput(
                &raw_stream as *const *const _,
                device.id(),
                ptr::null(), // *inputDriverInfo, not needed for normal operation
                buffer_size as c_int,
                ptr::null(), // PmTimeProcPtr, a procedure that returns time in ms,
                ptr::null(), // time_info, a pointer passed to the time procedure
                0,
            ) //latency
        }));

        Ok(OutputPort {
            stream: raw_stream,
            _context: context,
            device,
        })
    }

    /// Write a single `MidiEvent`.
    /// Returns an `Error::PortMidi(_)` if something went wrong.
    pub fn write_event<T: Into<MidiEvent>>(&mut self, midi_event: T) -> Result<()> {
        self.write_events(vec![midi_event])
    }

    /// Write a buffer of midi events to the output port.
    /// Returns an `Error::PortMidi(_)` if something went wrong.
    pub fn write_events<T: Into<MidiEvent>>(&mut self, midi_events: Vec<T>) -> Result<()> {
        let events: Vec<ffi::PmEvent> = midi_events
            .into_iter()
            .map(|event| event.into().into())
            .collect();
        Result::from(unsafe { ffi::Pm_Write(self.stream, events.as_ptr(), events.len() as c_int) })
    }

    /// Write a single `MidiMessage`.
    /// Returns an `Error::PortMidi(_)` if something went wrong.
    pub fn write_message<T: Into<MidiMessage>>(&mut self, midi_message: T) -> Result<()> {
        Result::from(unsafe { ffi::Pm_WriteShort(self.stream, 0, midi_message.into().into()) })
    }

    /// Returns the `DeviceInfo` of the Midi device that owns this port.
    pub fn device(&self) -> DeviceInfo {
        return self.device.clone();
    }

    // Write arbitrarily long EOX-terminated data
    pub fn write_sysex(&self, timestamp: ffi::PmTimestamp, msg: &[u8]) -> Result<()> {
        // Sysex writes MUST be EOX-terminated
        if Some(&ffi::MIDI_EOX) != msg.last() {
            Err(Error::Invalid)
        } else {
            Result::from(unsafe { ffi::Pm_WriteSysEx(self.stream, timestamp, msg.as_ptr()) })
        }
    }
}
impl<'a> Drop for OutputPort<'a> {
    fn drop(&mut self) {
        if let Err(err) = Result::from(unsafe { ffi::Pm_Close(self.stream) }) {
            println!("{}", err);
        }
    }
}
unsafe impl<'a> Send for OutputPort<'a> {}
<|endoftext|>


<commit_before>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The Gamma and derived distributions.

use core::num::Float;

use {Rng, Open01};
use super::normal::StandardNormal;
use super::{IndependentSample, Sample, Exp};

/// The Gamma distribution `Gamma(shape, scale)` distribution.
///
/// The density function of this distribution is
///
/// ```text
/// f(x) =  x^(k - 1) * exp(-x / θ) / (Γ(k) * θ^k)
/// ```
///
/// where `Γ` is the Gamma function, `k` is the shape and `θ` is the
/// scale and both `k` and `θ` are strictly positive.
///
/// The algorithm used is that described by Marsaglia & Tsang 2000[1],
/// falling back to directly sampling from an Exponential for `shape
/// == 1`, and using the boosting technique described in [1] for
/// `shape < 1`.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{IndependentSample, Gamma};
///
/// let gamma = Gamma::new(2.0, 5.0);
/// let v = gamma.ind_sample(&mut rand::task_rng());
/// println!("{} is from a Gamma(2, 5) distribution", v);
/// ```
///
/// [1]: George Marsaglia and Wai Wan Tsang. 2000. "A Simple Method
/// for Generating Gamma Variables" *ACM Trans. Math. Softw.* 26, 3
/// (September 2000),
/// 363-372. DOI:[10.1145/358407.358414](http://doi.acm.org/10.1145/358407.358414)
pub struct Gamma {
    repr: GammaRepr,
}

enum GammaRepr {
    Large(GammaLargeShape),
    One(Exp),
    Small(GammaSmallShape)
}

// These two helpers could be made public, but saving the
// match-on-Gamma-enum branch from using them directly (e.g. if one
// knows that the shape is always > 1) doesn't appear to be much
// faster.

/// Gamma distribution where the shape parameter is less than 1.
///
/// Note, samples from this require a compulsory floating-point `pow`
/// call, which makes it significantly slower than sampling from a
/// gamma distribution where the shape parameter is greater than or
/// equal to 1.
///
/// See `Gamma` for sampling from a Gamma distribution with general
/// shape parameters.
struct GammaSmallShape {
    inv_shape: f64,
    large_shape: GammaLargeShape
}

/// Gamma distribution where the shape parameter is larger than 1.
///
/// See `Gamma` for sampling from a Gamma distribution with general
/// shape parameters.
struct GammaLargeShape {
    shape: f64,
    scale: f64,
    c: f64,
    d: f64
}

impl Gamma {
    /// Construct an object representing the `Gamma(shape, scale)`
    /// distribution.
    ///
    /// Fails if `shape <= 0` or `scale <= 0`.
    pub fn new(shape: f64, scale: f64) -> Gamma {
        assert!(shape > 0.0, "Gamma::new called with shape <= 0");
        assert!(scale > 0.0, "Gamma::new called with scale <= 0");

        let repr = match shape {
            1.0        => One(Exp::new(1.0 / scale)),
            0.0 .. 1.0 => Small(GammaSmallShape::new_raw(shape, scale)),
            _          => Large(GammaLargeShape::new_raw(shape, scale))
        };
        Gamma { repr: repr }
    }
}

impl GammaSmallShape {
    fn new_raw(shape: f64, scale: f64) -> GammaSmallShape {
        GammaSmallShape {
            inv_shape: 1. / shape,
            large_shape: GammaLargeShape::new_raw(shape + 1.0, scale)
        }
    }
}

impl GammaLargeShape {
    fn new_raw(shape: f64, scale: f64) -> GammaLargeShape {
        let d = shape - 1. / 3.;
        GammaLargeShape {
            shape: shape,
            scale: scale,
            c: 1. / (9. * d).sqrt(),
            d: d
        }
    }
}

impl Sample<f64> for Gamma {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl Sample<f64> for GammaSmallShape {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl Sample<f64> for GammaLargeShape {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}

impl IndependentSample<f64> for Gamma {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        match self.repr {
            Small(ref g) => g.ind_sample(rng),
            One(ref g) => g.ind_sample(rng),
            Large(ref g) => g.ind_sample(rng),
        }
    }
}
impl IndependentSample<f64> for GammaSmallShape {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        let Open01(u) = rng.gen::<Open01<f64>>();

        self.large_shape.ind_sample(rng) * u.powf(self.inv_shape)
    }
}
impl IndependentSample<f64> for GammaLargeShape {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        loop {
            let StandardNormal(x) = rng.gen::<StandardNormal>();
            let v_cbrt = 1.0 + self.c * x;
            if v_cbrt <= 0.0 { // a^3 <= 0 iff a <= 0
                continue
            }

            let v = v_cbrt * v_cbrt * v_cbrt;
            let Open01(u) = rng.gen::<Open01<f64>>();

            let x_sqr = x * x;
            if u < 1.0 - 0.0331 * x_sqr * x_sqr ||
                u.ln() < 0.5 * x_sqr + self.d * (1.0 - v + v.ln()) {
                return self.d * v * self.scale
            }
        }
    }
}

/// The chi-squared distribution `χ²(k)`, where `k` is the degrees of
/// freedom.
///
/// For `k > 0` integral, this distribution is the sum of the squares
/// of `k` independent standard normal random variables. For other
/// `k`, this uses the equivalent characterisation `χ²(k) = Gamma(k/2,
/// 2)`.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{ChiSquared, IndependentSample};
///
/// let chi = ChiSquared::new(11.0);
/// let v = chi.ind_sample(&mut rand::task_rng());
/// println!("{} is from a χ²(11) distribution", v)
/// ```
pub struct ChiSquared {
    repr: ChiSquaredRepr,
}

enum ChiSquaredRepr {
    // k == 1, Gamma(alpha, ..) is particularly slow for alpha < 1,
    // e.g. when alpha = 1/2 as it would be for this case, so special-
    // casing and using the definition of N(0,1)^2 is faster.
    DoFExactlyOne,
    DoFAnythingElse(Gamma),
}

impl ChiSquared {
    /// Create a new chi-squared distribution with degrees-of-freedom
    /// `k`. Fails if `k < 0`.
    pub fn new(k: f64) -> ChiSquared {
        let repr = if k == 1.0 {
            DoFExactlyOne
        } else {
            assert!(k > 0.0, "ChiSquared::new called with `k` < 0");
            DoFAnythingElse(Gamma::new(0.5 * k, 2.0))
        };
        ChiSquared { repr: repr }
    }
}
impl Sample<f64> for ChiSquared {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for ChiSquared {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        match self.repr {
            DoFExactlyOne => {
                // k == 1 => N(0,1)^2
                let StandardNormal(norm) = rng.gen::<StandardNormal>();
                norm * norm
            }
            DoFAnythingElse(ref g) => g.ind_sample(rng)
        }
    }
}

/// The Fisher F distribution `F(m, n)`.
///
/// This distribution is equivalent to the ratio of two normalised
/// chi-squared distributions, that is, `F(m,n) = (χ²(m)/m) /
/// (χ²(n)/n)`.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{FisherF, IndependentSample};
///
/// let f = FisherF::new(2.0, 32.0);
/// let v = f.ind_sample(&mut rand::task_rng());
/// println!("{} is from an F(2, 32) distribution", v)
/// ```
pub struct FisherF {
    numer: ChiSquared,
    denom: ChiSquared,
    // denom_dof / numer_dof so that this can just be a straight
    // multiplication, rather than a division.
    dof_ratio: f64,
}

impl FisherF {
    /// Create a new `FisherF` distribution, with the given
    /// parameter. Fails if either `m` or `n` are not positive.
    pub fn new(m: f64, n: f64) -> FisherF {
        assert!(m > 0.0, "FisherF::new called with `m < 0`");
        assert!(n > 0.0, "FisherF::new called with `n < 0`");

        FisherF {
            numer: ChiSquared::new(m),
            denom: ChiSquared::new(n),
            dof_ratio: n / m
        }
    }
}
impl Sample<f64> for FisherF {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for FisherF {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        self.numer.ind_sample(rng) / self.denom.ind_sample(rng) * self.dof_ratio
    }
}

/// The Student t distribution, `t(nu)`, where `nu` is the degrees of
/// freedom.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{StudentT, IndependentSample};
///
/// let t = StudentT::new(11.0);
/// let v = t.ind_sample(&mut rand::task_rng());
/// println!("{} is from a t(11) distribution", v)
/// ```
pub struct StudentT {
    chi: ChiSquared,
    dof: f64
}

impl StudentT {
    /// Create a new Student t distribution with `n` degrees of
    /// freedom. Fails if `n <= 0`.
    pub fn new(n: f64) -> StudentT {
        assert!(n > 0.0, "StudentT::new called with `n <= 0`");
        StudentT {
            chi: ChiSquared::new(n),
            dof: n
        }
    }
}
impl Sample<f64> for StudentT {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for StudentT {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        let StandardNormal(norm) = rng.gen::<StandardNormal>();
        norm * (self.dof / self.chi.ind_sample(rng)).sqrt()
    }
}

#[cfg(test)]
mod test {
    use std::prelude::*;

    use distributions::{Sample, IndependentSample};
    use super::{ChiSquared, StudentT, FisherF};

    #[test]
    fn test_chi_squared_one() {
        let mut chi = ChiSquared::new(1.0);
        let mut rng = ::test::rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    fn test_chi_squared_small() {
        let mut chi = ChiSquared::new(0.5);
        let mut rng = ::test::rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    fn test_chi_squared_large() {
        let mut chi = ChiSquared::new(30.0);
        let mut rng = ::test::rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    #[should_fail]
    fn test_chi_squared_invalid_dof() {
        ChiSquared::new(-1.0);
    }

    #[test]
    fn test_f() {
        let mut f = FisherF::new(2.0, 32.0);
        let mut rng = ::test::rng();
        for _ in range(0, 1000) {
            f.sample(&mut rng);
            f.ind_sample(&mut rng);
        }
    }

    #[test]
    fn test_t() {
        let mut t = StudentT::new(11.0);
        let mut rng = ::test::rng();
        for _ in range(0, 1000) {
            t.sample(&mut rng);
            t.ind_sample(&mut rng);
        }
    }
}

#[cfg(test)]
mod bench {
    extern crate test;
    use std::prelude::*;
    use self::test::Bencher;
    use std::mem::size_of;
    use distributions::IndependentSample;
    use super::Gamma;

    #[bench]
    fn bench_gamma_large_shape(b: &mut Bencher) {
        let gamma = Gamma::new(10., 1.0);
        let mut rng = ::test::weak_rng();

        b.iter(|| {
            for _ in range(0, ::RAND_BENCH_N) {
                gamma.ind_sample(&mut rng);
            }
        });
        b.bytes = size_of::<f64>() as u64 * ::RAND_BENCH_N;
    }

    #[bench]
    fn bench_gamma_small_shape(b: &mut Bencher) {
        let gamma = Gamma::new(0.1, 1.0);
        let mut rng = ::test::weak_rng();

        b.iter(|| {
            for _ in range(0, ::RAND_BENCH_N) {
                gamma.ind_sample(&mut rng);
            }
        });
        b.bytes = size_of::<f64>() as u64 * ::RAND_BENCH_N;
    }
}
<commit_msg>Remove the dead code identified by the new lint<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The Gamma and derived distributions.

use core::num::Float;

use {Rng, Open01};
use super::normal::StandardNormal;
use super::{IndependentSample, Sample, Exp};

/// The Gamma distribution `Gamma(shape, scale)` distribution.
///
/// The density function of this distribution is
///
/// ```text
/// f(x) =  x^(k - 1) * exp(-x / θ) / (Γ(k) * θ^k)
/// ```
///
/// where `Γ` is the Gamma function, `k` is the shape and `θ` is the
/// scale and both `k` and `θ` are strictly positive.
///
/// The algorithm used is that described by Marsaglia & Tsang 2000[1],
/// falling back to directly sampling from an Exponential for `shape
/// == 1`, and using the boosting technique described in [1] for
/// `shape < 1`.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{IndependentSample, Gamma};
///
/// let gamma = Gamma::new(2.0, 5.0);
/// let v = gamma.ind_sample(&mut rand::task_rng());
/// println!("{} is from a Gamma(2, 5) distribution", v);
/// ```
///
/// [1]: George Marsaglia and Wai Wan Tsang. 2000. "A Simple Method
/// for Generating Gamma Variables" *ACM Trans. Math. Softw.* 26, 3
/// (September 2000),
/// 363-372. DOI:[10.1145/358407.358414](http://doi.acm.org/10.1145/358407.358414)
pub struct Gamma {
    repr: GammaRepr,
}

enum GammaRepr {
    Large(GammaLargeShape),
    One(Exp),
    Small(GammaSmallShape)
}

// These two helpers could be made public, but saving the
// match-on-Gamma-enum branch from using them directly (e.g. if one
// knows that the shape is always > 1) doesn't appear to be much
// faster.

/// Gamma distribution where the shape parameter is less than 1.
///
/// Note, samples from this require a compulsory floating-point `pow`
/// call, which makes it significantly slower than sampling from a
/// gamma distribution where the shape parameter is greater than or
/// equal to 1.
///
/// See `Gamma` for sampling from a Gamma distribution with general
/// shape parameters.
struct GammaSmallShape {
    inv_shape: f64,
    large_shape: GammaLargeShape
}

/// Gamma distribution where the shape parameter is larger than 1.
///
/// See `Gamma` for sampling from a Gamma distribution with general
/// shape parameters.
struct GammaLargeShape {
    scale: f64,
    c: f64,
    d: f64
}

impl Gamma {
    /// Construct an object representing the `Gamma(shape, scale)`
    /// distribution.
    ///
    /// Fails if `shape <= 0` or `scale <= 0`.
    pub fn new(shape: f64, scale: f64) -> Gamma {
        assert!(shape > 0.0, "Gamma::new called with shape <= 0");
        assert!(scale > 0.0, "Gamma::new called with scale <= 0");

        let repr = match shape {
            1.0        => One(Exp::new(1.0 / scale)),
            0.0 .. 1.0 => Small(GammaSmallShape::new_raw(shape, scale)),
            _          => Large(GammaLargeShape::new_raw(shape, scale))
        };
        Gamma { repr: repr }
    }
}

impl GammaSmallShape {
    fn new_raw(shape: f64, scale: f64) -> GammaSmallShape {
        GammaSmallShape {
            inv_shape: 1. / shape,
            large_shape: GammaLargeShape::new_raw(shape + 1.0, scale)
        }
    }
}

impl GammaLargeShape {
    fn new_raw(shape: f64, scale: f64) -> GammaLargeShape {
        let d = shape - 1. / 3.;
        GammaLargeShape {
            scale: scale,
            c: 1. / (9. * d).sqrt(),
            d: d
        }
    }
}

impl Sample<f64> for Gamma {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl Sample<f64> for GammaSmallShape {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl Sample<f64> for GammaLargeShape {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}

impl IndependentSample<f64> for Gamma {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        match self.repr {
            Small(ref g) => g.ind_sample(rng),
            One(ref g) => g.ind_sample(rng),
            Large(ref g) => g.ind_sample(rng),
        }
    }
}
impl IndependentSample<f64> for GammaSmallShape {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        let Open01(u) = rng.gen::<Open01<f64>>();

        self.large_shape.ind_sample(rng) * u.powf(self.inv_shape)
    }
}
impl IndependentSample<f64> for GammaLargeShape {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        loop {
            let StandardNormal(x) = rng.gen::<StandardNormal>();
            let v_cbrt = 1.0 + self.c * x;
            if v_cbrt <= 0.0 { // a^3 <= 0 iff a <= 0
                continue
            }

            let v = v_cbrt * v_cbrt * v_cbrt;
            let Open01(u) = rng.gen::<Open01<f64>>();

            let x_sqr = x * x;
            if u < 1.0 - 0.0331 * x_sqr * x_sqr ||
                u.ln() < 0.5 * x_sqr + self.d * (1.0 - v + v.ln()) {
                return self.d * v * self.scale
            }
        }
    }
}

/// The chi-squared distribution `χ²(k)`, where `k` is the degrees of
/// freedom.
///
/// For `k > 0` integral, this distribution is the sum of the squares
/// of `k` independent standard normal random variables. For other
/// `k`, this uses the equivalent characterisation `χ²(k) = Gamma(k/2,
/// 2)`.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{ChiSquared, IndependentSample};
///
/// let chi = ChiSquared::new(11.0);
/// let v = chi.ind_sample(&mut rand::task_rng());
/// println!("{} is from a χ²(11) distribution", v)
/// ```
pub struct ChiSquared {
    repr: ChiSquaredRepr,
}

enum ChiSquaredRepr {
    // k == 1, Gamma(alpha, ..) is particularly slow for alpha < 1,
    // e.g. when alpha = 1/2 as it would be for this case, so special-
    // casing and using the definition of N(0,1)^2 is faster.
    DoFExactlyOne,
    DoFAnythingElse(Gamma),
}

impl ChiSquared {
    /// Create a new chi-squared distribution with degrees-of-freedom
    /// `k`. Fails if `k < 0`.
    pub fn new(k: f64) -> ChiSquared {
        let repr = if k == 1.0 {
            DoFExactlyOne
        } else {
            assert!(k > 0.0, "ChiSquared::new called with `k` < 0");
            DoFAnythingElse(Gamma::new(0.5 * k, 2.0))
        };
        ChiSquared { repr: repr }
    }
}
impl Sample<f64> for ChiSquared {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for ChiSquared {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        match self.repr {
            DoFExactlyOne => {
                // k == 1 => N(0,1)^2
                let StandardNormal(norm) = rng.gen::<StandardNormal>();
                norm * norm
            }
            DoFAnythingElse(ref g) => g.ind_sample(rng)
        }
    }
}

/// The Fisher F distribution `F(m, n)`.
///
/// This distribution is equivalent to the ratio of two normalised
/// chi-squared distributions, that is, `F(m,n) = (χ²(m)/m) /
/// (χ²(n)/n)`.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{FisherF, IndependentSample};
///
/// let f = FisherF::new(2.0, 32.0);
/// let v = f.ind_sample(&mut rand::task_rng());
/// println!("{} is from an F(2, 32) distribution", v)
/// ```
pub struct FisherF {
    numer: ChiSquared,
    denom: ChiSquared,
    // denom_dof / numer_dof so that this can just be a straight
    // multiplication, rather than a division.
    dof_ratio: f64,
}

impl FisherF {
    /// Create a new `FisherF` distribution, with the given
    /// parameter. Fails if either `m` or `n` are not positive.
    pub fn new(m: f64, n: f64) -> FisherF {
        assert!(m > 0.0, "FisherF::new called with `m < 0`");
        assert!(n > 0.0, "FisherF::new called with `n < 0`");

        FisherF {
            numer: ChiSquared::new(m),
            denom: ChiSquared::new(n),
            dof_ratio: n / m
        }
    }
}
impl Sample<f64> for FisherF {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for FisherF {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        self.numer.ind_sample(rng) / self.denom.ind_sample(rng) * self.dof_ratio
    }
}

/// The Student t distribution, `t(nu)`, where `nu` is the degrees of
/// freedom.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{StudentT, IndependentSample};
///
/// let t = StudentT::new(11.0);
/// let v = t.ind_sample(&mut rand::task_rng());
/// println!("{} is from a t(11) distribution", v)
/// ```
pub struct StudentT {
    chi: ChiSquared,
    dof: f64
}

impl StudentT {
    /// Create a new Student t distribution with `n` degrees of
    /// freedom. Fails if `n <= 0`.
    pub fn new(n: f64) -> StudentT {
        assert!(n > 0.0, "StudentT::new called with `n <= 0`");
        StudentT {
            chi: ChiSquared::new(n),
            dof: n
        }
    }
}
impl Sample<f64> for StudentT {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for StudentT {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        let StandardNormal(norm) = rng.gen::<StandardNormal>();
        norm * (self.dof / self.chi.ind_sample(rng)).sqrt()
    }
}

#[cfg(test)]
mod test {
    use std::prelude::*;

    use distributions::{Sample, IndependentSample};
    use super::{ChiSquared, StudentT, FisherF};

    #[test]
    fn test_chi_squared_one() {
        let mut chi = ChiSquared::new(1.0);
        let mut rng = ::test::rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    fn test_chi_squared_small() {
        let mut chi = ChiSquared::new(0.5);
        let mut rng = ::test::rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    fn test_chi_squared_large() {
        let mut chi = ChiSquared::new(30.0);
        let mut rng = ::test::rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    #[should_fail]
    fn test_chi_squared_invalid_dof() {
        ChiSquared::new(-1.0);
    }

    #[test]
    fn test_f() {
        let mut f = FisherF::new(2.0, 32.0);
        let mut rng = ::test::rng();
        for _ in range(0, 1000) {
            f.sample(&mut rng);
            f.ind_sample(&mut rng);
        }
    }

    #[test]
    fn test_t() {
        let mut t = StudentT::new(11.0);
        let mut rng = ::test::rng();
        for _ in range(0, 1000) {
            t.sample(&mut rng);
            t.ind_sample(&mut rng);
        }
    }
}

#[cfg(test)]
mod bench {
    extern crate test;
    use std::prelude::*;
    use self::test::Bencher;
    use std::mem::size_of;
    use distributions::IndependentSample;
    use super::Gamma;

    #[bench]
    fn bench_gamma_large_shape(b: &mut Bencher) {
        let gamma = Gamma::new(10., 1.0);
        let mut rng = ::test::weak_rng();

        b.iter(|| {
            for _ in range(0, ::RAND_BENCH_N) {
                gamma.ind_sample(&mut rng);
            }
        });
        b.bytes = size_of::<f64>() as u64 * ::RAND_BENCH_N;
    }

    #[bench]
    fn bench_gamma_small_shape(b: &mut Bencher) {
        let gamma = Gamma::new(0.1, 1.0);
        let mut rng = ::test::weak_rng();

        b.iter(|| {
            for _ in range(0, ::RAND_BENCH_N) {
                gamma.ind_sample(&mut rng);
            }
        });
        b.bytes = size_of::<f64>() as u64 * ::RAND_BENCH_N;
    }
}
<|endoftext|>


<commit_before>//! Guards
//!
//! See: [6.6 Guards](http://erlang.org/doc/apps/erts/absform.html#id88356)
use ast;
use ast::literal;
use ast::common;

pub type Tuple = common::Tuple<Guard>;
pub type Cons = common::Cons<Guard>;
pub type Binary = common::Binary<Guard>;
pub type UnaryOp = common::UnaryOp<Guard>;
pub type BinaryOp = common::BinaryOp<Guard>;
pub type Record = common::Record<Guard>;
pub type RecordIndex = common::RecordIndex<Guard>;
pub type Map = common::Map<Guard>;
pub type LocalCall = common::LocalCall<Guard>;
pub type RemoteCall = common::RemoteCall<Guard>;

#[derive(Debug,Clone)]
pub struct OrGuard {
    and_guards: Vec<Guard>,
}
impl OrGuard {
    pub fn new(and_guards: Vec<Guard>) -> Self {
        OrGuard { and_guards: and_guards }
    }
}

#[derive(Debug,Clone)]
pub enum Guard {
    Integer(Box<literal::Integer>),
    Float(Box<literal::Float>),
    String(Box<literal::Str>),
    Char(Box<literal::Char>),
    Atom(Box<literal::Atom>),
    Var(Box<common::Var>),
    Tuple(Box<Tuple>),
    Nil(Box<common::Nil>),
    Cons(Box<Cons>),
    Binary(Box<Binary>),
    UnaryOp(Box<UnaryOp>),
    BinaryOp(Box<BinaryOp>),
    Record(Box<Record>),
    RecordIndex(Box<RecordIndex>),
    LocalCall(Box<LocalCall>),
    RemoteCall(Box<RemoteCall>),
}
impl_from!(Guard::Integer(literal::Integer));
impl_from!(Guard::Float(literal::Float));
impl_from!(Guard::String(literal::Str));
impl_from!(Guard::Char(literal::Char));
impl_from!(Guard::Atom(literal::Atom));
impl_from!(Guard::Var(common::Var));
impl_from!(Guard::Tuple(Tuple));
impl_from!(Guard::Nil(common::Nil));
impl_from!(Guard::Cons(Cons));
impl_from!(Guard::Binary(Binary));
impl_from!(Guard::UnaryOp(UnaryOp));
impl_from!(Guard::BinaryOp(BinaryOp));
impl_from!(Guard::Record(Record));
impl_from!(Guard::RecordIndex(RecordIndex));
impl_from!(Guard::LocalCall(LocalCall));
impl_from!(Guard::RemoteCall(RemoteCall));
impl ast::Node for Guard {
    fn line(&self) -> ast::LineNum {
        match *self {
            Guard::Integer(ref x) => x.line(),
            Guard::Float(ref x) => x.line(),
            Guard::String(ref x) => x.line(),
            Guard::Char(ref x) => x.line(),
            Guard::Atom(ref x) => x.line(),
            Guard::Var(ref x) => x.line(),
            Guard::Tuple(ref x) => x.line(),
            Guard::Nil(ref x) => x.line(),
            Guard::Cons(ref x) => x.line(),
            Guard::Binary(ref x) => x.line(),
            Guard::UnaryOp(ref x) => x.line(),
            Guard::BinaryOp(ref x) => x.line(),
            Guard::Record(ref x) => x.line(),
            Guard::RecordIndex(ref x) => x.line(),
            Guard::LocalCall(ref x) => x.line(),
            Guard::RemoteCall(ref x) => x.line(),
        }
    }
}
<commit_msg>Export and_guards field<commit_after>//! Guards
//!
//! See: [6.6 Guards](http://erlang.org/doc/apps/erts/absform.html#id88356)
use ast;
use ast::literal;
use ast::common;

pub type Tuple = common::Tuple<Guard>;
pub type Cons = common::Cons<Guard>;
pub type Binary = common::Binary<Guard>;
pub type UnaryOp = common::UnaryOp<Guard>;
pub type BinaryOp = common::BinaryOp<Guard>;
pub type Record = common::Record<Guard>;
pub type RecordIndex = common::RecordIndex<Guard>;
pub type Map = common::Map<Guard>;
pub type LocalCall = common::LocalCall<Guard>;
pub type RemoteCall = common::RemoteCall<Guard>;

#[derive(Debug,Clone)]
pub struct OrGuard {
    pub and_guards: Vec<Guard>,
}
impl OrGuard {
    pub fn new(and_guards: Vec<Guard>) -> Self {
        OrGuard { and_guards: and_guards }
    }
}

#[derive(Debug,Clone)]
pub enum Guard {
    Integer(Box<literal::Integer>),
    Float(Box<literal::Float>),
    String(Box<literal::Str>),
    Char(Box<literal::Char>),
    Atom(Box<literal::Atom>),
    Var(Box<common::Var>),
    Tuple(Box<Tuple>),
    Nil(Box<common::Nil>),
    Cons(Box<Cons>),
    Binary(Box<Binary>),
    UnaryOp(Box<UnaryOp>),
    BinaryOp(Box<BinaryOp>),
    Record(Box<Record>),
    RecordIndex(Box<RecordIndex>),
    LocalCall(Box<LocalCall>),
    RemoteCall(Box<RemoteCall>),
}
impl_from!(Guard::Integer(literal::Integer));
impl_from!(Guard::Float(literal::Float));
impl_from!(Guard::String(literal::Str));
impl_from!(Guard::Char(literal::Char));
impl_from!(Guard::Atom(literal::Atom));
impl_from!(Guard::Var(common::Var));
impl_from!(Guard::Tuple(Tuple));
impl_from!(Guard::Nil(common::Nil));
impl_from!(Guard::Cons(Cons));
impl_from!(Guard::Binary(Binary));
impl_from!(Guard::UnaryOp(UnaryOp));
impl_from!(Guard::BinaryOp(BinaryOp));
impl_from!(Guard::Record(Record));
impl_from!(Guard::RecordIndex(RecordIndex));
impl_from!(Guard::LocalCall(LocalCall));
impl_from!(Guard::RemoteCall(RemoteCall));
impl ast::Node for Guard {
    fn line(&self) -> ast::LineNum {
        match *self {
            Guard::Integer(ref x) => x.line(),
            Guard::Float(ref x) => x.line(),
            Guard::String(ref x) => x.line(),
            Guard::Char(ref x) => x.line(),
            Guard::Atom(ref x) => x.line(),
            Guard::Var(ref x) => x.line(),
            Guard::Tuple(ref x) => x.line(),
            Guard::Nil(ref x) => x.line(),
            Guard::Cons(ref x) => x.line(),
            Guard::Binary(ref x) => x.line(),
            Guard::UnaryOp(ref x) => x.line(),
            Guard::BinaryOp(ref x) => x.line(),
            Guard::Record(ref x) => x.line(),
            Guard::RecordIndex(ref x) => x.line(),
            Guard::LocalCall(ref x) => x.line(),
            Guard::RemoteCall(ref x) => x.line(),
        }
    }
}
<|endoftext|>


<commit_before>use rustc::lint::Context;
use rustc::middle::const_eval::lookup_const_by_id;
use syntax::ast::*;
use syntax::ptr::P;

pub enum FloatWidth {
    Fw32,
    Fw64,
    FwAny
}

impl From<FloatTy> for FloatWidth {
    fn from(ty: FloatTy) -> FloatWidth {
        match ty {
            TyF32 => Fw32,
            TyF64 => Fw64,
        }
    }
}

#[derive(PartialEq, Eq, Debug, Clone)]
pub struct Constant {
    constant: ConstantVariant,
    needed_resolution: bool
}

impl Constant {
    fn new(variant: ConstantVariant) -> Constant {
        Constant { constant: variant, needed_resolution: false }
    }

    fn new_resolved(variant: ConstantVariant) -> Constant {
        Constant { constant: variant, needed_resolution: true }
    }
}

/// a Lit_-like enum to fold constant `Expr`s into
#[derive(PartialEq, Eq, Debug, Clone)]
pub enum ConstantVariant {
    /// a String "abc"
    ConstantStr(&'static str, StrStyle),
    /// a Binary String b"abc"
    ConstantBinary(Rc<Vec<u8>>),
    /// a single byte b'a'
    ConstantByte(u8),
    /// a single char 'a'
    ConstantChar(char),
    /// an integer
    ConstantInt(u64, LitIntType),
    /// a float with given type
    ConstantFloat(Cow<'static, str>, FloatWidth),
    /// true or false
    ConstantBool(bool),
    /// an array of constants
    ConstantVec(Vec<Constant>),
    /// also an array, but with only one constant, repeated N times
    ConstantRepeat(Constant, usize),
    /// a tuple of constants
    ConstantTuple(Vec<Constant>),
}

impl ConstantVariant {
    /// convert to u64 if possible
    ///
    /// # panics
    ///
    /// if the constant could not be converted to u64 losslessly
    fn as_u64(&self) -> u64 {
        if let &ConstantInt(val, _) = self {
            val // TODO we may want to check the sign if any
        } else {
            panic!("Could not convert a {:?} to u64");
        }
    }
}

/// simple constant folding: Insert an expression, get a constant or none.
pub fn constant(cx: &Context, e: &Expr) -> Option<Constant> {
    match e {
        &ExprParen(ref inner) => constant(cx, inner),
        &ExprPath(_, _) => fetch_path(cx, e),
        &ExprBlock(ref block) => constant_block(cx, inner),
        &ExprIf(ref cond, ref then, ref otherwise) =>
            constant_if(cx, cond, then, otherwise),
        &ExprLit(ref lit) => Some(lit_to_constant(lit)),
        &ExprVec(ref vec) => constant_vec(cx, vec),
        &ExprTup(ref tup) => constant_tup(cx, tup),
        &ExprRepeat(ref value, ref number) =>
            constant_binop_apply(cx, value, number,|v, n| Constant {
                constant: ConstantRepeat(v, n.constant.as_u64()),
                needed_resolution: v.needed_resolution || n.needed_resolution
            }),
        &ExprUnary(op, ref operand) => constant(cx, operand).and_then(
            |o| match op {
                UnNot =>
                    if let ConstantBool(b) = o.variant {
                        Some(Constant{
                            needed_resolution: o.needed_resolution,
                            constant: ConstantBool(!b),
                        })
                    } else { None },
                UnNeg => constant_negate(o),
                UnUniq | UnDeref => o,
            }),
        &ExprBinary(op, ref left, ref right) =>
            constant_binop(op, left, right),
        //TODO: add other expressions
        _ => None,
    }
}

fn lit_to_constant(lit: &Lit_) -> Constant {
    match lit {
        &LitStr(ref is, style) => Constant::new(ConstantStr(&*is, style)),
        &LitBinary(ref blob) => Constant::new(ConstantBinary(blob.clone())),
        &LitByte(b) => Constant::new(ConstantByte(b)),
        &LitChar(c) => Constant::new(ConstantChar(c)),
        &LitInt(value, ty) => Constant::new(ConstantInt(value, ty)),
        &LitFloat(ref is, ty) =>
            Constant::new(ConstantFloat(Cow::Borrowed(&*is), ty.into())),
        &LitFloatUnsuffixed(InternedString) =>
            Constant::new(ConstantFloat(Cow::Borrowed(&*is), FwAny)),
        &LitBool(b) => Constant::new(ConstantBool(b)),
    }
}

/// create `Some(ConstantVec(..))` of all constants, unless there is any
/// non-constant part
fn constant_vec(cx: &Context, vec: &[&Expr]) -> Option<Constant> {
    let mut parts = Vec::new();
    let mut resolved = false;
    for opt_part in vec {
        match constant(cx, opt_part) {
            Some(ref p) => {
                resolved |= p.needed_resolution;
                parts.push(p)
            },
            None => { return None; },
        }
    }
    Some(Constant {
        constant: ConstantVec(parts),
        needed_resolution: resolved
    })
}

fn constant_tup(cx: &Context, tup: &[&Expr]) -> Option<Constant> {
    let mut parts = Vec::new();
    let mut resolved = false;
    for opt_part in vec {
        match constant(cx, opt_part) {
            Some(ref p) => {
                resolved |= p.needed_resolution;
                parts.push(p)
            },
            None => { return None; },
        }
    }
    Some(Constant {
        constant: ConstantTuple(parts),
        needed_resolution: resolved
    })
}

/// lookup a possibly constant expression from a ExprPath
fn fetch_path(cx: &Context, e: &Expr) -> Option<Constant> {
    if let Some(&PathResolution { base_def: DefConst(id), ..}) =
            cx.tcx.def_map.borrow().get(&e.id) {
        lookup_const_by_id(cx.tcx, id, None).map(
            |l| Constant::new_resolved(constant(cx, l).constant))
    } else { None }
}

/// A block can only yield a constant if it only has one constant expression
fn constant_block(cx: &Context, block: &Block) -> Option<Constant> {
    if block.stmts.is_empty() {
        block.expr.map(|b| constant(cx, b))
    } else { None }
}

fn constant_if(cx: &Context, cond: &Expr, then: &Expr, otherwise: &Expr) ->
        Option<Constant> {
    if let Some(Constant{ constant: ConstantBool(b), needed_resolution: res }) =
            constant(cx, cond) {
        let part = constant(cx, if b { then } else { otherwise });
        Some(Constant {
            constant: part.constant,
            needed_resolution: res || part.needed_resolution,
        })
    } else { None }
}

fn constant_negate(o: Constant) -> Option<Constant> {
    Some(Constant{
        needed_resolution: o.needed_resolution,
        constant: match o.constant {
            &ConstantInt(value, ty) =>
                ConstantInt(value, match ty {
                    SignedIntLit(ity, sign) =>
                        SignedIntLit(ity, neg_sign(sign)),
                    UnsuffixedIntLit(sign) => UnsuffixedIntLit(neg_sign(sign)),
                    _ => { return None; },
                }),
            &LitFloat(ref is, ref ty) => ConstantFloat(neg_float_str(is), ty),
            _ => { return None; },
        }
    })
}

fn neg_sign(s: Sign) -> Sign {
    match s {
        Sign::Plus => Sign::Minus,
        Sign::Minus => Sign::Plus,
    }
}

fn neg_float_str(s: &InternedString) -> Cow<'static, str> {
    if s.startsWith('-') {
        Cow::Borrowed(s[1..])
    } else {
        Cow::Owned(format!("-{}", &*s))
    }
}

fn constant_binop(cx: &Context, op: BinOp, left: &Expr, right: &Expr)
        -> Option<Constant> {
    match op.node {
        //BiAdd,
        //BiSub,
        //BiMul,
        //BiDiv,
        //BiRem,
        BiAnd => constant_short_circuit(cx, left, right, false),
        BiOr => constant_short_circuit(cx, left, right, true),
        //BiBitXor,
        //BiBitAnd,
        //BiBitOr,
        //BiShl,
        //BiShr,
        //BiEq,
        //BiLt,
        //BiLe,
        //BiNe,
        //BiGe,
        //BiGt,
        _ => None,
    }
}

fn constant_binop_apply<F>(cx: &Context, left: &Expr, right: &Expr, op: F)
        -> Option<Constant>
where F: FnMut(ConstantVariant, ConstantVariant) -> Option<ConstantVariant> {
    constant(cx, left).and_then(|l| constant(cx, right).and_then(
        |r| Constant {
            needed_resolution: l.needed_resolution || r.needed_resolution,
            constant: op(l.constant, r.constant)
        }))
}

fn constant_short_circuit(cx: &Context, left: &Expr, right: &Expr, b: bool) ->
        Option<Constant> {
    let leftconst = constant(cx, left);
    if let ConstantBool(lbool) = leftconst.constant {
        if l == b {
            Some(leftconst)
        } else {
            let rightconst = constant(cx, right);
            if let ConstantBool(rbool) = rightconst.constant {
                Some(Constant {
                    constant: rightconst.constant,
                    needed_resolution: leftconst.needed_resolution ||
                                       rightconst.needed_resolution,
                })
            } else { None }
        }
    } else { None }
}
<commit_msg>initial addition and subtraction for bytes and ints<commit_after>use rustc::lint::Context;
use rustc::middle::const_eval::lookup_const_by_id;
use syntax::ast::*;
use syntax::ptr::P;

pub enum FloatWidth {
    Fw32,
    Fw64,
    FwAny
}

impl From<FloatTy> for FloatWidth {
    fn from(ty: FloatTy) -> FloatWidth {
        match ty {
            TyF32 => Fw32,
            TyF64 => Fw64,
        }
    }
}

#[derive(PartialEq, Eq, Debug, Clone)]
pub struct Constant {
    constant: ConstantVariant,
    needed_resolution: bool
}

impl Constant {
    fn new(variant: ConstantVariant) -> Constant {
        Constant { constant: variant, needed_resolution: false }
    }

    fn new_resolved(variant: ConstantVariant) -> Constant {
        Constant { constant: variant, needed_resolution: true }
    }
}

/// a Lit_-like enum to fold constant `Expr`s into
#[derive(PartialEq, Eq, Debug, Clone)]
pub enum ConstantVariant {
    /// a String "abc"
    ConstantStr(&'static str, StrStyle),
    /// a Binary String b"abc"
    ConstantBinary(Rc<Vec<u8>>),
    /// a single byte b'a'
    ConstantByte(u8),
    /// a single char 'a'
    ConstantChar(char),
    /// an integer
    ConstantInt(u64, LitIntType),
    /// a float with given type
    ConstantFloat(Cow<'static, str>, FloatWidth),
    /// true or false
    ConstantBool(bool),
    /// an array of constants
    ConstantVec(Vec<Constant>),
    /// also an array, but with only one constant, repeated N times
    ConstantRepeat(Constant, usize),
    /// a tuple of constants
    ConstantTuple(Vec<Constant>),
}

impl ConstantVariant {
    /// convert to u64 if possible
    ///
    /// # panics
    ///
    /// if the constant could not be converted to u64 losslessly
    fn as_u64(&self) -> u64 {
        if let &ConstantInt(val, _) = self {
            val // TODO we may want to check the sign if any
        } else {
            panic!("Could not convert a {:?} to u64");
        }
    }
}

/// simple constant folding: Insert an expression, get a constant or none.
pub fn constant(cx: &Context, e: &Expr) -> Option<Constant> {
    match e {
        &ExprParen(ref inner) => constant(cx, inner),
        &ExprPath(_, _) => fetch_path(cx, e),
        &ExprBlock(ref block) => constant_block(cx, inner),
        &ExprIf(ref cond, ref then, ref otherwise) =>
            constant_if(cx, cond, then, otherwise),
        &ExprLit(ref lit) => Some(lit_to_constant(lit)),
        &ExprVec(ref vec) => constant_vec(cx, vec),
        &ExprTup(ref tup) => constant_tup(cx, tup),
        &ExprRepeat(ref value, ref number) =>
            constant_binop_apply(cx, value, number,|v, n| Constant {
                constant: ConstantRepeat(v, n.constant.as_u64()),
                needed_resolution: v.needed_resolution || n.needed_resolution
            }),
        &ExprUnary(op, ref operand) => constant(cx, operand).and_then(
            |o| match op {
                UnNot =>
                    if let ConstantBool(b) = o.variant {
                        Some(Constant{
                            needed_resolution: o.needed_resolution,
                            constant: ConstantBool(!b),
                        })
                    } else { None },
                UnNeg => constant_negate(o),
                UnUniq | UnDeref => o,
            }),
        &ExprBinary(op, ref left, ref right) =>
            constant_binop(op, left, right),
        //TODO: add other expressions
        _ => None,
    }
}

fn lit_to_constant(lit: &Lit_) -> Constant {
    match lit {
        &LitStr(ref is, style) => Constant::new(ConstantStr(&*is, style)),
        &LitBinary(ref blob) => Constant::new(ConstantBinary(blob.clone())),
        &LitByte(b) => Constant::new(ConstantByte(b)),
        &LitChar(c) => Constant::new(ConstantChar(c)),
        &LitInt(value, ty) => Constant::new(ConstantInt(value, ty)),
        &LitFloat(ref is, ty) =>
            Constant::new(ConstantFloat(Cow::Borrowed(&*is), ty.into())),
        &LitFloatUnsuffixed(InternedString) =>
            Constant::new(ConstantFloat(Cow::Borrowed(&*is), FwAny)),
        &LitBool(b) => Constant::new(ConstantBool(b)),
    }
}

/// create `Some(ConstantVec(..))` of all constants, unless there is any
/// non-constant part
fn constant_vec(cx: &Context, vec: &[&Expr]) -> Option<Constant> {
    let mut parts = Vec::new();
    let mut resolved = false;
    for opt_part in vec {
        match constant(cx, opt_part) {
            Some(ref p) => {
                resolved |= p.needed_resolution;
                parts.push(p)
            },
            None => { return None; },
        }
    }
    Some(Constant {
        constant: ConstantVec(parts),
        needed_resolution: resolved
    })
}

fn constant_tup(cx: &Context, tup: &[&Expr]) -> Option<Constant> {
    let mut parts = Vec::new();
    let mut resolved = false;
    for opt_part in vec {
        match constant(cx, opt_part) {
            Some(ref p) => {
                resolved |= p.needed_resolution;
                parts.push(p)
            },
            None => { return None; },
        }
    }
    Some(Constant {
        constant: ConstantTuple(parts),
        needed_resolution: resolved
    })
}

/// lookup a possibly constant expression from a ExprPath
fn fetch_path(cx: &Context, e: &Expr) -> Option<Constant> {
    if let Some(&PathResolution { base_def: DefConst(id), ..}) =
            cx.tcx.def_map.borrow().get(&e.id) {
        lookup_const_by_id(cx.tcx, id, None).map(
            |l| Constant::new_resolved(constant(cx, l).constant))
    } else { None }
}

/// A block can only yield a constant if it only has one constant expression
fn constant_block(cx: &Context, block: &Block) -> Option<Constant> {
    if block.stmts.is_empty() {
        block.expr.map(|b| constant(cx, b))
    } else { None }
}

fn constant_if(cx: &Context, cond: &Expr, then: &Expr, otherwise: &Expr) ->
        Option<Constant> {
    if let Some(Constant{ constant: ConstantBool(b), needed_resolution: res }) =
            constant(cx, cond) {
        let part = constant(cx, if b { then } else { otherwise });
        Some(Constant {
            constant: part.constant,
            needed_resolution: res || part.needed_resolution,
        })
    } else { None }
}

fn constant_negate(o: Constant) -> Option<Constant> {
    Some(Constant{
        needed_resolution: o.needed_resolution,
        constant: match o.constant {
            &ConstantInt(value, ty) =>
                ConstantInt(value, match ty {
                    SignedIntLit(ity, sign) =>
                        SignedIntLit(ity, neg_sign(sign)),
                    UnsuffixedIntLit(sign) => UnsuffixedIntLit(neg_sign(sign)),
                    _ => { return None; },
                }),
            &LitFloat(ref is, ref ty) => ConstantFloat(neg_float_str(is), ty),
            _ => { return None; },
        }
    })
}

fn neg_sign(s: Sign) -> Sign {
    match s {
        Sign::Plus => Sign::Minus,
        Sign::Minus => Sign::Plus,
    }
}

fn neg_float_str(s: &InternedString) -> Cow<'static, str> {
    if s.startsWith('-') {
        Cow::Borrowed(s[1..])
    } else {
        Cow::Owned(format!("-{}", &*s))
    }
}

fn is_negative(ty: LitIntType) -> bool {
    match ty {
        SignedIntLit(_, sign) | UnsuffixedIntLit(sign) => sign == Minus,
        UnsignedIntLit(_) => false,
    }
}

fn unify_int_type(l: LitIntType, r: LitIntType, s: Sign) -> Option(LitIntType) {
    match (l, r) {
        (SignedIntLit(lty, _), SignedIntLit(rty, _)) => if lty == rty {
            Some(SignedIntLit(lty, s)) } else { None },
        (UnsignedIntLit(lty), UnsignedIntLit(rty)) =>
            if Sign == Plus && lty == rty {
                Some(UnsignedIntLit(lty))
            } else { None },
        (UnsuffixedIntLit(_), UnsuffixedIntLit(_)) => UnsuffixedIntLit(s),
        (SignedIntLit(lty, _), UnsuffixedIntLit(_)) => SignedIntLit(lty, s),
        (UnsignedIntLit(lty), UnsuffixedIntLit(rs)) => if rs == Plus {
            Some(UnsignedIntLit(lty)) } else { None },
        (UnsuffixedIntLit(_), SignedIntLit(rty, _)) => SignedIntLit(rty, s),
        (UnsuffixedIntLit(ls), UnsignedIntLit(rty)) => if ls == Plus {
            Some(UnsignedIntLit(rty)) } else { None },
        _ => None,
    }
}

fn constant_binop(cx: &Context, op: BinOp, left: &Expr, right: &Expr)
        -> Option<Constant> {
    match op.node {
        BiAdd => constant_binop_apply(cx, left, right, |l, r|
            match (l, r) {
                (ConstantByte(l8), ConstantByte(r8)) =>
                    l8.checked_add(r8).map(|v| ConstantByte(v)),
                (ConstantInt(l64, lty), ConstantInt(r64, rty)) => {
                    let (ln, rn) = (is_negative(lty), is_negative(rty));
                    if ln == rn {
                        unify_int_type(lty, rty, if ln { Minus } else { Plus })
                            .and_then(|ty| l64.checked_add(r64).map(
                                |v| ConstantInt(v, ty)))
                    } else {
                        if ln {
                            add_neg_int(r64, rty, l64, lty)
                        } else {
                            add_neg_int(l64, lty, r64, rty)
                        }
                    }
                },
                // TODO: float
                _ => None
            }),
        BiSub => constant_binop_apply(cx, left, right, |l, r|
            match (l, r) {
                (ConstantByte(l8), ConstantByte(r8)) => if r8 > l8 {
                    None } else { Some(ConstantByte(l8 - r8)) },
                (ConstantInt(l64, lty), ConstantInt(r64, rty)) => {
                    let (ln, rn) = (is_negative(lty), is_negative(rty));
                    match (ln, rn) {
                        (false, false) => sub_int(l64, lty, r64, rty, r64 > l64),
                        (true, true) => sub_int(l64, lty, r64, rty, l64 > r64),
                        (true, false) => unify_int_type(lty, rty, Minus)
                            .and_then(|ty| l64.checked_add(r64).map(
                                |v| ConstantInt(v, ty))),
                        (false, true) => unify_int_type(lty, rty, Plus)
                            .and_then(|ty| l64.checked_add(r64).map(
                                |v| ConstantInt(v, ty))),
                    }
                },
                _ => None,
            }),
        //BiMul,
        //BiDiv,
        //BiRem,
        BiAnd => constant_short_circuit(cx, left, right, false),
        BiOr => constant_short_circuit(cx, left, right, true),
        //BiBitXor,
        //BiBitAnd,
        //BiBitOr,
        //BiShl,
        //BiShr,
        //BiEq,
        //BiLt,
        //BiLe,
        //BiNe,
        //BiGe,
        //BiGt,
        _ => None,
    }
}

fn add_neg_int(pos: u64, pty: LitIntType, neg: u64, nty: LitIntType) ->
        Some(Constant) {
    if neg > pos {
        unify_int_type(nty, pty, Minus).map(|ty| ConstantInt(neg - pos, ty))
    } else {
        unify_int_type(nty, pty, Plus).map(|ty| ConstantInt(pos - neg, ty))
    }
}

fn sub_int(l: u64, lty: LitIntType, r: u64, rty: LitIntType, neg: Bool) ->
        Option<Constant> {
     unify_int_type(lty, rty, if neg { Minus } else { Plus }).and_then(
        |ty| l64.checked_sub(r64).map(|v| ConstantInt(v, ty)))
}

fn constant_binop_apply<F>(cx: &Context, left: &Expr, right: &Expr, op: F)
        -> Option<Constant>
where F: FnMut(ConstantVariant, ConstantVariant) -> Option<ConstantVariant> {
    constant(cx, left).and_then(|l| constant(cx, right).and_then(
        |r| Constant {
            needed_resolution: l.needed_resolution || r.needed_resolution,
            constant: op(l.constant, r.constant)
        }))
}

fn constant_short_circuit(cx: &Context, left: &Expr, right: &Expr, b: bool) ->
        Option<Constant> {
    let leftconst = constant(cx, left);
    if let ConstantBool(lbool) = leftconst.constant {
        if l == b {
            Some(leftconst)
        } else {
            let rightconst = constant(cx, right);
            if let ConstantBool(rbool) = rightconst.constant {
                Some(Constant {
                    constant: rightconst.constant,
                    needed_resolution: leftconst.needed_resolution ||
                                       rightconst.needed_resolution,
                })
            } else { None }
        }
    } else { None }
}
<|endoftext|>