<commit_before>pub mod search;

use std::thread;
use std::cmp::min;
use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver};
use std::time::SystemTime;
use std::num::Wrapping;
use basetypes::*;
use uci::{UciEngine, UciEngineFactory, EngineReply, OptionName, OptionDescription};
use position::Position;
use chess_move::*;
use tt::*;

const VERSION: &'static str = "0.1";
const MAX_DEPTH: u8 = 126;

/// Implements `UciEngine` trait.
pub struct Engine {
    tt: Arc<TranspositionTable>,
    reply_queue: VecDeque<EngineReply>,
    position: Position,
    is_thinking: bool,
    is_pondering: bool,

    // Tells the engine if it will be allowed to ponder. This option
    // is needed because the engine might change its time management
    // algorithm when pondering is allowed.
    pondering_is_allowed: bool,

    // A channel for sending commands to the search thread.
    commands: Sender<search::Command>,

    // A channel for receiving reports from the search thread.
    reports: Receiver<search::Report>,

    // Search status info
    search_id: usize,
    thinking_since: SystemTime,
    no_reports_since: SystemTime,
    current_depth: u8,
    searched_nodes: NodeCount,
    searched_time: u64, // milliseconds
    nps: u64, // nodes per second
    mangled_pv: bool, // `true` if the primary variation is imperfect
    stop_when: TimeManagement,
}

impl Engine {
    /// Creates a new instance.
    ///
    /// `tt_size_mb` is the preferred size of the transposition
    /// table in Mbytes.
    pub fn new(tt_size_mb: usize) -> Engine {
        let mut tt = TranspositionTable::new();
        tt.resize(tt_size_mb);
        let tt = Arc::new(tt);
        let (commands_tx, commands_rx) = channel();
        let (reports_tx, reports_rx) = channel();
        let tt2 = tt.clone();
        thread::spawn(move || {
            search::run_deepening(tt2, commands_rx, reports_tx);
        });

        Engine {
            tt: tt,
            reply_queue: VecDeque::new(),
            position: Position::from_fen("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 \
                                          1")
                          .ok()
                          .unwrap(),
            is_thinking: false,
            is_pondering: false,
            pondering_is_allowed: false,
            commands: commands_tx,
            reports: reports_rx,
            search_id: 0,
            thinking_since: SystemTime::now(),
            no_reports_since: SystemTime::now(),
            current_depth: 0,
            searched_nodes: 0,
            searched_time: 0,
            nps: 0,
            mangled_pv: false,
            stop_when: TimeManagement::Infinite,
        }
    }
}

impl UciEngine for Engine {
    fn set_option(&mut self, name: &str, value: &str) {
        match name {
            "Ponder" => {
                self.pondering_is_allowed = value == "true";
            }

            // An invalid option. Notice that we do not support
            // re-sizing of the transposition table once the engine
            // had started.
            _ => (),
        }
    }

    fn new_game(&mut self) {
        // Clearing the transposition table would not change anything,
        // so we do nothing.
    }

    fn position(&mut self, fen: &str, moves: &mut Iterator<Item = &str>) {
        if !self.is_thinking {
            if let Ok(p) = Position::from_history(fen, moves) {
                self.position = p;
                self.tt.new_search();
            }
        }
    }

    #[allow(unused_variables)]
    fn go(&mut self,
          searchmoves: Option<Vec<String>>,
          ponder: bool,
          wtime: Option<u64>,
          btime: Option<u64>,
          winc: Option<u64>,
          binc: Option<u64>,
          movestogo: Option<u64>,
          depth: Option<u64>,
          nodes: Option<u64>,
          mate: Option<u64>,
          movetime: Option<u64>,
          infinite: bool) {
        if !self.is_thinking {
            // Note: We ignore "searchmoves" and "mate" parameters.

            self.is_thinking = true;
            self.is_pondering = ponder;
            self.search_id = (Wrapping(self.search_id) + Wrapping(1)).0;
            self.thinking_since = SystemTime::now();
            self.no_reports_since = SystemTime::now();
            self.current_depth = 0;
            self.searched_nodes = 0;
            self.searched_time = 0;
            self.mangled_pv = false;

            // Figure out when we should stop thinking.
            let (time, inc) = if self.position.board().to_move() == WHITE {
                (wtime, winc.unwrap_or(0))
            } else {
                (btime, binc.unwrap_or(0))
            };
            self.stop_when = if infinite {
                TimeManagement::Infinite
            } else if movetime.is_some() {
                TimeManagement::MoveTime(movetime.unwrap())
            } else if nodes.is_some() {
                TimeManagement::Nodes(nodes.unwrap())
            } else if depth.is_some() {
                TimeManagement::Depth(depth.unwrap() as u8)
            } else {
                let time = time.unwrap_or(0);
                let movestogo = movestogo.unwrap_or(40);
                let movetime = (time + inc * movestogo) / movestogo;
                TimeManagement::MoveTimeHint(min(movetime, time / 2))
            };

            // Start deepening search to the maximum possible depth.
            self.commands
                .send(search::Command::Search {
                    search_id: self.search_id,
                    position: self.position.clone(),
                    depth: MAX_DEPTH,
                    lower_bound: -20000,
                    upper_bound: 20000,
                })
                .unwrap();
        }
    }

    fn ponder_hit(&mut self) {
        if self.is_thinking && self.is_pondering {
            self.is_pondering = false;
        }
    }

    fn stop(&mut self) {
        if self.is_thinking {
            self.commands.send(search::Command::Stop).unwrap();

            // TODO: send ponder move as well.
            let best_move = self.get_best_move();
            self.reply_queue.push_back(EngineReply::BestMove {
                best_move: best_move,
                ponder_move: None,
            });

            self.is_thinking = false;
        }
    }

    fn is_thinking(&self) -> bool {
        self.is_thinking
    }

    fn get_reply(&mut self) -> Option<EngineReply> {
        // Check if we should stop thinking.
        if self.is_thinking && !self.is_pondering {
            if match self.stop_when {
                TimeManagement::MoveTime(t) => self.searched_time >= t,
                TimeManagement::MoveTimeHint(t) => self.searched_time >= t,
                TimeManagement::Nodes(n) => self.searched_nodes >= n,
                TimeManagement::Depth(d) => self.current_depth > d,
                TimeManagement::Infinite => false,
            } {
                self.stop();
            }
        }

        // Check the reports queue.
        while let Ok(report) = self.reports.try_recv() {
            if self.is_thinking {
                match report {
                    search::Report::Progress { search_id, searched_nodes, depth }
                        if search_id == self.search_id => {
                        // Register search progress.
                        self.register_progress(depth, searched_nodes);
                        if self.no_reports_since.elapsed().unwrap().as_secs() > 20 {
                            if self.mangled_pv {
                                self.report_pv(depth - 1);
                            } else {
                                self.report_progress();
                            }
                        }
                    }
                    search::Report::Done { search_id, .. } if search_id == self.search_id => {
                        // Terminate the search (unless pondering of infinite).
                        if !self.is_pondering {
                            self.stop_when = match self.stop_when {
                                TimeManagement::Infinite => TimeManagement::Infinite,
                                _ => TimeManagement::MoveTime(0),
                            };
                        }
                    }
                    // We may still receive stale reports from already
                    // stopped searches.
                    _ => (),
                }
            }
        }

        self.reply_queue.pop_front()
    }
}

impl Engine {
    // A helper method.
    fn register_progress(&mut self, depth: u8, searched_nodes: NodeCount) {
        let thinking_duration = self.thinking_since.elapsed().unwrap();
        self.searched_nodes = searched_nodes;
        self.searched_time = 1000 * thinking_duration.as_secs() +
                             (thinking_duration.subsec_nanos() / 1000000) as u64;
        self.nps = 1000 * (self.nps + self.searched_nodes) / (1000 + self.searched_time);
        if self.current_depth < depth {
            self.current_depth = depth;
            self.report_pv(depth - 1);
            if !self.mangled_pv {
                self.report_progress();
            }
        }
    }

    // A helper method. It extracts the primary variation (PV) from
    // the transposition table (TT) and sends it to the GUI.
    fn report_pv(&mut self, depth: u8) {
        let mut prev_move = None;
        let mut p = self.position.clone();

        // Extract the PV, the value, and the bound from the TT.
        let mut pv = Vec::new();
        let mut value = -20000;
        let mut bound = BOUND_LOWER;
        while let Some(entry) = self.tt.probe(p.hash()) {
            if pv.len() < depth as usize && entry.bound() != BOUND_NONE {
                if let Some(m) = prev_move {
                    pv.push(m);
                }
                if pv.len() & 1 == 0 {
                    value = entry.value();
                    bound = entry.bound();
                } else {
                    // In this case the value stored in `entry` is
                    // from other side's perspective.
                    value = -entry.value();
                    bound = match entry.bound() {
                        BOUND_UPPER => BOUND_LOWER,
                        BOUND_LOWER => BOUND_UPPER,
                        x => x,
                    };
                }
                if bound == BOUND_EXACT {
                    if let Some(m) = p.try_move_digest(entry.move16()) {
                        if p.do_move(m) && !p.is_repeated() {
                            // Try to extend the PV.
                            prev_move = Some(m);
                            continue;
                        }
                    }
                }
            }
            break;
        }
        self.mangled_pv = bound != BOUND_EXACT || pv.len() < depth as usize;

        // Send the extracted info to the GUI.
        let value_suffix = match bound {
            BOUND_EXACT => "",
            BOUND_UPPER => " upperbound",
            BOUND_LOWER => " lowerbound",
            _ => panic!("unexpected bound type"),
        };
        let mut pv_string = String::new();
        for m in pv {
            pv_string.push(' ');
            pv_string.push_str(&m.notation());
        }
        self.reply_queue.push_back(EngineReply::Info(vec![
            ("depth".to_string(), format!("{}", depth)),
            ("score".to_string(), format!("cp {}{}", value, value_suffix)),
            ("time".to_string(), format!("{}", self.searched_time)),
            ("nodes".to_string(), format!("{}", self.searched_nodes)),
            ("nps".to_string(), format!("{}", self.nps)),
            ("pv".to_string(), pv_string),
        ]));
        self.no_reports_since = SystemTime::now();
    }

    // A helper method. It reports the depth, the node count, and
    // nodes per second to the GUI.
    fn report_progress(&mut self) {
        self.reply_queue.push_back(EngineReply::Info(vec![
            ("depth".to_string(), format!("{}", self.current_depth)),
            ("nodes".to_string(), format!("{}", self.searched_nodes)),
            ("nps".to_string(), format!("{}", self.nps)),
        ]));
        self.no_reports_since = SystemTime::now();
    }

