<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::cmp;
use std::fmt::Write;
use std::mem;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::{self, Receiver, TryRecvError};
use std::sync::{Arc, Mutex, RwLock};
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};

use glium::glutin::EventsLoopProxy;

use crate::model::ttable::{Score, TTable};
use crate::model::{Board, Move, Outcome};

const AI_MOVE_DELAY: Duration = Duration::from_millis(300);

const NEG_INFINITY: i16 = -0x7000;
const LOSE: i16 = -0x4000;
// Small contempt factor to discourage draws
const DRAW: i16 = 1;
const ASPIRATION_WIDTH: i16 = 51;

pub struct AI {
    status: Status,
    ttable: Arc<Mutex<TTable>>,
    pub debug_info: Arc<RwLock<String>>,
}

enum Status {
    Idle,
    // Either the AI thread is running, or there is a move waiting to be received
    Thinking {
        move_recv: Receiver<Move>,
        // We store and load this atomic with Ordering::Relaxed. It *should* be fine because it
        // doesn't interact with any other atomics--all we want to do is tell the AI thread to stop
        // searching for a move
        stop_signal: Arc<AtomicBool>,
        handle: JoinHandle<()>,
    },
}

impl AI {
    pub fn new() -> Self {
        Self {
            status: Status::Idle,
            ttable: Arc::new(Mutex::new(TTable::new())),
            debug_info: Arc::new(RwLock::new(String::new())),
        }
    }

    pub fn is_idle(&self) -> bool {
        match self.status {
            Status::Idle => true,
            Status::Thinking { .. } => false,
        }
    }

    pub fn stop(&mut self) {
        if let Status::Thinking {
            ref stop_signal, ..
        } = self.status
        {
            stop_signal.store(true, Ordering::Relaxed);
        }
        // Unconditionally assign because without NLL, we can't put this in the if let block above
        self.status = Status::Idle;
    }

    pub fn try_recv(&mut self) -> Option<Move> {
        use self::TryRecvError::*;

        let result;
        match self.status {
            Status::Idle => result = None,
            Status::Thinking { ref move_recv, .. } => match move_recv.try_recv() {
                Ok(mv) => result = Some(mv),
                Err(Empty) => result = None,
                Err(Disconnected) => panic!("Tried to receive move from disconnected sender"),
            },
        }

        // We can't set status in the Ok(mv) arm above without NLL
        if result.is_some() {
            self.status = Status::Idle;
        }
        result
    }

    pub fn think(
        &mut self,
        board: Board,
        board_list: Vec<Board>,
        depth: u8,
        events_proxy: EventsLoopProxy,
        delay: bool,
        ply_count: u64,
    ) {
        assert_ne!(depth, 0);

        let prev_status = mem::replace(&mut self.status, Status::Idle);

        let (move_sender, move_recv) = mpsc::channel();
        let stop_signal = Arc::new(AtomicBool::new(false));
        let stop_signal_clone = stop_signal.clone();

        let ttable_mutex = self.ttable.clone();
        let debug_info = self.debug_info.clone();

        let handle = thread::spawn(move || {
            let start = Instant::now();

            if let Status::Thinking {
                stop_signal,
                handle,
                ..
            } = prev_status
            {
                stop_signal.store(true, Ordering::Relaxed);
                handle
                    .join()
                    .expect("Old AI thread panicked when new AI thread joined on it");
            }

            // If the previous AI thread was sent the stop signal, but hasn't received it yet, we
            // will block here until it finishes. We won't have joined on its handle above because
            // stop throws away its Status enum.
            let mut ttable = match ttable_mutex.lock() {
                Ok(table) => table,
                Err(_poison_error) => panic!("Transposition table mutex is poisoned"),
            };

            if let Ok(mut debug_info) = debug_info.write() {
                debug_info.clear();
                writeln!(
                    debug_info,
                    "Turn {}: {:?} (Depth {})",
                    ply_count / 2 + 1,
                    board.turn,
                    depth
                )
                .unwrap();
            }

            if let SearchResult::Move(mv) = search_root(
                depth,
                board,
                board_list,
                &mut ttable,
                &stop_signal_clone,
                &debug_info,
                &events_proxy,
            ) {
                if stop_signal_clone.load(Ordering::Relaxed) {
                    return;
                }

                let elapsed = Instant::now() - start;
                if delay && elapsed < AI_MOVE_DELAY {
                    thread::sleep(AI_MOVE_DELAY - elapsed);
                }

                move_sender.send(mv).expect("AI failed to send Move");
                events_proxy
                    .wakeup()
                    .expect("Failed to wake up events loop");
            }
        });

        self.status = Status::Thinking {
            move_recv,
            stop_signal,
            handle,
        };
    }
}

enum SearchResult {
    Move(Move),
    Stopped,
}

fn search_root(
    depth: u8,
    board: Board,
    board_list: Vec<Board>,
    ttable: &mut TTable,
    stop_signal: &Arc<AtomicBool>,
    debug_info: &Arc<RwLock<String>>,
    events_proxy: &EventsLoopProxy,
) -> SearchResult {
    ttable.inc_age();

    // Only take positions after the last irreversible move
    let mut board_list: Vec<_> = board_list
        .into_iter()
        .rev()
        .take_while(|b| b.vitals == board.vitals)
        .collect();
    board_list.reverse();

    let mut moves: Vec<(Move, i16)> = board
        .generate_moves()
        .map(|mv| (mv, NEG_INFINITY))
        .collect();

    if moves.is_empty() {
        panic!("AI has no moves");
    }

    let mut pv = None;
    let mut iter_score = evaluate(&board);
    for depth in 0..depth {
        if stop_signal.load(Ordering::Relaxed) {
            return SearchResult::Stopped;
        }

        // Aspiration window search loop
        let mut asp_width = ASPIRATION_WIDTH;
        loop {
            let mut max_score = iter_score - asp_width;
            for pair in &mut moves {
                let mut new_board = board;
                new_board.apply_move(&pair.0);

                let mut new_pv = vec![];

                let score = -alphabeta_negamax(
                    &new_board,
                    &mut board_list,
                    &mut new_pv,
                    -(iter_score + asp_width),
                    -max_score,
                    depth,
                    ttable,
                );

                if score > max_score {
                    max_score = score;
                    pv = Some(new_pv);
                }
                pair.1 = score;
            }

            if max_score == iter_score + asp_width || max_score == iter_score - asp_width {
                // True score lies outside the window, so we search this position again
                asp_width *= 2;
            } else {
                break;
            }
        }

        moves.sort_by(|&(_, a), &(_, b)| b.cmp(&a));
        iter_score = moves[0].1;

        if let Ok(mut debug_info) = debug_info.write() {
            writeln!(debug_info, "\nDepth {}: {:>6}", depth, moves[0].1).unwrap();
            if let Some(ref mut pv) = pv {
                pv.push(moves[0].0);
                for mv in pv.iter().rev() {
                    writeln!(debug_info, "    {}", mv).unwrap();
                }
            }
            events_proxy
                .wakeup()
                .expect("Failed to wake up events loop");
        }
    }
    SearchResult::Move(moves[0].0)
}

fn alphabeta_negamax(
    board: &Board,
    // This list does not include the current board
    mut board_list: &mut Vec<Board>,
    pv: &mut Vec<Move>,
    mut alpha: i16,
    mut beta: i16,
    depth: u8,
    ttable: &mut TTable,
) -> i16 {
    let mut set_pv = move |score, new_pv| {
        if score > alpha && score < beta {
            *pv = new_pv;
        }
    };
    let set_ttable = |ttable: &mut TTable, score| {
        ttable.set(board.zobrist, score, depth as i8);
    };

    use self::Outcome::*;
    match board.outcome() {
        DrawStalemate | DrawInsufficientMaterial => {
            // This is safe to do because Board does not detect draws by threefold repetition
            set_ttable(ttable, Score::Exact(DRAW));
            set_pv(DRAW, vec![]);
            return DRAW;
        }
        Win(color) => {
            assert_ne!(color, board.turn);
            // Weight score by depth to encourage shorter wins. The shorter the win, the greater
            // `depth` will be, and so the larger the score will be. This also encourages the AI to
            // prolong a loss.
            let score = LOSE - i16::from(depth);
            set_ttable(ttable, Score::Exact(score));
            set_pv(score, vec![]);
            return score;
        }
        InProgress => {}
        DrawThreefoldRepetition => unreachable!(),
    }

    if board_list.len() >= 8 && board_list.iter().filter(|&&b| b == *board).count() >= 2 {
        set_pv(DRAW, vec![]);
        return DRAW;
    }

    if depth == 0 {
        let score = quiescence_search(board, alpha, beta, depth as i8, ttable);
        set_pv(score, vec![]);
        return score;
    }

    {
        match ttable.get(board.zobrist, depth as i8) {
            Some(Score::Exact(score)) => {
                // This will cut the PV short
                // TODO: Store the best move in the table and get the PV from that?
                set_pv(score, vec![]);
                return score;
            }
            Some(Score::Beta(score)) => {
                if score >= beta {
                    return score;
                }
                beta = score;
            }
            None => {}
        }
    }

    let mut best_score = NEG_INFINITY;
    let mut best_move = None;

    let mut new_pv = vec![];
    for mv in board.generate_moves() {
        let mut new_board = *board;
        new_board.apply_move(&mv);

        board_list.push(*board);
        let score = -alphabeta_negamax(
            &new_board,
            &mut board_list,
            &mut new_pv,
            -beta,
            -alpha,
            depth - 1,
            ttable,
        );
        board_list.pop();

        best_score = cmp::max(score, best_score);

        if score >= beta {
            set_ttable(ttable, Score::Beta(score));
            return beta;
        } else if score > alpha {
            alpha = score;
            best_move = Some(mv);
        }
    }
    set_ttable(ttable, Score::Exact(best_score));
    if let Some(mv) = best_move {
        new_pv.push(mv);
        set_pv(alpha, new_pv);
    }
    alpha
}

fn quiescence_search(
    board: &Board,
    mut alpha: i16,
    mut beta: i16,
    depth: i8,
    ttable: &mut TTable,
) -> i16 {
    let stand_pat = evaluate(board);
    if stand_pat >= beta {
        return beta;
    } else if stand_pat + 200 < alpha {
        // Delta pruning: we bet that no capture move will raise alpha by more than 200 centipieces
        // We don't switch this off in the endgame because a draw by insufficient material is very
        // unlikely (both sides must be unable to exchange and each have only 1 piece left)
        return alpha;
    } else if alpha < stand_pat {
        alpha = stand_pat;
    }

    match ttable.get(board.zobrist, depth) {
        Some(Score::Exact(score)) => {
            return score;
        }
        Some(Score::Beta(score)) => {
            if score >= beta {
                return score;
            }
            beta = score;
        }
        None => {}
    };

    let set_ttable = |ttable: &mut TTable, score| {
        ttable.set(board.zobrist, score, depth);
    };

    for mv in board.generate_captures() {
        let mut new_board = *board;
        new_board.apply_move(&mv);

        let score = -quiescence_search(&new_board, -beta, -alpha, depth - 1, ttable);

        if score >= beta {
            set_ttable(ttable, Score::Beta(score));
            return beta;
        } else if score > alpha {
            alpha = score;
        }
    }
    set_ttable(ttable, Score::Exact(alpha));
    alpha
}

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

    // If it's two hexes to exchange, then a piece is 100 and a hex is 50. If it's one hex, then we
    // halve the value of a piece so that both are 50. We could instead up the value of a hex to
    // 100, but this way we don't need to change the width of the aspiration window.
    let hex_factor = i16::from(board.hexes_to_exchange);
    let wp = hex_factor * 50 * i16::from(board.pieces(White));
    let bp = hex_factor * 50 * 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),
    }
}
<commit_msg>Fix NLL workarounds<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::cmp;
use std::fmt::Write;
use std::mem;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::{self, Receiver, TryRecvError};
use std::sync::{Arc, Mutex, RwLock};
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};

