<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 mut 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_2",
                    Some("lichess_merge"),
                    lichess_merge,
                    Some(KeyPrefix::SIZE),
                    8 * 1024,
                    3,
                ),
                column_family(
                    "lichess_game_2",
                    Some("lichess_game_merge"),
                    lichess_game_merge,
                    None,
                    4 * 1024,
                    0,
                ),
                // Player database (also shares lichess_game)
                column_family(
                    "player_2",
                    Some("player_merge"),
                    player_merge,
                    Some(KeyPrefix::SIZE),
                    8 * 1024,
                    3,
                ),
                column_family("player_status_2", 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_2").expect("cf lichess"),
            cf_lichess_game: self
                .inner
                .cf_handle("lichess_game_2")
                .expect("cf lichess_game_2"),

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

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>further reduce bloom filter size for large cfs<commit_after>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 mut 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_2",
                    Some("lichess_merge"),
                    lichess_merge,
                    Some(KeyPrefix::SIZE),
                    16 * 1024,
                    2,
                ),
                column_family(
                    "lichess_game_2",
                    Some("lichess_game_merge"),
                    lichess_game_merge,
                    None,
                    4 * 1024,
                    0,
                ),
                // Player database (also shares lichess_game)
                column_family(
                    "player_2",
                    Some("player_merge"),
                    player_merge,
                    Some(KeyPrefix::SIZE),
                    16 * 1024,
                    2,
                ),
                column_family("player_status_2", 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_2").expect("cf lichess"),
            cf_lichess_game: self
                .inner
                .cf_handle("lichess_game_2")
                .expect("cf lichess_game_2"),

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

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>pub type ID = [u16;12];

pub const NullID: ID = [0,0,0,0,0,0,0,0,0,0,0,0];

mod experiment{
    use std::iter::Iterator;
    #[derive(Copy, Clone, Show)]
    pub struct ID{
        data: [u8;15],
        end: u8
    }

    impl ID{
        pub fn new() -> ID{
            ID{
                data: [0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
                end: 0
            }
        }
        ///push last component id back
        #[inline(always)]
        pub fn push(&mut self, i: u16){
            if self.end > 14{
                panic!("to many ids");
            }
            if i >= 32768{
                panic!("value given is to big to be stored, limit the number of components used");
            }
            if i < 128{
                self.data[self.end as usize] = i as u8;
                self.end += 1;
            }else{
                self.data[self.end as usize] = (i >> 1) as u8;
                self.data[(self.end +1) as usize] = (i >> 6) as u8;
                self.end += 2;
            }
        }

        ///get an iterator over all methods
        #[inline(always)]
        pub fn iter<'a>(&'a self) -> ID_Iterator<'a>{
            ID_Iterator{
                id: self,
                pos: 0
            }
        }
    }

    pub struct ID_Iterator<'a>{
        id: &'a ID,
        pos: u8,
    }
    impl<'a> Iterator for ID_Iterator<'a>{
        type Item = u16;
        fn next(&mut self) -> Option<u16>{
            None
        }
    }

    /*#[test]
    fn test_ID(){
        let mut id = ID::new();
        id.push(127);
        id.push(12);
        id.push(1023);
        assert!(id.end == 4, "end has wrong size");
        assert!(id.data[0] == 0b01111111, "wrong number {}", 0b01111111);
        assert!(id.data[1] == 12);
        assert!(id.data[2] == 0b11111111, "wront it is {}", id.data[2]);
        assert!(id.data[3] == 0b11100000, "wront it is {}", id.data[3]);
    }*/
}
<commit_msg>add comments<commit_after>/// The id-type used internaly to identify each widget.
/// it is an array consistiong of 12 ids of the elements an their children
/// if the id == 0, the element has no id assigned. Thats why ids should start
/// from 1
pub type ID = [u16;12];

///The null id
pub const NullID: ID = [0,0,0,0,0,0,0,0,0,0,0,0];

mod experiment{
    use std::iter::Iterator;
    #[derive(Copy, Clone, Show)]
    pub struct ID{
        data: [u8;15],
        end: u8
    }

    impl ID{
        pub fn new() -> ID{
            ID{
                data: [0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
                end: 0
            }
        }
        ///push last component id back
        #[inline(always)]
        pub fn push(&mut self, i: u16){
            if self.end > 14{
                panic!("to many ids");
            }
            if i >= 32768{
                panic!("value given is to big to be stored, limit the number of components used");
            }
            if i < 128{
                self.data[self.end as usize] = i as u8;
                self.end += 1;
            }else{
                self.data[self.end as usize] = (i >> 1) as u8;
                self.data[(self.end +1) as usize] = (i >> 6) as u8;
                self.end += 2;
            }
        }

        ///get an iterator over all methods
        #[inline(always)]
        pub fn iter<'a>(&'a self) -> ID_Iterator<'a>{
            ID_Iterator{
                id: self,
                pos: 0
            }
        }
    }

    pub struct ID_Iterator<'a>{
        id: &'a ID,
        pos: u8,
    }
    impl<'a> Iterator for ID_Iterator<'a>{
        type Item = u16;
        fn next(&mut self) -> Option<u16>{
            None
        }
    }

    /*#[test]
    fn test_ID(){
        let mut id = ID::new();
        id.push(127);
        id.push(12);
        id.push(1023);
        assert!(id.end == 4, "end has wrong size");
        assert!(id.data[0] == 0b01111111, "wrong number {}", 0b01111111);
        assert!(id.data[1] == 12);
        assert!(id.data[2] == 0b11111111, "wront it is {}", id.data[2]);
        assert!(id.data[3] == 0b11100000, "wront it is {}", id.data[3]);
    }*/
}
<|endoftext|>


<commit_before>//! Defining things related to `ag::op::Op`.
//!
use crate::ndarray_ext::{NdArrayView, NdArrayViewMut};
use crate::smallvec::SmallVec;
use crate::tensor::{Tensor, TensorInternal};
use crate::{Float, NdArray};
use std::any::type_name;
use std::fmt;
use std::marker::PhantomData;
use std::mem;

// Properties for op's `compute` method.
// Actual number of inout/output nodes are around 1~2 in most cases.
pub(crate) const NUM_MAX_OUTPUT: usize = 2;
pub(crate) const NUM_MAX_INPUT: usize = 4;

pub(crate) type InputArray<T> = SmallVec<[T; NUM_MAX_INPUT]>;
pub(crate) type OutputArray<T> = SmallVec<[T; NUM_MAX_OUTPUT]>;

pub(crate) type ComputeResult<'v, T> = Result<crate::ArrRepr<'v, T>, OpError>;

/// Error in `Op`'s computation.
#[derive(Clone, Debug, PartialEq)]
pub enum OpError {
    NdArrayError(String, ndarray::ShapeError),
    IncompatibleShape(String),
    TypeUnsupported(String),
    InvalidDims(String),
    OutOfBounds(String),
}

impl std::error::Error for OpError {}

impl fmt::Display for OpError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            OpError::NdArrayError(pref, e) => write!(f, "{}: ", pref).and_then(|()| e.fmt(f)),
            OpError::IncompatibleShape(s) => write!(f, "{}: ", s),
            OpError::TypeUnsupported(s) => write!(f, "{}: ", s),
            OpError::InvalidDims(s) => write!(f, "{}: ", s),
            OpError::OutOfBounds(s) => write!(f, "{}: ", s),
        }
    }
}

/// Sequence of an op's compute result.
///
/// Op can have multiple output NdArrays.
pub(crate) type Results<'v, T> = OutputArray<Option<ComputeResult<'v, T>>>;

