<commit_before>/*
 * Copyright (C) 2017-2018 Ryan Huang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use std::sync::mpsc::{self, Receiver, Sender};
use std::thread::{self, JoinHandle};

use model::{Board, Move, Outcome};

const INFINITY: i32 = 2_147_483_647;
const NEG_INFINITY: i32 = -2_147_483_647;
const LOSE: i32 = -1_073_741_824;
const DRAW: i32 = 0;

pub struct AIHandle {
    pub move_receiver: Receiver<Option<Move>>,
    pub stop_sender: Sender<()>,
    handle: JoinHandle<()>,
}

pub fn ai_move(board: Board, depth: u32, prev_handle: Option<AIHandle>) -> AIHandle {
    assert!(depth != 0);

    let (move_sender, move_receiver) = mpsc::channel();
    let (stop_sender, stop_receiver) = mpsc::channel();

    let handle = thread::spawn(move || {
        if let Some(prev_handle) = prev_handle {
            // If send returns an error, the other thread has already terminated
            if prev_handle.stop_sender.send(()).is_ok() {
                prev_handle.handle.join().expect(
                    "Previous AI thread panicked while new AI thread was waiting for it to finish",
                );
            }
        }

        let mut max_score = NEG_INFINITY;
        let mut best_move = None;
        for mv in board.generate_moves() {
            if stop_receiver.try_recv().is_ok() {
                return;
            }

            let mut new_board = board;
            new_board.apply_move(&mv);

            let score = -alphabeta_negamax(&new_board, NEG_INFINITY, INFINITY, depth - 1);
            if score > max_score {
                max_score = score;
                best_move = Some(mv);
            }
        }
        move_sender.send(best_move).expect("AI failed to send Move");
    });

    AIHandle {
        move_receiver,
        stop_sender,
        handle,
    }
}

fn alphabeta_negamax(board: &Board, mut alpha: i32, beta: i32, depth: u32) -> i32 {
    match board.outcome() {
        Outcome::Draw => return DRAW,
        Outcome::Win(color) => {
            assert!(color != board.turn());
            // TODO: weight by depth to encourage shorter wins
            return LOSE;
        }
        _ => {}
    }

    if depth == 0 {
        evaluate(board)
    } else {
        let moves = board.generate_moves();
        for mv in moves {
            let mut new_board = *board;
            new_board.apply_move(&mv);

            let score = -alphabeta_negamax(&new_board, -beta, -alpha, depth - 1);
            if score >= beta {
                return beta;
            } else if score > alpha {
                alpha = score;
            }
        }
        alpha
    }
}

fn evaluate(board: &Board) -> i32 {
    use model::Color::*;

    let wp = 100 * i32::from(board.pieces(White));
    let bp = 100 * i32::from(board.pieces(Black));
    let wh = 50 * i32::from(board.hexes(White));
    let bh = 50 * i32::from(board.hexes(Black));

    match board.turn() {
        White => (wp + wh) - (bp + bh),
        Black => (bp + bh) - (wp + wh),
    }
}

pub fn perft(board: &Board, depth: u32) -> u32 {
    if depth == 0 {
        1
    } else {
        let mut sum = 0;
        for mv in board.generate_moves() {
            let mut new_board = *board;
            new_board.apply_move(&mv);
            sum += perft(&new_board, depth - 1);
        }
        sum
    }
}
<commit_msg>Change AI score type to i16<commit_after>/*
 * Copyright (C) 2017-2018 Ryan Huang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use std::sync::mpsc::{self, Receiver, Sender};
use std::thread::{self, JoinHandle};

use model::{Board, Move, Outcome};

const INFINITY: i16 = 0x7FFF;
const NEG_INFINITY: i16 = -0x7FFF;
const LOSE: i16 = -0x4000;
const DRAW: i16 = 0;

pub struct AIHandle {
    pub move_receiver: Receiver<Option<Move>>,
    pub stop_sender: Sender<()>,
    handle: JoinHandle<()>,
}

pub fn ai_move(board: Board, depth: u32, prev_handle: Option<AIHandle>) -> AIHandle {
    assert!(depth != 0);

    let (move_sender, move_receiver) = mpsc::channel();
    let (stop_sender, stop_receiver) = mpsc::channel();

    let handle = thread::spawn(move || {
        if let Some(prev_handle) = prev_handle {
            // If send returns an error, the other thread has already terminated
            if prev_handle.stop_sender.send(()).is_ok() {
                prev_handle.handle.join().expect(
                    "Previous AI thread panicked while new AI thread was waiting for it to finish",
                );
            }
        }

        let mut max_score = NEG_INFINITY;
        let mut best_move = None;
        for mv in board.generate_moves() {
            if stop_receiver.try_recv().is_ok() {
                return;
            }

            let mut new_board = board;
            new_board.apply_move(&mv);

            let score = -alphabeta_negamax(&new_board, NEG_INFINITY, INFINITY, depth - 1);
            if score > max_score {
                max_score = score;
                best_move = Some(mv);
            }
        }
        move_sender.send(best_move).expect("AI failed to send Move");
    });

    AIHandle {
        move_receiver,
        stop_sender,
        handle,
    }
}

fn alphabeta_negamax(board: &Board, mut alpha: i16, beta: i16, depth: u32) -> i16 {
    match board.outcome() {
        Outcome::Draw => return DRAW,
        Outcome::Win(color) => {
            assert!(color != board.turn());
            // TODO: weight by depth to encourage shorter wins
            return LOSE;
        }
        _ => {}
    }

    if depth == 0 {
        evaluate(board)
    } else {
        let moves = board.generate_moves();
        for mv in moves {
            let mut new_board = *board;
            new_board.apply_move(&mv);

            let score = -alphabeta_negamax(&new_board, -beta, -alpha, depth - 1);
            if score >= beta {
                return beta;
            } else if score > alpha {
                alpha = score;
            }
        }
        alpha
    }
}

fn evaluate(board: &Board) -> i16 {
    use model::Color::*;

    let wp = 100 * i16::from(board.pieces(White));
    let bp = 100 * i16::from(board.pieces(Black));
    let wh = 50 * i16::from(board.hexes(White));
    let bh = 50 * i16::from(board.hexes(Black));

    match board.turn() {
        White => (wp + wh) - (bp + bh),
        Black => (bp + bh) - (wp + wh),
    }
}

pub fn perft(board: &Board, depth: u32) -> u32 {
    if depth == 0 {
        1
    } else {
        let mut sum = 0;
        for mv in board.generate_moves() {
            let mut new_board = *board;
            new_board.apply_move(&mv);
            sum += perft(&new_board, depth - 1);
        }
        sum
    }
}
<|endoftext|>


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

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

use crate::model::{
    GameId, GameInfo, Key, KeyPrefix, LichessEntry, MastersEntry, MastersGame, Month,
    PersonalEntry, PersonalStatus, UserId, Year,
};

#[derive(Debug)]
pub struct Database {
    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 = DBWithThreadMode::open_cf_descriptors(
            &db_opts,
            path,
            vec![
                column_family(
                    "lichess",
                    Some("lichess merge"),
                    lichess_merge,
                    Some(KeyPrefix::SIZE),
                    4 * 1024,
                    3,
                ),
                column_family(
                    "personal",
                    Some("personal merge"),
                    personal_merge,
                    Some(KeyPrefix::SIZE),
                    4 * 1024,
                    3,
                ),
                column_family("game", Some("game merge"), game_merge, None, 1024, 0),
                column_family("player", None, void_merge, None, 1024, 0),
                column_family(
                    "masters",
                    Some("masters merge"),
                    masters_merge,
                    Some(KeyPrefix::SIZE),
                    4 * 1024,
                    3,
                ),
                column_family("masters_game", None, void_merge, None, 4 * 1024, 0),
            ],
        )?;

        Ok(Database { inner })
    }

    pub fn queryable(&self) -> QueryableDatabase<'_> {
        QueryableDatabase {
            db: &self.inner,
            cf_lichess: self.inner.cf_handle("lichess").expect("cf lichess"),
            cf_personal: self.inner.cf_handle("personal").expect("cf personal"),
            cf_game: self.inner.cf_handle("game").expect("cf game"),
            cf_player: self.inner.cf_handle("player").expect("cf player"),
            cf_masters: self.inner.cf_handle("masters").expect("cf masters"),
            cf_masters_game: self
                .inner
                .cf_handle("masters_game")
                .expect("cf masters_game"),
        }
    }
}

pub struct QueryableDatabase<'a> {
    db: &'a DB,
    cf_lichess: &'a ColumnFamily,
    cf_personal: &'a ColumnFamily,
    cf_game: &'a ColumnFamily,
    cf_player: &'a ColumnFamily,
    cf_masters: &'a ColumnFamily,
    cf_masters_game: &'a ColumnFamily,
}

impl QueryableDatabase<'_> {
    pub fn db_property(&self, name: &str) -> Result<Option<String>, rocksdb::Error> {
        self.db.property_value(name)
    }

    pub fn game_property(&self, name: &str) -> Result<Option<String>, rocksdb::Error> {
        self.db.property_value_cf(self.cf_game, name)
    }

    pub fn personal_property(&self, name: &str) -> Result<Option<String>, rocksdb::Error> {
        self.db.property_value_cf(self.cf_personal, name)
    }

    pub fn get_game_info(&self, id: GameId) -> Result<Option<GameInfo>, rocksdb::Error> {
        Ok(self.db.get_cf(self.cf_game, id.to_bytes())?.map(|buf| {
            let mut cursor = Cursor::new(buf);
            GameInfo::read(&mut cursor).expect("deserialize game info")
        }))
    }

    pub fn get_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
            .db
            .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 get_personal(
        &self,
        key: &KeyPrefix,
        since: Month,
        until: Month,
    ) -> Result<PersonalEntry, 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
            .db
            .iterator_cf_opt(self.cf_personal, opt, IteratorMode::Start);

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

        Ok(entry)
    }

    pub fn get_player_status(&self, id: &UserId) -> Result<Option<PersonalStatus>, rocksdb::Error> {
        Ok(self.db.get_cf(self.cf_player, id.as_str())?.map(|buf| {
            let mut cursor = Cursor::new(buf);
            PersonalStatus::read(&mut cursor).expect("deserialize status")
        }))
    }

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

    pub fn has_masters_game(&self, id: GameId) -> Result<bool, rocksdb::Error> {
        self.db
            .get_cf(self.cf_masters_game, id.to_bytes())
            .map(|maybe_entry| maybe_entry.is_some())
    }

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

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

    pub fn get_masters(
        &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
            .db
            .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) -> Batch<'_> {
        Batch {
            queryable: self,
            batch: WriteBatch::default(),
        }
    }
}

pub struct Batch<'a> {
    queryable: &'a QueryableDatabase<'a>,
    batch: WriteBatch,
}

impl Batch<'_> {
    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.queryable.cf_lichess,
            key.into_bytes(),
            cursor.into_inner(),
        );
    }

    pub fn merge_personal(&mut self, key: Key, entry: PersonalEntry) {
        let mut cursor = Cursor::new(Vec::with_capacity(PersonalEntry::SIZE_HINT));
        entry.write(&mut cursor).expect("serialize personal entry");
        self.batch.merge_cf(
            self.queryable.cf_personal,
            key.into_bytes(),
            cursor.into_inner(),
        );
    }

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

    pub fn merge_masters(&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.queryable.cf_masters,
            key.into_bytes(),
            cursor.into_inner(),
        );
    }

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

    pub fn write(self) -> Result<(), rocksdb::Error> {
        self.queryable.db.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 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<GameInfo> = 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 = GameInfo::read(&mut cursor).expect("read for game merge");
        if let Some(old_info) = info {
            new_info.indexed_personal.white |= old_info.indexed_personal.white;
            new_info.indexed_personal.black |= old_info.indexed_personal.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 game");
        cursor.into_inner()
    })
}

fn personal_merge(
    _key: &[u8],
    existing: Option<&[u8]>,
    operands: &mut MergeOperands,
) -> Option<Vec<u8>> {
    let mut entry = PersonalEntry::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 personal merge");
        size_hint += op.len();
    }
    let mut cursor = Cursor::new(Vec::with_capacity(size_hint));
    entry.write(&mut cursor).expect("write personal 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 operator only used to satisfy type checker")
}
<commit_msg>larger block sizes<commit_after>use std::{io::Cursor, path::Path};

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

use crate::model::{
    GameId, GameInfo, Key, KeyPrefix, LichessEntry, MastersEntry, MastersGame, Month,
    PersonalEntry, PersonalStatus, UserId, Year,
};

#[derive(Debug)]
pub struct Database {
    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 = DBWithThreadMode::open_cf_descriptors(
            &db_opts,
            path,
            vec![
                column_family(
                    "lichess",
                    Some("lichess merge"),
                    lichess_merge,
                    Some(KeyPrefix::SIZE),
                    8 * 1024,
                    3,
                ),
                column_family(
                    "personal",
                    Some("personal merge"),
                    personal_merge,
                    Some(KeyPrefix::SIZE),
                    8 * 1024,
                    3,
                ),
                column_family("game", Some("game merge"), game_merge, None, 4 * 1024, 0),
                column_family("player", None, void_merge, None, 4 * 1024, 0),
                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),
            ],
        )?;

        Ok(Database { inner })
    }

    pub fn queryable(&self) -> QueryableDatabase<'_> {
        QueryableDatabase {
            db: &self.inner,
            cf_lichess: self.inner.cf_handle("lichess").expect("cf lichess"),
            cf_personal: self.inner.cf_handle("personal").expect("cf personal"),
            cf_game: self.inner.cf_handle("game").expect("cf game"),
            cf_player: self.inner.cf_handle("player").expect("cf player"),
            cf_masters: self.inner.cf_handle("masters").expect("cf masters"),
            cf_masters_game: self
                .inner
                .cf_handle("masters_game")
                .expect("cf masters_game"),
        }
    }
}

pub struct QueryableDatabase<'a> {
    db: &'a DB,
    cf_lichess: &'a ColumnFamily,
    cf_personal: &'a ColumnFamily,
    cf_game: &'a ColumnFamily,
    cf_player: &'a ColumnFamily,
    cf_masters: &'a ColumnFamily,
    cf_masters_game: &'a ColumnFamily,
}

impl QueryableDatabase<'_> {
    pub fn db_property(&self, name: &str) -> Result<Option<String>, rocksdb::Error> {
        self.db.property_value(name)
    }

    pub fn game_property(&self, name: &str) -> Result<Option<String>, rocksdb::Error> {
        self.db.property_value_cf(self.cf_game, name)
    }

    pub fn personal_property(&self, name: &str) -> Result<Option<String>, rocksdb::Error> {
        self.db.property_value_cf(self.cf_personal, name)
    }

    pub fn get_game_info(&self, id: GameId) -> Result<Option<GameInfo>, rocksdb::Error> {
        Ok(self.db.get_cf(self.cf_game, id.to_bytes())?.map(|buf| {
            let mut cursor = Cursor::new(buf);
            GameInfo::read(&mut cursor).expect("deserialize game info")
        }))
    }

    pub fn get_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
            .db
            .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 get_personal(
        &self,
        key: &KeyPrefix,
        since: Month,
        until: Month,
    ) -> Result<PersonalEntry, 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
            .db
            .iterator_cf_opt(self.cf_personal, opt, IteratorMode::Start);

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

        Ok(entry)
    }

    pub fn get_player_status(&self, id: &UserId) -> Result<Option<PersonalStatus>, rocksdb::Error> {
        Ok(self.db.get_cf(self.cf_player, id.as_str())?.map(|buf| {
            let mut cursor = Cursor::new(buf);
            PersonalStatus::read(&mut cursor).expect("deserialize status")
        }))
    }

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

    pub fn has_masters_game(&self, id: GameId) -> Result<bool, rocksdb::Error> {
        self.db
            .get_cf(self.cf_masters_game, id.to_bytes())
            .map(|maybe_entry| maybe_entry.is_some())
    }

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

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

    pub fn get_masters(
        &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
            .db
            .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) -> Batch<'_> {
        Batch {
            queryable: self,
            batch: WriteBatch::default(),
        }
    }
}

pub struct Batch<'a> {
    queryable: &'a QueryableDatabase<'a>,
    batch: WriteBatch,
}

impl Batch<'_> {
    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.queryable.cf_lichess,
            key.into_bytes(),
            cursor.into_inner(),
        );
    }

    pub fn merge_personal(&mut self, key: Key, entry: PersonalEntry) {
        let mut cursor = Cursor::new(Vec::with_capacity(PersonalEntry::SIZE_HINT));
        entry.write(&mut cursor).expect("serialize personal entry");
        self.batch.merge_cf(
            self.queryable.cf_personal,
            key.into_bytes(),
            cursor.into_inner(),
        );
    }

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

    pub fn merge_masters(&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.queryable.cf_masters,
            key.into_bytes(),
            cursor.into_inner(),
        );
    }

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

    pub fn write(self) -> Result<(), rocksdb::Error> {
        self.queryable.db.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 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<GameInfo> = 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 = GameInfo::read(&mut cursor).expect("read for game merge");
        if let Some(old_info) = info {
            new_info.indexed_personal.white |= old_info.indexed_personal.white;
            new_info.indexed_personal.black |= old_info.indexed_personal.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 game");
        cursor.into_inner()
    })
}

fn personal_merge(
    _key: &[u8],
    existing: Option<&[u8]>,
    operands: &mut MergeOperands,
) -> Option<Vec<u8>> {
    let mut entry = PersonalEntry::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 personal merge");
        size_hint += op.len();
    }
    let mut cursor = Cursor::new(Vec::with_capacity(size_hint));
    entry.write(&mut cursor).expect("write personal 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 operator only used to satisfy type checker")
}
<|endoftext|>


<commit_before>use std::fmt::Show;
use std::mem;
use std::sync::Arc;

use pg;
use pg::{PostgresConnection, PostgresStatement, PostgresResult};
use pg::types::ToSql;
use r2d2::{mod, LoggingErrorHandler};
use r2d2_postgres::PostgresPoolManager;
use conduit::{Request, Response};
use conduit_middleware::Middleware;

use app::{App, RequestApp};
use util::{CargoResult, LazyCell, internal};

pub type Pool = r2d2::Pool<pg::PostgresConnection,
                           pg::error::PostgresConnectError,
                           PostgresPoolManager,
                           LoggingErrorHandler>;
type PooledConnnection<'a> =
        r2d2::PooledConnection<'a,
                               pg::PostgresConnection,
                               pg::error::PostgresConnectError,
                               PostgresPoolManager,
                               LoggingErrorHandler>;

pub fn pool(url: &str, config: r2d2::Config) -> Pool {
    let mgr = PostgresPoolManager::new(url, pg::NoSsl);
    r2d2::Pool::new(config, mgr, LoggingErrorHandler).unwrap()
}

pub struct TransactionMiddleware;

pub struct Transaction {
    // fields are destructed top-to-bottom so ensure we destroy them in the
    // right order.
    tx: LazyCell<pg::PostgresTransaction<'static>>,
    slot: LazyCell<PooledConnnection<'static>>,

    // Keep a handle to the app which keeps a handle to the database to ensure
    // that this `'static` is indeed at least a little more accurate (in that
    // it's alive for the period of time this `Transaction` is alive.
    app: Arc<App>,
}

impl Transaction {
    pub fn new(app: Arc<App>) -> Transaction {
        Transaction {
            app: app,
            slot: LazyCell::new(),
            tx: LazyCell::new(),
        }
    }

    pub fn conn(&self) -> CargoResult<&PostgresConnection> {
        // Here we want to tie the lifetime of a single connection the lifetime
        // of this request. Currently the lifetime of a connection is tied to
        // the lifetime of the pool from which it came from, which is the
        // mismatch.
        //
        // The unsafety here is frobbing lifetimes to ensure that they work out.
        // Additionally, any extension in a Request needs to live for the static
        // lifetime (yay!).
        //
        // To solve these problems, the private `Transaction` extension stores a
        // reference both to the pool (to keep it alive) as well as a connection
        // transmuted to the 'static lifetime. This allows us to allocate a
        // connection up front and then repeatedly hand it out.
        unsafe {
            if !self.slot.filled() {
                let conn = try!(self.app.database.get().map_err(|()| {
                    internal("failed to get a database connection")
                }));
                let conn: PooledConnnection = conn;
                self.slot.fill(mem::transmute::<_, PooledConnnection<'static>>(conn));
            }
        }
        Ok(&**self.slot.borrow().unwrap())
    }

    fn tx<'a>(&'a self) -> CargoResult<&'a pg::PostgresTransaction<'a>> {
        // Similar to above, the transaction for this request is actually tied
        // to the connection in the request itself, not 'static. We transmute it
        // to static as its paired with the inner connection to achieve the
        // desired effect.
        unsafe {
            if !self.tx.filled() {
                let conn = try!(self.conn());
                let t = try!(conn.transaction());
                let t = mem::transmute::<_, pg::PostgresTransaction<'static>>(t);
                self.tx.fill(t);
            }
        }
        let tx = self.tx.borrow();
        let tx: &'a pg::PostgresTransaction<'static> = tx.unwrap();
        Ok(tx)
    }
}

impl Middleware for TransactionMiddleware {
    fn before(&self, req: &mut Request) -> Result<(), Box<Show + 'static>> {
        if !req.extensions().contains::<Transaction>() {
            let app = req.app().clone();
            req.mut_extensions().insert(Transaction::new(app));
        }
        Ok(())
    }

    fn after(&self, req: &mut Request, res: Result<Response, Box<Show + 'static>>)
             -> Result<Response, Box<Show + 'static>> {
        match res {
            Ok(ref res) if res.status.val0() == 200 => {
                let tx = req.extensions().find::<Transaction>()
                            .expect("Transaction not present in request");
                match tx.tx.borrow() {
                    Some(tx) => {
                        if req.app().config.env != ::Test {
                            tx.set_commit();
                        }
                    }
                    None => {}
                }
            }
            _ => {}
        }
        return res;
    }
}

pub trait RequestTransaction<'a> {
    /// Return the lazily initialized postgres connection for this request.
    ///
    /// The connection will live for the lifetime of the request.
    fn db_conn(self) -> CargoResult<&'a PostgresConnection>;

    /// Return the lazily initialized postgres transaction for this request.
    ///
    /// The transaction will live for the duration of the request, and it will
    /// only be set to commit() if a successful response code of 200 is seen.
    fn tx(self) -> CargoResult<&'a pg::PostgresTransaction<'a>>;
}

impl<'a> RequestTransaction<'a> for &'a Request + 'a {
    fn db_conn(self) -> CargoResult<&'a PostgresConnection> {
        self.extensions().find::<Transaction>()
            .expect("Transaction not present in request")
            .conn()
    }

    fn tx(self) -> CargoResult<&'a pg::PostgresTransaction<'a>> {
        self.extensions().find::<Transaction>()
            .expect("Transaction not present in request")
            .tx()
    }
}

pub trait Connection {
    fn prepare<'a>(&'a self, query: &str) -> PostgresResult<PostgresStatement<'a>>;
    fn execute(&self, query: &str, params: &[&ToSql]) -> PostgresResult<uint>;
}

impl Connection for pg::PostgresConnection {
    fn prepare<'a>(&'a self, query: &str) -> PostgresResult<PostgresStatement<'a>> {
        self.prepare(query)
    }
    fn execute(&self, query: &str, params: &[&ToSql]) -> PostgresResult<uint> {
        self.execute(query, params)
    }
}
//
// impl Connection for pg::pool::PooledPostgresConnection {
//     fn prepare<'a>(&'a self, query: &str) -> PostgresResult<PostgresStatement<'a>> {
//         self.prepare(query)
//     }
//     fn execute(&self, query: &str, params: &[&ToSql]) -> PostgresResult<uint> {
//         self.execute(query, params)
//     }
// }

impl<'a> Connection for pg::PostgresTransaction<'a> {
    fn prepare<'a>(&'a self, query: &str) -> PostgresResult<PostgresStatement<'a>> {
        self.prepare(query)
    }
    fn execute(&self, query: &str, params: &[&ToSql]) -> PostgresResult<uint> {
        self.execute(query, params)
    }
}
<commit_msg>Commit all successful responses<commit_after>use std::fmt::Show;
use std::mem;
use std::sync::Arc;

use pg;
use pg::{PostgresConnection, PostgresStatement, PostgresResult};
use pg::types::ToSql;
use r2d2::{mod, LoggingErrorHandler};
use r2d2_postgres::PostgresPoolManager;
use conduit::{Request, Response};
use conduit_middleware::Middleware;

use app::{App, RequestApp};
use util::{CargoResult, LazyCell, internal};

pub type Pool = r2d2::Pool<pg::PostgresConnection,
                           pg::error::PostgresConnectError,
                           PostgresPoolManager,
                           LoggingErrorHandler>;
type PooledConnnection<'a> =
        r2d2::PooledConnection<'a,
                               pg::PostgresConnection,
                               pg::error::PostgresConnectError,
                               PostgresPoolManager,
                               LoggingErrorHandler>;

pub fn pool(url: &str, config: r2d2::Config) -> Pool {
    let mgr = PostgresPoolManager::new(url, pg::NoSsl);
    r2d2::Pool::new(config, mgr, LoggingErrorHandler).unwrap()
}

pub struct TransactionMiddleware;

pub struct Transaction {
    // fields are destructed top-to-bottom so ensure we destroy them in the
    // right order.
    tx: LazyCell<pg::PostgresTransaction<'static>>,
    slot: LazyCell<PooledConnnection<'static>>,

    // Keep a handle to the app which keeps a handle to the database to ensure
    // that this `'static` is indeed at least a little more accurate (in that
    // it's alive for the period of time this `Transaction` is alive.
    app: Arc<App>,
}

impl Transaction {
    pub fn new(app: Arc<App>) -> Transaction {
        Transaction {
            app: app,
            slot: LazyCell::new(),
            tx: LazyCell::new(),
        }
    }

    pub fn conn(&self) -> CargoResult<&PostgresConnection> {
        // Here we want to tie the lifetime of a single connection the lifetime
        // of this request. Currently the lifetime of a connection is tied to
        // the lifetime of the pool from which it came from, which is the
        // mismatch.
        //
        // The unsafety here is frobbing lifetimes to ensure that they work out.
        // Additionally, any extension in a Request needs to live for the static
        // lifetime (yay!).
        //
        // To solve these problems, the private `Transaction` extension stores a
        // reference both to the pool (to keep it alive) as well as a connection
        // transmuted to the 'static lifetime. This allows us to allocate a
        // connection up front and then repeatedly hand it out.
        unsafe {
            if !self.slot.filled() {
                let conn = try!(self.app.database.get().map_err(|()| {
                    internal("failed to get a database connection")
                }));
                let conn: PooledConnnection = conn;
                self.slot.fill(mem::transmute::<_, PooledConnnection<'static>>(conn));
            }
        }
        Ok(&**self.slot.borrow().unwrap())
    }

    fn tx<'a>(&'a self) -> CargoResult<&'a pg::PostgresTransaction<'a>> {
        // Similar to above, the transaction for this request is actually tied
        // to the connection in the request itself, not 'static. We transmute it
        // to static as its paired with the inner connection to achieve the
        // desired effect.
        unsafe {
            if !self.tx.filled() {
                let conn = try!(self.conn());
                let t = try!(conn.transaction());
                let t = mem::transmute::<_, pg::PostgresTransaction<'static>>(t);
                self.tx.fill(t);
            }
        }
        let tx = self.tx.borrow();
        let tx: &'a pg::PostgresTransaction<'static> = tx.unwrap();
        Ok(tx)
    }
}

impl Middleware for TransactionMiddleware {
    fn before(&self, req: &mut Request) -> Result<(), Box<Show + 'static>> {
        if !req.extensions().contains::<Transaction>() {
            let app = req.app().clone();
            req.mut_extensions().insert(Transaction::new(app));
        }
        Ok(())
    }

    fn after(&self, req: &mut Request, res: Result<Response, Box<Show + 'static>>)
             -> Result<Response, Box<Show + 'static>> {
        if res.is_ok() {
            let tx = req.extensions().find::<Transaction>()
                        .expect("Transaction not present in request");
            match tx.tx.borrow() {
                Some(tx) => {
                    if req.app().config.env != ::Test {
                        tx.set_commit();
                    }
                }
                None => {}
            }
        }
        return res;
    }
}

pub trait RequestTransaction<'a> {
    /// Return the lazily initialized postgres connection for this request.
    ///
    /// The connection will live for the lifetime of the request.
    fn db_conn(self) -> CargoResult<&'a PostgresConnection>;

    /// Return the lazily initialized postgres transaction for this request.
    ///
    /// The transaction will live for the duration of the request, and it will
    /// only be set to commit() if a successful response code of 200 is seen.
    fn tx(self) -> CargoResult<&'a pg::PostgresTransaction<'a>>;
}

impl<'a> RequestTransaction<'a> for &'a Request + 'a {
    fn db_conn(self) -> CargoResult<&'a PostgresConnection> {
        self.extensions().find::<Transaction>()
            .expect("Transaction not present in request")
            .conn()
    }

    fn tx(self) -> CargoResult<&'a pg::PostgresTransaction<'a>> {
        self.extensions().find::<Transaction>()
            .expect("Transaction not present in request")
            .tx()
    }
}

pub trait Connection {
    fn prepare<'a>(&'a self, query: &str) -> PostgresResult<PostgresStatement<'a>>;
    fn execute(&self, query: &str, params: &[&ToSql]) -> PostgresResult<uint>;
}

impl Connection for pg::PostgresConnection {
    fn prepare<'a>(&'a self, query: &str) -> PostgresResult<PostgresStatement<'a>> {
        self.prepare(query)
    }
    fn execute(&self, query: &str, params: &[&ToSql]) -> PostgresResult<uint> {
        self.execute(query, params)
    }
}
//
// impl Connection for pg::pool::PooledPostgresConnection {
//     fn prepare<'a>(&'a self, query: &str) -> PostgresResult<PostgresStatement<'a>> {
//         self.prepare(query)
//     }
//     fn execute(&self, query: &str, params: &[&ToSql]) -> PostgresResult<uint> {
//         self.execute(query, params)
//     }
// }

impl<'a> Connection for pg::PostgresTransaction<'a> {
    fn prepare<'a>(&'a self, query: &str) -> PostgresResult<PostgresStatement<'a>> {
        self.prepare(query)
    }
    fn execute(&self, query: &str, params: &[&ToSql]) -> PostgresResult<uint> {
        self.execute(query, params)
    }
}
<|endoftext|>


<commit_before>use std::str;
use std::cell::RefCell;
use std::time::Instant;
use std::collections::BTreeMap;
use std::net::SocketAddr;
use mio::net::UdpSocket;
use std::io::{self,Write};
use sozu_command::proxy::{FilteredData,MetricsData};

mod network_drain;
mod local_drain;

use self::network_drain::NetworkDrain;
use self::local_drain::LocalDrain;

thread_local! {
  pub static METRICS: RefCell<Aggregator> = RefCell::new(Aggregator::new(String::from("sozu")));
}

#[derive(Debug,Clone,PartialEq)]
pub enum MetricData {
  Gauge(usize),
  GaugeAdd(i64),
  Count(i64),
  Time(usize),
}

impl MetricData {
  fn is_time(&self) -> bool {
    match self {
      &MetricData::Time(_) => true,
      _ => false,
    }
  }

  fn update(&mut self, key: &'static str, m: MetricData) -> bool {
    match (self, m) {
      (&mut MetricData::Gauge(ref mut v1), MetricData::Gauge(v2)) => {
        let changed = *v1 != v2;
        *v1 = v2;
        changed
      },
      (&mut MetricData::Gauge(ref mut v1), MetricData::GaugeAdd(v2)) => {
        debug_assert!(*v1 as i64 + v2 >= 0, "metric {} underflow: previous value: {}, adding: {}", key, v1, v2);
        let changed = v2 != 0;
        let res = *v1 as i64 + v2;
        *v1 = if res >= 0 {
          res as usize
        } else {
          error!("metric {} underflow: previous value: {}, adding: {}", key, v1, v2);
          0
        };

        changed
      },
      (&mut MetricData::Count(ref mut v1), MetricData::Count(v2)) => {
        let changed = v2 != 0;
        *v1 += v2;
        changed
      },
      (s,m) => panic!("tried to update metric {} of value {:?} with an incompatible metric: {:?}", key, s, m)
    }
  }
}

#[derive(Debug,Clone)]
pub struct StoredMetricData {
  last_sent: Instant,
  updated:   bool,
  data:      MetricData,
}

impl StoredMetricData {
  pub fn new(last_sent: Instant, data: MetricData) -> StoredMetricData {
    StoredMetricData {
      last_sent,
      updated: true,
      data: if let MetricData::GaugeAdd(v) = data {
        MetricData::Gauge(v as usize)
      } else { data }
    }
  }

  pub fn update(&mut self, key: &'static str, m: MetricData) {
    let updated = self.data.update(key, m);
    if ! self.updated {
      self.updated = updated;
    }
  }
}

pub fn setup<O: Into<String>>(metrics_host: &SocketAddr, origin: O, use_tagged_metrics: bool, prefix: Option<String>) {
  let metrics_socket = udp_bind();

  debug!("setting up metrics: local address = {:#?}", metrics_socket.local_addr());

  METRICS.with(|metrics| {
    if let Some(p) = prefix {
      (*metrics.borrow_mut()).set_up_prefix(p);
    }
    (*metrics.borrow_mut()).set_up_remote(metrics_socket, metrics_host.clone());
    (*metrics.borrow_mut()).set_up_origin(origin.into());
    (*metrics.borrow_mut()).set_up_tagged_metrics(use_tagged_metrics);
  });
}

pub trait Subscriber {
  fn receive_metric(&mut self, label: &'static str, app_id: Option<&str>, backend_id: Option<&str>, metric: MetricData);
}

pub struct Aggregator {
  prefix:  String,
  network: Option<NetworkDrain>,
  local:   LocalDrain,
}

impl Aggregator {
  pub fn new(prefix: String) -> Aggregator {
    Aggregator {
      prefix: prefix.clone(),
      network: None,
      local: LocalDrain::new(prefix),
    }
  }

  pub fn set_up_prefix(&mut self, prefix: String) {
    self.prefix = prefix;
  }

  pub fn set_up_remote(&mut self, socket: UdpSocket, addr: SocketAddr) {
    self.network = Some(NetworkDrain::new(self.prefix.clone(), socket, addr));
  }

  pub fn set_up_origin(&mut self, origin: String) {
    self.network.as_mut().map(|n| n.origin = origin);
  }

  pub fn set_up_tagged_metrics(&mut self, tagged: bool) {
    self.network.as_mut().map(|n| n.use_tagged_metrics = tagged);
  }

  pub fn socket(&self) -> Option<&UdpSocket> {
    self.network.as_ref().map(|n| &n.remote.get_ref().socket)
  }

  pub fn count_add(&mut self, key: &'static str, count_value: i64) {
    self.receive_metric(key, None, None, MetricData::Count(count_value));
  }

  pub fn set_gauge(&mut self, key: &'static str, gauge_value: usize) {
    self.receive_metric(key, None, None, MetricData::Gauge(gauge_value));
  }

  pub fn gauge_add(&mut self, key: &'static str, gauge_value: i64) {
    self.receive_metric(key, None, None, MetricData::GaugeAdd(gauge_value));
  }

  pub fn writable(&mut self) {
    if let Some(ref mut net) = self.network.as_mut() {
      net.writable();
    }
  }

  pub fn send_data(&mut self) {
    if let Some(ref mut net) = self.network.as_mut() {
      net.send_data();
    }
  }

  pub fn dump_metrics_data(&mut self) -> MetricsData {
    self.local.dump_metrics_data()
  }

  pub fn dump_process_data(&mut self) -> BTreeMap<String, FilteredData> {
    self.local.dump_process_data()
  }
}

impl Subscriber for Aggregator {
  fn receive_metric(&mut self, label: &'static str, app_id: Option<&str>, backend_id: Option<&str>, metric: MetricData) {
    if let Some(ref mut net) = self.network.as_mut() {
      net.receive_metric(label, app_id, backend_id, metric.clone());
    }
    self.local.receive_metric(label, app_id, backend_id, metric.clone());
  }
}

#[derive(Debug,Clone,PartialEq)]
pub struct MetricLine {
  label:      &'static str,
  app_id:     Option<String>,
  backend_id: Option<String>,
  /// in milliseconds
  duration:   usize,
}

pub struct MetricSocket {
  pub addr:   SocketAddr,
  pub socket: UdpSocket,
}

impl Write for MetricSocket {
  fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
    self.socket.send_to(buf, &self.addr)
  }

  fn flush(&mut self) -> io::Result<()> {
    Ok(())
  }
}

pub fn udp_bind() -> UdpSocket {
  UdpSocket::bind(&("0.0.0.0:0".parse().unwrap())).expect("could not parse address")
}

#[macro_export]
macro_rules! count (
  ($key:expr, $value: expr) => ({
    let v = $value;
    $crate::metrics::METRICS.with(|metrics| {
      (*metrics.borrow_mut()).count_add($key, v);
    });
  })
);

#[macro_export]
macro_rules! incr (
  ($key:expr) => (count!($key, 1););
  ($key:expr, $app_id:expr, $backend_id:expr) => {
    use $crate::metrics::Subscriber;

    $crate::metrics::METRICS.with(|metrics| {
      (*metrics.borrow_mut()).receive_metric($key, $app_id, $backend_id, $crate::metrics::MetricData::Count(1));
    });
  }
);

#[macro_export]
macro_rules! decr (
  ($key:expr) => (count!($key, -1);)
);

#[macro_export]
macro_rules! gauge (
  ($key:expr, $value: expr) => ({
    let v = $value;
    $crate::metrics::METRICS.with(|metrics| {
      (*metrics.borrow_mut()).set_gauge($key, v);
    });
  })
);

#[macro_export]
macro_rules! gauge_add (
  ($key:expr, $value: expr) => ({
    let v = $value;
    $crate::metrics::METRICS.with(|metrics| {
      (*metrics.borrow_mut()).gauge_add($key, v);
    });
  });
  ($key:expr, $value:expr, $app_id:expr, $backend_id:expr) => {
    use $crate::metrics::Subscriber;
    let v = $value;

    $crate::metrics::METRICS.with(|metrics| {
      (*metrics.borrow_mut()).receive_metric($key, $app_id, $backend_id, $crate::metrics::MetricData::GaugeAdd(v));
    });
  }
);

#[macro_export]
macro_rules! time (
  ($key:expr, $value: expr) => {
    use $crate::metrics::{MetricData,Subscriber};
    let v = $value;
    $crate::metrics::METRICS.with(|metrics| {
      let ref mut m = *metrics.borrow_mut();

      m.receive_metric($key, None, None, MetricData::Time(v as usize));
    });
  };
  ($key:expr, $app_id:expr, $value: expr) => {
    use $crate::metrics::{MetricData,Subscriber};
    let v = $value;
    $crate::metrics::METRICS.with(|metrics| {
      let ref mut m = *metrics.borrow_mut();
      let app: &str = $app_id;

      m.receive_metric($key, Some(app), None, MetricData::Time(v as usize));
    });
  }
);

#[macro_export]
macro_rules! record_backend_metrics (
  ($app_id:expr, $backend_id:expr, $response_time: expr, $backend_connection_time: expr, $bin: expr, $bout: expr) => {
    use $crate::metrics::{MetricData,Subscriber};
    $crate::metrics::METRICS.with(|metrics| {
      let ref mut m = *metrics.borrow_mut();
      let app_id: &str = $app_id.as_str();
      let backend_id: &str = $backend_id;

      m.receive_metric("bytes_in", Some(app_id), Some(backend_id), MetricData::Count($bin as i64));
      m.receive_metric("bytes_out", Some(app_id), Some(backend_id), MetricData::Count($bout as i64));
      m.receive_metric("response_time", Some(app_id), Some(backend_id), MetricData::Time($response_time as usize));
      if let Some(t) = $backend_connection_time {
        m.receive_metric("connection_time", Some(app_id), Some(backend_id), MetricData::Time(t.num_milliseconds() as usize));
      }
    });
  }
);

<commit_msg>fix conflicts in metrics macro<commit_after>use std::str;
use std::cell::RefCell;
use std::time::Instant;
use std::collections::BTreeMap;
use std::net::SocketAddr;
use mio::net::UdpSocket;
use std::io::{self,Write};
use sozu_command::proxy::{FilteredData,MetricsData};

mod network_drain;
mod local_drain;

use self::network_drain::NetworkDrain;
use self::local_drain::LocalDrain;

thread_local! {
  pub static METRICS: RefCell<Aggregator> = RefCell::new(Aggregator::new(String::from("sozu")));
}

#[derive(Debug,Clone,PartialEq)]
pub enum MetricData {
  Gauge(usize),
  GaugeAdd(i64),
  Count(i64),
  Time(usize),
}

impl MetricData {
  fn is_time(&self) -> bool {
    match self {
      &MetricData::Time(_) => true,
      _ => false,
    }
  }

  fn update(&mut self, key: &'static str, m: MetricData) -> bool {
    match (self, m) {
      (&mut MetricData::Gauge(ref mut v1), MetricData::Gauge(v2)) => {
        let changed = *v1 != v2;
        *v1 = v2;
        changed
      },
      (&mut MetricData::Gauge(ref mut v1), MetricData::GaugeAdd(v2)) => {
        debug_assert!(*v1 as i64 + v2 >= 0, "metric {} underflow: previous value: {}, adding: {}", key, v1, v2);
        let changed = v2 != 0;
        let res = *v1 as i64 + v2;
        *v1 = if res >= 0 {
          res as usize
        } else {
          error!("metric {} underflow: previous value: {}, adding: {}", key, v1, v2);
          0
        };

        changed
      },
      (&mut MetricData::Count(ref mut v1), MetricData::Count(v2)) => {
        let changed = v2 != 0;
        *v1 += v2;
        changed
      },
      (s,m) => panic!("tried to update metric {} of value {:?} with an incompatible metric: {:?}", key, s, m)
    }
  }
}

#[derive(Debug,Clone)]
pub struct StoredMetricData {
  last_sent: Instant,
  updated:   bool,
  data:      MetricData,
}

impl StoredMetricData {
  pub fn new(last_sent: Instant, data: MetricData) -> StoredMetricData {
    StoredMetricData {
      last_sent,
      updated: true,
      data: if let MetricData::GaugeAdd(v) = data {
        MetricData::Gauge(v as usize)
      } else { data }
    }
  }

  pub fn update(&mut self, key: &'static str, m: MetricData) {
    let updated = self.data.update(key, m);
    if ! self.updated {
      self.updated = updated;
    }
  }
}

pub fn setup<O: Into<String>>(metrics_host: &SocketAddr, origin: O, use_tagged_metrics: bool, prefix: Option<String>) {
  let metrics_socket = udp_bind();

  debug!("setting up metrics: local address = {:#?}", metrics_socket.local_addr());

  METRICS.with(|metrics| {
    if let Some(p) = prefix {
      (*metrics.borrow_mut()).set_up_prefix(p);
    }
    (*metrics.borrow_mut()).set_up_remote(metrics_socket, metrics_host.clone());
    (*metrics.borrow_mut()).set_up_origin(origin.into());
    (*metrics.borrow_mut()).set_up_tagged_metrics(use_tagged_metrics);
  });
}

pub trait Subscriber {
  fn receive_metric(&mut self, label: &'static str, app_id: Option<&str>, backend_id: Option<&str>, metric: MetricData);
}

pub struct Aggregator {
  prefix:  String,
  network: Option<NetworkDrain>,
  local:   LocalDrain,
}

impl Aggregator {
  pub fn new(prefix: String) -> Aggregator {
    Aggregator {
      prefix: prefix.clone(),
      network: None,
      local: LocalDrain::new(prefix),
    }
  }

  pub fn set_up_prefix(&mut self, prefix: String) {
    self.prefix = prefix;
  }

  pub fn set_up_remote(&mut self, socket: UdpSocket, addr: SocketAddr) {
    self.network = Some(NetworkDrain::new(self.prefix.clone(), socket, addr));
  }

  pub fn set_up_origin(&mut self, origin: String) {
    self.network.as_mut().map(|n| n.origin = origin);
  }

  pub fn set_up_tagged_metrics(&mut self, tagged: bool) {
    self.network.as_mut().map(|n| n.use_tagged_metrics = tagged);
  }

  pub fn socket(&self) -> Option<&UdpSocket> {
    self.network.as_ref().map(|n| &n.remote.get_ref().socket)
  }

  pub fn count_add(&mut self, key: &'static str, count_value: i64) {
    self.receive_metric(key, None, None, MetricData::Count(count_value));
  }

  pub fn set_gauge(&mut self, key: &'static str, gauge_value: usize) {
    self.receive_metric(key, None, None, MetricData::Gauge(gauge_value));
  }

  pub fn gauge_add(&mut self, key: &'static str, gauge_value: i64) {
    self.receive_metric(key, None, None, MetricData::GaugeAdd(gauge_value));
  }

  pub fn writable(&mut self) {
    if let Some(ref mut net) = self.network.as_mut() {
      net.writable();
    }
  }

  pub fn send_data(&mut self) {
    if let Some(ref mut net) = self.network.as_mut() {
      net.send_data();
    }
  }

  pub fn dump_metrics_data(&mut self) -> MetricsData {
    self.local.dump_metrics_data()
  }

  pub fn dump_process_data(&mut self) -> BTreeMap<String, FilteredData> {
    self.local.dump_process_data()
  }
}

impl Subscriber for Aggregator {
  fn receive_metric(&mut self, label: &'static str, app_id: Option<&str>, backend_id: Option<&str>, metric: MetricData) {
    if let Some(ref mut net) = self.network.as_mut() {
      net.receive_metric(label, app_id, backend_id, metric.clone());
    }
    self.local.receive_metric(label, app_id, backend_id, metric.clone());
  }
}

#[derive(Debug,Clone,PartialEq)]
pub struct MetricLine {
  label:      &'static str,
  app_id:     Option<String>,
  backend_id: Option<String>,
  /// in milliseconds
  duration:   usize,
}

pub struct MetricSocket {
  pub addr:   SocketAddr,
  pub socket: UdpSocket,
}

impl Write for MetricSocket {
  fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
    self.socket.send_to(buf, &self.addr)
  }

  fn flush(&mut self) -> io::Result<()> {
    Ok(())
  }
}

pub fn udp_bind() -> UdpSocket {
  UdpSocket::bind(&("0.0.0.0:0".parse().unwrap())).expect("could not parse address")
}

#[macro_export]
macro_rules! count (
  ($key:expr, $value: expr) => ({
    let v = $value;
    $crate::metrics::METRICS.with(|metrics| {
      (*metrics.borrow_mut()).count_add($key, v);
    });
  })
);

#[macro_export]
macro_rules! incr (
  ($key:expr) => (count!($key, 1););
  ($key:expr, $app_id:expr, $backend_id:expr) => {
    use $crate::metrics::Subscriber;

    $crate::metrics::METRICS.with(|metrics| {
      (*metrics.borrow_mut()).receive_metric($key, $app_id, $backend_id, $crate::metrics::MetricData::Count(1));
    });
  }
);

#[macro_export]
macro_rules! decr (
  ($key:expr) => (count!($key, -1);)
);

#[macro_export]
macro_rules! gauge (
  ($key:expr, $value: expr) => ({
    let v = $value;
    $crate::metrics::METRICS.with(|metrics| {
      (*metrics.borrow_mut()).set_gauge($key, v);
    });
  })
);

#[macro_export]
macro_rules! gauge_add (
  ($key:expr, $value: expr) => ({
    let v = $value;
    $crate::metrics::METRICS.with(|metrics| {
      (*metrics.borrow_mut()).gauge_add($key, v);
    });
  });
  ($key:expr, $value:expr, $app_id:expr, $backend_id:expr) => {
    use $crate::metrics::Subscriber;
    let v = $value;

    $crate::metrics::METRICS.with(|metrics| {
      (*metrics.borrow_mut()).receive_metric($key, $app_id, $backend_id, $crate::metrics::MetricData::GaugeAdd(v));
    });
  }
);

#[macro_export]
macro_rules! time (
  ($key:expr, $value: expr) => {
    use $crate::metrics::{MetricData,Subscriber};
    let v = $value;
    $crate::metrics::METRICS.with(|metrics| {
      let ref mut m = *metrics.borrow_mut();

      m.receive_metric($key, None, None, MetricData::Time(v as usize));
    });
  };
  ($key:expr, $app_id:expr, $value: expr) => ({
    use $crate::metrics::{MetricData,Subscriber};
    let v = $value;
    $crate::metrics::METRICS.with(|metrics| {
      let ref mut m = *metrics.borrow_mut();
      let app: &str = $app_id;

      m.receive_metric($key, Some(app), None, MetricData::Time(v as usize));
    });
  })
);

#[macro_export]
macro_rules! record_backend_metrics (
  ($app_id:expr, $backend_id:expr, $response_time: expr, $backend_connection_time: expr, $bin: expr, $bout: expr) => {
    use $crate::metrics::{MetricData,Subscriber};
    $crate::metrics::METRICS.with(|metrics| {
      let ref mut m = *metrics.borrow_mut();
      let app_id: &str = $app_id.as_str();
      let backend_id: &str = $backend_id;

      m.receive_metric("bytes_in", Some(app_id), Some(backend_id), MetricData::Count($bin as i64));
      m.receive_metric("bytes_out", Some(app_id), Some(backend_id), MetricData::Count($bout as i64));
      m.receive_metric("response_time", Some(app_id), Some(backend_id), MetricData::Time($response_time as usize));
      if let Some(t) = $backend_connection_time {
        m.receive_metric("connection_time", Some(app_id), Some(backend_id), MetricData::Time(t.num_milliseconds() as usize));
      }
    });
  }
);

<|endoftext|>


<commit_before><commit_msg>Add more benchmarks for rand_xoshiro RNGs<commit_after><|endoftext|>


<commit_before>extern crate tajo;

use std::{i8,i16};
use std::rc::Rc;
use std::cell::RefCell;

use tajo::schema::*;
use tajo::eval::{Eval, MapEval};
use tajo::eval::interpreter::*;
use tajo::rows::*;
use tajo::types::*;

pub fn make_test_schema() -> Schema {
  let mut columns = Vec::new();
  columns.push(Column::new("c0".to_string(), Ty::Bool));      // 0
  columns.push(Column::new("c1".to_string(), Ty::Int1));      // 1
  columns.push(Column::new("c2".to_string(), Ty::Int2));      // 2
  columns.push(Column::new("c3".to_string(), Ty::Int4));      // 3
  columns.push(Column::new("c4".to_string(), Ty::Int8));      // 4
  columns.push(Column::new("c5".to_string(), Ty::Float4));    // 5
  columns.push(Column::new("c6".to_string(), Ty::Float8));    // 6
  columns.push(Column::new("c7".to_string(), Ty::Date));      // 7
  columns.push(Column::new("c8".to_string(), Ty::Time));      // 8
  columns.push(Column::new("c9".to_string(), Ty::Timestamp)); // 9

  columns.push(Column::new("c10".to_string(), Ty::Text));     // 10
  // columns.push(Column::new_with_len("c12".to_string(), Ty::VARCHAR, 4));
  // columns.push(Column::new_with_len("c13".to_string(), Ty::CHAR, 4));

  Schema::new(columns)
}

pub fn fill_vector_block(rowblock: Rc<RefCell<RowBlock>>) {
  for i in (0..1024) {
    rowblock.put_int1(1, i, (i % (i8::MAX - 1) as usize) as i8);
    rowblock.put_int2(2, i, (i % (i16::MAX -1) as usize) as i16);
    rowblock.put_int4(3, i, i as i32);
    rowblock.put_int8(4, i, i as i64);
    rowblock.put_float4(5, i, i as f32);
    rowblock.put_float8(6, i, i as f64);
    rowblock.put_date(7, i, i as i32);
    rowblock.put_time(8, i, i as i64);
    rowblock.put_timestamp(9, i, i as i64);

    rowblock.put_text(10, i, &"Rust");
  }
}

#[test]
fn test_field() {
  let schema = make_test_schema();

  let mut field: Box<MapEval> = Box::new(Field::new(&Column::new("c3".to_string(), Ty::Int4)));
  let mut r: Rc<RefCell<RowBlock>> = Rc::new(RefCell::new(AllocatedVecRowBlock::new(schema.clone())));
  fill_vector_block(&mut *r);

  assert!(field.bind(&schema).is_ok());
  field.eval(&*r);
}<commit_msg>Fixed test APIs.<commit_after>extern crate tajo;

use std::{i8,i16};
use std::rc::Rc;
use std::cell::RefCell;

use tajo::schema::*;
use tajo::eval::{Eval, MapEval};
use tajo::eval::interpreter::*;
use tajo::rows::*;
use tajo::types::*;

pub fn make_test_schema() -> Schema {
  let mut columns = Vec::new();
  columns.push(Column::new("c0".to_string(), Ty::Bool));      // 0
  columns.push(Column::new("c1".to_string(), Ty::Int1));      // 1
  columns.push(Column::new("c2".to_string(), Ty::Int2));      // 2
  columns.push(Column::new("c3".to_string(), Ty::Int4));      // 3
  columns.push(Column::new("c4".to_string(), Ty::Int8));      // 4
  columns.push(Column::new("c5".to_string(), Ty::Float4));    // 5
  columns.push(Column::new("c6".to_string(), Ty::Float8));    // 6
  columns.push(Column::new("c7".to_string(), Ty::Date));      // 7
  columns.push(Column::new("c8".to_string(), Ty::Time));      // 8
  columns.push(Column::new("c9".to_string(), Ty::Timestamp)); // 9

  columns.push(Column::new("c10".to_string(), Ty::Text));     // 10
  // columns.push(Column::new_with_len("c12".to_string(), Ty::VARCHAR, 4));
  // columns.push(Column::new_with_len("c13".to_string(), Ty::CHAR, 4));

  Schema::new(columns)
}

pub fn fill_vector_block(rowblock: &mut RowBlock) {
  for i in (0..1024) {
    rowblock.put_int1(1, i, (i % (i8::MAX - 1) as usize) as i8);
    rowblock.put_int2(2, i, (i % (i16::MAX -1) as usize) as i16);
    rowblock.put_int4(3, i, i as i32);
    rowblock.put_int8(4, i, i as i64);
    rowblock.put_float4(5, i, i as f32);
    rowblock.put_float8(6, i, i as f64);
    rowblock.put_date(7, i, i as i32);
    rowblock.put_time(8, i, i as i64);
    rowblock.put_timestamp(9, i, i as i64);

    rowblock.put_text(10, i, &"Rust");
  }
}

#[test]
fn test_field() {
  let schema = make_test_schema();

  let mut field: Box<MapEval> = Box::new(Field::new(&Column::new("c3".to_string(), Ty::Int4)));
  let mut r: Box<RowBlock> = Box::new(AllocatedVecRowBlock::new(schema.clone()));
  fill_vector_block(&mut *r);

  assert!(field.bind(&schema).is_ok());
  field.eval(&*r);
}<|endoftext|>


<commit_before>use tensor::Tensor;
use std::ops::{Add, Sub, Mul, Div, Rem, BitAnd, BitOr, BitXor};

/*
macro_rules! add_impl {
    ($trait_name:ident, $fname:ident, $($t:ty)*) => ($(
        // mut T <op> T
        impl $trait_name<Tensor<$t>> for Tensor<$t> {
            type Output = Tensor<$t>;
            fn $fname(mut self, rhs: Tensor<$t>) -> Tensor<$t> {
                if rhs.is_scalar() {
                    let v = rhs[0];
                    for i in 0..self.size() {
                        self.data[i] = self.data[i].$fname(v);
                    }
                    self
                } else if self.is_scalar() {
                    let mut t: Tensor<$t> = Tensor::empty(&rhs.shape);
                    let v = self[0];
                    for i in 0..t.size() {
                        t.data[i] = v.$fname(rhs.data[i]);
                    }
                    t
                } else {
                    assert_eq!(self.shape, rhs.shape);
                    for i in 0..self.size() {
                        self.data[i] = self.data[i].$fname(rhs.data[i]);
                    }
                    self
                }
            }
        }

        // mut T <op> &T
        impl<'a> $trait_name<&'a Tensor<$t>> for Tensor<$t> {
            type Output = Tensor<$t>;
            fn $fname(mut self, rhs: &Tensor<$t>) -> Tensor<$t> {
                //println!("$fname T + &T");
                if rhs.is_scalar() {
                    let v = rhs[0];
                    for i in 0..self.size() {
                        self.data[i] = self.data[i].$fname(v);
                    }
                    self
                } else if self.is_scalar() {
                    let mut t: Tensor<$t> = Tensor::empty(&rhs.shape);
                    let v = self[0];
                    for i in 0..t.size() {
                        t.data[i] = v.$fname(rhs.data[i]);
                    }
                    t
                } else {
                    assert_eq!(self.shape, rhs.shape);
                    for i in 0..self.size() {
                        self.data[i] = self.data[i].$fname(rhs.data[i]);
                    }
                    self
                }
            }
        }

        // &T <op> &T
        impl<'a> $trait_name for &'a Tensor<$t> {
            type Output = Tensor<$t>;
            fn $fname(self, rhs: &Tensor<$t>) -> Tensor<$t> {
                //println!("$fname &T + &T");
                if rhs.is_scalar() {
                    let mut t = self.clone();
                    let v = rhs[0];
                    for i in 0..t.size() {
                        t.data[i] = t.data[i].$fname(v);
                    }
                    t
                } else if self.is_scalar() {
                    let mut t = Tensor::empty(&rhs.shape);
                    let v = self[0];
                    for i in 0..t.size() {
                        t.data[i] = v.$fname(rhs.data[i]);
                    }
                    t
                } else {
                    assert_eq!(self.shape, rhs.shape);
                    let mut t = self.clone();
                    for i in 0..t.size() {
                        t.data[i] = t.data[i].$fname(rhs.data[i]);
                    }
                    t
                }
            }
        }

        // T <op> S
        impl $trait_name<$t> for Tensor<$t> {
            type Output = Tensor<$t>;
            fn $fname(mut self, rhs: $t) -> Tensor<$t> {
                for i in 0..self.size() {
                    self.data[i] = self.data[i].$fname(rhs);
                }
                self
            }
        }
    )*)
}

// TODO: It would be better if these relied only on the traits, so that custom-made tensort types
// could be used.
//add_impl! { Add, add, usize u8 u16 u32 u64 isize i8 i16 i32 i64 f32 f64 }
add_impl! { Sub, sub, usize u8 u16 u32 u64 isize i8 i16 i32 i64 f32 f64 }
add_impl! { Mul, mul, usize u8 u16 u32 u64 isize i8 i16 i32 i64 f32 f64 }
add_impl! { Div, div, usize u8 u16 u32 u64 isize i8 i16 i32 i64 f32 f64 }
add_impl! { Rem, rem, usize u8 u16 u32 u64 isize i8 i16 i32 i64 }

add_impl! { BitAnd, bitand, usize u8 u16 u32 u64 isize i8 i16 i32 i64 bool }
add_impl! { BitOr, bitor, usize u8 u16 u32 u64 isize i8 i16 i32 i64 bool }
add_impl! { BitXor, bitxor, usize u8 u16 u32 u64 isize i8 i16 i32 i64 bool }
*/