use glium::glutin::EventsLoopProxy;

use crate::model::ttable::{Score, TTable};
use crate::model::{Board, Move, Outcome};

const AI_MOVE_DELAY: Duration = Duration::from_millis(300);

const NEG_INFINITY: i16 = -0x7000;
const LOSE: i16 = -0x4000;
// Small contempt factor to discourage draws
const DRAW: i16 = 1;
const ASPIRATION_WIDTH: i16 = 51;

pub struct AI {
    status: Status,
    ttable: Arc<Mutex<TTable>>,
    pub debug_info: Arc<RwLock<String>>,
}

enum Status {
    Idle,
    // Either the AI thread is running, or there is a move waiting to be received
    Thinking {
        move_recv: Receiver<Move>,
        // We store and load this atomic with Ordering::Relaxed. It *should* be fine because it
        // doesn't interact with any other atomics--all we want to do is tell the AI thread to stop
        // searching for a move
        stop_signal: Arc<AtomicBool>,
        handle: JoinHandle<()>,
    },
}

impl AI {
    pub fn new() -> Self {
        Self {
            status: Status::Idle,
            ttable: Arc::new(Mutex::new(TTable::new())),
            debug_info: Arc::new(RwLock::new(String::new())),
        }
    }

    pub fn is_idle(&self) -> bool {
        match self.status {
            Status::Idle => true,
            Status::Thinking { .. } => false,
        }
    }

    pub fn stop(&mut self) {
        if let Status::Thinking {
            ref stop_signal, ..
        } = self.status
        {
            stop_signal.store(true, Ordering::Relaxed);
            self.status = Status::Idle;
        }
    }

    pub fn try_recv(&mut self) -> Option<Move> {
        use self::TryRecvError::*;
        match self.status {
            Status::Idle => None,
            Status::Thinking { ref move_recv, .. } => match move_recv.try_recv() {
                Ok(mv) => {
                    self.status = Status::Idle;
                    Some(mv)
                }
                Err(Empty) => None,
                Err(Disconnected) => panic!("Tried to receive move from disconnected sender"),
            },
        }
    }

    pub fn think(
        &mut self,
        board: Board,
        board_list: Vec<Board>,
        depth: u8,
        events_proxy: EventsLoopProxy,
        delay: bool,
        ply_count: u64,
    ) {
        assert_ne!(depth, 0);

        let prev_status = mem::replace(&mut self.status, Status::Idle);

        let (move_sender, move_recv) = mpsc::channel();
        let stop_signal = Arc::new(AtomicBool::new(false));
        let stop_signal_clone = stop_signal.clone();

        let ttable_mutex = self.ttable.clone();
        let debug_info = self.debug_info.clone();

        let handle = thread::spawn(move || {
            let start = Instant::now();

            if let Status::Thinking {
                stop_signal,
                handle,
                ..
            } = prev_status
            {
                stop_signal.store(true, Ordering::Relaxed);
                handle
                    .join()
                    .expect("Old AI thread panicked when new AI thread joined on it");
            }

            // If the previous AI thread was sent the stop signal, but hasn't received it yet, we
            // will block here until it finishes. We won't have joined on its handle above because
            // stop throws away its Status enum.
            let mut ttable = match ttable_mutex.lock() {
                Ok(table) => table,
                Err(_poison_error) => panic!("Transposition table mutex is poisoned"),
            };

            if let Ok(mut debug_info) = debug_info.write() {
                debug_info.clear();
                writeln!(
                    debug_info,
                    "Turn {}: {:?} (Depth {})",
                    ply_count / 2 + 1,
                    board.turn,
                    depth
                )
                .unwrap();
            }

            if let SearchResult::Move(mv) = search_root(
                depth,
                board,
                board_list,
                &mut ttable,
                &stop_signal_clone,
                &debug_info,
                &events_proxy,
            ) {
                if stop_signal_clone.load(Ordering::Relaxed) {
                    return;
                }

                let elapsed = Instant::now() - start;
                if delay && elapsed < AI_MOVE_DELAY {
                    thread::sleep(AI_MOVE_DELAY - elapsed);
                }

                move_sender.send(mv).expect("AI failed to send Move");
                events_proxy
                    .wakeup()
                    .expect("Failed to wake up events loop");
            }
        });

        self.status = Status::Thinking {
            move_recv,
            stop_signal,
            handle,
        };
    }
}

enum SearchResult {
    Move(Move),
    Stopped,
}

fn search_root(
    depth: u8,
    board: Board,
    board_list: Vec<Board>,
    ttable: &mut TTable,
    stop_signal: &Arc<AtomicBool>,
    debug_info: &Arc<RwLock<String>>,
    events_proxy: &EventsLoopProxy,
) -> SearchResult {
    ttable.inc_age();

    // Only take positions after the last irreversible move
    let mut board_list: Vec<_> = board_list
        .into_iter()
        .rev()
        .take_while(|b| b.vitals == board.vitals)
        .collect();
    board_list.reverse();

    let mut moves: Vec<(Move, i16)> = board
        .generate_moves()
        .map(|mv| (mv, NEG_INFINITY))
        .collect();

    if moves.is_empty() {
        panic!("AI has no moves");
    }

    let mut pv = None;
    let mut iter_score = evaluate(&board);
    for depth in 0..depth {
        if stop_signal.load(Ordering::Relaxed) {
            return SearchResult::Stopped;
        }

        // Aspiration window search loop
        let mut asp_width = ASPIRATION_WIDTH;
        loop {
            let mut max_score = iter_score - asp_width;
            for pair in &mut moves {
                let mut new_board = board;
                new_board.apply_move(&pair.0);

                let mut new_pv = vec![];

                let score = -alphabeta_negamax(
                    &new_board,
                    &mut board_list,
                    &mut new_pv,
                    -(iter_score + asp_width),
                    -max_score,
                    depth,
                    ttable,
                );

                if score > max_score {
                    max_score = score;
                    pv = Some(new_pv);
                }
                pair.1 = score;
            }

            if max_score == iter_score + asp_width || max_score == iter_score - asp_width {
                // True score lies outside the window, so we search this position again
                asp_width *= 2;
            } else {
                break;
            }
        }

        moves.sort_by(|&(_, a), &(_, b)| b.cmp(&a));
        iter_score = moves[0].1;

        if let Ok(mut debug_info) = debug_info.write() {
            writeln!(debug_info, "\nDepth {}: {:>6}", depth, moves[0].1).unwrap();
            if let Some(ref mut pv) = pv {
                pv.push(moves[0].0);
                for mv in pv.iter().rev() {
                    writeln!(debug_info, "    {}", mv).unwrap();
                }
            }
            events_proxy
                .wakeup()
                .expect("Failed to wake up events loop");
        }
    }
    SearchResult::Move(moves[0].0)
}

fn alphabeta_negamax(
    board: &Board,
    // This list does not include the current board
    mut board_list: &mut Vec<Board>,
    pv: &mut Vec<Move>,
    mut alpha: i16,
    mut beta: i16,
    depth: u8,
    ttable: &mut TTable,
) -> i16 {
    let mut set_pv = move |score, new_pv| {
        if score > alpha && score < beta {
            *pv = new_pv;
        }
    };
    let set_ttable = |ttable: &mut TTable, score| {
        ttable.set(board.zobrist, score, depth as i8);
    };

    use self::Outcome::*;
    match board.outcome() {
        DrawStalemate | DrawInsufficientMaterial => {
            // This is safe to do because Board does not detect draws by threefold repetition
            set_ttable(ttable, Score::Exact(DRAW));
            set_pv(DRAW, vec![]);
            return DRAW;
        }
        Win(color) => {
            assert_ne!(color, board.turn);
            // Weight score by depth to encourage shorter wins. The shorter the win, the greater
            // `depth` will be, and so the larger the score will be. This also encourages the AI to
            // prolong a loss.
            let score = LOSE - i16::from(depth);
            set_ttable(ttable, Score::Exact(score));
            set_pv(score, vec![]);
            return score;
        }
        InProgress => {}
        DrawThreefoldRepetition => unreachable!(),
    }

    if board_list.len() >= 8 && board_list.iter().filter(|&&b| b == *board).count() >= 2 {
        set_pv(DRAW, vec![]);
        return DRAW;
    }

    if depth == 0 {
        let score = quiescence_search(board, alpha, beta, depth as i8, ttable);
        set_pv(score, vec![]);
        return score;
    }

    {
        match ttable.get(board.zobrist, depth as i8) {
            Some(Score::Exact(score)) => {
                // This will cut the PV short
                // TODO: Store the best move in the table and get the PV from that?
                set_pv(score, vec![]);
                return score;
            }
            Some(Score::Beta(score)) => {
                if score >= beta {
                    return score;
                }
                beta = score;
            }
            None => {}
        }
    }

    let mut best_score = NEG_INFINITY;
    let mut best_move = None;

    let mut new_pv = vec![];
    for mv in board.generate_moves() {
        let mut new_board = *board;
        new_board.apply_move(&mv);

        board_list.push(*board);
        let score = -alphabeta_negamax(
            &new_board,
            &mut board_list,
            &mut new_pv,
            -beta,
            -alpha,
            depth - 1,
            ttable,
        );
        board_list.pop();

        best_score = cmp::max(score, best_score);

        if score >= beta {
            set_ttable(ttable, Score::Beta(score));
            return beta;
        } else if score > alpha {
            alpha = score;
            best_move = Some(mv);
        }
    }
    set_ttable(ttable, Score::Exact(best_score));
    if let Some(mv) = best_move {
        new_pv.push(mv);
        set_pv(alpha, new_pv);
    }
    alpha
}