    fn get_best_move(&mut self) -> String {
        let mut m = match self.tt.probe(self.position.hash()) {
            Some(entry) => entry.move16(),
            None => 0,
        };
        if m == 0 {
            // Pick the first legal move.
            let mut first_legal_move = Move::invalid();
            let mut v = MoveStack::new();
            self.position.generate_moves(&mut v);
            while let Some(x) = v.pop() {
                if self.position.do_move(x) {
                    self.position.undo_move();
                    first_legal_move = x;
                    break;
                }
            }
            m = first_legal_move.digest();
        }
        if m != 0 {
            let move_type = move_type(m);
            let orig_square = orig_square(m);
            let dest_square = dest_square(m);
            let promoted_piece = match aux_data(m) {
                0 => "q",
                1 => "r",
                2 => "b",
                3 => "n",
                _ => panic!("invalid promoted piece code"),
            };
            format!("{}{}{}",
                    notation(orig_square),
                    notation(dest_square),
                    if move_type == MOVE_PROMOTION {
                        promoted_piece
                    } else {
                        ""
                    })

        } else {
            "0000".to_string()
        }
    }
}

/// Implements `UciEngineFactory` trait.
pub struct EngineFactory;

impl UciEngineFactory<Engine> for EngineFactory {
    fn name(&self) -> String {
        format!("Alcibiades {}", VERSION)
    }

    fn author(&self) -> String {
        "Evgeni Pandurski".to_string()
    }

    fn options(&self) -> Vec<(OptionName, OptionDescription)> {
        vec![
            // TODO: Calculate a sane limit for the hash size.
            ("Hash".to_string(), OptionDescription::Spin { min: 1, max: 2048, default: 16 }),
            ("Ponder".to_string(), OptionDescription::Check { default: false }),
        ]
    }

    fn create(&self, hash_size_mb: Option<usize>) -> Engine {
        Engine::new(hash_size_mb.unwrap_or(16))
    }
}

enum TimeManagement {
    MoveTime(u64),
    MoveTimeHint(u64),
    Nodes(NodeCount),
    Depth(u8),
    Infinite,
}
<commit_msg>more beautifications<commit_after>pub mod search;

use std::thread;
use std::cmp::min;
use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver};
use std::time::SystemTime;
use std::num::Wrapping;
use basetypes::*;
use uci::{UciEngine, UciEngineFactory, EngineReply, OptionName, OptionDescription};
use position::Position;
use chess_move::*;
use tt::*;

const VERSION: &'static str = "0.1";
const MAX_DEPTH: u8 = 126;

/// Implements `UciEngine` trait.
pub struct Engine {
    tt: Arc<TranspositionTable>,
    reply_queue: VecDeque<EngineReply>,
    position: Position,
    is_thinking: bool,
    is_pondering: bool,

    // Tells the engine if it will be allowed to ponder. This option
    // is needed because the engine might change its time management
    // algorithm when pondering is allowed.
    pondering_is_allowed: bool,

    // A channel for sending commands to the search thread.
    commands: Sender<search::Command>,

    // A channel for receiving reports from the search thread.
    reports: Receiver<search::Report>,

    // Search status info
    search_id: usize,
    thinking_since: SystemTime,
    silent_since: SystemTime,
    current_depth: u8,
    searched_nodes: NodeCount,
    searched_time: u64, // milliseconds
    nps: u64, // nodes per second
    mangled_pv: bool, // `true` if the primary variation is imperfect
    stop_when: TimeManagement,
}

impl Engine {
    /// Creates a new instance.
    ///
    /// `tt_size_mb` is the preferred size of the transposition
    /// table in Mbytes.
    pub fn new(tt_size_mb: usize) -> Engine {
        let mut tt = TranspositionTable::new();
        tt.resize(tt_size_mb);
        let tt = Arc::new(tt);
        let (commands_tx, commands_rx) = channel();
        let (reports_tx, reports_rx) = channel();
        let tt2 = tt.clone();
        thread::spawn(move || {
            search::run_deepening(tt2, commands_rx, reports_tx);
        });

        Engine {
            tt: tt,
            reply_queue: VecDeque::new(),
            position: Position::from_fen("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 \
                                          1")
                          .ok()
                          .unwrap(),
            is_thinking: false,
            is_pondering: false,
            pondering_is_allowed: false,
            commands: commands_tx,
            reports: reports_rx,
            search_id: 0,
            thinking_since: SystemTime::now(),
            silent_since: SystemTime::now(),
            current_depth: 0,
            searched_nodes: 0,
            searched_time: 0,
            nps: 0,
            mangled_pv: false,
            stop_when: TimeManagement::Infinite,
        }
    }
}

impl UciEngine for Engine {
    fn set_option(&mut self, name: &str, value: &str) {
        match name {
            "Ponder" => {
                self.pondering_is_allowed = value == "true";
            }

            // An invalid option. Notice that we do not support
            // re-sizing of the transposition table once the engine
            // had started.
            _ => (),
        }
    }

    fn new_game(&mut self) {
        // Clearing the transposition table would not change anything,
        // so we do nothing.
    }

    fn position(&mut self, fen: &str, moves: &mut Iterator<Item = &str>) {
        if !self.is_thinking {
            if let Ok(p) = Position::from_history(fen, moves) {
                self.position = p;
                self.tt.new_search();
            }
        }
    }

    #[allow(unused_variables)]
    fn go(&mut self,
          searchmoves: Option<Vec<String>>,
          ponder: bool,
          wtime: Option<u64>,
          btime: Option<u64>,
          winc: Option<u64>,
          binc: Option<u64>,
          movestogo: Option<u64>,
          depth: Option<u64>,
          nodes: Option<u64>,
          mate: Option<u64>,
          movetime: Option<u64>,
          infinite: bool) {
        if !self.is_thinking {
            // Note: We ignore "searchmoves" and "mate" parameters.

            self.is_thinking = true;
            self.is_pondering = ponder;
            self.search_id = (Wrapping(self.search_id) + Wrapping(1)).0;
            self.thinking_since = SystemTime::now();
            self.silent_since = SystemTime::now();
            self.current_depth = 0;
            self.searched_nodes = 0;
            self.searched_time = 0;
            self.mangled_pv = false;

            // Figure out when we should stop thinking.
            let (time, inc) = if self.position.board().to_move() == WHITE {
                (wtime, winc.unwrap_or(0))
            } else {
                (btime, binc.unwrap_or(0))
            };
            self.stop_when = if infinite {
                TimeManagement::Infinite
            } else if movetime.is_some() {
                TimeManagement::MoveTime(movetime.unwrap())
            } else if nodes.is_some() {
                TimeManagement::Nodes(nodes.unwrap())
            } else if depth.is_some() {
                TimeManagement::Depth(depth.unwrap() as u8)
            } else {
                let time = time.unwrap_or(0);
                let movestogo = movestogo.unwrap_or(40);
                let movetime = (time + inc * movestogo) / movestogo;
                TimeManagement::MoveTimeHint(min(movetime, time / 2))
            };

            // Start deepening search to the maximum possible depth.
            self.commands
                .send(search::Command::Search {
                    search_id: self.search_id,
                    position: self.position.clone(),
                    depth: MAX_DEPTH,
                    lower_bound: -20000,
                    upper_bound: 20000,
                })
                .unwrap();
        }
    }

    fn ponder_hit(&mut self) {
        if self.is_thinking && self.is_pondering {
            self.is_pondering = false;
        }
    }

    fn stop(&mut self) {
        if self.is_thinking {
            self.commands.send(search::Command::Stop).unwrap();

            // TODO: send ponder move as well.
            let best_move = self.get_best_move();
            self.reply_queue.push_back(EngineReply::BestMove {
                best_move: best_move,
                ponder_move: None,
            });

            self.is_thinking = false;
        }
    }

    fn is_thinking(&self) -> bool {
        self.is_thinking
    }

    fn get_reply(&mut self) -> Option<EngineReply> {
        // Check if we should stop thinking.
        if self.is_thinking && !self.is_pondering {
            if match self.stop_when {
                TimeManagement::MoveTime(t) => self.searched_time >= t,
                TimeManagement::MoveTimeHint(t) => self.searched_time >= t,
                TimeManagement::Nodes(n) => self.searched_nodes >= n,
                TimeManagement::Depth(d) => self.current_depth > d,
                TimeManagement::Infinite => false,
            } {
                self.stop();
            }
        }

        // Check the reports queue.
        while let Ok(report) = self.reports.try_recv() {
            if self.is_thinking {
                match report {
                    search::Report::Progress { search_id, depth, searched_nodes }
                        if search_id == self.search_id => {
                        // Register search progress.
                        self.register_progress(depth, searched_nodes);
                        if self.silent_since.elapsed().unwrap().as_secs() > 20 {
                            if self.mangled_pv {
                                self.report_pv(depth - 1);
                            } else {
                                self.report_progress();
                            }
                        }
                    }
                    search::Report::Done { search_id, .. } if search_id == self.search_id => {
                        // Terminate the search (unless pondering of infinite).
                        if !self.is_pondering {
                            self.stop_when = match self.stop_when {
                                TimeManagement::Infinite => TimeManagement::Infinite,
                                _ => TimeManagement::MoveTime(0),
                            };
                        }
                    }
                    // We may still receive stale reports from already
                    // stopped searches.
                    _ => (),
                }
            }
        }

        self.reply_queue.pop_front()
    }
}

impl Engine {
    // A helper method. It updates the engine search status info and
    // makes sure that the PV is sent to the GUI for every newly
    // reached depth.
    fn register_progress(&mut self, depth: u8, searched_nodes: NodeCount) {
        let thinking_duration = self.thinking_since.elapsed().unwrap();
        self.searched_time = 1000 * thinking_duration.as_secs() +
                             (thinking_duration.subsec_nanos() / 1000000) as u64;
        self.searched_nodes = searched_nodes;
        self.nps = 1000 * (self.nps + self.searched_nodes) / (1000 + self.searched_time);
        if self.current_depth < depth {
            self.current_depth = depth;
            self.report_pv(depth - 1);
            if !self.mangled_pv {
                self.report_progress();
            }
        }
    }