macro_rules! add_impl {
    ($trait_name:ident, $func_name:ident) => (
        // T <op> T
        impl<T: Copy + $trait_name<Output=T>> $trait_name for Tensor<T> {
            type Output = Tensor<T>;
            fn $func_name(mut self, rhs: Self::Output) -> Self::Output {
                if rhs.is_scalar() {
                    let v: T = rhs.scalar_value() as T;
                    for i in 0..self.size() {
                        self.data[i] = self.data[i].$func_name(v);
                    }
                    self
                } else if self.is_scalar() {
                    let mut t = Tensor::empty(&rhs.shape);
                    let v = self.scalar_value();
                    for i in 0..t.size() {
                        t.data[i] = v.$func_name(rhs.data[i]);
                    }
                    t
                } else {
                    assert_eq!(self.shape, rhs.shape);
                    for i in 0..self.size() {
                        self.data[i] = self.data[i].$func_name(rhs.data[i]);
                    }
                    self
                }
            }
        }
        // T <op> &T
        impl<'a, T: Copy + $trait_name<Output=T>> $trait_name<&'a Tensor<T>> for Tensor<T> {
            type Output = Tensor<T>;
            fn $func_name(mut self, rhs: &Self::Output) -> Self::Output {
                //println!("add T + &T");
                if rhs.is_scalar() {
                    let v = rhs[0];
                    for i in 0..self.size() {
                        self.data[i] = self.data[i].$func_name(v);
                    }
                    self
                } else if self.is_scalar() {
                    let mut t: Tensor<T> = Tensor::empty(&rhs.shape);
                    let v = self[0];
                    for i in 0..t.size() {
                        t.data[i] = v.$func_name(rhs.data[i]);
                    }
                    t
                } else {
                    assert_eq!(self.shape, rhs.shape);
                    for i in 0..self.size() {
                        self.data[i] = self.data[i].$func_name(rhs.data[i]);
                    }
                    self
                }
            }
        }

        // &T <op> &T
        impl<'a, T: Copy + $trait_name<Output=T>> $trait_name<&'a Tensor<T>> for &'a Tensor<T> {
            type Output = Tensor<T>;
            fn $func_name(self, rhs: &Self::Output) -> Self::Output {
                //println!("$fname &T + &T");
                if rhs.is_scalar() {
                    let mut t = self.clone();
                    let v = rhs[0];
                    for i in 0..t.size() {
                        t.data[i] = t.data[i].$func_name(v);
                    }
                    t
                } else if self.is_scalar() {
                    let mut t = Tensor::empty(&rhs.shape);
                    let v = self[0];
                    for i in 0..t.size() {
                        t.data[i] = v.$func_name(rhs.data[i]);
                    }
                    t
                } else {
                    assert_eq!(self.shape, rhs.shape);
                    let mut t = self.clone();
                    for i in 0..t.size() {
                        t.data[i] = t.data[i].$func_name(rhs.data[i]);
                    }
                    t
                }
            }
        }

        // T <op> S
        impl<T: Copy + $trait_name<Output=T>> $trait_name<T> for Tensor<T> {
            type Output = Tensor<T>;
            fn $func_name(mut self, rhs: T) -> Self::Output {
                for i in 0..self.size() {
                    self.data[i] = self.data[i].$func_name(rhs);
                }
                self
            }
        }
    )
}