fn quiescence_search(
    board: &Board,
    mut alpha: i16,
    mut beta: i16,
    depth: i8,
    ttable: &mut TTable,
) -> i16 {
    let stand_pat = evaluate(board);
    if stand_pat >= beta {
        return beta;
    } else if stand_pat + 200 < alpha {
        // Delta pruning: we bet that no capture move will raise alpha by more than 200 centipieces
        // We don't switch this off in the endgame because a draw by insufficient material is very
        // unlikely (both sides must be unable to exchange and each have only 1 piece left)
        return alpha;
    } else if alpha < stand_pat {
        alpha = stand_pat;
    }

    match ttable.get(board.zobrist, depth) {
        Some(Score::Exact(score)) => {
            return score;
        }
        Some(Score::Beta(score)) => {
            if score >= beta {
                return score;
            }
            beta = score;
        }
        None => {}
    };

    let set_ttable = |ttable: &mut TTable, score| {
        ttable.set(board.zobrist, score, depth);
    };

    for mv in board.generate_captures() {
        let mut new_board = *board;
        new_board.apply_move(&mv);

        let score = -quiescence_search(&new_board, -beta, -alpha, depth - 1, ttable);

        if score >= beta {
            set_ttable(ttable, Score::Beta(score));
            return beta;
        } else if score > alpha {
            alpha = score;
        }
    }
    set_ttable(ttable, Score::Exact(alpha));
    alpha
}

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

    // If it's two hexes to exchange, then a piece is 100 and a hex is 50. If it's one hex, then we
    // halve the value of a piece so that both are 50. We could instead up the value of a hex to
    // 100, but this way we don't need to change the width of the aspiration window.
    let hex_factor = i16::from(board.hexes_to_exchange);
    let wp = hex_factor * 50 * i16::from(board.pieces(White));
    let bp = hex_factor * 50 * 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),
    }
}
<|endoftext|>


<commit_before>use crate::model::{
    AnnoLichess, GameId, GameInfo, PersonalEntry, PersonalKey, PersonalKeyPrefix, PersonalStatus,
    UserId,
};
use rocksdb::{
    BlockBasedIndexType, BlockBasedOptions, ColumnFamily, ColumnFamilyDescriptor, DBWithThreadMode,
    Direction, IteratorMode, MergeOperands, Options, ReadOptions, SliceTransform, DB,
};
use std::{io::Cursor, path::Path};

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

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 personal_opts = Options::default();
        personal_opts.set_merge_operator_associative("personal merge", personal_merge);
        personal_opts
            .set_prefix_extractor(SliceTransform::create_fixed_prefix(PersonalKeyPrefix::SIZE));
        let mut personal_block_opts = BlockBasedOptions::default();
        personal_block_opts.set_index_type(BlockBasedIndexType::HashSearch);
        personal_block_opts.set_block_size(4 * 1024);
        personal_block_opts.set_bloom_filter(8, true);
        personal_opts.set_block_based_table_factory(&personal_block_opts);

        let mut game_opts = Options::default();
        game_opts.set_merge_operator_associative("game merge", game_merge);
        game_opts.set_prefix_extractor(SliceTransform::create_noop());
        let mut game_block_opts = BlockBasedOptions::default();
        game_block_opts.set_index_type(BlockBasedIndexType::HashSearch);
        game_block_opts.set_block_size(1024);
        game_opts.set_block_based_table_factory(&game_block_opts);

        let mut player_opts = Options::default();
        player_opts.set_prefix_extractor(SliceTransform::create_noop());
        let mut player_block_opts = BlockBasedOptions::default();
        player_block_opts.set_index_type(BlockBasedIndexType::HashSearch);
        player_block_opts.set_block_size(1024);
        player_opts.set_block_based_table_factory(&player_block_opts);

        let inner = DBWithThreadMode::open_cf_descriptors(
            &db_opts,
            path,
            vec![
                ColumnFamilyDescriptor::new("personal", personal_opts),
                ColumnFamilyDescriptor::new("game", game_opts),
                ColumnFamilyDescriptor::new("player", player_opts),
            ],
        )?;

        Ok(Database { inner })
    }

    pub fn queryable(&self) -> QueryableDatabase<'_> {
        QueryableDatabase {
            db: &self.inner,
            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"),
        }
    }
}

pub struct QueryableDatabase<'a> {
    pub db: &'a DB,
    pub cf_personal: &'a ColumnFamily,
    pub cf_game: &'a ColumnFamily,
    pub cf_player: &'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 merge_game_info(&self, id: GameId, info: GameInfo) -> Result<(), rocksdb::Error> {
        let mut cursor = Cursor::new(Vec::with_capacity(GameInfo::SIZE_HINT));
        info.write(&mut cursor).expect("serialize game info");
        self.db
            .merge_cf(self.cf_game, id.to_bytes(), cursor.into_inner())
    }

    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 merge_personal(
        &self,
        key: PersonalKey,
        entry: PersonalEntry,
    ) -> Result<(), rocksdb::Error> {
        let mut cursor = Cursor::new(Vec::with_capacity(PersonalEntry::SIZE_HINT));
        entry.write(&mut cursor).expect("serialize personal entry");
        self.db
            .merge_cf(self.cf_personal, key.into_bytes(), cursor.into_inner())
    }

    pub fn get_personal(
        &self,
        key: PersonalKeyPrefix,
        since: AnnoLichess,
    ) -> Result<PersonalEntry, rocksdb::Error> {
        let mut entry = PersonalEntry::default();

        let mut opt = ReadOptions::default();
        opt.set_prefix_same_as_start(true);
        let iterator = self.db.iterator_cf_opt(
            self.cf_personal,
            opt,
            IteratorMode::From(&key.with_year(since).into_bytes(), Direction::Forward),
        );

        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 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 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")
        }))
    }
}

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.white |= old_info.indexed.white;
            new_info.indexed.black |= old_info.indexed.black;
        }
        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())
}
<commit_msg>allow permanently flagging games as unrated<commit_after>use crate::model::{
    AnnoLichess, GameId, GameInfo, PersonalEntry, PersonalKey, PersonalKeyPrefix, PersonalStatus,
    UserId,
};
use rocksdb::{
    BlockBasedIndexType, BlockBasedOptions, ColumnFamily, ColumnFamilyDescriptor, DBWithThreadMode,
    Direction, IteratorMode, MergeOperands, Options, ReadOptions, SliceTransform, DB,
};
use std::{io::Cursor, path::Path};

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

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 personal_opts = Options::default();
        personal_opts.set_merge_operator_associative("personal merge", personal_merge);
        personal_opts
            .set_prefix_extractor(SliceTransform::create_fixed_prefix(PersonalKeyPrefix::SIZE));
        let mut personal_block_opts = BlockBasedOptions::default();
        personal_block_opts.set_index_type(BlockBasedIndexType::HashSearch);
        personal_block_opts.set_block_size(4 * 1024);
        personal_block_opts.set_bloom_filter(8, true);
        personal_opts.set_block_based_table_factory(&personal_block_opts);

        let mut game_opts = Options::default();
        game_opts.set_merge_operator_associative("game merge", game_merge);
        game_opts.set_prefix_extractor(SliceTransform::create_noop());
        let mut game_block_opts = BlockBasedOptions::default();
        game_block_opts.set_index_type(BlockBasedIndexType::HashSearch);
        game_block_opts.set_block_size(1024);
        game_opts.set_block_based_table_factory(&game_block_opts);

        let mut player_opts = Options::default();
        player_opts.set_prefix_extractor(SliceTransform::create_noop());
        let mut player_block_opts = BlockBasedOptions::default();
        player_block_opts.set_index_type(BlockBasedIndexType::HashSearch);
        player_block_opts.set_block_size(1024);
        player_opts.set_block_based_table_factory(&player_block_opts);

        let inner = DBWithThreadMode::open_cf_descriptors(
            &db_opts,
            path,
            vec![
                ColumnFamilyDescriptor::new("personal", personal_opts),
                ColumnFamilyDescriptor::new("game", game_opts),
                ColumnFamilyDescriptor::new("player", player_opts),
            ],
        )?;

        Ok(Database { inner })
    }

    pub fn queryable(&self) -> QueryableDatabase<'_> {
        QueryableDatabase {
            db: &self.inner,
            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"),
        }
    }
}

pub struct QueryableDatabase<'a> {
    pub db: &'a DB,
    pub cf_personal: &'a ColumnFamily,
    pub cf_game: &'a ColumnFamily,
    pub cf_player: &'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 merge_game_info(&self, id: GameId, info: GameInfo) -> Result<(), rocksdb::Error> {
        let mut cursor = Cursor::new(Vec::with_capacity(GameInfo::SIZE_HINT));
        info.write(&mut cursor).expect("serialize game info");
        self.db
            .merge_cf(self.cf_game, id.to_bytes(), cursor.into_inner())
    }

    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 merge_personal(
        &self,
        key: PersonalKey,
        entry: PersonalEntry,
    ) -> Result<(), rocksdb::Error> {
        let mut cursor = Cursor::new(Vec::with_capacity(PersonalEntry::SIZE_HINT));
        entry.write(&mut cursor).expect("serialize personal entry");
        self.db
            .merge_cf(self.cf_personal, key.into_bytes(), cursor.into_inner())
    }

    pub fn get_personal(
        &self,
        key: PersonalKeyPrefix,
        since: AnnoLichess,
    ) -> Result<PersonalEntry, rocksdb::Error> {
        let mut entry = PersonalEntry::default();

        let mut opt = ReadOptions::default();
        opt.set_prefix_same_as_start(true);
        let iterator = self.db.iterator_cf_opt(
            self.cf_personal,
            opt,
            IteratorMode::From(&key.with_year(since).into_bytes(), Direction::Forward),
        );

        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 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 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")
        }))
    }
}

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.white |= old_info.indexed.white;
            new_info.indexed.black |= old_info.indexed.black;
            new_info.rated &= old_info.rated;
        }
        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())
}
<|endoftext|>


<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]

use std::{collections::HashSet, fmt, time::Duration};

use structopt::StructOpt;
use tokio::time;

use crate::{
    cluster::Cluster,
    experiments::{Context, Experiment, ExperimentParam},
    instance::Instance,
    tx_emitter::EmitJobRequest,
};
use async_trait::async_trait;
use diem_logger::info;
use std::time::Instant;

const EXPERIMENT_DURATION_TIMEOUT_SECS: u64 = 600;
const STATE_SYNC_COMMITTED_COUNTER_NAME: &str = "diem_state_sync_version.synced";

#[derive(StructOpt, Debug)]
pub struct StateSyncPerformanceParams {
    emit_transactions_duration_secs: u64,
}

impl StateSyncPerformanceParams {
    pub fn new(emit_transactions_duration_secs: u64) -> Self {
        Self {
            emit_transactions_duration_secs,
        }
    }
}

pub struct StateSyncPerformance {
    params: StateSyncPerformanceParams,
    fullnode_instance: Instance,
    validator_instance: Instance,
}

impl ExperimentParam for StateSyncPerformanceParams {
    type E = StateSyncPerformance;

    fn build(self, cluster: &Cluster) -> Self::E {
        let validator_instance = cluster.random_validator_instance();
        let fullnode_instance = cluster.random_fullnode_instance();
        Self::E {
            params: self,
            fullnode_instance,
            validator_instance,
        }
    }
}