    // A helper method. It extracts the primary variation (PV) from
    // the transposition table (TT) and sends it to the GUI.
    fn report_pv(&mut self, depth: u8) {
        let mut prev_move = None;
        let mut p = self.position.clone();

        // Extract the PV, the value, and the bound from the TT.
        let mut pv = Vec::new();
        let mut value = -20000;
        let mut bound = BOUND_LOWER;
        while let Some(entry) = self.tt.probe(p.hash()) {
            if pv.len() < depth as usize && entry.bound() != BOUND_NONE {
                if let Some(m) = prev_move {
                    pv.push(m);
                }
                if pv.len() & 1 == 0 {
                    value = entry.value();
                    bound = entry.bound();
                } else {
                    // In this case the value stored in `entry` is
                    // from other side's perspective.
                    value = -entry.value();
                    bound = match entry.bound() {
                        BOUND_UPPER => BOUND_LOWER,
                        BOUND_LOWER => BOUND_UPPER,
                        x => x,
                    };
                }
                if bound == BOUND_EXACT {
                    if let Some(m) = p.try_move_digest(entry.move16()) {
                        if p.do_move(m) && !p.is_repeated() {
                            // Try to extend the PV.
                            prev_move = Some(m);
                            continue;
                        }
                    }
                }
            }
            break;
        }
        self.mangled_pv = bound != BOUND_EXACT || pv.len() < depth as usize;

        // Send the extracted info to the GUI.
        let value_suffix = match bound {
            BOUND_EXACT => "",
            BOUND_UPPER => " upperbound",
            BOUND_LOWER => " lowerbound",
            _ => panic!("unexpected bound type"),
        };
        let mut pv_string = String::new();
        for m in pv {
            pv_string.push(' ');
            pv_string.push_str(&m.notation());
        }
        self.reply_queue.push_back(EngineReply::Info(vec![
            ("depth".to_string(), format!("{}", depth)),
            ("score".to_string(), format!("cp {}{}", value, value_suffix)),
            ("time".to_string(), format!("{}", self.searched_time)),
            ("nodes".to_string(), format!("{}", self.searched_nodes)),
            ("nps".to_string(), format!("{}", self.nps)),
            ("pv".to_string(), pv_string),
        ]));
        self.silent_since = SystemTime::now();
    }

    // A helper method. It reports the current depth, the node count,
    // and the nodes per second to the GUI.
    fn report_progress(&mut self) {
        self.reply_queue.push_back(EngineReply::Info(vec![
            ("depth".to_string(), format!("{}", self.current_depth)),
            ("nodes".to_string(), format!("{}", self.searched_nodes)),
            ("nps".to_string(), format!("{}", self.nps)),
        ]));
        self.silent_since = SystemTime::now();
    }

    fn get_best_move(&mut self) -> String {
        let mut m = match self.tt.probe(self.position.hash()) {
            Some(entry) => entry.move16(),
            None => 0,
        };
        if m == 0 {
            // Pick the first legal move.
            let mut first_legal_move = Move::invalid();
            let mut v = MoveStack::new();
            self.position.generate_moves(&mut v);
            while let Some(x) = v.pop() {
                if self.position.do_move(x) {
                    self.position.undo_move();
                    first_legal_move = x;
                    break;
                }
            }
            m = first_legal_move.digest();
        }
        if m != 0 {
            let move_type = move_type(m);
            let orig_square = orig_square(m);
            let dest_square = dest_square(m);
            let promoted_piece = match aux_data(m) {
                0 => "q",
                1 => "r",
                2 => "b",
                3 => "n",
                _ => panic!("invalid promoted piece code"),
            };
            format!("{}{}{}",
                    notation(orig_square),
                    notation(dest_square),
                    if move_type == MOVE_PROMOTION {
                        promoted_piece
                    } else {
                        ""
                    })

        } else {
            "0000".to_string()
        }
    }
}

/// Implements `UciEngineFactory` trait.
pub struct EngineFactory;

impl UciEngineFactory<Engine> for EngineFactory {
    fn name(&self) -> String {
        format!("Alcibiades {}", VERSION)
    }

    fn author(&self) -> String {
        "Evgeni Pandurski".to_string()
    }

    fn options(&self) -> Vec<(OptionName, OptionDescription)> {
        vec![
            // TODO: Calculate a sane limit for the hash size.
            ("Hash".to_string(), OptionDescription::Spin { min: 1, max: 2048, default: 16 }),
            ("Ponder".to_string(), OptionDescription::Check { default: false }),
        ]
    }

    fn create(&self, hash_size_mb: Option<usize>) -> Engine {
        Engine::new(hash_size_mb.unwrap_or(16))
    }
}

enum TimeManagement {
    MoveTime(u64),
    MoveTimeHint(u64),
    Nodes(NodeCount),
    Depth(u8),
    Infinite,
}
<|endoftext|>


<commit_before>//! Facilities for implementing chess engines.

pub mod time_manager;

use std::collections::VecDeque;
use std::ops::Deref;
use std::sync::Arc;
use std::time::{SystemTime, Duration};
use std::cmp::min;
use std::io;
use uci::*;
use search::*;
use self::time_manager::*;

const NAME: &'static str = "Alcibiades";
const AUTHOR: &'static str = "Evgeni Pandurski";
const VERSION: &'static str = "0.1";
const START_FEN: &'static str = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 1";

struct SearchStatus {
    pub done: bool,
    pub depth: Depth,
    pub searched_nodes: u64,

    // The duration of the search in milliseconds.
    pub duration_millis: u64,
}

impl Default for SearchStatus {
    fn default() -> Self {
        SearchStatus {
            done: false,
            depth: 0,
            searched_nodes: 0,
            duration_millis: 0,
        }
    }
}

enum PlayWhen {
    TimeManagement(TimeManager), // Stop when the time manager says so.
    MoveTime(u64), // Stop after the given number of milliseconds.
    Nodes(u64), // Stop when the given number of nodes has been searched.
    Depth(Depth), // Stop when the given search depth has been completed.
    Never, // An infinite search.
}

struct Engine<S: SearchExecutor<ReportData = Vec<Variation>>> {
    tt: Arc<S::HashTable>,
    position: S::SearchNode,
    searcher: S,
    queue: VecDeque<EngineReply>,

    // The starting time of the current/last search.
    started_at: SystemTime,

    // The status of the current/last search.
    status: SearchStatus,

    // Nodes per second statistics.
    nps_stats: (u64, u64, u64),

    // Helps the engine decide when to show periodic progress reports.
    silent_since: SystemTime,

    // Whether the engine is thinking in pondering mode at the moment.
    is_pondering: bool,

    // Tells the engine when it must stop thinking and play the best move.
    play_when: PlayWhen,

    // Tells the engine if it will be allowed to ponder. This option
    // is needed because the engine might change its time management
    // algorithm when pondering is allowed.
    pondering_is_allowed: bool,
}

impl<S: SearchExecutor<ReportData = Vec<Variation>>> UciEngine for Engine<S> {
    fn name() -> String {
        format!("{} {}", NAME, VERSION)
    }

    fn author() -> String {
        AUTHOR.to_string()
    }

    fn options() -> Vec<(String, OptionDescription)> {
        // Add up all suported options.
        let mut options = vec![
            ("Hash".to_string(), OptionDescription::Spin { min: 1, max: 64 * 1024, default: 16 }),
            ("Clear Hash".to_string(), OptionDescription::Button),
            ("Ponder".to_string(), OptionDescription::Check { default: false }),
        ];
        options.extend(S::options());

        // Remove duplicated options.
        options.sort_by(|a, b| a.0.cmp(&b.0));
        let mut options_dedup = vec![];
        let mut prev_name = String::from("");
        for o in options.drain(..) {
            if o.0 == prev_name {
                continue;
            }
            prev_name = o.0.clone();
            options_dedup.push(o);
        }
        options_dedup
    }

    fn new(tt_size_mb: Option<usize>) -> Engine<S> {
        let tt = Arc::new(S::HashTable::new(tt_size_mb));
        Engine {
            tt: tt.clone(),
            position: S::SearchNode::from_history(START_FEN, &mut vec![].into_iter()).ok().unwrap(),
            searcher: S::new(tt),
            queue: VecDeque::new(),
            started_at: SystemTime::now(),
            status: SearchStatus { done: true, ..Default::default() },
            nps_stats: (0, 0, 0),
            silent_since: SystemTime::now(),
            is_pondering: false,
            play_when: PlayWhen::Never,
            pondering_is_allowed: false,
        }
    }

    fn set_option(&mut self, name: &str, value: &str) {
        match name {
            "Ponder" => {
                // TODO: Can we remove this?
                self.pondering_is_allowed = value == "true";
            }
            "Hash" => {
                // We do not support re-sizing of the transposition
                // table once the engine has been started.
                return;
            }
            "Clear Hash" => {
                self.tt.clear();
            }
            _ => (),
        }
        S::set_option(name, value);
    }

    fn new_game(&mut self) {
        self.tt.clear();
    }

    fn position(&mut self, fen: &str, moves: &mut Iterator<Item = &str>) {
        if let Ok(p) = S::SearchNode::from_history(fen, moves) {
            self.position = p;
        }
    }

    fn go(&mut self, params: GoParams) {
        self.terminate();

        // TODO: What should we do with "mate"?
        let GoParams { mut searchmoves,
                       ponder,
                       wtime,
                       btime,
                       winc,
                       binc,
                       movestogo,
                       depth,
                       nodes,
                       movetime,
                       infinite,
                       .. } = params;

        // Validate `searchmoves`.
        let searchmoves = {
            let mut moves = vec![];
            let legal_moves = self.position.legal_moves();
            if !searchmoves.is_empty() {
                searchmoves.sort();
                for m in legal_moves.iter() {
                    if searchmoves.binary_search(&m.notation()).is_ok() {
                        moves.push(*m);
                    }
                }
            };
            if moves.is_empty() {
                legal_moves
            } else {
                moves
            }
        };

        // Start a new search.
        self.tt.new_search();
        self.started_at = SystemTime::now();
        self.status = Default::default();
        self.nps_stats = (self.nps_stats.0, 0, 0);
        self.silent_since = SystemTime::now();
        self.is_pondering = ponder;
        self.play_when = if infinite {
            PlayWhen::Never
        } else if movetime.is_some() {
            PlayWhen::MoveTime(movetime.unwrap())
        } else if nodes.is_some() {
            PlayWhen::Nodes(nodes.unwrap())
        } else if depth.is_some() {
            PlayWhen::Depth(min(depth.unwrap(), DEPTH_MAX as u64) as Depth)
        } else {
            PlayWhen::TimeManagement(TimeManager::new(self.position.board().to_move,
                                                      self.pondering_is_allowed,
                                                      wtime,
                                                      btime,
                                                      winc,
                                                      binc,
                                                      movestogo))
        };
        self.searcher.start_search(SearchParams {
            search_id: 0,
            position: self.position.clone(),
            depth: DEPTH_MAX,
            lower_bound: VALUE_MIN,
            upper_bound: VALUE_MAX,
            searchmoves: searchmoves,
        });
    }