/// Operation trait. `Tensor` wraps trait-object of this.
///
/// # Implementing differentiable operations
///
/// Many of well-known ops are pre-defined in `Graph`, but you can also
/// implement custom ops by hand.
///
/// ```
/// use ndarray;
/// use autograd as ag;
/// use ag::Graph;
///
/// type NdArray<T: ag::Float> = ndarray::Array<T, ndarray::IxDyn>;
///
/// // Implements `Op` trait for `Sigmoid`.
/// struct Sigmoid;
///
/// impl<T: ag::Float> ag::op::Op<T> for Sigmoid {
///     fn compute(
///         &self,
///         ctx: &mut ag::op::ComputeContext<T>,
///     ) {
///         let x: &ag::NdArrayView<_> = &ctx.input(0);
///         // Use `ndarray::Array::mapv` for element-wise computation.
///         let half = T::from(0.5).unwrap();
///         let y = x.mapv(move |a| ((a * half).tanh() * half) + half);
///         ctx.append_output(y);
///     }
///
///     fn grad(&self, ctx: &mut ag::op::GradientContext<T>) {
///         // Symbolic gradient of the input of Sigmoid
///         let gy = ctx.output_grad();
///         let y = ctx.output();
///         let gx = gy * (y - ctx.graph().square(y));
///         ctx.append_input_grad(Some(gx));
///     }
/// }
///
///  use ag::tensor::Input;
///
/// // Symbolic `sigmoid` function for end-user.
/// fn sigmoid<'graph, F: ag::Float>(x: &ag::Tensor<'graph, F>, g: &'graph ag::Graph<F>)
/// -> ag::Tensor<'graph, F> {
///     ag::Tensor::builder()
///            .set_inputs(&[Input::new(x)])
///            .build(g, Sigmoid)
/// }
/// ```
pub trait Op<F: Float> {
    /// Name of this op
    fn name(&self) -> &str {
        type_name::<Self>()
    }

    /// Runs this op with `ComputeContext`.
    fn compute(&self, ctx: &mut ComputeContext<F>);

    /// Returns symbolic gradients for input nodes by use of output's gradients etc.
    fn grad(&self, ctx: &mut GradientContext<F>);
}

pub(crate) struct DummyOp<F: Float> {
    pub phantom: PhantomData<F>,
}

impl<F: Float> DummyOp<F> {
    #[allow(dead_code)]
    pub(crate) fn new() -> Self {
        DummyOp {
            phantom: PhantomData,
        }
    }
}

impl<F: Float> Op<F> for DummyOp<F> {
    fn compute(&self, _: &mut ComputeContext<F>) {}
    fn grad(&self, _: &mut GradientContext<F>) {}
}

/// Wrapper object of NdArrayView/NdArrayViewMut which is fed to `Op::compute`
///
/// Used in `OpComputeContext`.
pub(crate) enum OpInput<'v, T: Float> {
    /// Read-only view
    RO(Option<NdArrayView<'v, T>>),
    /// Read-write view
    RW(Option<NdArrayViewMut<'v, T>>),
}

impl<'v, T: Float> OpInput<'v, T> {
    #[inline]
    /// Make a read-only input array
    pub fn new(x: NdArrayView<'v, T>) -> Self {
        OpInput::RO(Some(x))
    }

    #[inline]
    /// Make a read/write input array
    pub fn new_mut(x: NdArrayViewMut<'v, T>) -> Self {
        OpInput::RW(Some(x))
    }
}

/// Context of an `Op`'s computation phase.
///
/// # Example
///
/// ```
/// use autograd as ag;
///
/// // Implementing `Op` trait for `Sigmoid`.
/// struct Sigmoid;
///
/// impl<T: ag::Float> ag::op::Op<T> for Sigmoid {
///     fn compute(
///         &self,
///         ctx: &mut ag::op::ComputeContext<T>,
///     ) {
///         // Getting the first input array.
///         let x: &ag::NdArrayView<_> = &ctx.input(0);
///         let half = T::from(0.5).unwrap();
///         let y = x.mapv(move |a| ((a * half).tanh() * half) + half);
///         // Put the computed result.
///         ctx.append_output(y);
///     }
///
///     fn grad(&self, ctx: &mut ag::op::GradientContext<T>) { /* ... */ }
/// }
/// ```
pub struct ComputeContext<'t, 'v, T: Float> {
    node: &'t TensorInternal<T>,
    // Input arrays
    xs: InputArray<OpInput<'v, T>>,
    // Output arrays
    ys: Results<'v, T>,
}

impl<'g, 't, 'v, T: Float> ComputeContext<'t, 'v, T> {
    pub(crate) fn extract_outputs(self) -> Results<'v, T> {
        self.ys
    }

    #[inline]
    pub(crate) fn new(node: &'t TensorInternal<T>, xs: InputArray<OpInput<'v, T>>) -> Self {
        ComputeContext {
            node,
            xs,
            ys: OutputArray::new(),
        }
    }

    /// Grabs the `i` th input array as a *read-only* array view.
    ///
    /// Calling `input(i)` more than once causes panic.
    #[inline]
    pub fn input(&mut self, i: usize) -> NdArrayView<'v, T> {
        let x = match self.xs.get_mut(i) {
            Some(x) => x,
            None => panic!("Bad op impl: input index out of range."),
        };
        match x {
            OpInput::RO(ref mut a) => match a.take() {
                Some(ret) => ret,
                None => panic!(
                    "Bad op impl of {}: input({})/input_mut({}) cannot be called twice",
                    self.node.get_op().name(),
                    i,
                    i
                ),
            },
            _ => {
                panic!(
                    "Bad op impl of {}: cannot perform immutable borrowing for input({})",
                    self.node.get_op().name(),
                    i
                );
            }
        }
    }

    /// Grabs the `i` th input array as a *read-write* array view.
    ///
    /// Calling `input_mut(i)` more than once causes panic.
    #[inline]
    pub fn input_mut(&mut self, i: usize) -> NdArrayViewMut<'v, T> {
        let x = match self.xs.get_mut(i) {
            Some(x) => x,
            None => panic!("Bad op impl: {}'s input doesn't exist.", i),
        };
        match x {
            OpInput::RW(ref mut a) => match a.take() {
                Some(ret) => ret,
                None => panic!(
                    "Bad op impl of {}: input({})/input_mut({}) cannot be called twice",
                    self.node.get_op().name(),
                    i,
                    i
                ),
            },
            _ => {
                panic!(
                    "Bad op impl of {}: cannot perform mutable borrowing for input({})",
                    self.node.get_op().name(),
                    i
                );
            }
        }
    }

    /// Appends an `ndarray::ArrayView` to the back of the output list of the current op.
    ///
    /// NOTE: Implementor of `Op::compute` must not forget to call `append_*` as many as the number of its output in `Op::compute`, otherwise panic occurs.
    #[inline]
    pub fn append_output_view(&mut self, y: NdArrayView<'v, T>) {
        self.ys.push(Some(Ok(crate::ArrRepr::View(y))));
    }

    /// Appends an ndarray to the back of the output list of the current op.
    ///
    /// NOTE: Implementor of `Op::compute` must not forget to call `append_*` as many as the number of its output in `Op::compute`, otherwise panic occurs.
    #[inline]
    pub fn append_output(&mut self, y: NdArray<T>) {
        self.ys.push(Some(Ok(crate::ArrRepr::Owned(y))));
    }

    /// Appends an empty result to the back of the output list of the current op.
    ///
    /// For example, this is used for gradient-descent-optimizers.
    ///
    /// NOTE: Implementor of `Op::compute` must not forget to call `append_*` as many as the number of its output in `Op::compute`, otherwise panic occurs.
    #[inline]
    pub fn append_empty_output(&mut self) {
        self.ys.push(None);
    }

    /// Marks this op's compute session as failed.
    ///
    /// NOTE: After called this function, Op::compute should return early.
    #[inline]
    pub fn set_error(&mut self, y: OpError) {
        self.ys.push(Some(Err(y)));
    }

    /// Returns a number of input arrays.
    #[inline]
    pub fn num_inputs(&self) -> usize {
        self.xs.len()
    }
}