#[async_trait]
impl Experiment for StateSyncPerformance {
    fn affected_validators(&self) -> HashSet<String> {
        let mut result = HashSet::new();
        result.insert(self.validator_instance.peer_name().clone());
        result
    }

    async fn run(&mut self, context: &mut Context<'_>) -> anyhow::Result<()> {
        // Stop the fullnode and clear all data so that it falls behind
        info!("Stopping the fullnode: {}", self.fullnode_instance);
        self.fullnode_instance.stop().await?;
        self.fullnode_instance.clean_data().await?;

        // Execute and commit transactions on the validators for the specified duration
        let emit_transactions_duration_secs = self.params.emit_transactions_duration_secs;
        info!(
            "Executing transactions for {} seconds",
            emit_transactions_duration_secs
        );
        let emit_job_request = EmitJobRequest::for_instances(
            context.cluster.validator_instances().to_vec(),
            context.global_emit_job_request,
            0,
            0,
        );
        let _ = context
            .tx_emitter
            .emit_txn_for(
                Duration::from_secs(emit_transactions_duration_secs),
                emit_job_request,
            )
            .await?;

        // Read the validator synced version
        let validator_synced_version = self.read_validator_synced_version();
        info!(
            "The validator is now synced at version: {}",
            validator_synced_version
        );

        // Restart the fullnode so that it starts state syncing to catch up
        info!(
            "Waiting for the fullnode to wake up: {}",
            self.fullnode_instance
        );
        self.fullnode_instance.start().await?;
        self.fullnode_instance
            .wait_json_rpc(Instant::now() + Duration::from_secs(120))
            .await?;

        // Wait for the fullnode to catch up to the expected version
        info!(
            "The fullnode is now up. Waiting for it to state sync to the expected version: {}",
            validator_synced_version
        );
        let start_instant = Instant::now();
        while self.read_fullnode_synced_version() < validator_synced_version {
            time::sleep(Duration::from_secs(1)).await;
        }
        info!(
            "The fullnode has caught up to version: {}",
            validator_synced_version
        );

        // Calculate the state sync throughput
        let time_to_state_sync = start_instant.elapsed();
        let state_sync_throughput =
            (validator_synced_version * 1000.0) / time_to_state_sync.as_millis() as f64;
        let state_sync_throughput_message = format!(
            "State sync throughput : {:?} txn/sec",
            state_sync_throughput,
        );
        info!("Time to state sync {:?}", time_to_state_sync);

        // Display the state sync throughput and report the results
        info!("{}", state_sync_throughput_message);
        context.report.report_text(state_sync_throughput_message);
        context
            .report
            .report_metric(self, "state_sync_throughput", state_sync_throughput as f64);

        Ok(())
    }

    fn deadline(&self) -> Duration {
        Duration::from_secs(EXPERIMENT_DURATION_TIMEOUT_SECS)
    }
}

impl StateSyncPerformance {
    fn read_fullnode_synced_version(&self) -> f64 {
        Self::read_synced_counter(&self.fullnode_instance)
    }

    fn read_validator_synced_version(&self) -> f64 {
        Self::read_synced_counter(&self.validator_instance)
    }

    // Reads the state sync "synced counter" for the given instance. If no
    // counter is found, returns zero.
    fn read_synced_counter(instance: &Instance) -> f64 {
        instance
            .counter(STATE_SYNC_COMMITTED_COUNTER_NAME)
            .unwrap_or(0.0)
    }
}

impl fmt::Display for StateSyncPerformance {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "StateSyncPerformance")
    }
}
<commit_msg>[Cluster Test] Update the state sync benchmark to use reasonable precision for throughput numbers.<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]

use std::{collections::HashSet, fmt, time::Duration};

use structopt::StructOpt;
use tokio::time;

use crate::{
    cluster::Cluster,
    experiments::{Context, Experiment, ExperimentParam},
    instance::Instance,
    tx_emitter::EmitJobRequest,
};
use async_trait::async_trait;
use diem_logger::info;
use std::time::Instant;

const EXPERIMENT_DURATION_TIMEOUT_SECS: u64 = 600;
const STATE_SYNC_COMMITTED_COUNTER_NAME: &str = "diem_state_sync_version.synced";

#[derive(StructOpt, Debug)]
pub struct StateSyncPerformanceParams {
    emit_transactions_duration_secs: u64,
}

impl StateSyncPerformanceParams {
    pub fn new(emit_transactions_duration_secs: u64) -> Self {
        Self {
            emit_transactions_duration_secs,
        }
    }
}

pub struct StateSyncPerformance {
    params: StateSyncPerformanceParams,
    fullnode_instance: Instance,
    validator_instance: Instance,
}

impl ExperimentParam for StateSyncPerformanceParams {
    type E = StateSyncPerformance;

    fn build(self, cluster: &Cluster) -> Self::E {
        let validator_instance = cluster.random_validator_instance();
        let fullnode_instance = cluster.random_fullnode_instance();
        Self::E {
            params: self,
            fullnode_instance,
            validator_instance,
        }
    }
}

#[async_trait]
impl Experiment for StateSyncPerformance {
    fn affected_validators(&self) -> HashSet<String> {
        let mut result = HashSet::new();
        result.insert(self.validator_instance.peer_name().clone());
        result
    }

    async fn run(&mut self, context: &mut Context<'_>) -> anyhow::Result<()> {
        // Stop the fullnode and clear all data so that it falls behind
        info!("Stopping the fullnode: {}", self.fullnode_instance);
        self.fullnode_instance.stop().await?;
        self.fullnode_instance.clean_data().await?;

        // Execute and commit transactions on the validators for the specified duration
        let emit_transactions_duration_secs = self.params.emit_transactions_duration_secs;
        info!(
            "Executing transactions for {} seconds",
            emit_transactions_duration_secs
        );
        let emit_job_request = EmitJobRequest::for_instances(
            context.cluster.validator_instances().to_vec(),
            context.global_emit_job_request,
            0,
            0,
        );
        let _ = context
            .tx_emitter
            .emit_txn_for(
                Duration::from_secs(emit_transactions_duration_secs),
                emit_job_request,
            )
            .await?;

        // Read the validator synced version
        let validator_synced_version = self.read_validator_synced_version();
        info!(
            "The validator is now synced at version: {}",
            validator_synced_version
        );

        // Restart the fullnode so that it starts state syncing to catch up
        info!(
            "Waiting for the fullnode to wake up: {}",
            self.fullnode_instance
        );
        self.fullnode_instance.start().await?;
        self.fullnode_instance
            .wait_json_rpc(Instant::now() + Duration::from_secs(120))
            .await?;

        // Wait for the fullnode to catch up to the expected version
        info!(
            "The fullnode is now up. Waiting for it to state sync to the expected version: {}",
            validator_synced_version
        );
        let start_instant = Instant::now();
        while self.read_fullnode_synced_version() < validator_synced_version {
            time::sleep(Duration::from_secs(1)).await;
        }
        info!(
            "The fullnode has caught up to version: {}",
            validator_synced_version
        );

        // Calculate the state sync throughput
        let time_to_state_sync = start_instant.elapsed();
        let state_sync_throughput = validator_synced_version as u64 / time_to_state_sync.as_secs();
        let state_sync_throughput_message =
            format!("State sync throughput : {} txn/sec", state_sync_throughput,);
        info!("Time to state sync {:?}", time_to_state_sync);

        // Display the state sync throughput and report the results
        info!("{}", state_sync_throughput_message);
        context.report.report_text(state_sync_throughput_message);
        context
            .report
            .report_metric(self, "state_sync_throughput", state_sync_throughput as f64);

        Ok(())
    }

    fn deadline(&self) -> Duration {
        Duration::from_secs(EXPERIMENT_DURATION_TIMEOUT_SECS)
    }
}

impl StateSyncPerformance {
    fn read_fullnode_synced_version(&self) -> f64 {
        Self::read_synced_counter(&self.fullnode_instance)
    }

    fn read_validator_synced_version(&self) -> f64 {
        Self::read_synced_counter(&self.validator_instance)
    }

    // Reads the state sync "synced counter" for the given instance. If no
    // counter is found, returns zero.
    fn read_synced_counter(instance: &Instance) -> f64 {
        instance
            .counter(STATE_SYNC_COMMITTED_COUNTER_NAME)
            .unwrap_or(0.0)
    }
}

impl fmt::Display for StateSyncPerformance {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "StateSyncPerformance")
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Guillaume Pinot (@TeXitoi) <texitoi@texitoi.eu>,
// Kevin Knapp (@kbknapp) <kbknapp@gmail.com>, and
// Andrew Hobden (@hoverbear) <andrew@hoverbear.org>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
//
// This work was derived from Structopt (https://github.com/TeXitoi/structopt)
// commit#ea76fa1b1b273e65e3b0b1046643715b49bec51f which is licensed under the
// MIT/Apache 2.0 license.

mod utils;

use clap::Clap;
use utils::*;

#[derive(Clap, PartialEq, Debug)]
enum Opt {
    /// Fetch stuff from GitHub
    Fetch {
        #[clap(long)]
        all: bool,
        #[clap(short, long)]
        /// Overwrite local branches.
        force: bool,
        repo: String,
    },

    Add {
        #[clap(short, long)]
        interactive: bool,
        #[clap(short, long)]
        verbose: bool,
    },
}

#[test]
fn test_fetch() {
    assert_eq!(
        Opt::Fetch {
            all: true,
            force: false,
            repo: "origin".to_string()
        },
        Opt::parse_from(&["test", "fetch", "--all", "origin"])
    );
    assert_eq!(
        Opt::Fetch {
            all: false,
            force: true,
            repo: "origin".to_string()
        },
        Opt::parse_from(&["test", "fetch", "-f", "origin"])
    );
}

#[test]
fn test_add() {
    assert_eq!(
        Opt::Add {
            interactive: false,
            verbose: false
        },
        Opt::parse_from(&["test", "add"])
    );
    assert_eq!(
        Opt::Add {
            interactive: true,
            verbose: true
        },
        Opt::parse_from(&["test", "add", "-i", "-v"])
    );
}

#[test]
fn test_no_parse() {
    let result = Opt::try_parse_from(&["test", "badcmd", "-i", "-v"]);
    assert!(result.is_err());

    let result = Opt::try_parse_from(&["test", "add", "--badoption"]);
    assert!(result.is_err());

    let result = Opt::try_parse_from(&["test"]);
    assert!(result.is_err());
}

#[derive(Clap, PartialEq, Debug)]
enum Opt2 {
    DoSomething { arg: String },
}

#[test]
/// This test is specifically to make sure that hyphenated subcommands get
/// processed correctly.
fn test_hyphenated_subcommands() {
    assert_eq!(
        Opt2::DoSomething {
            arg: "blah".to_string()
        },
        Opt2::parse_from(&["test", "do-something", "blah"])
    );
}

#[derive(Clap, PartialEq, Debug)]
enum Opt3 {
    Add,
    Init,
    Fetch,
}