    fn ponder_hit(&mut self) {
        if self.status.done {
            self.queue_best_move();
        } else {
            self.is_pondering = false;
        }
    }

    fn stop(&mut self) {
        self.terminate();
        self.queue_best_move();
    }

    fn wait_for_reply(&mut self, duration: Duration) -> Option<EngineReply> {
        if self.queue.is_empty() {
            let done = self.status.done;

            // Wait for the search thread to do some work, and
            // hopefully update the status. (We must do this even when
            // the search is done -- in that case the next line will
            // just yield the CPU to another process.)
            self.wait_status_update(duration);

            // See if we must play now.
            if !done {
                let &SearchStatus { done, depth, searched_nodes, duration_millis, .. } =
                    &self.status;
                if !self.is_pondering &&
                   match self.play_when {
                    PlayWhen::TimeManagement(ref tm) => done || tm.must_play(),
                    PlayWhen::MoveTime(t) => done || duration_millis >= t,
                    PlayWhen::Nodes(n) => done || searched_nodes >= n,
                    PlayWhen::Depth(d) => done || depth >= d,
                    PlayWhen::Never => false,
                } {
                    self.stop();
                }
            }
        }

        self.queue.pop_front()
    }

    fn exit(&mut self) {
        self.terminate();
    }
}

impl<S: SearchExecutor<ReportData = Vec<Variation>>> Engine<S> {
    fn queue_progress_report(&mut self) {
        let SearchStatus { ref depth, ref searched_nodes, .. } = self.status;
        self.queue.push_back(EngineReply::Info(vec![
            InfoItem { info_type: "depth".to_string(), data: format!("{}", depth) },
            InfoItem { info_type: "nodes".to_string(), data: format!("{}", searched_nodes) },
            InfoItem { info_type: "nps".to_string(), data: format!("{}", self.nps_stats.0) },
        ]));
        self.silent_since = SystemTime::now();
    }

    fn queue_pv(&mut self, variations: &Vec<Variation>) {
        let SearchStatus { ref depth, ref searched_nodes, ref duration_millis, .. } = self.status;
        for (i, &Variation { ref moves, value, bound }) in variations.iter().enumerate() {
            let bound_suffix = match bound {
                BOUND_EXACT => "",
                BOUND_UPPER => " upperbound",
                BOUND_LOWER => " lowerbound",
                _ => panic!("unexpected bound type"),
            };
            let mut moves_string = String::new();
            for m in moves {
                moves_string.push_str(&m.notation());
                moves_string.push(' ');
            }
            self.queue.push_back(EngineReply::Info(vec![
                InfoItem { info_type: "depth".to_string(), data: format!("{}", depth) },
                InfoItem { info_type: "multipv".to_string(), data: format!("{}", i + 1) },
                InfoItem { info_type: "score".to_string(), data: format!("cp {}{}", value, bound_suffix) },
                InfoItem { info_type: "time".to_string(), data: format!("{}", duration_millis) },
                InfoItem { info_type: "nodes".to_string(), data: format!("{}", searched_nodes) },
                InfoItem { info_type: "nps".to_string(), data: format!("{}", self.nps_stats.0) },
                InfoItem { info_type: "pv".to_string(), data: moves_string },
            ]));
        }
        if !variations.is_empty() {
            self.silent_since = SystemTime::now();
        }
    }

    fn queue_best_move(&mut self) {
        let pv = extract_pv(self.tt.deref(), &self.position, 2);
        let best_move = if let Some(m) = pv.moves.get(0) {
            m.notation()
        } else {
            // If there is no best move, pick the first legal move.
            self.position.legal_moves().get(0).map_or("0000".to_string(), |m| m.notation())
        };
        self.queue.push_back(EngineReply::BestMove {
            best_move: best_move,
            ponder_move: pv.moves.get(1).map(|m| m.notation()),
        });
        self.silent_since = SystemTime::now();
    }

    fn terminate(&mut self) {
        self.searcher.terminate_search();
        while !self.status.done {
            self.wait_status_update(Duration::from_millis(1000));
        }
    }

    fn wait_status_update(&mut self, duration: Duration) {
        self.searcher.wait_report(duration);
        while let Ok(r) = self.searcher.try_recv_report() {
            self.process_report(r)
        }
    }

    fn process_report(&mut self, report: SearchReport<Vec<Variation>>) {
        let duration_millis = {
            let d = self.started_at.elapsed().unwrap();
            1000 * d.as_secs() + (d.subsec_nanos() / 1_000_000) as u64
        };
        self.status = SearchStatus {
            done: report.done,
            depth: report.depth,
            searched_nodes: report.searched_nodes,
            duration_millis: duration_millis,
        };

        // Update `self.nps_stats` each 1000 milliseconds.
        let elapsed_millis = duration_millis - self.nps_stats.2;
        if elapsed_millis >= 1000 {
            let nodes = report.searched_nodes - self.nps_stats.1;
            self.nps_stats = (1000 * nodes / elapsed_millis, report.searched_nodes, duration_millis)
        }

        // Inform the time manager.
        if let PlayWhen::TimeManagement(ref mut tm) = self.play_when {
            tm.update_status(&report);
        }

        // Show the primary variations provided with the report (if any).
        self.queue_pv(&report.data);

        // If nothing has happened for a while, show a progress report.
        if self.silent_since.elapsed().unwrap().as_secs() > 10 {
            self.queue_progress_report();
        }
    }
}

/// Serves UCI commands until a "quit" command is received.
///
/// The current thread will block until the UCI session is closed.
///
/// Returns `Err` if the handshake was unsuccessful, or if an IO error
/// occurred.
pub fn run_server<S: SearchExecutor<ReportData = Vec<Variation>>>() -> io::Result<()> {
    run_engine::<Engine<S>>()
}
<commit_msg>beautify<commit_after>//! Facilities for implementing chess engines.

pub mod time_manager;

use std::collections::VecDeque;
use std::ops::Deref;
use std::sync::Arc;
use std::time::{SystemTime, Duration};
use std::cmp::min;
use std::io;
use uci::*;
use search::*;
use self::time_manager::*;

const NAME: &'static str = "Alcibiades";
const AUTHOR: &'static str = "Evgeni Pandurski";
const VERSION: &'static str = "0.1";
const START_FEN: &'static str = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 1";

struct SearchStatus {
    pub done: bool,
    pub depth: Depth,
    pub searched_nodes: u64,

    // The duration of the search in milliseconds.
    pub duration_millis: u64,
}

impl Default for SearchStatus {
    fn default() -> Self {
        SearchStatus {
            done: false,
            depth: 0,
            searched_nodes: 0,
            duration_millis: 0,
        }
    }
}

enum PlayWhen {
    TimeManagement(TimeManager), // Stop when the time manager says so.
    MoveTime(u64), // Stop after the given number of milliseconds.
    Nodes(u64), // Stop when the given number of nodes has been searched.
    Depth(Depth), // Stop when the given search depth has been completed.
    Never, // An infinite search.
}

struct Engine<T: SearchExecutor<ReportData = Vec<Variation>>> {
    tt: Arc<T::HashTable>,
    position: T::SearchNode,
    searcher: T,
    queue: VecDeque<EngineReply>,

    // The starting time of the current/last search.
    started_at: SystemTime,

    // The status of the current/last search.
    status: SearchStatus,

    // Nodes per second statistics.
    nps_stats: (u64, u64, u64),

    // Helps the engine decide when to show periodic progress reports.
    silent_since: SystemTime,

    // Whether the engine is thinking in pondering mode at the moment.
    is_pondering: bool,

    // Tells the engine when it must stop thinking and play the best move.
    play_when: PlayWhen,

    // Tells the engine if it will be allowed to ponder. This option
    // is needed because the engine might change its time management
    // algorithm when pondering is allowed.
    pondering_is_allowed: bool,
}

impl<T: SearchExecutor<ReportData = Vec<Variation>>> UciEngine for Engine<T> {
    fn name() -> String {
        format!("{} {}", NAME, VERSION)
    }

    fn author() -> String {
        AUTHOR.to_string()
    }

    fn options() -> Vec<(String, OptionDescription)> {
        // Add up all suported options.
        let mut options = vec![
            ("Hash".to_string(), OptionDescription::Spin { min: 1, max: 64 * 1024, default: 16 }),
            ("Clear Hash".to_string(), OptionDescription::Button),
            ("Ponder".to_string(), OptionDescription::Check { default: false }),
        ];
        options.extend(T::options());

        // Remove duplicated options.
        options.sort_by(|a, b| a.0.cmp(&b.0));
        let mut options_dedup = vec![];
        let mut prev_name = String::from("");
        for o in options.drain(..) {
            if o.0 == prev_name {
                continue;
            }
            prev_name = o.0.clone();
            options_dedup.push(o);
        }
        options_dedup
    }

    fn new(tt_size_mb: Option<usize>) -> Engine<T> {
        let tt = Arc::new(T::HashTable::new(tt_size_mb));
        Engine {
            tt: tt.clone(),
            position: T::SearchNode::from_history(START_FEN, &mut vec![].into_iter()).ok().unwrap(),
            searcher: T::new(tt),
            queue: VecDeque::new(),
            started_at: SystemTime::now(),
            status: SearchStatus { done: true, ..Default::default() },
            nps_stats: (0, 0, 0),
            silent_since: SystemTime::now(),
            is_pondering: false,
            play_when: PlayWhen::Never,
            pondering_is_allowed: false,
        }
    }

    fn set_option(&mut self, name: &str, value: &str) {
        match name {
            "Ponder" => {
                // TODO: Can we remove this?
                self.pondering_is_allowed = value == "true";
            }
            "Hash" => {
                // We do not support re-sizing of the transposition
                // table once the engine has been started.
                return;
            }
            "Clear Hash" => {
                self.tt.clear();
            }
            _ => (),
        }
        T::set_option(name, value);
    }

    fn new_game(&mut self) {
        self.tt.clear();
    }

    fn position(&mut self, fen: &str, moves: &mut Iterator<Item = &str>) {
        if let Ok(p) = T::SearchNode::from_history(fen, moves) {
            self.position = p;
        }
    }