/// Context of an `Op`'s gradient propagation phase.
///
/// This is passed to an `Op` through `Op::grad`.
/// `Op::grad` should provide the gradients of its inputs by calling `GradientContext::set_input_grads`.
///
/// Use `graph()` to access `Graph` object for tensor computations.
///
/// ```
/// use autograd as ag;
///
/// struct Sigmoid;
///
/// impl<T: ag::Float> ag::op::Op<T> for Sigmoid {
///     fn compute(&self, ctx: &mut ag::op::ComputeContext<T>) { /* ... */ }
///
///     fn grad(&self, ctx: &mut ag::op::GradientContext<T>) {
///         // Symbolic gradient of the input of Sigmoid
///         let gy = ctx.output_grad();
///         // Symbolic output tensor
///         let y = ctx.output();
///         // `Tensor` computations
///         let gx = gy * (y - ctx.graph().square(y));
///         // Propagates input's gradient.
///         ctx.append_input_grad(Some(gx));
///     }
/// }
/// ```
pub struct GradientContext<'g, T: Float> {
    gy: Tensor<'g, T>,
    y: Tensor<'g, T>,
    graph: &'g crate::graph::Graph<T>,
    gxs: InputArray<Option<Tensor<'g, T>>>,
}

impl<'g, T: Float> GradientContext<'g, T> {
    pub(crate) fn new(
        gy: Tensor<'g, T>,
        y: Tensor<'g, T>,
        graph: &'g crate::graph::Graph<T>,
    ) -> Self {
        GradientContext {
            gy,
            y,
            graph,
            gxs: InputArray::new(),
        }
    }

    // Call Op::grad and return `gxs`
    pub(crate) fn extract_input_grads(mut self) -> InputArray<Option<Tensor<'g, T>>> {
        let id = self.y.id;
        unsafe {
            // steal op
            let stolen = mem::replace(&mut self.graph().access_inner_mut(id).op, None).unwrap();
            // call Op::grad
            stolen.grad(&mut self);
            // restore
            mem::swap(&mut self.graph().access_inner_mut(id).op, &mut Some(stolen));
            debug_assert!(
                !self.gxs.is_empty(),
                "Bad Op impl: GradientContext::set_input_grads was not called"
            );
            self.gxs
        }
    }

    /// Returns the symbolic gradient of the op's output.
    #[inline]
    pub fn output_grad(&self) -> Tensor<'g, T> {
        self.gy
    }

    /// Grabs the symbolic output of the op.
    #[inline]
    pub fn output(&self) -> Tensor<'g, T> {
        self.y
    }

    /// Grabs the `i` th symbolic input.
    #[inline]
    pub fn input(&self, i: usize) -> Tensor<'g, T> {
        return self
            .y
            .input_tensor(i, self.graph)
            .expect("bad Op::grad impl");
    }

    /// Returns the number of inputs.
    #[inline]
    pub fn num_inputs(&self) -> usize {
        unsafe { self.y.inner().in_edges.len() }
    }

    /// Returns a graph object that is usable for tensor computations in the context.
    #[inline]
    pub fn graph(&self) -> &'g crate::graph::Graph<T> {
        self.graph
    }

    /// Back-propagates the input's gradient.
    ///
    /// Appends the given tensor to the back of the input-gradient-list.
    /// `None` argument indicates that the `Op`'s input doesn't have gradient.
    /// Note that `Op::grad` must call this function as many as `num_inputs()`.
    #[inline]
    pub fn append_input_grad(&mut self, gx: Option<Tensor<'g, T>>) {
        self.gxs.push(gx);
    }
}
<commit_msg>doc fix<commit_after>//! Defining things related to `ag::op::Op`.
//!
use crate::ndarray_ext::{NdArrayView, NdArrayViewMut};
use crate::smallvec::SmallVec;
use crate::tensor::{Tensor, TensorInternal};
use crate::{Float, NdArray};
use std::any::type_name;
use std::fmt;
use std::marker::PhantomData;
use std::mem;

// Properties for op's `compute` method.
// Actual number of inout/output nodes are around 1~2 in most cases.
pub(crate) const NUM_MAX_OUTPUT: usize = 2;
pub(crate) const NUM_MAX_INPUT: usize = 4;

pub(crate) type InputArray<T> = SmallVec<[T; NUM_MAX_INPUT]>;
pub(crate) type OutputArray<T> = SmallVec<[T; NUM_MAX_OUTPUT]>;

pub(crate) type ComputeResult<'v, T> = Result<crate::ArrRepr<'v, T>, OpError>;

/// Error in `Op`'s computation.
#[derive(Clone, Debug, PartialEq)]
pub enum OpError {
    NdArrayError(String, ndarray::ShapeError),
    IncompatibleShape(String),
    TypeUnsupported(String),
    InvalidDims(String),
    OutOfBounds(String),
}

impl std::error::Error for OpError {}

impl fmt::Display for OpError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            OpError::NdArrayError(pref, e) => write!(f, "{}: ", pref).and_then(|()| e.fmt(f)),
            OpError::IncompatibleShape(s) => write!(f, "{}: ", s),
            OpError::TypeUnsupported(s) => write!(f, "{}: ", s),
            OpError::InvalidDims(s) => write!(f, "{}: ", s),
            OpError::OutOfBounds(s) => write!(f, "{}: ", s),
        }
    }
}

/// Sequence of an op's compute result.
///
/// Op can have multiple output NdArrays.
pub(crate) type Results<'v, T> = OutputArray<Option<ComputeResult<'v, T>>>;

/// Operation trait. `Tensor` wraps trait-object of this.
///
/// # Implementing differentiable operations
///
/// Many of well-known ops are pre-defined in `Graph`, but you can also
/// implement custom ops by hand.
///
/// ```
/// use ndarray;
/// use autograd as ag;
/// use ag::Graph;
///
/// type NdArray<T: ag::Float> = ndarray::Array<T, ndarray::IxDyn>;
///
/// // Implements `Op` trait for `Sigmoid`.
/// struct Sigmoid;
///
/// impl<T: ag::Float> ag::op::Op<T> for Sigmoid {
///     fn compute(
///         &self,
///         ctx: &mut ag::op::ComputeContext<T>,
///     ) {
///         let x: &ag::NdArrayView<_> = &ctx.input(0);
///         // Use `ndarray::Array::mapv` for element-wise computation.
///         let half = T::from(0.5).unwrap();
///         let y = x.mapv(move |a| ((a * half).tanh() * half) + half);
///         ctx.append_output(y);
///     }
///
///     fn grad(&self, ctx: &mut ag::op::GradientContext<T>) {
///         // Symbolic gradient of the input of Sigmoid
///         let gy = ctx.output_grad();
///         let y = ctx.output();
///         let gx = gy * (y - ctx.graph().square(y));
///         ctx.append_input_grad(Some(gx));
///     }
/// }
///
///  use ag::tensor::Input;
///
/// // Symbolic `sigmoid` function for end-user.
/// fn sigmoid<'graph, F: ag::Float>(x: &ag::Tensor<'graph, F>, g: &'graph ag::Graph<F>)
/// -> ag::Tensor<'graph, F> {
///     ag::Tensor::builder()
///            .set_inputs(&[Input::new(x)])
///            .build(g, Sigmoid)
/// }
/// ```
pub trait Op<F: Float> {
    /// Name of this op
    fn name(&self) -> &str {
        type_name::<Self>()
    }

    /// Runs this op with `ComputeContext`.
    fn compute(&self, ctx: &mut ComputeContext<F>);

    /// Returns symbolic gradients for input nodes by use of output's gradients etc.
    fn grad(&self, ctx: &mut GradientContext<F>);
}

pub(crate) struct DummyOp<F: Float> {
    pub phantom: PhantomData<F>,
}

impl<F: Float> DummyOp<F> {
    #[allow(dead_code)]
    pub(crate) fn new() -> Self {
        DummyOp {
            phantom: PhantomData,
        }
    }
}

impl<F: Float> Op<F> for DummyOp<F> {
    fn compute(&self, _: &mut ComputeContext<F>) {}
    fn grad(&self, _: &mut GradientContext<F>) {}
}