#[test]
fn test_null_commands() {
    assert_eq!(Opt3::Add, Opt3::parse_from(&["test", "add"]));
    assert_eq!(Opt3::Init, Opt3::parse_from(&["test", "init"]));
    assert_eq!(Opt3::Fetch, Opt3::parse_from(&["test", "fetch"]));
}

#[derive(Clap, PartialEq, Debug)]
#[clap(about = "Not shown")]
struct Add {
    file: String,
}
/// Not shown
#[derive(Clap, PartialEq, Debug)]
struct Fetch {
    remote: String,
}
#[derive(Clap, PartialEq, Debug)]
enum Opt4 {
    // Not shown
    /// Add a file
    Add(Add),
    Init,
    /// download history from remote
    Fetch(Fetch),
}

#[test]
fn test_tuple_commands() {
    assert_eq!(
        Opt4::Add(Add {
            file: "f".to_string()
        }),
        Opt4::parse_from(&["test", "add", "f"])
    );
    assert_eq!(Opt4::Init, Opt4::parse_from(&["test", "init"]));
    assert_eq!(
        Opt4::Fetch(Fetch {
            remote: "origin".to_string()
        }),
        Opt4::parse_from(&["test", "fetch", "origin"])
    );

    let output = get_long_help::<Opt4>();

    assert!(output.contains("download history from remote"));
    assert!(output.contains("Add a file"));
    assert!(!output.contains("Not shown"));
}

#[test]
fn external_subcommand() {
    #[derive(Debug, PartialEq, Clap)]
    struct Opt {
        #[clap(global = true, long)]
        global: bool,
        #[clap(subcommand)]
        sub: Subcommands,
    }

    #[derive(Debug, PartialEq, Clap)]
    enum Subcommands {
        Add,
        Remove,
        Global(GlobalCmd),
        #[clap(external_subcommand)]
        Other(Vec<String>),
    }

    #[derive(Debug, PartialEq, Clap)]
    struct GlobalCmd {
        #[clap(from_global)]
        global: bool,
    }

    assert_eq!(
        Opt::parse_from(&["test", "add"]),
        Opt {
            global: false,
            sub: Subcommands::Add
        }
    );

    assert_eq!(
        Opt::parse_from(&["test", "remove"]),
        Opt {
            global: false,
            sub: Subcommands::Remove
        }
    );

    assert_eq!(
        Opt::parse_from(&["test", "global"]),
        Opt {
            global: false,
            sub: Subcommands::Global(GlobalCmd { global: false })
        }
    );

    assert_eq!(
        Opt::parse_from(&["test", "global", "--global"]),
        Opt {
            global: true,
            sub: Subcommands::Global(GlobalCmd { global: true })
        }
    );

    assert!(Opt::try_parse_from(&["test"]).is_err());

    assert_eq!(
        Opt::try_parse_from(&["test", "git", "status"]).unwrap(),
        Opt {
            global: false,
            sub: Subcommands::Other(vec!["git".into(), "status".into()])
        }
    );
}

#[test]
fn external_subcommand_os_string() {
    use std::ffi::OsString;

    #[derive(Debug, PartialEq, Clap)]
    struct Opt {
        #[clap(subcommand)]
        sub: Subcommands,
    }

    #[derive(Debug, PartialEq, Clap)]
    enum Subcommands {
        #[clap(external_subcommand)]
        Other(Vec<OsString>),
    }

    assert_eq!(
        Opt::try_parse_from(&["test", "git", "status"]).unwrap(),
        Opt {
            sub: Subcommands::Other(vec!["git".into(), "status".into()])
        }
    );

    assert!(Opt::try_parse_from(&["test"]).is_err());
}

#[test]
fn external_subcommand_optional() {
    #[derive(Debug, PartialEq, Clap)]
    struct Opt {
        #[clap(subcommand)]
        sub: Option<Subcommands>,
    }

    #[derive(Debug, PartialEq, Clap)]
    enum Subcommands {
        #[clap(external_subcommand)]
        Other(Vec<String>),
    }

    assert_eq!(
        Opt::try_parse_from(&["test", "git", "status"]).unwrap(),
        Opt {
            sub: Some(Subcommands::Other(vec!["git".into(), "status".into()]))
        }
    );

    assert_eq!(Opt::try_parse_from(&["test"]).unwrap(), Opt { sub: None });
}

// #[test]
// #[ignore] // FIXME (@CreepySkeleton)
// fn enum_in_enum_subsubcommand() {
//     #[derive(Clap, Debug, PartialEq)]
//     pub enum Opt {
//         Daemon(DaemonCommand),
//     }

//     #[derive(Clap, Debug, PartialEq)]
//     pub enum DaemonCommand {
//         Start,
//         Stop,
//     }

//     let result = Opt::try_parse_from(&["test"]);
//     assert!(result.is_err());

//     let result = Opt::try_parse_from(&["test", "daemon"]);
//     assert!(result.is_err());

//     let result = Opt::parse_from(&["test", "daemon", "start"]);
//     assert_eq!(Opt::Daemon(DaemonCommand::Start), result);
// }

// WTF??? It tests that `flatten` is essentially a synonym
// for `subcommand`. It's not documented and I doubt it was supposed to
// work this way.
// #[test]
// #[ignore]
// fn flatten_enum() {
//     #[derive(Clap, Debug, PartialEq)]
//     struct Opt {
//         #[clap(flatten)]
//         sub_cmd: SubCmd,
//     }
//     #[derive(Clap, Debug, PartialEq)]
//     enum SubCmd {
//         Foo,
//         Bar,
//     }

//     assert!(Opt::try_parse_from(&["test"]).is_err());
//     assert_eq!(
//         Opt::parse_from(&["test", "foo"]),
//         Opt {
//             sub_cmd: SubCmd::Foo
//         }
//     );
// }
<commit_msg>Create a separate test for from_global<commit_after>// Copyright 2018 Guillaume Pinot (@TeXitoi) <texitoi@texitoi.eu>,
// Kevin Knapp (@kbknapp) <kbknapp@gmail.com>, and
// Andrew Hobden (@hoverbear) <andrew@hoverbear.org>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
//
// This work was derived from Structopt (https://github.com/TeXitoi/structopt)
// commit#ea76fa1b1b273e65e3b0b1046643715b49bec51f which is licensed under the
// MIT/Apache 2.0 license.

mod utils;

use clap::Clap;
use utils::*;

#[derive(Clap, PartialEq, Debug)]
enum Opt {
    /// Fetch stuff from GitHub
    Fetch {
        #[clap(long)]
        all: bool,
        #[clap(short, long)]
        /// Overwrite local branches.
        force: bool,
        repo: String,
    },

    Add {
        #[clap(short, long)]
        interactive: bool,
        #[clap(short, long)]
        verbose: bool,
    },
}

#[test]
fn test_fetch() {
    assert_eq!(
        Opt::Fetch {
            all: true,
            force: false,
            repo: "origin".to_string()
        },
        Opt::parse_from(&["test", "fetch", "--all", "origin"])
    );
    assert_eq!(
        Opt::Fetch {
            all: false,
            force: true,
            repo: "origin".to_string()
        },
        Opt::parse_from(&["test", "fetch", "-f", "origin"])
    );
}

#[test]
fn test_add() {
    assert_eq!(
        Opt::Add {
            interactive: false,
            verbose: false
        },
        Opt::parse_from(&["test", "add"])
    );
    assert_eq!(
        Opt::Add {
            interactive: true,
            verbose: true
        },
        Opt::parse_from(&["test", "add", "-i", "-v"])
    );
}

#[test]
fn test_no_parse() {
    let result = Opt::try_parse_from(&["test", "badcmd", "-i", "-v"]);
    assert!(result.is_err());

    let result = Opt::try_parse_from(&["test", "add", "--badoption"]);
    assert!(result.is_err());

    let result = Opt::try_parse_from(&["test"]);
    assert!(result.is_err());
}

#[derive(Clap, PartialEq, Debug)]
enum Opt2 {
    DoSomething { arg: String },
}

#[test]
/// This test is specifically to make sure that hyphenated subcommands get
/// processed correctly.
fn test_hyphenated_subcommands() {
    assert_eq!(
        Opt2::DoSomething {
            arg: "blah".to_string()
        },
        Opt2::parse_from(&["test", "do-something", "blah"])
    );
}

#[derive(Clap, PartialEq, Debug)]
enum Opt3 {
    Add,
    Init,
    Fetch,
}

#[test]
fn test_null_commands() {
    assert_eq!(Opt3::Add, Opt3::parse_from(&["test", "add"]));
    assert_eq!(Opt3::Init, Opt3::parse_from(&["test", "init"]));
    assert_eq!(Opt3::Fetch, Opt3::parse_from(&["test", "fetch"]));
}

#[derive(Clap, PartialEq, Debug)]
#[clap(about = "Not shown")]
struct Add {
    file: String,
}
/// Not shown
#[derive(Clap, PartialEq, Debug)]
struct Fetch {
    remote: String,
}
#[derive(Clap, PartialEq, Debug)]
enum Opt4 {
    // Not shown
    /// Add a file
    Add(Add),
    Init,
    /// download history from remote
    Fetch(Fetch),
}

#[test]
fn test_tuple_commands() {
    assert_eq!(
        Opt4::Add(Add {
            file: "f".to_string()
        }),
        Opt4::parse_from(&["test", "add", "f"])
    );
    assert_eq!(Opt4::Init, Opt4::parse_from(&["test", "init"]));
    assert_eq!(
        Opt4::Fetch(Fetch {
            remote: "origin".to_string()
        }),
        Opt4::parse_from(&["test", "fetch", "origin"])
    );

    let output = get_long_help::<Opt4>();

    assert!(output.contains("download history from remote"));
    assert!(output.contains("Add a file"));
    assert!(!output.contains("Not shown"));
}

#[test]
fn global_passed_down() {
    #[derive(Debug, PartialEq, Clap)]
    struct Opt {
        #[clap(global = true, long)]
        global: bool,
        #[clap(subcommand)]
        sub: Subcommands,
    }

    #[derive(Debug, PartialEq, Clap)]
    enum Subcommands {
        Add,
        Global(GlobalCmd),
    }

    #[derive(Debug, PartialEq, Clap)]
    struct GlobalCmd {
        #[clap(from_global)]
        global: bool,
    }

    assert_eq!(
        Opt::parse_from(&["test", "global"]),
        Opt {
            global: false,
            sub: Subcommands::Global(GlobalCmd { global: false })
        }
    );

    assert_eq!(
        Opt::parse_from(&["test", "global", "--global"]),
        Opt {
            global: true,
            sub: Subcommands::Global(GlobalCmd { global: true })
        }
    );
}