// Any operation supported on T should be supported on Tensor<T>, as long as T supports Copy
add_impl!(Add, add);
add_impl!(Sub, sub);
add_impl!(Mul, mul);
add_impl!(Div, div);
add_impl!(Rem, rem);

add_impl!(BitAnd, bitand);
add_impl!(BitOr, bitor);
add_impl!(BitXor, bitxor);

<commit_msg>Added Neg trait<commit_after>use tensor::Tensor;
use std::ops::{Add, Sub, Mul, Div, Rem, Neg, BitAnd, BitOr, BitXor};

macro_rules! add_impl {
    ($trait_name:ident, $func_name:ident) => (
        // T <op> T
        impl<T: Copy + $trait_name<Output=T>> $trait_name for Tensor<T> {
            type Output = Tensor<T>;
            fn $func_name(mut self, rhs: Self::Output) -> Self::Output {
                if rhs.is_scalar() {
                    let v: T = rhs.scalar_value() as T;
                    for i in 0..self.size() {
                        self.data[i] = self.data[i].$func_name(v);
                    }
                    self
                } else if self.is_scalar() {
                    let mut t = Tensor::empty(&rhs.shape);
                    let v = self.scalar_value();
                    for i in 0..t.size() {
                        t.data[i] = v.$func_name(rhs.data[i]);
                    }
                    t
                } else {
                    assert_eq!(self.shape, rhs.shape);
                    for i in 0..self.size() {
                        self.data[i] = self.data[i].$func_name(rhs.data[i]);
                    }
                    self
                }
            }
        }
        // T <op> &T
        impl<'a, T: Copy + $trait_name<Output=T>> $trait_name<&'a Tensor<T>> for Tensor<T> {
            type Output = Tensor<T>;
            fn $func_name(mut self, rhs: &Self::Output) -> Self::Output {
                //println!("add T + &T");
                if rhs.is_scalar() {
                    let v = rhs[0];
                    for i in 0..self.size() {
                        self.data[i] = self.data[i].$func_name(v);
                    }
                    self
                } else if self.is_scalar() {
                    let mut t: Tensor<T> = Tensor::empty(&rhs.shape);
                    let v = self[0];
                    for i in 0..t.size() {
                        t.data[i] = v.$func_name(rhs.data[i]);
                    }
                    t
                } else {
                    assert_eq!(self.shape, rhs.shape);
                    for i in 0..self.size() {
                        self.data[i] = self.data[i].$func_name(rhs.data[i]);
                    }
                    self
                }
            }
        }

        // &T <op> &T
        impl<'a, T: Copy + $trait_name<Output=T>> $trait_name<&'a Tensor<T>> for &'a Tensor<T> {
            type Output = Tensor<T>;
            fn $func_name(self, rhs: &Self::Output) -> Self::Output {
                //println!("$fname &T + &T");
                if rhs.is_scalar() {
                    let mut t = self.clone();
                    let v = rhs[0];
                    for i in 0..t.size() {
                        t.data[i] = t.data[i].$func_name(v);
                    }
                    t
                } else if self.is_scalar() {
                    let mut t = Tensor::empty(&rhs.shape);
                    let v = self[0];
                    for i in 0..t.size() {
                        t.data[i] = v.$func_name(rhs.data[i]);
                    }
                    t
                } else {
                    assert_eq!(self.shape, rhs.shape);
                    let mut t = self.clone();
                    for i in 0..t.size() {
                        t.data[i] = t.data[i].$func_name(rhs.data[i]);
                    }
                    t
                }
            }
        }

        // T <op> S
        impl<T: Copy + $trait_name<Output=T>> $trait_name<T> for Tensor<T> {
            type Output = Tensor<T>;
            fn $func_name(mut self, rhs: T) -> Self::Output {
                for i in 0..self.size() {
                    self.data[i] = self.data[i].$func_name(rhs);
                }
                self
            }
        }
    )
}