    fn go(&mut self, mut params: GoParams) {
        self.terminate();

        // Validate `params.searchmoves`.
        //
        // TODO: What should we do with "mate"?
        let searchmoves = {
            let mut moves = vec![];
            let legal_moves = self.position.legal_moves();
            if !params.searchmoves.is_empty() {
                params.searchmoves.sort();
                for m in legal_moves.iter() {
                    if params.searchmoves.binary_search(&m.notation()).is_ok() {
                        moves.push(*m);
                    }
                }
            };
            if moves.is_empty() {
                legal_moves
            } else {
                moves
            }
        };

        // Start a new search.
        self.tt.new_search();
        self.started_at = SystemTime::now();
        self.status = Default::default();
        self.nps_stats = (self.nps_stats.0, 0, 0);
        self.silent_since = self.started_at;
        self.is_pondering = params.ponder;
        self.play_when = if params.infinite {
            PlayWhen::Never
        } else if params.movetime.is_some() {
            PlayWhen::MoveTime(params.movetime.unwrap())
        } else if params.nodes.is_some() {
            PlayWhen::Nodes(params.nodes.unwrap())
        } else if params.depth.is_some() {
            PlayWhen::Depth(min(params.depth.unwrap(), DEPTH_MAX as u64) as Depth)
        } else {
            PlayWhen::TimeManagement(TimeManager::new(self.position.board().to_move,
                                                      self.pondering_is_allowed,
                                                      params.wtime,
                                                      params.btime,
                                                      params.winc,
                                                      params.binc,
                                                      params.movestogo))
        };
        self.searcher.start_search(SearchParams {
            search_id: 0,
            position: self.position.clone(),
            depth: DEPTH_MAX,
            lower_bound: VALUE_MIN,
            upper_bound: VALUE_MAX,
            searchmoves: searchmoves,
        });
    }

    fn ponder_hit(&mut self) {
        if self.status.done {
            self.queue_best_move();
        } else {
            self.is_pondering = false;
        }
    }

    fn stop(&mut self) {
        self.terminate();
        self.queue_best_move();
    }

    fn wait_for_reply(&mut self, duration: Duration) -> Option<EngineReply> {
        if self.queue.is_empty() {
            let is_thinking = !self.status.done;

            // Wait for the search thread to do some work, and
            // hopefully update the status. (We must do this even when
            // the engine is not thinking -- in that case the next
            // line will just yield the CPU to another process.)
            self.wait_status_update(duration);

            // See if we must stop thinking and play.
            if is_thinking && !self.is_pondering &&
               match self.play_when {
                PlayWhen::TimeManagement(ref tm) => self.status.done || tm.must_play(),
                PlayWhen::MoveTime(t) => self.status.done || self.status.duration_millis >= t,
                PlayWhen::Nodes(n) => self.status.done || self.status.searched_nodes >= n,
                PlayWhen::Depth(d) => self.status.done || self.status.depth >= d,
                PlayWhen::Never => false,
            } {
                self.stop();
            }
        }

        self.queue.pop_front()
    }

    fn exit(&mut self) {
        self.terminate();
    }
}

impl<T: SearchExecutor<ReportData = Vec<Variation>>> Engine<T> {
    fn queue_progress_info(&mut self) {
        let SearchStatus { ref depth, ref searched_nodes, ref duration_millis, .. } = self.status;
        self.queue.push_back(EngineReply::Info(vec![
            InfoItem { info_type: "depth".to_string(), data: format!("{}", depth) },
            InfoItem { info_type: "time".to_string(), data: format!("{}", duration_millis) },
            InfoItem { info_type: "nodes".to_string(), data: format!("{}", searched_nodes) },
            InfoItem { info_type: "nps".to_string(), data: format!("{}", self.nps_stats.0) },
        ]));
        self.silent_since = SystemTime::now();
    }

    fn queue_pv(&mut self, variations: &Vec<Variation>) {
        let SearchStatus { ref depth, ref searched_nodes, ref duration_millis, .. } = self.status;
        for (i, &Variation { ref moves, value, bound }) in variations.iter().enumerate() {
            let bound_suffix = match bound {
                BOUND_EXACT => "",
                BOUND_UPPER => " upperbound",
                BOUND_LOWER => " lowerbound",
                _ => panic!("unexpected bound type"),
            };
            let mut pv = String::new();
            for m in moves {
                pv.push_str(&m.notation());
                pv.push(' ');
            }
            self.queue.push_back(EngineReply::Info(vec![
                InfoItem { info_type: "depth".to_string(), data: format!("{}", depth) },
                InfoItem { info_type: "multipv".to_string(), data: format!("{}", i + 1) },
                InfoItem { info_type: "score".to_string(), data: format!("cp {}{}", value, bound_suffix) },
                InfoItem { info_type: "time".to_string(), data: format!("{}", duration_millis) },
                InfoItem { info_type: "nodes".to_string(), data: format!("{}", searched_nodes) },
                InfoItem { info_type: "nps".to_string(), data: format!("{}", self.nps_stats.0) },
                InfoItem { info_type: "pv".to_string(), data: pv },
            ]));
        }
        if !variations.is_empty() {
            self.silent_since = SystemTime::now();
        }
    }

    fn queue_best_move(&mut self) {
        let pv = extract_pv(self.tt.deref(), &self.position, 2);
        let best_move = if let Some(m) = pv.moves.get(0) {
            m.notation()
        } else {
            // If there is no best move, pick the first legal move.
            self.position.legal_moves().get(0).map_or("0000".to_string(), |m| m.notation())
        };
        self.queue.push_back(EngineReply::BestMove {
            best_move: best_move,
            ponder_move: pv.moves.get(1).map(|m| m.notation()),
        });
        self.silent_since = SystemTime::now();
    }

    fn terminate(&mut self) {
        self.searcher.terminate_search();
        while !self.status.done {
            self.wait_status_update(Duration::from_millis(1000));
        }
    }

    fn wait_status_update(&mut self, duration: Duration) {
        self.searcher.wait_report(duration);
        while let Ok(r) = self.searcher.try_recv_report() {
            self.process_report(r)
        }
    }

    fn process_report(&mut self, report: SearchReport<Vec<Variation>>) {
        let duration_millis = {
            let d = self.started_at.elapsed().unwrap();
            1000 * d.as_secs() + (d.subsec_nanos() / 1_000_000) as u64
        };
        self.status = SearchStatus {
            done: report.done,
            depth: report.depth,
            searched_nodes: report.searched_nodes,
            duration_millis: duration_millis,
        };

        // Update `self.nps_stats` each 1000 milliseconds.
        let elapsed_millis = duration_millis - self.nps_stats.2;
        if elapsed_millis >= 1000 {
            let nodes = report.searched_nodes - self.nps_stats.1;
            self.nps_stats = (1000 * nodes / elapsed_millis, report.searched_nodes, duration_millis)
        }

        // Inform the time manager.
        if let PlayWhen::TimeManagement(ref mut tm) = self.play_when {
            tm.update_status(&report);
        }

        // Show the primary variations provided with the report (if any).
        self.queue_pv(&report.data);

        // If nothing has happened for a while, show progress info.
        if self.silent_since.elapsed().unwrap().as_secs() > 10 {
            self.queue_progress_info();
        }
    }
}

/// Serves UCI commands until a "quit" command is received.
///
/// The current thread will block until the UCI session is closed.
///
/// Returns `Err` if the handshake was unsuccessful, or if an IO error
/// occurred.
pub fn run_server<T: SearchExecutor<ReportData = Vec<Variation>>>() -> io::Result<()> {
    run_engine::<Engine<T>>()
}
<|endoftext|>


<commit_before>// Copyright 2013-2015, The Gtk-rs Project Developers.
// See the COPYRIGHT file at the top-level directory of this distribution.
// Licensed under the MIT license, see the LICENSE file or <http://opensource.org/licenses/MIT>

//! OpenGL context

use glib::translate::*;
use glib::types;
use cursor::Cursor;
use display::Display;
use object::Object;
use screen::Screen;
use window::Window;
use ffi;

/// GLContext is an object representing the platform-specific OpenGL drawing context.
///
/// GLContexts are created for a GdkWindow using gdk_window_create_gl_context(), and the context
/// will match the GdkVisual of the window.
///
/// A GLContext is not tied to any particular normal framebuffer. For instance, it cannot draw to
/// the GdkWindow back buffer. The GDK repaint system is in full control of the painting to that.
/// Instead, you can create render buffers or textures and use gdk_cairo_draw_from_gl() in the draw
/// function of your widget to draw them. Then GDK will handle the integration of your rendering
/// with that of other widgets.
///
/// Support for GLContext is platform-specific, context creation can fail, returning None context.
///
/// A GLContext has to be made "current" in order to start using it, otherwise any OpenGL call will
/// be ignored.
///
/// ##Creating a new OpenGL context
///
/// In order to create a new GLContext instance you need a GdkWindow, which you typically get during
/// the realize call of a widget.
///
/// A GLContext is not realized until either gdk_gl_context_make_current(), or until it is realized
/// using gdk_gl_context_realize(). It is possible to specify details of the GL context like the
/// OpenGL version to be used, or whether the GL context should have extra state validation enabled
/// after calling gdk_window_create_gl_context() by calling gdk_gl_context_realize(). If the
/// realization fails you have the option to change the settings of the GLContext and try again.
///
/// ##Using a GLContext
///
/// You will need to make the GLContext the current context before issuing OpenGL calls; the system
/// sends OpenGL commands to whichever context is current. It is possible to have multiple contexts,
/// so you always need to ensure that the one which you want to draw with is the current one before
/// issuing commands:
///
/// ```Rust,ignore
/// GLContext::make_current(context);
/// ```
///
/// You can now perform your drawing using OpenGL commands.
///
/// You can check which GLContext is the current one by using gdk_gl_context_get_current(); you can
/// also unset any GLContext that is currently set by calling gdk_gl_context_clear_current().
pub type GLContext = Object<ffi::GdkGLContext>;

impl GLContext {
    /// Retrieves the Display the context is created for
    pub fn get_display(&self) -> Option<Display> {
        unsafe { from_glib_none(ffi::gdk_gl_context_get_display(self.to_glib_none().0)) }
    }

    /// Retrieves the GdkWindow used by self.
    pub fn get_window(&self) -> Option<Window> {
        unsafe { from_glib_none(ffi::gdk_gl_context_get_window(self.to_glib_none().0)) }
    }

    /// Retrieves the GLContext that this context share data with.
    pub fn get_shared_context(&self) -> Option<Window> {
        unsafe { from_glib_none(ffi::gdk_gl_context_get_shared_context(self.to_glib_none().0)) }
    }

    /// Retrieves the OpenGL version of the context.
    ///
    /// The context must be realized prior to calling this function.
    pub fn get_version(&self) -> (i32, i32) {
        let mut major = 0;
        let mut minor = 0;

        unsafe { ffi::gdk_gl_context_get_version(self.to_glib_none().0, &mut major, &mut minor) };
        (major, minor)
    }