#[test]
fn external_subcommand() {
    #[derive(Debug, PartialEq, Clap)]
    struct Opt {
        #[clap(subcommand)]
        sub: Subcommands,
    }

    #[derive(Debug, PartialEq, Clap)]
    enum Subcommands {
        Add,
        Remove,
        #[clap(external_subcommand)]
        Other(Vec<String>),
    }

    assert_eq!(
        Opt::parse_from(&["test", "add"]),
        Opt {
            sub: Subcommands::Add
        }
    );

    assert_eq!(
        Opt::parse_from(&["test", "remove"]),
        Opt {
            sub: Subcommands::Remove
        }
    );

    assert!(Opt::try_parse_from(&["test"]).is_err());

    assert_eq!(
        Opt::try_parse_from(&["test", "git", "status"]).unwrap(),
        Opt {
            sub: Subcommands::Other(vec!["git".into(), "status".into()])
        }
    );
}

#[test]
fn external_subcommand_os_string() {
    use std::ffi::OsString;

    #[derive(Debug, PartialEq, Clap)]
    struct Opt {
        #[clap(subcommand)]
        sub: Subcommands,
    }

    #[derive(Debug, PartialEq, Clap)]
    enum Subcommands {
        #[clap(external_subcommand)]
        Other(Vec<OsString>),
    }

    assert_eq!(
        Opt::try_parse_from(&["test", "git", "status"]).unwrap(),
        Opt {
            sub: Subcommands::Other(vec!["git".into(), "status".into()])
        }
    );

    assert!(Opt::try_parse_from(&["test"]).is_err());
}

#[test]
fn external_subcommand_optional() {
    #[derive(Debug, PartialEq, Clap)]
    struct Opt {
        #[clap(subcommand)]
        sub: Option<Subcommands>,
    }

    #[derive(Debug, PartialEq, Clap)]
    enum Subcommands {
        #[clap(external_subcommand)]
        Other(Vec<String>),
    }

    assert_eq!(
        Opt::try_parse_from(&["test", "git", "status"]).unwrap(),
        Opt {
            sub: Some(Subcommands::Other(vec!["git".into(), "status".into()]))
        }
    );

    assert_eq!(Opt::try_parse_from(&["test"]).unwrap(), Opt { sub: None });
}

// #[test]
// #[ignore] // FIXME (@CreepySkeleton)
// fn enum_in_enum_subsubcommand() {
//     #[derive(Clap, Debug, PartialEq)]
//     pub enum Opt {
//         Daemon(DaemonCommand),
//     }

//     #[derive(Clap, Debug, PartialEq)]
//     pub enum DaemonCommand {
//         Start,
//         Stop,
//     }

//     let result = Opt::try_parse_from(&["test"]);
//     assert!(result.is_err());

//     let result = Opt::try_parse_from(&["test", "daemon"]);
//     assert!(result.is_err());

//     let result = Opt::parse_from(&["test", "daemon", "start"]);
//     assert_eq!(Opt::Daemon(DaemonCommand::Start), result);
// }

// WTF??? It tests that `flatten` is essentially a synonym
// for `subcommand`. It's not documented and I doubt it was supposed to
// work this way.
// #[test]
// #[ignore]
// fn flatten_enum() {
//     #[derive(Clap, Debug, PartialEq)]
//     struct Opt {
//         #[clap(flatten)]
//         sub_cmd: SubCmd,
//     }
//     #[derive(Clap, Debug, PartialEq)]
//     enum SubCmd {
//         Foo,
//         Bar,
//     }

//     assert!(Opt::try_parse_from(&["test"]).is_err());
//     assert_eq!(
//         Opt::parse_from(&["test", "foo"]),
//         Opt {
//             sub_cmd: SubCmd::Foo
//         }
//     );
// }
<|endoftext|>


<commit_before>use std;
use std::collections::HashMap;
use std::cell::RefCell;
use typed_arena::Arena;
use llvm_sys::prelude::*;
use llvm_sys::core::*;

pub struct TypeContext<'t> {
    backing_store: Arena<TypeVariant<'t>>,
    type_references: RefCell<HashMap<TypeVariant<'t>, &'t TypeVariant<'t>>>,
}

impl<'t> TypeContext<'t> {
    pub fn new() -> Self {
        TypeContext {
            backing_store: Arena::new(),
            type_references: RefCell::new(HashMap::new()),
        }
    }

    fn get(&'t self, variant: TypeVariant<'t>) -> &'t TypeVariant<'t> {
        if let Some(var) = self.type_references.borrow().get(&variant) {
            return var;
        }

        let var = self.backing_store.alloc(variant);
        self.type_references.borrow_mut().insert(variant, var);
        var
    }
}

#[derive(Copy, Clone)]
pub struct Type<'t> {
    pub variant: &'t TypeVariant<'t>,
}

impl<'t> PartialEq for Type<'t> {
    fn eq(&self, rhs: &Self) -> bool {
        self.variant as *const _ == rhs.variant as *const _
    }
}

impl<'t> Eq for Type<'t> { }

impl<'t> std::hash::Hash for Type<'t> {
    fn hash<H>(&self, state: &mut H) where H: std::hash::Hasher {
        (self.variant as *const _).hash(state);
    }
}

impl<'t> std::fmt::Debug for Type<'t> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        match *self.variant {
            TypeVariant::SInt(size) => write!(f, "SInt({:?})", size),
            TypeVariant::UInt(size) => write!(f, "SInt({:?})", size),
            TypeVariant::Bool => write!(f, "Bool"),
            TypeVariant::Unit => write!(f, "Unit"),
            TypeVariant::Diverging => write!(f, "Diverging"),
            TypeVariant::Reference(inner) =>
                write!(f, "Ref({:?})", inner),
            TypeVariant::Infer(i) => write!(f, "Infer({:?})", i),
            TypeVariant::InferInt(i) => write!(f, "InferInt({:?})", i),
        }
    }
}

// -- constructors --
impl<'t> Type<'t> {
    pub fn infer(ctxt: &'t TypeContext<'t>) -> Self {
        Type {
            variant: ctxt.get(TypeVariant::Infer(None)),
        }
    }

    pub fn infer_int(ctxt: &'t TypeContext<'t>) -> Self {
        Type {
            variant: ctxt.get(TypeVariant::InferInt(None)),
        }
    }

    pub fn sint(int: Int, ctxt: &'t TypeContext<'t>) -> Self {
        Type {
            variant: ctxt.get(TypeVariant::SInt(int)),
        }
    }
    pub fn uint(int: Int, ctxt: &'t TypeContext<'t>) -> Self {
        Type {
            variant: ctxt.get(TypeVariant::UInt(int)),
        }
    }
    pub fn bool(ctxt: &'t TypeContext<'t>) -> Self {
        Type {
            variant: ctxt.get(TypeVariant::Bool),
        }
    }
    pub fn unit(ctxt: &'t TypeContext<'t>) -> Self {
        Type {
            variant: ctxt.get(TypeVariant::Unit),
        }
    }
    pub fn diverging(ctxt: &'t TypeContext<'t>) -> Self {
        Type {
            variant: ctxt.get(TypeVariant::Diverging),
        }
    }

    pub fn ref_(ty: Type<'t>, ctxt: &'t TypeContext<'t>) -> Self {
        Type {
            variant: ctxt.get(TypeVariant::Reference(ty))
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum TypeVariant<'t> {
    SInt(Int),
    UInt(Int),
    Bool,
    Unit,

    Diverging,

    Reference(Type<'t>),

    Infer(Option<u32>),
    InferInt(Option<u32>),
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum Int {
    I8,
    I16,
    I32,
    I64,
}

#[derive(Debug, PartialEq, Eq)]
pub struct Function<'t> {
    input: Vec<Type<'t>>,
    output: Type<'t>,
}

impl<'t> Function<'t> {
    pub fn new(input: Vec<Type<'t>>, output: Type<'t>) -> Self {
        Function {
            input: input,
            output: output,
        }
    }

    pub fn input(&self) -> &[Type<'t>] {
        &self.input
    }

    pub fn output(&self) -> Type<'t> {
        self.output
    }

    pub fn to_llvm(&self) -> LLVMTypeRef {
        unsafe {
            let mut args = self.input.iter().map(|a| a.to_llvm())
                .collect::<Vec<_>>();
            LLVMFunctionType(self.output.to_llvm_ret(), args.as_mut_ptr(),
                args.len() as u32, false as LLVMBool)
        }
    }
}

impl<'t> Type<'t> {
    pub fn to_llvm(&self) -> LLVMTypeRef {
        unsafe {
            match *self.variant {
                TypeVariant::SInt(ref size) | TypeVariant::UInt(ref size)
                    => LLVMIntType(size.size()),
                TypeVariant::Bool => LLVMInt1Type(),
                TypeVariant::Unit => LLVMStructType(std::ptr::null_mut(), 0,
                        false as LLVMBool),
                TypeVariant::Reference(inner) =>
                    LLVMPointerType(inner.to_llvm(), 0),
                TypeVariant::Diverging
                    => panic!("ICE: Attempted to get the LLVM type of \
                        Diverging"),
                TypeVariant::Infer(_) | TypeVariant::InferInt(_) =>
                    panic!("ICE: Attempted to get the LLVM type of an \
                        inference variable: {:?}", self),
            }
        }
    }

    pub fn to_llvm_ret(&self) -> LLVMTypeRef {
        unsafe {
            match *self.variant {
                TypeVariant::SInt(ref size) | TypeVariant::UInt(ref size) =>
                    LLVMIntType(size.size()),
                TypeVariant::Bool => LLVMInt1Type(),
                TypeVariant::Unit => LLVMVoidType(),
                TypeVariant::Reference(inner) =>
                    LLVMPointerType(inner.to_llvm(), 0),
                TypeVariant::Diverging => LLVMVoidType(),
                TypeVariant::Infer(_) | TypeVariant::InferInt(_) =>
                    panic!("ICE: Attempted to get the LLVM return type of an \
                        inference variable"),
            }
        }
    }

    pub fn is_final_type(&self) -> bool {
        match *self.variant {
            TypeVariant::SInt(_) | TypeVariant::UInt(_) | TypeVariant::Bool
            | TypeVariant::Unit | TypeVariant::Diverging
                => true,
            TypeVariant::Reference(inner) => inner.is_final_type(),
            TypeVariant::Infer(_) | TypeVariant::InferInt(_) => false,
        }
    }

    pub fn generate_inference_id(&mut self, uf: &mut UnionFind<'t>,
            ctxt: &'t TypeContext<'t>) {
        self.variant = self.get_inference_type(uf, ctxt);
    }

    fn get_inference_type(&self, uf: &mut UnionFind<'t>,
            ctxt: &'t TypeContext<'t>) -> &'t TypeVariant<'t> {
        match *self.variant {
            TypeVariant::Infer(None) => {
                ctxt.get(TypeVariant::Infer(Some(uf.next_id())))
            }
            TypeVariant::InferInt(None) => {
                ctxt.get(TypeVariant::InferInt(Some(uf.next_id())))
            }
            TypeVariant::Reference(inner) => {
                ctxt.get(TypeVariant::Reference(
                    Type {
                        variant: inner.get_inference_type(uf, ctxt),
                    }
                ))
            }
            ref t @ TypeVariant::SInt(_) | ref t @ TypeVariant::UInt(_)
            | ref t @ TypeVariant::Bool | ref t @ TypeVariant::Diverging
            | ref t @ TypeVariant::Unit | ref t @ TypeVariant::Infer(Some(_))
            | ref t @ TypeVariant::InferInt(Some(_)) => t,
        }
    }

    pub fn finalize(&mut self, uf: &mut UnionFind<'t>,
            ctxt: &'t TypeContext<'t>) -> Result<(), ()> {
        *self = match self.get_final_ty(uf, ctxt) {
            Some(t) => t,
            None => return Err(())
        };
        Ok(())
    }

    fn get_final_ty(&self, uf: &mut UnionFind<'t>, ctxt: &'t TypeContext<'t>)
            -> Option<Type<'t>> {
        match *self.variant {
            TypeVariant::SInt(_) | TypeVariant::UInt(_) | TypeVariant::Bool
            | TypeVariant::Unit | TypeVariant::Diverging => {
                Some(*self)
            }
            TypeVariant::Reference(inner) => {
                match inner.get_final_ty(uf, ctxt) {
                    Some(inner) => Some(Type::ref_(inner, ctxt)),
                    None => None,
                }
            }
            TypeVariant::Infer(_) | TypeVariant::InferInt(_) => {
                match uf.resolve(*self) {
                    Some(t) => t.get_final_ty(uf, ctxt),
                    None => return None,
                }
            }
        }
    }
}

impl<'t> std::fmt::Display for Type<'t> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        let s = match *self.variant {
            TypeVariant::SInt(Int::I8) => "s8",
            TypeVariant::SInt(Int::I16) => "s16",
            TypeVariant::SInt(Int::I32) => "s32",
            TypeVariant::SInt(Int::I64) => "s64",
            TypeVariant::UInt(Int::I8) => "u8",
            TypeVariant::UInt(Int::I16) => "u16",
            TypeVariant::UInt(Int::I32) => "u32",
            TypeVariant::UInt(Int::I64) => "u64",
            TypeVariant::Bool => "bool",
            TypeVariant::Unit => "()",
            TypeVariant::Diverging => "!",
            TypeVariant::Reference(inner) => return write!(f, "&{}", inner),
            TypeVariant::Infer(_) | TypeVariant::InferInt(_) => "_",
        };
        write!(f, "{}", s)
    }
}