// Any operation supported on T should be supported on Tensor<T>, as long as T supports Copy
add_impl!(Add, add);
add_impl!(Sub, sub);
add_impl!(Mul, mul);
add_impl!(Div, div);
add_impl!(Rem, rem);

add_impl!(BitAnd, bitand);
add_impl!(BitOr, bitor);
add_impl!(BitXor, bitxor);

// -T
impl<T: Copy + Neg<Output=T>> Neg for Tensor<T> {
    type Output = Tensor<T>;
    fn neg(mut self) -> Self::Output {
        for i in 0..self.size() {
            self.data[i] = -self.data[i];
        }
        self
    }
}

// -&T
impl<'a, T: Copy + Neg<Output=T>> Neg for &'a Tensor<T> {
    type Output = Tensor<T>;
    fn neg(self) -> Self::Output {
        let mut t = Tensor::empty(&self.shape);
        for i in 0..self.size() {
            t.data[i] = -self.data[i];
        }
        t
    }
}
<|endoftext|>


<commit_before>use tensor::Tensor;
use std::ops::{Add, Sub, Mul, Div, Rem};

macro_rules! add_impl {
    ($trait_name:ident, $fname:ident, $($t:ty)*) => ($(
        // mut T <op> T
        impl $trait_name<Tensor<$t>> for Tensor<$t> {
            type Output = Tensor<$t>;
            fn $fname(mut self, rhs: Tensor<$t>) -> Tensor<$t> {
                if rhs.is_scalar() {
                    let v = rhs[0];
                    for i in 0..self.size() {
                        self.data[i] = self.data[i].$fname(v);
                    }
                    self
                } else if self.is_scalar() {
                    let mut t: Tensor<$t> = Tensor::zeros(&rhs.shape);
                    let v = self[0];
                    for i in 0..t.size() {
                        t.data[i] = v.$fname(rhs.data[i]);
                    }
                    t
                } else {
                    assert_eq!(self.shape, rhs.shape);
                    for i in 0..self.size() {
                        self.data[i] = self.data[i].$fname(rhs.data[i]);
                    }
                    self
                }
            }
        }

        // mut T <op> &T
        impl<'a> $trait_name<&'a Tensor<$t>> for Tensor<$t> {
            type Output = Tensor<$t>;
            fn $fname(mut self, rhs: &Tensor<$t>) -> Tensor<$t> {
                //println!("$fname T + &T");
                if rhs.is_scalar() {
                    let v = rhs[0];
                    for i in 0..self.size() {
                        self.data[i] = self.data[i].$fname(v);
                    }
                    self
                } else if self.is_scalar() {
                    let mut t: Tensor<$t> = Tensor::zeros(&rhs.shape);
                    let v = self[0];
                    for i in 0..t.size() {
                        t.data[i] = v.$fname(rhs.data[i]);
                    }
                    t
                } else {
                    assert_eq!(self.shape, rhs.shape);
                    for i in 0..self.size() {
                        self.data[i] = self.data[i].$fname(rhs.data[i]);
                    }
                    self
                }
            }
        }

        // &T <op> &T
        impl<'a> $trait_name for &'a Tensor<$t> {
            type Output = Tensor<$t>;
            fn $fname(self, rhs: &Tensor<$t>) -> Tensor<$t> {
                //println!("$fname &T + &T");
                if rhs.is_scalar() {
                    let mut t: Tensor<$t> = self.clone();
                    let v = rhs[0];
                    for i in 0..t.size() {
                        t.data[i] = t.data[i].$fname(v);
                    }
                    t
                } else if self.is_scalar() {
                    let mut t: Tensor<$t> = Tensor::zeros(&rhs.shape);
                    let v = self[0];
                    for i in 0..t.size() {
                        t.data[i] = v.$fname(rhs.data[i]);
                    }
                    t
                } else {
                    assert_eq!(self.shape, rhs.shape);
                    let mut t: Tensor<$t> = self.clone();
                    for i in 0..t.size() {
                        t.data[i] = t.data[i].$fname(rhs.data[i]);
                    }
                    t
                }
            }
        }

        // T <op> S
        impl $trait_name<$t> for Tensor<$t> {
            type Output = Tensor<$t>;
            fn $fname(mut self, rhs: $t) -> Tensor<$t> {
                for i in 0..self.size() {
                    self.data[i] = self.data[i].$fname(rhs);
                }
                self
            }
        }
    )*)
}