    /// Sets the major and minor version of OpenGL to request.
    ///
    /// Setting major and minor to zero will use the default values.
    ///
    /// The GLContext must not be realized or made current prior to calling this function.
    pub fn set_required_version(&self, major: i32, minor: i32) {
        unsafe { ffi::gdk_gl_context_set_required_version(self.to_glib_none().0, major, minor) }
    }

    /// Sets whether the GLContext should perform extra validations and run time checking.
    /// This is useful during development, but has additional overhead.
    ///
    /// The GLContext must not be realized or made current prior to calling this function.
    pub fn set_debug_enabled(&self, enabled: bool) {
        unsafe { ffi::gdk_gl_context_set_debug_enabled(self.to_glib_none().0, enabled.to_glib()) }
    }

    /// Retrieves the value set using set_debug_enabled().
    pub fn get_debug_enabled(&self) -> bool {
        unsafe { from_glib(ffi::gdk_gl_context_get_debug_enabled(self.to_glib_none().0)) }
    }

    /// Sets whether the GLContext should be forward compatible.
    ///
    /// Forward compatibile contexts must not support OpenGL functionality that has been
    /// marked as deprecated in the requested version; non-forward compatible contexts, on
    /// the other hand, must support both deprecated and non deprecated functionality.
    ///
    /// The GLContext must not be realized or made current prior to calling this function.
    pub fn set_forward_compatible(&self, compatible: bool) {
        unsafe { ffi::gdk_gl_context_set_forward_compatible(self.to_glib_none().0, compatible.to_glib()) }
    }

    /// Retrieves the value set using set_forward_compatible().
    pub fn get_forward_compatible(&self) -> bool {
        unsafe { from_glib(ffi::gdk_gl_context_get_forward_compatible(self.to_glib_none().0)) }
    }

    /*
    /// Realizes the given GLContext.
    ///
    /// It is safe to call this function on a realized GLContext.
    pub fn realize(&self, *mut *mut glib::error) -> bool;
    */

    /// Makes self the current context.
    pub fn make_current(&self) {
        unsafe { ffi::gdk_gl_context_make_current(self.to_glib_none().0) }
    }

    /// Retrieves the current GLContext.
    pub fn get_current() -> Option<GLContext> {
        unsafe { from_glib_none(ffi::gdk_gl_context_get_current()) }
    }

    /// Clears the current GLContext.
    ///
    /// Any OpenGL call after this function returns will be ignored until GLContext::make_current()
    /// is called.
    pub fn clear_current() {
        unsafe { ffi::gdk_gl_context_clear_current() }
    }
}
<commit_msg>Add missing trait impl on GLContext<commit_after>// Copyright 2013-2015, The Gtk-rs Project Developers.
// See the COPYRIGHT file at the top-level directory of this distribution.
// Licensed under the MIT license, see the LICENSE file or <http://opensource.org/licenses/MIT>

//! OpenGL context

use glib::translate::*;
use glib::types;
use display::Display;
use object::Object;
use window::Window;
use ffi;

/// GLContext is an object representing the platform-specific OpenGL drawing context.
///
/// GLContexts are created for a GdkWindow using gdk_window_create_gl_context(), and the context
/// will match the GdkVisual of the window.
///
/// A GLContext is not tied to any particular normal framebuffer. For instance, it cannot draw to
/// the GdkWindow back buffer. The GDK repaint system is in full control of the painting to that.
/// Instead, you can create render buffers or textures and use gdk_cairo_draw_from_gl() in the draw
/// function of your widget to draw them. Then GDK will handle the integration of your rendering
/// with that of other widgets.
///
/// Support for GLContext is platform-specific, context creation can fail, returning None context.
///
/// A GLContext has to be made "current" in order to start using it, otherwise any OpenGL call will
/// be ignored.
///
/// ##Creating a new OpenGL context
///
/// In order to create a new GLContext instance you need a GdkWindow, which you typically get during
/// the realize call of a widget.
///
/// A GLContext is not realized until either gdk_gl_context_make_current(), or until it is realized
/// using gdk_gl_context_realize(). It is possible to specify details of the GL context like the
/// OpenGL version to be used, or whether the GL context should have extra state validation enabled
/// after calling gdk_window_create_gl_context() by calling gdk_gl_context_realize(). If the
/// realization fails you have the option to change the settings of the GLContext and try again.
///
/// ##Using a GLContext
///
/// You will need to make the GLContext the current context before issuing OpenGL calls; the system
/// sends OpenGL commands to whichever context is current. It is possible to have multiple contexts,
/// so you always need to ensure that the one which you want to draw with is the current one before
/// issuing commands:
///
/// ```Rust,ignore
/// GLContext::make_current(context);
/// ```
///
/// You can now perform your drawing using OpenGL commands.
///
/// You can check which GLContext is the current one by using gdk_gl_context_get_current(); you can
/// also unset any GLContext that is currently set by calling gdk_gl_context_clear_current().
pub type GLContext = Object<ffi::GdkGLContext>;

impl types::StaticType for GLContext {
    fn static_type() -> types::Type { unsafe { from_glib(ffi::gdk_gl_context_get_type()) } }
}

impl GLContext {
    /// Retrieves the Display the context is created for
    pub fn get_display(&self) -> Option<Display> {
        unsafe { from_glib_none(ffi::gdk_gl_context_get_display(self.to_glib_none().0)) }
    }

    /// Retrieves the GdkWindow used by self.
    pub fn get_window(&self) -> Option<Window> {
        unsafe { from_glib_none(ffi::gdk_gl_context_get_window(self.to_glib_none().0)) }
    }

    /// Retrieves the GLContext that this context share data with.
    pub fn get_shared_context(&self) -> Option<GLContext> {
        unsafe { from_glib_none(ffi::gdk_gl_context_get_shared_context(self.to_glib_none().0)) }
    }

    /// Retrieves the OpenGL version of the context.
    ///
    /// The context must be realized prior to calling this function.
    pub fn get_version(&self) -> (i32, i32) {
        let mut major = 0;
        let mut minor = 0;

        unsafe { ffi::gdk_gl_context_get_version(self.to_glib_none().0, &mut major, &mut minor) };
        (major, minor)
    }

    /// Sets the major and minor version of OpenGL to request.
    ///
    /// Setting major and minor to zero will use the default values.
    ///
    /// The GLContext must not be realized or made current prior to calling this function.
    pub fn set_required_version(&self, major: i32, minor: i32) {
        unsafe { ffi::gdk_gl_context_set_required_version(self.to_glib_none().0, major, minor) }
    }

    /// Sets whether the GLContext should perform extra validations and run time checking.
    /// This is useful during development, but has additional overhead.
    ///
    /// The GLContext must not be realized or made current prior to calling this function.
    pub fn set_debug_enabled(&self, enabled: bool) {
        unsafe { ffi::gdk_gl_context_set_debug_enabled(self.to_glib_none().0, enabled.to_glib()) }
    }

    /// Retrieves the value set using set_debug_enabled().
    pub fn get_debug_enabled(&self) -> bool {
        unsafe { from_glib(ffi::gdk_gl_context_get_debug_enabled(self.to_glib_none().0)) }
    }

    /// Sets whether the GLContext should be forward compatible.
    ///
    /// Forward compatibile contexts must not support OpenGL functionality that has been
    /// marked as deprecated in the requested version; non-forward compatible contexts, on
    /// the other hand, must support both deprecated and non deprecated functionality.
    ///
    /// The GLContext must not be realized or made current prior to calling this function.
    pub fn set_forward_compatible(&self, compatible: bool) {
        unsafe { ffi::gdk_gl_context_set_forward_compatible(self.to_glib_none().0, compatible.to_glib()) }
    }

    /// Retrieves the value set using set_forward_compatible().
    pub fn get_forward_compatible(&self) -> bool {
        unsafe { from_glib(ffi::gdk_gl_context_get_forward_compatible(self.to_glib_none().0)) }
    }

    /*
    /// Realizes the given GLContext.
    ///
    /// It is safe to call this function on a realized GLContext.
    pub fn realize(&self, *mut *mut glib::error) -> bool;
    */

    /// Makes self the current context.
    pub fn make_current(&self) {
        unsafe { ffi::gdk_gl_context_make_current(self.to_glib_none().0) }
    }

    /// Retrieves the current GLContext.
    pub fn get_current() -> Option<GLContext> {
        unsafe { from_glib_none(ffi::gdk_gl_context_get_current()) }
    }

    /// Clears the current GLContext.
    ///
    /// Any OpenGL call after this function returns will be ignored until GLContext::make_current()
    /// is called.
    pub fn clear_current() {
        unsafe { ffi::gdk_gl_context_clear_current() }
    }
}
<|endoftext|>


<commit_before>#![allow(unused_imports, unused_variables, dead_code)]
use std::fmt::Write as FmtWrite;
use std::io::Write;
use std::marker::PhantomData;
use std::str::FromStr;
use std::{cmp, fmt, io, mem};

use bytes::{BufMut, Bytes, BytesMut};
use http::header::{
    HeaderValue, ACCEPT_ENCODING, CONNECTION, CONTENT_LENGTH, DATE, TRANSFER_ENCODING,
};
use http::{HeaderMap, Method, StatusCode, Version};

use crate::body::BodyLength;
use crate::config::ServiceConfig;
use crate::header::ContentEncoding;
use crate::helpers;
use crate::message::{ConnectionType, RequestHead, ResponseHead};
use crate::request::Request;
use crate::response::Response;

const AVERAGE_HEADER_SIZE: usize = 30;

#[derive(Debug)]
pub(crate) struct MessageEncoder<T: MessageType> {
    pub length: BodyLength,
    pub te: TransferEncoding,
    _t: PhantomData<T>,
}

impl<T: MessageType> Default for MessageEncoder<T> {
    fn default() -> Self {
        MessageEncoder {
            length: BodyLength::None,
            te: TransferEncoding::empty(),
            _t: PhantomData,
        }
    }
}

pub(crate) trait MessageType: Sized {
    fn status(&self) -> Option<StatusCode>;

    fn connection_type(&self) -> Option<ConnectionType>;

    fn headers(&self) -> &HeaderMap;

    fn chunked(&self) -> bool;

    fn encode_status(&mut self, dst: &mut BytesMut) -> io::Result<()>;