impl Int {
    #[allow(dead_code)]
    pub fn shift_mask(&self) -> u64 {
        match *self {
            Int::I8 => (1 << 3) - 1,
            Int::I16 => (1 << 4) - 1,
            Int::I32 => (1 << 5) - 1,
            Int::I64 => (1 << 6) - 1,
        }
    }

    pub fn size(&self) -> u32 {
        match *self {
            Int::I8 => 8,
            Int::I16 => 16,
            Int::I32 => 32,
            Int::I64 => 64,
        }
    }
}

pub struct UnionFind<'t> {
    current_id: u32,
    group_parents: Vec<u32>,
    parents_ty: Vec<Option<Type<'t>>>,
}

impl<'t> UnionFind<'t> {
    pub fn new() -> Self {
        UnionFind {
            current_id: 0,
            group_parents: Vec::new(),
            parents_ty: Vec::new(),
        }
    }

    fn union(&mut self, a: u32, b: u32) {
        let b = self.find(b) as usize;
        self.group_parents[b] = self.find(a);
    }

    fn find(&self, mut n: u32) -> u32 {
        while self.group_parents[n as usize] != n {
            n = self.group_parents[n as usize];
        }
        n
    }

    pub fn unify(&mut self, a: Type<'t>, b: Type<'t>) -> Result<(), ()> {
        let a_actual = self.resolve(a);
        let b_actual = self.resolve(b);
        match (a_actual, b_actual) {
            (Some(a), Some(b)) => {
                if a.is_final_type() && b.is_final_type() && a == b {
                    Ok(())
                } else {
                    match (*a.variant, *b.variant) {
                        (TypeVariant::Reference(lhs),
                                TypeVariant::Reference(rhs)) => {
                            self.unify(lhs, rhs)
                        }
                        _ => Err(())
                    }
                }
            }
            (None, None) => {
                match (*a.variant, *b.variant) {
                    (TypeVariant::Infer(Some(lid)),
                            TypeVariant::Infer(Some(rid)))
                    | (TypeVariant::Infer(Some(lid)),
                            TypeVariant::InferInt(Some(rid)))
                    | (TypeVariant::InferInt(Some(lid)),
                            TypeVariant::Infer(Some(rid)))
                    | (TypeVariant::InferInt(Some(lid)),
                            TypeVariant::InferInt(Some(rid))) => {
                        self.union(lid, rid);
                        Ok(())
                    }
                    (l, r)
                        => panic!("actual ty isn't working: {:?}, {:?}", l, r)
                }
            }
            (Some(ty), None) => {
                match *b.variant {
                    TypeVariant::Infer(Some(id)) => {
                        let id = self.find(id) as usize;
                        self.parents_ty[id] = Some(ty);
                        Ok(())
                    }
                    TypeVariant::InferInt(Some(id)) => {
                        match *ty.variant {
                            TypeVariant::UInt(_) | TypeVariant::SInt(_) => {
                                let id = self.find(id) as usize;
                                self.parents_ty[id] = Some(ty);
                                Ok(())
                            }
                            _ => {
                                Err(())
                            }
                        }
                    }
                    t => panic!("actual ty isn't working: {:?}", t)
                }
            }
            (None, Some(ty)) => {
                match *a.variant {
                    TypeVariant::Infer(Some(id)) => {
                        let id = self.find(id) as usize;
                        self.parents_ty[id] = Some(ty);
                        Ok(())
                    }
                    TypeVariant::InferInt(Some(id)) => {
                        match *ty.variant {
                            TypeVariant::UInt(_) | TypeVariant::SInt(_) => {
                                let id = self.find(id) as usize;
                                self.parents_ty[id] = Some(ty);
                                Ok(())
                            }
                            _ => {
                                Err(())
                            }
                        }
                    }
                    t => panic!("actual ty isn't working: {:?}", t)
                }
            }
        }
    }

    pub fn resolve(&self, ty: Type<'t>) -> Option<Type<'t>> {
        match *ty.variant {
            TypeVariant::Infer(Some(id)) | TypeVariant::InferInt(Some(id)) => {
                self.parents_ty[self.find(id) as usize]
            }

            TypeVariant::Infer(None) | TypeVariant::InferInt(None) => {
                None
            }

            _ => Some(ty)
        }
    }

    fn next_id(&mut self) -> u32 {
        if self.current_id == u32::max_value() {
            panic!()
        } else {
            self.group_parents.push(self.current_id);
            self.parents_ty.push(None);
            self.current_id += 1;
            self.current_id - 1
        }
    }
}
<commit_msg>Add some better errors to unify<commit_after>use std;
use std::collections::HashMap;
use std::cell::RefCell;
use typed_arena::Arena;
use llvm_sys::prelude::*;
use llvm_sys::core::*;

pub struct TypeContext<'t> {
    backing_store: Arena<TypeVariant<'t>>,
    type_references: RefCell<HashMap<TypeVariant<'t>, &'t TypeVariant<'t>>>,
}

impl<'t> TypeContext<'t> {
    pub fn new() -> Self {
        TypeContext {
            backing_store: Arena::new(),
            type_references: RefCell::new(HashMap::new()),
        }
    }

    fn get(&'t self, variant: TypeVariant<'t>) -> &'t TypeVariant<'t> {
        if let Some(var) = self.type_references.borrow().get(&variant) {
            return var;
        }

        let var = self.backing_store.alloc(variant);
        self.type_references.borrow_mut().insert(variant, var);
        var
    }
}

#[derive(Copy, Clone)]
pub struct Type<'t> {
    pub variant: &'t TypeVariant<'t>,
}

impl<'t> PartialEq for Type<'t> {
    fn eq(&self, rhs: &Self) -> bool {
        self.variant as *const _ == rhs.variant as *const _
    }
}

impl<'t> Eq for Type<'t> { }

impl<'t> std::hash::Hash for Type<'t> {
    fn hash<H>(&self, state: &mut H) where H: std::hash::Hasher {
        (self.variant as *const _).hash(state);
    }
}

impl<'t> std::fmt::Debug for Type<'t> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        match *self.variant {
            TypeVariant::SInt(size) => write!(f, "SInt({:?})", size),
            TypeVariant::UInt(size) => write!(f, "SInt({:?})", size),
            TypeVariant::Bool => write!(f, "Bool"),
            TypeVariant::Unit => write!(f, "Unit"),
            TypeVariant::Diverging => write!(f, "Diverging"),
            TypeVariant::Reference(inner) =>
                write!(f, "Ref({:?})", inner),
            TypeVariant::Infer(i) => write!(f, "Infer({:?})", i),
            TypeVariant::InferInt(i) => write!(f, "InferInt({:?})", i),
        }
    }
}

// -- constructors --
impl<'t> Type<'t> {
    pub fn infer(ctxt: &'t TypeContext<'t>) -> Self {
        Type {
            variant: ctxt.get(TypeVariant::Infer(None)),
        }
    }

    pub fn infer_int(ctxt: &'t TypeContext<'t>) -> Self {
        Type {
            variant: ctxt.get(TypeVariant::InferInt(None)),
        }
    }

    pub fn sint(int: Int, ctxt: &'t TypeContext<'t>) -> Self {
        Type {
            variant: ctxt.get(TypeVariant::SInt(int)),
        }
    }
    pub fn uint(int: Int, ctxt: &'t TypeContext<'t>) -> Self {
        Type {
            variant: ctxt.get(TypeVariant::UInt(int)),
        }
    }
    pub fn bool(ctxt: &'t TypeContext<'t>) -> Self {
        Type {
            variant: ctxt.get(TypeVariant::Bool),
        }
    }
    pub fn unit(ctxt: &'t TypeContext<'t>) -> Self {
        Type {
            variant: ctxt.get(TypeVariant::Unit),
        }
    }
    pub fn diverging(ctxt: &'t TypeContext<'t>) -> Self {
        Type {
            variant: ctxt.get(TypeVariant::Diverging),
        }
    }