add_impl! { Add, add, usize u8 u16 u32 u64 isize i8 i16 i32 i64 f32 f64 }
add_impl! { Sub, sub, usize u8 u16 u32 u64 isize i8 i16 i32 i64 f32 f64 }
add_impl! { Mul, mul, usize u8 u16 u32 u64 isize i8 i16 i32 i64 f32 f64 }
add_impl! { Div, div, usize u8 u16 u32 u64 isize i8 i16 i32 i64 f32 f64 }
add_impl! { Rem, rem, usize u8 u16 u32 u64 isize i8 i16 i32 i64 }
<commit_msg>Added bitwise operations.<commit_after>use tensor::Tensor;
use std::ops::{Add, Sub, Mul, Div, Rem, BitAnd, BitOr, BitXor};

macro_rules! add_impl {
    ($trait_name:ident, $fname:ident, $($t:ty)*) => ($(
        // mut T <op> T
        impl $trait_name<Tensor<$t>> for Tensor<$t> {
            type Output = Tensor<$t>;
            fn $fname(mut self, rhs: Tensor<$t>) -> Tensor<$t> {
                if rhs.is_scalar() {
                    let v = rhs[0];
                    for i in 0..self.size() {
                        self.data[i] = self.data[i].$fname(v);
                    }
                    self
                } else if self.is_scalar() {
                    let mut t: Tensor<$t> = Tensor::empty(&rhs.shape);
                    let v = self[0];
                    for i in 0..t.size() {
                        t.data[i] = v.$fname(rhs.data[i]);
                    }
                    t
                } else {
                    assert_eq!(self.shape, rhs.shape);
                    for i in 0..self.size() {
                        self.data[i] = self.data[i].$fname(rhs.data[i]);
                    }
                    self
                }
            }
        }

        // mut T <op> &T
        impl<'a> $trait_name<&'a Tensor<$t>> for Tensor<$t> {
            type Output = Tensor<$t>;
            fn $fname(mut self, rhs: &Tensor<$t>) -> Tensor<$t> {
                //println!("$fname T + &T");
                if rhs.is_scalar() {
                    let v = rhs[0];
                    for i in 0..self.size() {
                        self.data[i] = self.data[i].$fname(v);
                    }
                    self
                } else if self.is_scalar() {
                    let mut t: Tensor<$t> = Tensor::empty(&rhs.shape);
                    let v = self[0];
                    for i in 0..t.size() {
                        t.data[i] = v.$fname(rhs.data[i]);
                    }
                    t
                } else {
                    assert_eq!(self.shape, rhs.shape);
                    for i in 0..self.size() {
                        self.data[i] = self.data[i].$fname(rhs.data[i]);
                    }
                    self
                }
            }
        }

        // &T <op> &T
        impl<'a> $trait_name for &'a Tensor<$t> {
            type Output = Tensor<$t>;
            fn $fname(self, rhs: &Tensor<$t>) -> Tensor<$t> {
                //println!("$fname &T + &T");
                if rhs.is_scalar() {
                    let mut t = self.clone();
                    let v = rhs[0];
                    for i in 0..t.size() {
                        t.data[i] = t.data[i].$fname(v);
                    }
                    t
                } else if self.is_scalar() {
                    let mut t = Tensor::empty(&rhs.shape);
                    let v = self[0];
                    for i in 0..t.size() {
                        t.data[i] = v.$fname(rhs.data[i]);
                    }
                    t
                } else {
                    assert_eq!(self.shape, rhs.shape);
                    let mut t = self.clone();
                    for i in 0..t.size() {
                        t.data[i] = t.data[i].$fname(rhs.data[i]);
                    }
                    t
                }
            }
        }

        // T <op> S
        impl $trait_name<$t> for Tensor<$t> {
            type Output = Tensor<$t>;
            fn $fname(mut self, rhs: $t) -> Tensor<$t> {
                for i in 0..self.size() {
                    self.data[i] = self.data[i].$fname(rhs);
                }
                self
            }
        }
    )*)
}

// TODO: It would be better if these relied only on the traits, so that custom-made tensort types
// could be used.
add_impl! { Add, add, usize u8 u16 u32 u64 isize i8 i16 i32 i64 f32 f64 }
add_impl! { Sub, sub, usize u8 u16 u32 u64 isize i8 i16 i32 i64 f32 f64 }
add_impl! { Mul, mul, usize u8 u16 u32 u64 isize i8 i16 i32 i64 f32 f64 }
add_impl! { Div, div, usize u8 u16 u32 u64 isize i8 i16 i32 i64 f32 f64 }
add_impl! { Rem, rem, usize u8 u16 u32 u64 isize i8 i16 i32 i64 }

add_impl! { BitAnd, bitand, usize u8 u16 u32 u64 isize i8 i16 i32 i64 bool }
add_impl! { BitOr, bitor, usize u8 u16 u32 u64 isize i8 i16 i32 i64 bool }
add_impl! { BitXor, bitxor, usize u8 u16 u32 u64 isize i8 i16 i32 i64 bool }
<|endoftext|>