    fn encode_headers(
        &mut self,
        dst: &mut BytesMut,
        version: Version,
        mut length: BodyLength,
        ctype: ConnectionType,
        config: &ServiceConfig,
    ) -> io::Result<()> {
        let mut skip_len = length != BodyLength::Stream;

        // Content length
        if let Some(status) = self.status() {
            match status {
                StatusCode::NO_CONTENT
                | StatusCode::CONTINUE
                | StatusCode::PROCESSING => length = BodyLength::None,
                StatusCode::SWITCHING_PROTOCOLS => {
                    skip_len = true;
                    length = BodyLength::Stream;
                }
                _ => (),
            }
        }
        match length {
            BodyLength::Stream => {
                if self.chunked() {
                    dst.extend_from_slice(b"\r\ntransfer-encoding: chunked\r\n")
                }
            }
            BodyLength::Empty => {
                dst.extend_from_slice(b"\r\ncontent-length: 0\r\n");
            }
            BodyLength::Sized(len) => helpers::write_content_length(len, dst),
            BodyLength::Sized64(len) => {
                dst.extend_from_slice(b"\r\ncontent-length: ");
                write!(dst.writer(), "{}", len)?;
                dst.extend_from_slice(b"\r\n");
            }
            BodyLength::None => dst.extend_from_slice(b"\r\n"),
        }

        // Connection
        match ctype {
            ConnectionType::Upgrade => dst.extend_from_slice(b"connection: upgrade\r\n"),
            ConnectionType::KeepAlive if version < Version::HTTP_11 => {
                dst.extend_from_slice(b"connection: keep-alive\r\n")
            }
            ConnectionType::Close if version >= Version::HTTP_11 => {
                dst.extend_from_slice(b"connection: close\r\n")
            }
            _ => (),
        }

        // write headers
        let mut pos = 0;
        let mut has_date = false;
        let mut remaining = dst.remaining_mut();
        let mut buf = unsafe { &mut *(dst.bytes_mut() as *mut [u8]) };
        for (key, value) in self.headers() {
            match *key {
                CONNECTION => continue,
                TRANSFER_ENCODING | CONTENT_LENGTH if skip_len => continue,
                DATE => {
                    has_date = true;
                }
                _ => (),
            }

            let v = value.as_ref();
            let k = key.as_str().as_bytes();
            let len = k.len() + v.len() + 4;
            if len > remaining {
                unsafe {
                    dst.advance_mut(pos);
                }
                pos = 0;
                dst.reserve(len);
                remaining = dst.remaining_mut();
                unsafe {
                    buf = &mut *(dst.bytes_mut() as *mut _);
                }
            }

            buf[pos..pos + k.len()].copy_from_slice(k);
            pos += k.len();
            buf[pos..pos + 2].copy_from_slice(b": ");
            pos += 2;
            buf[pos..pos + v.len()].copy_from_slice(v);
            pos += v.len();
            buf[pos..pos + 2].copy_from_slice(b"\r\n");
            pos += 2;
            remaining -= len;
        }
        unsafe {
            dst.advance_mut(pos);
        }

        // optimized date header, set_date writes \r\n
        if !has_date {
            config.set_date(dst);
        } else {
            // msg eof
            dst.extend_from_slice(b"\r\n");
        }

        Ok(())
    }
}

impl MessageType for Response<()> {
    fn status(&self) -> Option<StatusCode> {
        Some(self.head().status)
    }

    fn chunked(&self) -> bool {
        !self.head().no_chunking
    }

    fn connection_type(&self) -> Option<ConnectionType> {
        self.head().ctype
    }

    fn headers(&self) -> &HeaderMap {
        &self.head().headers
    }

    fn encode_status(&mut self, dst: &mut BytesMut) -> io::Result<()> {
        let head = self.head();
        let reason = head.reason().as_bytes();
        dst.reserve(256 + head.headers.len() * AVERAGE_HEADER_SIZE + reason.len());

        // status line
        helpers::write_status_line(head.version, head.status.as_u16(), dst);
        dst.extend_from_slice(reason);
        Ok(())
    }
}

impl MessageType for RequestHead {
    fn status(&self) -> Option<StatusCode> {
        None
    }

    fn connection_type(&self) -> Option<ConnectionType> {
        self.ctype
    }

    fn chunked(&self) -> bool {
        !self.no_chunking
    }

    fn headers(&self) -> &HeaderMap {
        &self.headers
    }

    fn encode_status(&mut self, dst: &mut BytesMut) -> io::Result<()> {
        write!(
            Writer(dst),
            "{} {} {}",
            self.method,
            self.uri.path_and_query().map(|u| u.as_str()).unwrap_or("/"),
            match self.version {
                Version::HTTP_09 => "HTTP/0.9",
                Version::HTTP_10 => "HTTP/1.0",
                Version::HTTP_11 => "HTTP/1.1",
                Version::HTTP_2 => "HTTP/2.0",
            }
        )
        .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
    }
}

impl<T: MessageType> MessageEncoder<T> {
    /// Encode message
    pub fn encode_chunk(&mut self, msg: &[u8], buf: &mut BytesMut) -> io::Result<bool> {
        self.te.encode(msg, buf)
    }

    /// Encode eof
    pub fn encode_eof(&mut self, buf: &mut BytesMut) -> io::Result<()> {
        self.te.encode_eof(buf)
    }

    pub fn encode(
        &mut self,
        dst: &mut BytesMut,
        message: &mut T,
        head: bool,
        version: Version,
        length: BodyLength,
        ctype: ConnectionType,
        config: &ServiceConfig,
    ) -> io::Result<()> {
        // transfer encoding
        if !head {
            self.te = match length {
                BodyLength::Empty => TransferEncoding::empty(),
                BodyLength::Sized(len) => TransferEncoding::length(len as u64),
                BodyLength::Sized64(len) => TransferEncoding::length(len),
                BodyLength::Stream => {
                    if message.chunked() {
                        TransferEncoding::chunked()
                    } else {
                        TransferEncoding::eof()
                    }
                }
                BodyLength::None => TransferEncoding::empty(),
            };
        } else {
            self.te = TransferEncoding::empty();
        }

        message.encode_status(dst)?;
        message.encode_headers(dst, version, length, ctype, config)
    }
}

/// Encoders to handle different Transfer-Encodings.
#[derive(Debug)]
pub(crate) struct TransferEncoding {
    kind: TransferEncodingKind,
}

#[derive(Debug, PartialEq, Clone)]
enum TransferEncodingKind {
    /// An Encoder for when Transfer-Encoding includes `chunked`.
    Chunked(bool),
    /// An Encoder for when Content-Length is set.
    ///
    /// Enforces that the body is not longer than the Content-Length header.
    Length(u64),
    /// An Encoder for when Content-Length is not known.
    ///
    /// Application decides when to stop writing.
    Eof,
}

impl TransferEncoding {
    #[inline]
    pub fn empty() -> TransferEncoding {
        TransferEncoding {
            kind: TransferEncodingKind::Length(0),
        }
    }

    #[inline]
    pub fn eof() -> TransferEncoding {
        TransferEncoding {
            kind: TransferEncodingKind::Eof,
        }
    }

    #[inline]
    pub fn chunked() -> TransferEncoding {
        TransferEncoding {
            kind: TransferEncodingKind::Chunked(false),
        }
    }

    #[inline]
    pub fn length(len: u64) -> TransferEncoding {
        TransferEncoding {
            kind: TransferEncodingKind::Length(len),
        }
    }

    /// Encode message. Return `EOF` state of encoder
    #[inline]
    pub fn encode(&mut self, msg: &[u8], buf: &mut BytesMut) -> io::Result<bool> {
        match self.kind {
            TransferEncodingKind::Eof => {
                let eof = msg.is_empty();
                buf.extend_from_slice(msg);
                Ok(eof)
            }
            TransferEncodingKind::Chunked(ref mut eof) => {
                if *eof {
                    return Ok(true);
                }

                if msg.is_empty() {
                    *eof = true;
                    buf.extend_from_slice(b"0\r\n\r\n");
                } else {
                    writeln!(Writer(buf), "{:X}\r", msg.len())
                        .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;

                    buf.reserve(msg.len() + 2);
                    buf.extend_from_slice(msg);
                    buf.extend_from_slice(b"\r\n");
                }
                Ok(*eof)
            }
            TransferEncodingKind::Length(ref mut remaining) => {
                if *remaining > 0 {
                    if msg.is_empty() {
                        return Ok(*remaining == 0);
                    }
                    let len = cmp::min(*remaining, msg.len() as u64);

                    buf.extend_from_slice(&msg[..len as usize]);

                    *remaining -= len as u64;
                    Ok(*remaining == 0)
                } else {
                    Ok(true)
                }
            }
        }
    }

    /// Encode eof. Return `EOF` state of encoder
    #[inline]
    pub fn encode_eof(&mut self, buf: &mut BytesMut) -> io::Result<()> {
        match self.kind {
            TransferEncodingKind::Eof => Ok(()),
            TransferEncodingKind::Length(rem) => {
                if rem != 0 {
                    Err(io::Error::new(io::ErrorKind::UnexpectedEof, ""))
                } else {
                    Ok(())
                }
            }
            TransferEncodingKind::Chunked(ref mut eof) => {
                if !*eof {
                    *eof = true;
                    buf.extend_from_slice(b"0\r\n\r\n");
                }
                Ok(())
            }
        }
    }
}

struct Writer<'a>(pub &'a mut BytesMut);

impl<'a> io::Write for Writer<'a> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.extend_from_slice(buf);
        Ok(buf.len())
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bytes::Bytes;

    #[test]
    fn test_chunked_te() {
        let mut bytes = BytesMut::new();
        let mut enc = TransferEncoding::chunked();
        {
            assert!(!enc.encode(b"test", &mut bytes).ok().unwrap());
            assert!(enc.encode(b"", &mut bytes).ok().unwrap());
        }
        assert_eq!(
            bytes.take().freeze(),
            Bytes::from_static(b"4\r\ntest\r\n0\r\n\r\n")
        );
    }
}
<commit_msg>do not skip content length is no chunking is selected<commit_after>#![allow(unused_imports, unused_variables, dead_code)]
use std::fmt::Write as FmtWrite;
use std::io::Write;
use std::marker::PhantomData;
use std::str::FromStr;
use std::{cmp, fmt, io, mem};

use bytes::{BufMut, Bytes, BytesMut};
use http::header::{
    HeaderValue, ACCEPT_ENCODING, CONNECTION, CONTENT_LENGTH, DATE, TRANSFER_ENCODING,
};
use http::{HeaderMap, Method, StatusCode, Version};

use crate::body::BodyLength;
use crate::config::ServiceConfig;
use crate::header::ContentEncoding;
use crate::helpers;
use crate::message::{ConnectionType, RequestHead, ResponseHead};
use crate::request::Request;
use crate::response::Response;

const AVERAGE_HEADER_SIZE: usize = 30;

#[derive(Debug)]
pub(crate) struct MessageEncoder<T: MessageType> {
    pub length: BodyLength,
    pub te: TransferEncoding,
    _t: PhantomData<T>,
}