    pub fn ref_(ty: Type<'t>, ctxt: &'t TypeContext<'t>) -> Self {
        Type {
            variant: ctxt.get(TypeVariant::Reference(ty))
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum TypeVariant<'t> {
    SInt(Int),
    UInt(Int),
    Bool,
    Unit,

    Diverging,

    Reference(Type<'t>),

    Infer(Option<u32>),
    InferInt(Option<u32>),
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum Int {
    I8,
    I16,
    I32,
    I64,
}

#[derive(Debug, PartialEq, Eq)]
pub struct Function<'t> {
    input: Vec<Type<'t>>,
    output: Type<'t>,
}

impl<'t> Function<'t> {
    pub fn new(input: Vec<Type<'t>>, output: Type<'t>) -> Self {
        Function {
            input: input,
            output: output,
        }
    }

    pub fn input(&self) -> &[Type<'t>] {
        &self.input
    }

    pub fn output(&self) -> Type<'t> {
        self.output
    }

    pub fn to_llvm(&self) -> LLVMTypeRef {
        unsafe {
            let mut args = self.input.iter().map(|a| a.to_llvm())
                .collect::<Vec<_>>();
            LLVMFunctionType(self.output.to_llvm_ret(), args.as_mut_ptr(),
                args.len() as u32, false as LLVMBool)
        }
    }
}

impl<'t> Type<'t> {
    pub fn to_llvm(&self) -> LLVMTypeRef {
        unsafe {
            match *self.variant {
                TypeVariant::SInt(ref size) | TypeVariant::UInt(ref size)
                    => LLVMIntType(size.size()),
                TypeVariant::Bool => LLVMInt1Type(),
                TypeVariant::Unit => LLVMStructType(std::ptr::null_mut(), 0,
                        false as LLVMBool),
                TypeVariant::Reference(inner) =>
                    LLVMPointerType(inner.to_llvm(), 0),
                TypeVariant::Diverging
                    => panic!("ICE: Attempted to get the LLVM type of \
                        Diverging"),
                TypeVariant::Infer(_) | TypeVariant::InferInt(_) =>
                    panic!("ICE: Attempted to get the LLVM type of an \
                        inference variable: {:?}", self),
            }
        }
    }

    pub fn to_llvm_ret(&self) -> LLVMTypeRef {
        unsafe {
            match *self.variant {
                TypeVariant::SInt(ref size) | TypeVariant::UInt(ref size) =>
                    LLVMIntType(size.size()),
                TypeVariant::Bool => LLVMInt1Type(),
                TypeVariant::Unit => LLVMVoidType(),
                TypeVariant::Reference(inner) =>
                    LLVMPointerType(inner.to_llvm(), 0),
                TypeVariant::Diverging => LLVMVoidType(),
                TypeVariant::Infer(_) | TypeVariant::InferInt(_) =>
                    panic!("ICE: Attempted to get the LLVM return type of an \
                        inference variable"),
            }
        }
    }

    pub fn is_final_type(&self) -> bool {
        match *self.variant {
            TypeVariant::SInt(_) | TypeVariant::UInt(_) | TypeVariant::Bool
            | TypeVariant::Unit | TypeVariant::Diverging
                => true,
            TypeVariant::Reference(inner) => inner.is_final_type(),
            TypeVariant::Infer(_) | TypeVariant::InferInt(_) => false,
        }
    }

    pub fn generate_inference_id(&mut self, uf: &mut UnionFind<'t>,
            ctxt: &'t TypeContext<'t>) {
        self.variant = self.get_inference_type(uf, ctxt);
    }

    fn get_inference_type(&self, uf: &mut UnionFind<'t>,
            ctxt: &'t TypeContext<'t>) -> &'t TypeVariant<'t> {
        match *self.variant {
            TypeVariant::Infer(None) => {
                ctxt.get(TypeVariant::Infer(Some(uf.next_id())))
            }
            TypeVariant::InferInt(None) => {
                ctxt.get(TypeVariant::InferInt(Some(uf.next_id())))
            }
            TypeVariant::Reference(inner) => {
                ctxt.get(TypeVariant::Reference(
                    Type {
                        variant: inner.get_inference_type(uf, ctxt),
                    }
                ))
            }
            ref t @ TypeVariant::SInt(_) | ref t @ TypeVariant::UInt(_)
            | ref t @ TypeVariant::Bool | ref t @ TypeVariant::Diverging
            | ref t @ TypeVariant::Unit | ref t @ TypeVariant::Infer(Some(_))
            | ref t @ TypeVariant::InferInt(Some(_)) => t,
        }
    }

    pub fn finalize(&mut self, uf: &mut UnionFind<'t>,
            ctxt: &'t TypeContext<'t>) -> Result<(), ()> {
        *self = match self.get_final_ty(uf, ctxt) {
            Some(t) => t,
            None => return Err(())
        };
        Ok(())
    }

    fn get_final_ty(&self, uf: &mut UnionFind<'t>, ctxt: &'t TypeContext<'t>)
            -> Option<Type<'t>> {
        match *self.variant {
            TypeVariant::SInt(_) | TypeVariant::UInt(_) | TypeVariant::Bool
            | TypeVariant::Unit | TypeVariant::Diverging => {
                Some(*self)
            }
            TypeVariant::Reference(inner) => {
                match inner.get_final_ty(uf, ctxt) {
                    Some(inner) => Some(Type::ref_(inner, ctxt)),
                    None => None,
                }
            }
            TypeVariant::Infer(_) | TypeVariant::InferInt(_) => {
                match uf.resolve(*self) {
                    Some(t) => t.get_final_ty(uf, ctxt),
                    None => return None,
                }
            }
        }
    }
}

impl<'t> std::fmt::Display for Type<'t> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        let s = match *self.variant {
            TypeVariant::SInt(Int::I8) => "s8",
            TypeVariant::SInt(Int::I16) => "s16",
            TypeVariant::SInt(Int::I32) => "s32",
            TypeVariant::SInt(Int::I64) => "s64",
            TypeVariant::UInt(Int::I8) => "u8",
            TypeVariant::UInt(Int::I16) => "u16",
            TypeVariant::UInt(Int::I32) => "u32",
            TypeVariant::UInt(Int::I64) => "u64",
            TypeVariant::Bool => "bool",
            TypeVariant::Unit => "()",
            TypeVariant::Diverging => "!",
            TypeVariant::Reference(inner) => return write!(f, "&{}", inner),
            TypeVariant::Infer(_) | TypeVariant::InferInt(_) => "_",
        };
        write!(f, "{}", s)
    }
}

impl Int {
    #[allow(dead_code)]
    pub fn shift_mask(&self) -> u64 {
        match *self {
            Int::I8 => (1 << 3) - 1,
            Int::I16 => (1 << 4) - 1,
            Int::I32 => (1 << 5) - 1,
            Int::I64 => (1 << 6) - 1,
        }
    }

    pub fn size(&self) -> u32 {
        match *self {
            Int::I8 => 8,
            Int::I16 => 16,
            Int::I32 => 32,
            Int::I64 => 64,
        }
    }
}

pub struct UnionFind<'t> {
    current_id: u32,
    group_parents: Vec<u32>,
    parents_ty: Vec<Option<Type<'t>>>,
}

impl<'t> UnionFind<'t> {
    pub fn new() -> Self {
        UnionFind {
            current_id: 0,
            group_parents: Vec::new(),
            parents_ty: Vec::new(),
        }
    }

    fn union(&mut self, a: u32, b: u32) {
        let b = self.find(b) as usize;
        self.group_parents[b] = self.find(a);
    }

    fn find(&self, mut n: u32) -> u32 {
        while self.group_parents[n as usize] != n {
            n = self.group_parents[n as usize];
        }
        n
    }

    pub fn unify(&mut self, a: Type<'t>, b: Type<'t>) -> Result<(), ()> {
        let a_actual = self.resolve(a);
        let b_actual = self.resolve(b);
        match (a_actual, b_actual) {
            (Some(a), Some(b)) => {
                if a.is_final_type() && b.is_final_type() && a == b {
                    Ok(())
                } else {
                    match (*a.variant, *b.variant) {
                        (TypeVariant::Reference(lhs),
                                TypeVariant::Reference(rhs)) => {
                            self.unify(lhs, rhs)
                        }
                        _ => Err(())
                    }
                }
            }
            (None, None) => {
                match (*a.variant, *b.variant) {
                    (TypeVariant::Infer(Some(lid)),
                            TypeVariant::Infer(Some(rid)))
                    | (TypeVariant::Infer(Some(lid)),
                            TypeVariant::InferInt(Some(rid)))
                    | (TypeVariant::InferInt(Some(lid)),
                            TypeVariant::Infer(Some(rid)))
                    | (TypeVariant::InferInt(Some(lid)),
                            TypeVariant::InferInt(Some(rid))) => {
                        self.union(lid, rid);
                        Ok(())
                    }
                    (lhs @ TypeVariant::Infer(None), rhs)
                    | (lhs @ TypeVariant::InferInt(None), rhs)
                    | (rhs, lhs @ TypeVariant::Infer(None))
                    | (lhs, rhs @ TypeVariant::InferInt(None)) =>
                        panic!("ICE: attempted to unify {:?} with {:?}",
                            lhs, rhs),
                    (l, r)
                        => panic!("actual ty isn't working: {:?}, {:?}", l, r)
                }
            }
            (Some(ty), None) => {
                match *b.variant {
                    TypeVariant::Infer(Some(id)) => {
                        let id = self.find(id) as usize;
                        self.parents_ty[id] = Some(ty);
                        Ok(())
                    }
                    TypeVariant::InferInt(Some(id)) => {
                        match *ty.variant {
                            TypeVariant::UInt(_) | TypeVariant::SInt(_) => {
                                let id = self.find(id) as usize;
                                self.parents_ty[id] = Some(ty);
                                Ok(())
                            }
                            _ => {
                                Err(())
                            }
                        }
                    }
                    t @ TypeVariant::Infer(None)
                    | t @ TypeVariant::InferInt(None) =>
                        panic!("ICE: attempted to unify {:?} with {:?}",
                            ty, t),
                    t => panic!("ICE: resolve isn't working: {:?}", t)
                }
            }
            (None, Some(ty)) => {
                match *a.variant {
                    TypeVariant::Infer(Some(id)) => {
                        let id = self.find(id) as usize;
                        self.parents_ty[id] = Some(ty);
                        Ok(())
                    }
                    TypeVariant::InferInt(Some(id)) => {
                        match *ty.variant {
                            TypeVariant::UInt(_) | TypeVariant::SInt(_) => {
                                let id = self.find(id) as usize;
                                self.parents_ty[id] = Some(ty);
                                Ok(())
                            }
                            _ => {
                                Err(())
                            }
                        }
                    }
                    t @ TypeVariant::Infer(None)
                    | t @ TypeVariant::InferInt(None) =>
                        panic!("ICE: attempted to unify {:?} with {:?}",
                            ty, t),
                    t => panic!("ICE: resolve isn't working: {:?}", t)
                }
            }
        }
    }

    pub fn resolve(&self, ty: Type<'t>) -> Option<Type<'t>> {
        match *ty.variant {
            TypeVariant::Infer(Some(id)) | TypeVariant::InferInt(Some(id)) => {
                self.parents_ty[self.find(id) as usize]
            }

            TypeVariant::Infer(None) | TypeVariant::InferInt(None) => {
                None
            }

            _ => Some(ty)
        }
    }

    fn next_id(&mut self) -> u32 {
        if self.current_id == u32::max_value() {
            panic!()
        } else {
            self.group_parents.push(self.current_id);
            self.parents_ty.push(None);
            self.current_id += 1;
            self.current_id - 1
        }
    }
}
<|endoftext|>