/// Wrapper object of NdArrayView/NdArrayViewMut which is fed to `Op::compute`
///
/// Used in `OpComputeContext`.
pub(crate) enum OpInput<'v, T: Float> {
    /// Read-only view
    RO(Option<NdArrayView<'v, T>>),
    /// Read-write view
    RW(Option<NdArrayViewMut<'v, T>>),
}

impl<'v, T: Float> OpInput<'v, T> {
    #[inline]
    /// Make a read-only input array
    pub fn new(x: NdArrayView<'v, T>) -> Self {
        OpInput::RO(Some(x))
    }

    #[inline]
    /// Make a read/write input array
    pub fn new_mut(x: NdArrayViewMut<'v, T>) -> Self {
        OpInput::RW(Some(x))
    }
}

/// Context of an `Op`'s computation phase.
///
/// # Example
///
/// ```
/// use autograd as ag;
///
/// // Implementing `Op` trait for `Sigmoid`.
/// struct Sigmoid;
///
/// impl<T: ag::Float> ag::op::Op<T> for Sigmoid {
///     fn compute(
///         &self,
///         ctx: &mut ag::op::ComputeContext<T>,
///     ) {
///         // Getting the first input array.
///         let x: &ag::NdArrayView<_> = &ctx.input(0);
///         let half = T::from(0.5).unwrap();
///         let y = x.mapv(move |a| ((a * half).tanh() * half) + half);
///         // Put the computed result.
///         ctx.append_output(y);
///     }
///
///     fn grad(&self, ctx: &mut ag::op::GradientContext<T>) { /* ... */ }
/// }
/// ```
pub struct ComputeContext<'t, 'v, T: Float> {
    node: &'t TensorInternal<T>,
    // Input arrays
    xs: InputArray<OpInput<'v, T>>,
    // Output arrays
    ys: Results<'v, T>,
}

impl<'g, 't, 'v, T: Float> ComputeContext<'t, 'v, T> {
    pub(crate) fn extract_outputs(self) -> Results<'v, T> {
        self.ys
    }

    #[inline]
    pub(crate) fn new(node: &'t TensorInternal<T>, xs: InputArray<OpInput<'v, T>>) -> Self {
        ComputeContext {
            node,
            xs,
            ys: OutputArray::new(),
        }
    }

    /// Grabs the `i` th input array as a *read-only* array view.
    ///
    /// Calling `input(i)` more than once causes panic.
    #[inline]
    pub fn input(&mut self, i: usize) -> NdArrayView<'v, T> {
        let x = match self.xs.get_mut(i) {
            Some(x) => x,
            None => panic!("Bad op impl: input index out of range."),
        };
        match x {
            OpInput::RO(ref mut a) => match a.take() {
                Some(ret) => ret,
                None => panic!(
                    "Bad op impl of {}: input({})/input_mut({}) cannot be called twice",
                    self.node.get_op().name(),
                    i,
                    i
                ),
            },
            _ => {
                panic!(
                    "Bad op impl of {}: cannot perform immutable borrowing for input({})",
                    self.node.get_op().name(),
                    i
                );
            }
        }
    }

    /// Grabs the `i` th input array as a *read-write* array view.
    ///
    /// Calling `input_mut(i)` more than once causes panic.
    #[inline]
    pub fn input_mut(&mut self, i: usize) -> NdArrayViewMut<'v, T> {
        let x = match self.xs.get_mut(i) {
            Some(x) => x,
            None => panic!("Bad op impl: {}'s input doesn't exist.", i),
        };
        match x {
            OpInput::RW(ref mut a) => match a.take() {
                Some(ret) => ret,
                None => panic!(
                    "Bad op impl of {}: input({})/input_mut({}) cannot be called twice",
                    self.node.get_op().name(),
                    i,
                    i
                ),
            },
            _ => {
                panic!(
                    "Bad op impl of {}: cannot perform mutable borrowing for input({})",
                    self.node.get_op().name(),
                    i
                );
            }
        }
    }

    /// Appends an `ndarray::ArrayView` to the back of the output list of the current op.
    ///
    /// NOTE: Implementor of `Op::compute` must not forget to call `append_*` as many as the number of its output in `Op::compute`, otherwise panic occurs.
    #[inline]
    pub fn append_output_view(&mut self, y: NdArrayView<'v, T>) {
        self.ys.push(Some(Ok(crate::ArrRepr::View(y))));
    }

    /// Appends an ndarray to the back of the output list of the current op.
    ///
    /// NOTE: Implementor of `Op::compute` must not forget to call `append_*` as many as the number of its output in `Op::compute`, otherwise panic occurs.
    #[inline]
    pub fn append_output(&mut self, y: NdArray<T>) {
        self.ys.push(Some(Ok(crate::ArrRepr::Owned(y))));
    }

    /// Appends an empty result to the back of the output list of the current op.
    ///
    /// For example, this is used for gradient-descent-optimizers.
    ///
    /// NOTE: Implementor of `Op::compute` must not forget to call `append_*` as many as the number of its output in `Op::compute`, otherwise panic occurs.
    #[inline]
    pub fn append_empty_output(&mut self) {
        self.ys.push(None);
    }

    /// Marks this op's compute session as failed.
    ///
    /// NOTE: After called this function, Op::compute should return early.
    #[inline]
    pub fn set_error(&mut self, y: OpError) {
        self.ys.push(Some(Err(y)));
    }

    /// Returns a number of input arrays.
    #[inline]
    pub fn num_inputs(&self) -> usize {
        self.xs.len()
    }
}

/// Context of an `Op`'s gradient propagation phase.
///
/// This is passed to an `Op` through `Op::grad`.
/// `Op::grad` should provide the gradients of its inputs by calling `GradientContext::append_input_grad`.
///
/// Use `graph()` to access `Graph` object for tensor computations.
///
/// ```
/// use autograd as ag;
///
/// struct Sigmoid;
///
/// impl<T: ag::Float> ag::op::Op<T> for Sigmoid {
///     fn compute(&self, ctx: &mut ag::op::ComputeContext<T>) { /* ... */ }
///
///     fn grad(&self, ctx: &mut ag::op::GradientContext<T>) {
///         // Symbolic gradient of the input of Sigmoid
///         let gy = ctx.output_grad();
///         // Symbolic output tensor
///         let y = ctx.output();
///         // `Tensor` computations
///         let gx = gy * (y - ctx.graph().square(y));
///         // Propagates input's gradient.
///         ctx.append_input_grad(Some(gx));
///     }
/// }
/// ```
pub struct GradientContext<'g, T: Float> {
    gy: Tensor<'g, T>,
    y: Tensor<'g, T>,
    graph: &'g crate::graph::Graph<T>,
    gxs: InputArray<Option<Tensor<'g, T>>>,
}

impl<'g, T: Float> GradientContext<'g, T> {
    pub(crate) fn new(
        gy: Tensor<'g, T>,
        y: Tensor<'g, T>,
        graph: &'g crate::graph::Graph<T>,
    ) -> Self {
        GradientContext {
            gy,
            y,
            graph,
            gxs: InputArray::new(),
        }
    }