impl<T: MessageType> Default for MessageEncoder<T> {
    fn default() -> Self {
        MessageEncoder {
            length: BodyLength::None,
            te: TransferEncoding::empty(),
            _t: PhantomData,
        }
    }
}

pub(crate) trait MessageType: Sized {
    fn status(&self) -> Option<StatusCode>;

    fn connection_type(&self) -> Option<ConnectionType>;

    fn headers(&self) -> &HeaderMap;

    fn chunked(&self) -> bool;

    fn encode_status(&mut self, dst: &mut BytesMut) -> io::Result<()>;

    fn encode_headers(
        &mut self,
        dst: &mut BytesMut,
        version: Version,
        mut length: BodyLength,
        ctype: ConnectionType,
        config: &ServiceConfig,
    ) -> io::Result<()> {
        let chunked = self.chunked();
        let mut skip_len = length != BodyLength::Stream;

        // Content length
        if let Some(status) = self.status() {
            match status {
                StatusCode::NO_CONTENT
                | StatusCode::CONTINUE
                | StatusCode::PROCESSING => length = BodyLength::None,
                StatusCode::SWITCHING_PROTOCOLS => {
                    skip_len = true;
                    length = BodyLength::Stream;
                }
                _ => (),
            }
        }
        match length {
            BodyLength::Stream => {
                if chunked {
                    dst.extend_from_slice(b"\r\ntransfer-encoding: chunked\r\n")
                } else {
                    skip_len = false;
                }
            }
            BodyLength::Empty => {
                dst.extend_from_slice(b"\r\ncontent-length: 0\r\n");
            }
            BodyLength::Sized(len) => helpers::write_content_length(len, dst),
            BodyLength::Sized64(len) => {
                dst.extend_from_slice(b"\r\ncontent-length: ");
                write!(dst.writer(), "{}", len)?;
                dst.extend_from_slice(b"\r\n");
            }
            BodyLength::None => dst.extend_from_slice(b"\r\n"),
        }

        // Connection
        match ctype {
            ConnectionType::Upgrade => dst.extend_from_slice(b"connection: upgrade\r\n"),
            ConnectionType::KeepAlive if version < Version::HTTP_11 => {
                dst.extend_from_slice(b"connection: keep-alive\r\n")
            }
            ConnectionType::Close if version >= Version::HTTP_11 => {
                dst.extend_from_slice(b"connection: close\r\n")
            }
            _ => (),
        }

        // write headers
        let mut pos = 0;
        let mut has_date = false;
        let mut remaining = dst.remaining_mut();
        let mut buf = unsafe { &mut *(dst.bytes_mut() as *mut [u8]) };
        for (key, value) in self.headers() {
            match *key {
                CONNECTION => continue,
                TRANSFER_ENCODING | CONTENT_LENGTH if skip_len => continue,
                DATE => {
                    has_date = true;
                }
                _ => (),
            }

            let v = value.as_ref();
            let k = key.as_str().as_bytes();
            let len = k.len() + v.len() + 4;
            if len > remaining {
                unsafe {
                    dst.advance_mut(pos);
                }
                pos = 0;
                dst.reserve(len);
                remaining = dst.remaining_mut();
                unsafe {
                    buf = &mut *(dst.bytes_mut() as *mut _);
                }
            }

            buf[pos..pos + k.len()].copy_from_slice(k);
            pos += k.len();
            buf[pos..pos + 2].copy_from_slice(b": ");
            pos += 2;
            buf[pos..pos + v.len()].copy_from_slice(v);
            pos += v.len();
            buf[pos..pos + 2].copy_from_slice(b"\r\n");
            pos += 2;
            remaining -= len;
        }
        unsafe {
            dst.advance_mut(pos);
        }

        // optimized date header, set_date writes \r\n
        if !has_date {
            config.set_date(dst);
        } else {
            // msg eof
            dst.extend_from_slice(b"\r\n");
        }

        Ok(())
    }
}

impl MessageType for Response<()> {
    fn status(&self) -> Option<StatusCode> {
        Some(self.head().status)
    }

    fn chunked(&self) -> bool {
        !self.head().no_chunking
    }

    fn connection_type(&self) -> Option<ConnectionType> {
        self.head().ctype
    }

    fn headers(&self) -> &HeaderMap {
        &self.head().headers
    }

    fn encode_status(&mut self, dst: &mut BytesMut) -> io::Result<()> {
        let head = self.head();
        let reason = head.reason().as_bytes();
        dst.reserve(256 + head.headers.len() * AVERAGE_HEADER_SIZE + reason.len());

        // status line
        helpers::write_status_line(head.version, head.status.as_u16(), dst);
        dst.extend_from_slice(reason);
        Ok(())
    }
}

impl MessageType for RequestHead {
    fn status(&self) -> Option<StatusCode> {
        None
    }

    fn connection_type(&self) -> Option<ConnectionType> {
        self.ctype
    }

    fn chunked(&self) -> bool {
        !self.no_chunking
    }

    fn headers(&self) -> &HeaderMap {
        &self.headers
    }

    fn encode_status(&mut self, dst: &mut BytesMut) -> io::Result<()> {
        write!(
            Writer(dst),
            "{} {} {}",
            self.method,
            self.uri.path_and_query().map(|u| u.as_str()).unwrap_or("/"),
            match self.version {
                Version::HTTP_09 => "HTTP/0.9",
                Version::HTTP_10 => "HTTP/1.0",
                Version::HTTP_11 => "HTTP/1.1",
                Version::HTTP_2 => "HTTP/2.0",
            }
        )
        .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
    }
}

impl<T: MessageType> MessageEncoder<T> {
    /// Encode message
    pub fn encode_chunk(&mut self, msg: &[u8], buf: &mut BytesMut) -> io::Result<bool> {
        self.te.encode(msg, buf)
    }

    /// Encode eof
    pub fn encode_eof(&mut self, buf: &mut BytesMut) -> io::Result<()> {
        self.te.encode_eof(buf)
    }

    pub fn encode(
        &mut self,
        dst: &mut BytesMut,
        message: &mut T,
        head: bool,
        version: Version,
        length: BodyLength,
        ctype: ConnectionType,
        config: &ServiceConfig,
    ) -> io::Result<()> {
        // transfer encoding
        if !head {
            self.te = match length {
                BodyLength::Empty => TransferEncoding::empty(),
                BodyLength::Sized(len) => TransferEncoding::length(len as u64),
                BodyLength::Sized64(len) => TransferEncoding::length(len),
                BodyLength::Stream => {
                    if message.chunked() {
                        TransferEncoding::chunked()
                    } else {
                        TransferEncoding::eof()
                    }
                }
                BodyLength::None => TransferEncoding::empty(),
            };
        } else {
            self.te = TransferEncoding::empty();
        }

        message.encode_status(dst)?;
        message.encode_headers(dst, version, length, ctype, config)
    }
}

/// Encoders to handle different Transfer-Encodings.
#[derive(Debug)]
pub(crate) struct TransferEncoding {
    kind: TransferEncodingKind,
}

#[derive(Debug, PartialEq, Clone)]
enum TransferEncodingKind {
    /// An Encoder for when Transfer-Encoding includes `chunked`.
    Chunked(bool),
    /// An Encoder for when Content-Length is set.
    ///
    /// Enforces that the body is not longer than the Content-Length header.
    Length(u64),
    /// An Encoder for when Content-Length is not known.
    ///
    /// Application decides when to stop writing.
    Eof,
}

impl TransferEncoding {
    #[inline]
    pub fn empty() -> TransferEncoding {
        TransferEncoding {
            kind: TransferEncodingKind::Length(0),
        }
    }

    #[inline]
    pub fn eof() -> TransferEncoding {
        TransferEncoding {
            kind: TransferEncodingKind::Eof,
        }
    }

    #[inline]
    pub fn chunked() -> TransferEncoding {
        TransferEncoding {
            kind: TransferEncodingKind::Chunked(false),
        }
    }

    #[inline]
    pub fn length(len: u64) -> TransferEncoding {
        TransferEncoding {
            kind: TransferEncodingKind::Length(len),
        }
    }

    /// Encode message. Return `EOF` state of encoder
    #[inline]
    pub fn encode(&mut self, msg: &[u8], buf: &mut BytesMut) -> io::Result<bool> {
        match self.kind {
            TransferEncodingKind::Eof => {
                let eof = msg.is_empty();
                buf.extend_from_slice(msg);
                Ok(eof)
            }
            TransferEncodingKind::Chunked(ref mut eof) => {
                if *eof {
                    return Ok(true);
                }

                if msg.is_empty() {
                    *eof = true;
                    buf.extend_from_slice(b"0\r\n\r\n");
                } else {
                    writeln!(Writer(buf), "{:X}\r", msg.len())
                        .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;

                    buf.reserve(msg.len() + 2);
                    buf.extend_from_slice(msg);
                    buf.extend_from_slice(b"\r\n");
                }
                Ok(*eof)
            }
            TransferEncodingKind::Length(ref mut remaining) => {
                if *remaining > 0 {
                    if msg.is_empty() {
                        return Ok(*remaining == 0);
                    }
                    let len = cmp::min(*remaining, msg.len() as u64);

                    buf.extend_from_slice(&msg[..len as usize]);

                    *remaining -= len as u64;
                    Ok(*remaining == 0)
                } else {
                    Ok(true)
                }
            }
        }
    }

    /// Encode eof. Return `EOF` state of encoder
    #[inline]
    pub fn encode_eof(&mut self, buf: &mut BytesMut) -> io::Result<()> {
        match self.kind {
            TransferEncodingKind::Eof => Ok(()),
            TransferEncodingKind::Length(rem) => {
                if rem != 0 {
                    Err(io::Error::new(io::ErrorKind::UnexpectedEof, ""))
                } else {
                    Ok(())
                }
            }
            TransferEncodingKind::Chunked(ref mut eof) => {
                if !*eof {
                    *eof = true;
                    buf.extend_from_slice(b"0\r\n\r\n");
                }
                Ok(())
            }
        }
    }
}

struct Writer<'a>(pub &'a mut BytesMut);

impl<'a> io::Write for Writer<'a> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.extend_from_slice(buf);
        Ok(buf.len())
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bytes::Bytes;

    #[test]
    fn test_chunked_te() {
        let mut bytes = BytesMut::new();
        let mut enc = TransferEncoding::chunked();
        {
            assert!(!enc.encode(b"test", &mut bytes).ok().unwrap());
            assert!(enc.encode(b"", &mut bytes).ok().unwrap());
        }
        assert_eq!(
            bytes.take().freeze(),
            Bytes::from_static(b"4\r\ntest\r\n0\r\n\r\n")
        );
    }
}
<|endoftext|>