    // Call Op::grad and return `gxs`
    pub(crate) fn extract_input_grads(mut self) -> InputArray<Option<Tensor<'g, T>>> {
        let id = self.y.id;
        unsafe {
            // steal op
            let stolen = mem::replace(&mut self.graph().access_inner_mut(id).op, None).unwrap();
            // call Op::grad
            stolen.grad(&mut self);
            // restore
            mem::swap(&mut self.graph().access_inner_mut(id).op, &mut Some(stolen));
            debug_assert!(
                !self.gxs.is_empty(),
                "Bad Op impl: GradientContext::append_input_grad was not called"
            );
            self.gxs
        }
    }

    /// Returns the symbolic gradient of the op's output.
    #[inline]
    pub fn output_grad(&self) -> Tensor<'g, T> {
        self.gy
    }

    /// Grabs the symbolic output of the op.
    #[inline]
    pub fn output(&self) -> Tensor<'g, T> {
        self.y
    }

    /// Grabs the `i` th symbolic input.
    #[inline]
    pub fn input(&self, i: usize) -> Tensor<'g, T> {
        return self
            .y
            .input_tensor(i, self.graph)
            .expect("bad Op::grad impl");
    }

    /// Returns the number of inputs.
    #[inline]
    pub fn num_inputs(&self) -> usize {
        unsafe { self.y.inner().in_edges.len() }
    }

    /// Returns a graph object that is usable for tensor computations in the context.
    #[inline]
    pub fn graph(&self) -> &'g crate::graph::Graph<T> {
        self.graph
    }

    /// Back-propagates the input's gradient.
    ///
    /// Appends the given tensor to the back of the input-gradient-list.
    /// `None` argument indicates that the `Op`'s input doesn't have gradient.
    /// Note that `Op::grad` must call this function as many as `num_inputs()`.
    #[inline]
    pub fn append_input_grad(&mut self, gx: Option<Tensor<'g, T>>) {
        self.gxs.push(gx);
    }
}
<|endoftext|>


<commit_before>use std::{env, fs};
use std::fs::{File, OpenOptions, Permissions};
use std::io::{Read, Result, Write};
#[cfg(unix)]
use std::os::unix::fs::PermissionsExt;
use std::path::{Path, PathBuf};

#[cfg(feature = "temp")]
use tempdir;

use FileSystem;
#[cfg(unix)]
use UnixFileSystem;
#[cfg(feature = "temp")]
use {TempDir, TempFileSystem};

/// Tracks a temporary directory that will be deleted once the struct goes out of scope.
///
/// This is a wrapper around a [`TempDir`].
///
/// [`TempDir`]: https://doc.rust-lang.org/tempdir/tempdir/struct.TempDir.html
#[cfg(feature = "temp")]
#[derive(Debug)]
pub struct OsTempDir(tempdir::TempDir);

#[cfg(feature = "temp")]
impl TempDir for OsTempDir {
    fn path(&self) -> &Path {
        self.0.path()
    }
}

/// An implementation of FileSystem that interacts with the actual operating system's file system.
///
/// This is primarily a wrapper for [`fs`] methods.
///
/// [`fs`]: https://doc.rust-lang.org/std/fs/index.html
#[derive(Clone, Debug, Default)]
pub struct OsFileSystem {}

impl OsFileSystem {
    pub fn new() -> Self {
        OsFileSystem {}
    }
}

impl FileSystem for OsFileSystem {
    fn current_dir(&self) -> Result<PathBuf> {
        env::current_dir()
    }

    fn set_current_dir<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        env::set_current_dir(path)
    }

    fn is_dir<P: AsRef<Path>>(&self, path: P) -> bool {
        path.as_ref().is_dir()
    }

    fn is_file<P: AsRef<Path>>(&self, path: P) -> bool {
        path.as_ref().is_file()
    }

    fn create_dir<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        fs::create_dir(path)
    }

    fn create_dir_all<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        fs::create_dir_all(path)
    }

    fn remove_dir<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        fs::remove_dir(path)
    }

    fn remove_dir_all<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        fs::remove_dir_all(path)
    }

    fn write_file<P, B>(&self, path: P, buf: B) -> Result<()>
        where P: AsRef<Path>,
              B: AsRef<[u8]>
    {
        let mut file = File::create(path)?;
        file.write_all(buf.as_ref())
    }

    fn read_file<P: AsRef<Path>>(&self, path: P) -> Result<Vec<u8>> {
        let mut contents = Vec::<u8>::new();
        let mut file = File::open(path)?;

        file.read_to_end(&mut contents)?;

        Ok(contents)
    }

    fn read_file_to_string<P: AsRef<Path>>(&self, path: P) -> Result<String> {
        let mut contents = String::new();
        let mut file = File::open(path)?;

        file.read_to_string(&mut contents)?;

        Ok(contents)
    }

    fn create_file<P, B>(&self, path: P, buf: B) -> Result<()>
        where P: AsRef<Path>,
              B: AsRef<[u8]>
    {
        let mut file = OpenOptions::new().write(true)
            .create_new(true)
            .open(path)?;

        file.write_all(buf.as_ref())
    }

    fn remove_file<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        fs::remove_file(path)
    }

    fn copy_file<P, Q>(&self, from: P, to: Q) -> Result<()>
        where P: AsRef<Path>,
              Q: AsRef<Path>
    {
        fs::copy(from, to).and(Ok(()))
    }

    fn rename<P, Q>(&self, from: P, to: Q) -> Result<()>
        where P: AsRef<Path>,
              Q: AsRef<Path>
    {
        fs::rename(from, to)
    }

    fn readonly<P: AsRef<Path>>(&self, path: P) -> Result<bool> {
        permissions(path.as_ref()).map(|p| p.readonly())
    }

    fn set_readonly<P: AsRef<Path>>(&self, path: P, readonly: bool) -> Result<()> {
        let mut permissions = permissions(path.as_ref())?;

        permissions.set_readonly(readonly);

        fs::set_permissions(path, permissions)
    }
}

#[cfg(unix)]
impl UnixFileSystem for OsFileSystem {
    fn mode<P: AsRef<Path>>(&self, path: P) -> Result<u32> {
        permissions(path.as_ref()).map(|p| p.mode())
    }

    fn set_mode<P: AsRef<Path>>(&self, path: P, mode: u32) -> Result<()> {
        let mut permissions = permissions(path.as_ref())?;

        permissions.set_mode(mode);

        fs::set_permissions(path, permissions)
    }
}

#[cfg(feature = "temp")]
impl TempFileSystem for OsFileSystem {
    type TempDir = OsTempDir;

    fn temp_dir<S: AsRef<str>>(&self, prefix: S) -> Result<Self::TempDir> {
        tempdir::TempDir::new(prefix.as_ref()).map(OsTempDir)
    }
}

fn permissions(path: &Path) -> Result<Permissions> {
    let metadata = fs::metadata(path)?;

    Ok(metadata.permissions())
}
<commit_msg>Fix clippy warning<commit_after>use std::{env, fs};
use std::fs::{File, OpenOptions, Permissions};
use std::io::{Read, Result, Write};
#[cfg(unix)]
use std::os::unix::fs::PermissionsExt;
use std::path::{Path, PathBuf};

#[cfg(feature = "temp")]
use tempdir;

use FileSystem;
#[cfg(unix)]
use UnixFileSystem;
#[cfg(feature = "temp")]
use {TempDir, TempFileSystem};

/// Tracks a temporary directory that will be deleted once the struct goes out of scope.
///
/// This is a wrapper around a [`TempDir`].
///
/// [`TempDir`]: https://doc.rust-lang.org/tempdir/tempdir/struct.TempDir.html
#[cfg(feature = "temp")]
#[derive(Debug)]
pub struct OsTempDir(tempdir::TempDir);

#[cfg(feature = "temp")]
impl TempDir for OsTempDir {
    fn path(&self) -> &Path {
        self.0.path()
    }
}

/// An implementation of `FileSystem` that interacts with the actual operating system's file system.
///
/// This is primarily a wrapper for [`fs`] methods.
///
/// [`fs`]: https://doc.rust-lang.org/std/fs/index.html
#[derive(Clone, Debug, Default)]
pub struct OsFileSystem {}

impl OsFileSystem {
    pub fn new() -> Self {
        OsFileSystem {}
    }
}

impl FileSystem for OsFileSystem {
    fn current_dir(&self) -> Result<PathBuf> {
        env::current_dir()
    }

    fn set_current_dir<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        env::set_current_dir(path)
    }

    fn is_dir<P: AsRef<Path>>(&self, path: P) -> bool {
        path.as_ref().is_dir()
    }

    fn is_file<P: AsRef<Path>>(&self, path: P) -> bool {
        path.as_ref().is_file()
    }

    fn create_dir<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        fs::create_dir(path)
    }

    fn create_dir_all<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        fs::create_dir_all(path)
    }

    fn remove_dir<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        fs::remove_dir(path)
    }

    fn remove_dir_all<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        fs::remove_dir_all(path)
    }

    fn write_file<P, B>(&self, path: P, buf: B) -> Result<()>
        where P: AsRef<Path>,
              B: AsRef<[u8]>
    {
        let mut file = File::create(path)?;
        file.write_all(buf.as_ref())
    }

    fn read_file<P: AsRef<Path>>(&self, path: P) -> Result<Vec<u8>> {
        let mut contents = Vec::<u8>::new();
        let mut file = File::open(path)?;

        file.read_to_end(&mut contents)?;

        Ok(contents)
    }

    fn read_file_to_string<P: AsRef<Path>>(&self, path: P) -> Result<String> {
        let mut contents = String::new();
        let mut file = File::open(path)?;

        file.read_to_string(&mut contents)?;

        Ok(contents)
    }

    fn create_file<P, B>(&self, path: P, buf: B) -> Result<()>
        where P: AsRef<Path>,
              B: AsRef<[u8]>
    {
        let mut file = OpenOptions::new().write(true)
            .create_new(true)
            .open(path)?;

        file.write_all(buf.as_ref())
    }

    fn remove_file<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        fs::remove_file(path)
    }

    fn copy_file<P, Q>(&self, from: P, to: Q) -> Result<()>
        where P: AsRef<Path>,
              Q: AsRef<Path>
    {
        fs::copy(from, to).and(Ok(()))
    }

    fn rename<P, Q>(&self, from: P, to: Q) -> Result<()>
        where P: AsRef<Path>,
              Q: AsRef<Path>
    {
        fs::rename(from, to)
    }

    fn readonly<P: AsRef<Path>>(&self, path: P) -> Result<bool> {
        permissions(path.as_ref()).map(|p| p.readonly())
    }

    fn set_readonly<P: AsRef<Path>>(&self, path: P, readonly: bool) -> Result<()> {
        let mut permissions = permissions(path.as_ref())?;

        permissions.set_readonly(readonly);

        fs::set_permissions(path, permissions)
    }
}

#[cfg(unix)]
impl UnixFileSystem for OsFileSystem {
    fn mode<P: AsRef<Path>>(&self, path: P) -> Result<u32> {
        permissions(path.as_ref()).map(|p| p.mode())
    }

    fn set_mode<P: AsRef<Path>>(&self, path: P, mode: u32) -> Result<()> {
        let mut permissions = permissions(path.as_ref())?;

        permissions.set_mode(mode);

        fs::set_permissions(path, permissions)
    }
}

#[cfg(feature = "temp")]
impl TempFileSystem for OsFileSystem {
    type TempDir = OsTempDir;

    fn temp_dir<S: AsRef<str>>(&self, prefix: S) -> Result<Self::TempDir> {
        tempdir::TempDir::new(prefix.as_ref()).map(OsTempDir)
    }
}

fn permissions(path: &Path) -> Result<Permissions> {
    let metadata = fs::metadata(path)?;

    Ok(metadata.permissions())
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

//! Volume Groups

use std::io::Result;
use std::io::Error;
use std::io::ErrorKind::Other;
use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::str::FromStr;
use std::borrow::Cow;

use time::now;
use nix::sys::utsname::uname;

use lv::{LV, Segment};
use pv::{PV, Device};
use pvlabel::PvHeader;
use parser::{LvmTextMap, Entry};
use lvmetad;
use dm::DM;
use util::{align_to, make_uuid};

/// A Volume Group.
#[derive(Debug, PartialEq, Clone)]
pub struct VG {
    /// Name.
    pub name: String,
    /// Uuid.
    pub id: String,
    /// The generation of metadata this VG represents.
    pub seqno: u64,
    /// Always "LVM2".
    pub format: String,
    /// Status.
    pub status: Vec<String>,
    /// Flags.
    pub flags: Vec<String>,
    /// Size of each extent, in 512-byte sectors.
    pub extent_size: u64,
    /// Maximum number of LVs, 0 means no limit.
    pub max_lv: u64,
    /// Maximum number of PVs, 0 means no limit.
    pub max_pv: u64,
    /// How many metadata copies (?)
    pub metadata_copies: u64,
    /// Physical Volumes within this volume group.
    pub pvs: BTreeMap<String, PV>,
    /// Logical Volumes within this volume group.
    pub lvs: BTreeMap<String, LV>,
}

impl VG {
    /// The total number of extents in use in the volume group.
    pub fn extents_in_use(&self) -> u64 {
        self.lvs
            .values()
            .map(|x| x.used_extents())
            .sum()
    }

    /// The total number of free extents in the volume group.
    pub fn extents_free(&self) -> u64 {
        self.extents() - self.extents_in_use()
    }

    /// The total number of extents in the volume group.
    pub fn extents(&self) -> u64 {
        self.pvs
            .values()
            .map(|x| x.pe_count)
            .sum()
    }

    // Recursively walk DM deps to see if device is present
    fn depends_on(dev: Device, dm_majors: &BTreeSet<u32>, dm: &DM) -> bool {
        if !dm_majors.contains(&dev.major) {
            return false;
        }

        match dm.list_deps(dev) {
            Err(_) => return false,
            Ok(dep_list) => {
                for d in dep_list {
                    if d == dev {
                        return true;
                    } else if Self::depends_on(d, dm_majors, dm) {
                        return true;
                    }
                }
            }
        }

        false
    }

    /// Add a non-affiliated PV to this VG.
    pub fn add_pv(&mut self, pvh: &PvHeader) -> Result<()> {
        // add_pv_to_vg
        // check pv is not on an LV from the vg:
        // 1) is pv's major a devicemapper major?
        // 2) equiv. of dev_manager_device_uses_vg()

        let dm_majors = DM::dm_majors();
        let dev = try!(Device::from_str(&pvh.dev_path.to_string_lossy()));
        {
            let dm = try!(DM::new(&self));
            if Self::depends_on(dev, &dm_majors, &dm) {
                return Err(Error::new(Other, "Dependency loops prohibited"));
            }
        }

        // check pv is not already in the VG or another VG
        // Does it have text metadata??
        if let Ok(metadata) = pvh.read_metadata() {

            // Find the textmap for the vg, among all the other stuff.
            // (It's the only textmap.)
            let mut vg_name = Cow::Borrowed("<unknown>");
            for (key, value) in metadata {
                match value {
                    Entry::TextMap(_) => {
                        vg_name = Cow::Owned(key);
                        break
                    },
                    _ => {}
                }
            }

            return Err(Error::new(Other, format!("PV already in VG {}", vg_name)));
        }

        let da = try!(pvh.data_areas.get(0)
                      .ok_or(Error::new(Other, "Could not find data area in PV")));

        // figure out how many extents fit in the PV's data area
        // pe_start aligned to extent size
        let pe_start = align_to(da.offset as usize, self.extent_size as usize) as u64;
        let mda1_size = match pvh.metadata_areas.get(1) {
            Some(pvarea) => pvarea.size,
            None => 0,
        };
        let area_size = pvh.size - pe_start - mda1_size;
        let pe_count = area_size / self.extent_size;

        let name = format!("pv{}", self.pvs.len());

        // make a PV and add it to self
        let pv = PV {
            name: name.clone(),
            id: pvh.uuid.clone(),
            device: dev,
            status: vec!["ALLOCATABLE".to_string()],
            flags: Vec::new(),
            dev_size: pvh.size,
            pe_start: pe_start,
            pe_count: pe_count,
        };

        self.pvs.insert(name, pv);

        self.commit()
    }

    /// Create a new linear logical volume in the volume group.
    pub fn new_linear_lv(&mut self, name: &str, extent_size: u64) -> Result<()> {
        if self.lvs.contains_key(name) {
            return Err(Error::new(Other, "LV already exists"));
        }

        let mut contig_area = None;
        for (pvname, areas) in self.free_areas() {
            for (start, len) in areas {
                if len >= extent_size {
                    contig_area = Some((pvname, start));
                    break;
                }
            }
        }

        // we don't support multiple segments yet
        let (pv_with_area, area_start) = match contig_area {
            None => return Err(Error::new(Other, "no contiguous area for new LV")),
            Some(x) => x,
        };

        let segment = Segment {
            name: "segment1".to_string(),
            start_extent: 0,
            extent_count: extent_size,
            ty: "striped".to_string(),
            stripes: vec![(pv_with_area, area_start)],
        };

        let mut lv = LV {
            name: name.to_string(),
            id: make_uuid(),
            status: vec!["READ".to_string(),
                         "WRITE".to_string(),
                         "VISIBLE".to_string()],
            flags: Vec::new(),
            creation_host: uname().nodename().to_string(),
            creation_time: now().to_timespec().sec,
            segments: vec![segment],
            device: None,
        };

        // poke dm and tell it about a new device
        {
            let dm = try!(DM::new(self));
            try!(dm.activate_device(&mut lv));
        }

        self.lvs.insert(name.to_string(), lv);

        self.commit()
    }

    /// Destroy a logical volume.
    pub fn lv_remove(&mut self, name: &str) -> Result<()> {
        match self.lvs.remove(name) {
            None => Err(Error::new(Other, "LV not found in VG")),
            Some(mut lv) => {
                {
                    let dm = try!(DM::new(self));
                    try!(dm.deactivate_device(&mut lv));
                }

                self.commit()
            },
        }
    }

    fn commit(&mut self) -> Result<()> {
        let map = self.clone().into();

        // TODO: atomicity of updating pvs, metad, dm
        for pv in self.pvs.values() {
            if let Some(path) = pv.device.path() {
                let mut pvheader = PvHeader::find_in_dev(&path)
                    .expect("could not find pvheader");

                try!(pvheader.write_metadata(&map));
            }
        }

        lvmetad::vg_update(&map)
    }

    // Returns e.g. {"pv0": {0: 45, 47: 100, 100: 200} }
    // This means extents 0-44 are used, 45 and 46 are not,
    // 47-99 are used, then 100-199 are used.
    //
    // Adjacent used areas are not merged.
    //
    // PVs with no used areas are not in the outer map at all.
    //
    fn used_areas(&self) -> BTreeMap<String, BTreeMap<u64, u64>> {
        let mut used_map = BTreeMap::new();

        for lv in self.lvs.values() {
            for seg in &lv.segments {
                for &(ref pvname, start) in &seg.stripes {
                    used_map.entry(pvname.to_string())
                        .or_insert(BTreeMap::new())
                        .insert(start as u64, seg.extent_count);
                }
            }
        }

        used_map
    }

    // Returns e.g. {"pv0": {45: 47, 200: 1000} }
    // (assuming pv0 has 1000 extents)
    //
    fn free_areas(&self) -> BTreeMap<String, BTreeMap<u64, u64>> {
        let mut free_map = BTreeMap::new();

        for (pvname, mut area_map) in self.used_areas() {

            // Insert an entry to mark the end of the PV so the fold works
            // correctly
            let pv = self.pvs.get(&pvname)
                .expect("area map name refers to nonexistent PV");
            area_map.insert(pv.pe_count, 0);

            area_map.iter()
                .fold(0, |prev_end, (start, len)| {
                    if prev_end < *start {
                        free_map.entry(pvname.clone())
                            .or_insert(BTreeMap::new())
                            .insert(prev_end, start-prev_end);
                    }
                    start + len
                });
        }

        // Also return completely-unused PVs
        for (pvname, pv) in &self.pvs {
            if !free_map.contains_key(pvname) {
                let mut map = BTreeMap::new();
                map.insert(0, pv.pe_count);
                free_map.insert(pvname.clone(), map);
            }
        }

        free_map
    }
}

impl From<VG> for LvmTextMap {
    fn from(vg: VG) -> Self {
        let mut map = LvmTextMap::new();

        map.insert("id".to_string(), Entry::String(vg.id));
        map.insert("seqno".to_string(),
                   Entry::Number(vg.seqno as i64 + 1));
        map.insert("format".to_string(), Entry::String(vg.format));

        map.insert("max_pv".to_string(), Entry::Number(0));
        map.insert("max_lv".to_string(), Entry::Number(0));

        map.insert("status".to_string(),
                   Entry::List(
                       Box::new(
                           vg.status
                               .into_iter()
                               .map(|x| Entry::String(x))
                               .collect())));

        map.insert("flags".to_string(),
                   Entry::List(
                       Box::new(
                           vg.flags
                               .into_iter()
                               .map(|x| Entry::String(x))
                               .collect())));

        map.insert("extent_size".to_string(),
                   Entry::Number(vg.extent_size as i64));
        map.insert("metadata_copies".to_string(),
                   Entry::Number(vg.metadata_copies as i64));
        map.insert("physical_volumes".to_string(),
                   Entry::TextMap(
                       Box::new(
                           vg.pvs
                               .into_iter()
                               .map(|(k, v)|
                                    (k, Entry::TextMap(Box::new(v.into()))))
                               .collect())));

        map.insert("logical_volumes".to_string(),
                   Entry::TextMap(
                       Box::new(
                           vg.lvs
                               .into_iter()
                               .map(|(k, v)|
                                    (k, Entry::TextMap(Box::new(v.into()))))
                               .collect())));

        let mut outer_map = LvmTextMap::new();

        outer_map.insert(vg.name, Entry::TextMap(Box::new(map)));

        outer_map
    }
}
<commit_msg>Use 'if let' instead of match in depends_on()<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

//! Volume Groups

use std::io::Result;
use std::io::Error;
use std::io::ErrorKind::Other;
use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::str::FromStr;
use std::borrow::Cow;

use time::now;
use nix::sys::utsname::uname;

use lv::{LV, Segment};
use pv::{PV, Device};
use pvlabel::PvHeader;
use parser::{LvmTextMap, Entry};
use lvmetad;
use dm::DM;
use util::{align_to, make_uuid};

/// A Volume Group.
#[derive(Debug, PartialEq, Clone)]
pub struct VG {
    /// Name.
    pub name: String,
    /// Uuid.
    pub id: String,
    /// The generation of metadata this VG represents.
    pub seqno: u64,
    /// Always "LVM2".
    pub format: String,
    /// Status.
    pub status: Vec<String>,
    /// Flags.
    pub flags: Vec<String>,
    /// Size of each extent, in 512-byte sectors.
    pub extent_size: u64,
    /// Maximum number of LVs, 0 means no limit.
    pub max_lv: u64,
    /// Maximum number of PVs, 0 means no limit.
    pub max_pv: u64,
    /// How many metadata copies (?)
    pub metadata_copies: u64,
    /// Physical Volumes within this volume group.
    pub pvs: BTreeMap<String, PV>,
    /// Logical Volumes within this volume group.
    pub lvs: BTreeMap<String, LV>,
}

impl VG {
    /// The total number of extents in use in the volume group.
    pub fn extents_in_use(&self) -> u64 {
        self.lvs
            .values()
            .map(|x| x.used_extents())
            .sum()
    }

    /// The total number of free extents in the volume group.
    pub fn extents_free(&self) -> u64 {
        self.extents() - self.extents_in_use()
    }

    /// The total number of extents in the volume group.
    pub fn extents(&self) -> u64 {
        self.pvs
            .values()
            .map(|x| x.pe_count)
            .sum()
    }

    // Recursively walk DM deps to see if device is present
    fn depends_on(dev: Device, dm_majors: &BTreeSet<u32>, dm: &DM) -> bool {
        if !dm_majors.contains(&dev.major) {
            return false;
        }

        if let Ok(dep_list) = dm.list_deps(dev) {
            for d in dep_list {
                if d == dev {
                    return true;
                } else if Self::depends_on(d, dm_majors, dm) {
                    return true;
                }
            }
        }

        false
    }

    /// Add a non-affiliated PV to this VG.
    pub fn add_pv(&mut self, pvh: &PvHeader) -> Result<()> {
        // add_pv_to_vg
        // check pv is not on an LV from the vg:
        // 1) is pv's major a devicemapper major?
        // 2) equiv. of dev_manager_device_uses_vg()

        let dm_majors = DM::dm_majors();
        let dev = try!(Device::from_str(&pvh.dev_path.to_string_lossy()));
        {
            let dm = try!(DM::new(&self));
            if Self::depends_on(dev, &dm_majors, &dm) {
                return Err(Error::new(Other, "Dependency loops prohibited"));
            }
        }

        // check pv is not already in the VG or another VG
        // Does it have text metadata??
        if let Ok(metadata) = pvh.read_metadata() {

            // Find the textmap for the vg, among all the other stuff.
            // (It's the only textmap.)
            let mut vg_name = Cow::Borrowed("<unknown>");
            for (key, value) in metadata {
                match value {
                    Entry::TextMap(_) => {
                        vg_name = Cow::Owned(key);
                        break
                    },
                    _ => {}
                }
            }

            return Err(Error::new(Other, format!("PV already in VG {}", vg_name)));
        }

        let da = try!(pvh.data_areas.get(0)
                      .ok_or(Error::new(Other, "Could not find data area in PV")));

        // figure out how many extents fit in the PV's data area
        // pe_start aligned to extent size
        let pe_start = align_to(da.offset as usize, self.extent_size as usize) as u64;
        let mda1_size = match pvh.metadata_areas.get(1) {
            Some(pvarea) => pvarea.size,
            None => 0,
        };
        let area_size = pvh.size - pe_start - mda1_size;
        let pe_count = area_size / self.extent_size;

        let name = format!("pv{}", self.pvs.len());

        // make a PV and add it to self
        let pv = PV {
            name: name.clone(),
            id: pvh.uuid.clone(),
            device: dev,
            status: vec!["ALLOCATABLE".to_string()],
            flags: Vec::new(),
            dev_size: pvh.size,
            pe_start: pe_start,
            pe_count: pe_count,
        };

        self.pvs.insert(name, pv);

        self.commit()
    }

    /// Create a new linear logical volume in the volume group.
    pub fn new_linear_lv(&mut self, name: &str, extent_size: u64) -> Result<()> {
        if self.lvs.contains_key(name) {
            return Err(Error::new(Other, "LV already exists"));
        }

        let mut contig_area = None;
        for (pvname, areas) in self.free_areas() {
            for (start, len) in areas {
                if len >= extent_size {
                    contig_area = Some((pvname, start));
                    break;
                }
            }
        }

        // we don't support multiple segments yet
        let (pv_with_area, area_start) = match contig_area {
            None => return Err(Error::new(Other, "no contiguous area for new LV")),
            Some(x) => x,
        };

        let segment = Segment {
            name: "segment1".to_string(),
            start_extent: 0,
            extent_count: extent_size,
            ty: "striped".to_string(),
            stripes: vec![(pv_with_area, area_start)],
        };

        let mut lv = LV {
            name: name.to_string(),
            id: make_uuid(),
            status: vec!["READ".to_string(),
                         "WRITE".to_string(),
                         "VISIBLE".to_string()],
            flags: Vec::new(),
            creation_host: uname().nodename().to_string(),
            creation_time: now().to_timespec().sec,
            segments: vec![segment],
            device: None,
        };

        // poke dm and tell it about a new device
        {
            let dm = try!(DM::new(self));
            try!(dm.activate_device(&mut lv));
        }

        self.lvs.insert(name.to_string(), lv);

        self.commit()
    }

    /// Destroy a logical volume.
    pub fn lv_remove(&mut self, name: &str) -> Result<()> {
        match self.lvs.remove(name) {
            None => Err(Error::new(Other, "LV not found in VG")),
            Some(mut lv) => {
                {
                    let dm = try!(DM::new(self));
                    try!(dm.deactivate_device(&mut lv));
                }

                self.commit()
            },
        }
    }

    fn commit(&mut self) -> Result<()> {
        let map = self.clone().into();

        // TODO: atomicity of updating pvs, metad, dm
        for pv in self.pvs.values() {
            if let Some(path) = pv.device.path() {
                let mut pvheader = PvHeader::find_in_dev(&path)
                    .expect("could not find pvheader");

                try!(pvheader.write_metadata(&map));
            }
        }

        lvmetad::vg_update(&map)
    }

    // Returns e.g. {"pv0": {0: 45, 47: 100, 100: 200} }
    // This means extents 0-44 are used, 45 and 46 are not,
    // 47-99 are used, then 100-199 are used.
    //
    // Adjacent used areas are not merged.
    //
    // PVs with no used areas are not in the outer map at all.
    //
    fn used_areas(&self) -> BTreeMap<String, BTreeMap<u64, u64>> {
        let mut used_map = BTreeMap::new();

        for lv in self.lvs.values() {
            for seg in &lv.segments {
                for &(ref pvname, start) in &seg.stripes {
                    used_map.entry(pvname.to_string())
                        .or_insert(BTreeMap::new())
                        .insert(start as u64, seg.extent_count);
                }
            }
        }

        used_map
    }

    // Returns e.g. {"pv0": {45: 47, 200: 1000} }
    // (assuming pv0 has 1000 extents)
    //
    fn free_areas(&self) -> BTreeMap<String, BTreeMap<u64, u64>> {
        let mut free_map = BTreeMap::new();

        for (pvname, mut area_map) in self.used_areas() {

            // Insert an entry to mark the end of the PV so the fold works
            // correctly
            let pv = self.pvs.get(&pvname)
                .expect("area map name refers to nonexistent PV");
            area_map.insert(pv.pe_count, 0);

            area_map.iter()
                .fold(0, |prev_end, (start, len)| {
                    if prev_end < *start {
                        free_map.entry(pvname.clone())
                            .or_insert(BTreeMap::new())
                            .insert(prev_end, start-prev_end);
                    }
                    start + len
                });
        }

        // Also return completely-unused PVs
        for (pvname, pv) in &self.pvs {
            if !free_map.contains_key(pvname) {
                let mut map = BTreeMap::new();
                map.insert(0, pv.pe_count);
                free_map.insert(pvname.clone(), map);
            }
        }

        free_map
    }
}

impl From<VG> for LvmTextMap {
    fn from(vg: VG) -> Self {
        let mut map = LvmTextMap::new();

        map.insert("id".to_string(), Entry::String(vg.id));
        map.insert("seqno".to_string(),
                   Entry::Number(vg.seqno as i64 + 1));
        map.insert("format".to_string(), Entry::String(vg.format));

        map.insert("max_pv".to_string(), Entry::Number(0));
        map.insert("max_lv".to_string(), Entry::Number(0));

        map.insert("status".to_string(),
                   Entry::List(
                       Box::new(
                           vg.status
                               .into_iter()
                               .map(|x| Entry::String(x))
                               .collect())));

        map.insert("flags".to_string(),
                   Entry::List(
                       Box::new(
                           vg.flags
                               .into_iter()
                               .map(|x| Entry::String(x))
                               .collect())));

        map.insert("extent_size".to_string(),
                   Entry::Number(vg.extent_size as i64));
        map.insert("metadata_copies".to_string(),
                   Entry::Number(vg.metadata_copies as i64));
        map.insert("physical_volumes".to_string(),
                   Entry::TextMap(
                       Box::new(
                           vg.pvs
                               .into_iter()
                               .map(|(k, v)|
                                    (k, Entry::TextMap(Box::new(v.into()))))
                               .collect())));

        map.insert("logical_volumes".to_string(),
                   Entry::TextMap(
                       Box::new(
                           vg.lvs
                               .into_iter()
                               .map(|(k, v)|
                                    (k, Entry::TextMap(Box::new(v.into()))))
                               .collect())));

        let mut outer_map = LvmTextMap::new();

        outer_map.insert(vg.name, Entry::TextMap(Box::new(map)));

        outer_map
    }
}
<|endoftext|>