<commit_before>//! Implements a generic chess engine.

use std::process;
use std::marker::PhantomData;
use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use std::time::{SystemTime, Duration};
use std::cmp::{min, max};
use uci::*;
use value::*;
use depth::*;
use hash_table::*;
use search_executor::{SearchParams, SearchReport, SearchExecutor};
use search_node::SearchNode;
use time_manager::{TimeManager, RemainingTime};

struct SearchStatus {
    pub done: bool,
    pub depth: Depth,
    pub value: Value,
    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,
            value: VALUE_UNKNOWN,
            searched_nodes: 0,
            duration_millis: 0,
        }
    }
}

enum PlayWhen<S, T>
    where S: SearchExecutor<ReportData = Vec<Variation>>,
          T: TimeManager<S>
{
    TimeManagement(T), // 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.
    Mate(i16), // Stop when a mate in the given number of moves is found.
    Never(PhantomData<S>), // An infinite search.
}

struct Engine<S, T>
    where S: SearchExecutor<ReportData = Vec<Variation>>,
          T: TimeManager<S>
{
    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<S, T>,
}

impl<S, T> UciEngine for Engine<S, T>
    where S: SearchExecutor<ReportData = Vec<Variation>>,
          T: TimeManager<S>
{
    fn name() -> &'static str {
        ENGINE.lock().unwrap().as_ref().unwrap().name
    }

    fn author() -> &'static str {
        ENGINE.lock().unwrap().as_ref().unwrap().author
    }

    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),
        ];
        options.extend(S::options());
        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<S, T> {
        const START_FEN: &'static str = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 1";
        let tt = Arc::new(S::HashTable::new(tt_size_mb));
        let started_at = SystemTime::now();
        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: started_at,
            status: SearchStatus { done: true, ..Default::default() },
            nps_stats: (0, 0, 0),
            silent_since: started_at,
            is_pondering: false,
            play_when: PlayWhen::Never(PhantomData),
        }
    }

    fn set_option(&mut self, name: &str, value: &str) {
        match name {
            "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);
        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) = S::SearchNode::from_history(fen, moves) {
            self.position = p;
        }
    }

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

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

        // Start a new search.
        let depth = params.depth.map_or(DEPTH_MAX, |x| min(x, DEPTH_MAX as u64) as Depth);
        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(PhantomData)
        } 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(depth)
        } else if params.mate.is_some() {
            PlayWhen::Mate(min(params.mate.unwrap(), (DEPTH_MAX + 1) as u64 / 2) as i16)
        } else {
            PlayWhen::TimeManagement(T::new(&self.position,
                                            RemainingTime {
                                                white_millis: params.wtime.unwrap_or(300_000),
                                                black_millis: params.btime.unwrap_or(300_000),
                                                winc_millis: params.winc.unwrap_or(0),
                                                binc_millis: params.binc.unwrap_or(0),
                                                movestogo: params.movestogo,
                                            }))
        };
        self.searcher.start_search(SearchParams {
            search_id: 0,
            position: self.position.clone(),
            depth: depth,
            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(&self.searcher)
                }
                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::Mate(m) => self.status.done || self.status.value > VALUE_MAX - 2 * m,
                PlayWhen::Never(_) => false,
            } {
                self.stop();
            }
        }

        self.queue.pop_front()
    }

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

impl<S, T> Engine<S, T>
    where S: SearchExecutor<ReportData = Vec<Variation>>,
          T: TimeManager<S>
{
    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) },
        ]));
    }

    fn queue_pv(&mut self, variations: &Vec<Variation>) {
        fn suffix(bound: BoundType) -> &'static str {
            match bound {
                BOUND_UPPER => " upperbound",
                BOUND_LOWER => " lowerbound",
                BOUND_EXACT => "",
                _ => panic!("unexpected bound type"),
            }
        }

        let SearchStatus { ref depth, ref searched_nodes, ref duration_millis, .. } = self.status;
        for (i, &Variation { ref moves, value, bound }) in variations.iter().enumerate() {
            let score = match value {
                v if bound & BOUND_UPPER != 0 && VALUE_MIN < v && v < VALUE_EVAL_MIN => {
                    format!("mate {}", (VALUE_MIN - v - 1) / 2)
                }
                v if bound & BOUND_LOWER != 0 && VALUE_EVAL_MAX < v && v < VALUE_MAX => {
                    format!("mate {}", (VALUE_MAX - v + 1) / 2)
                }
                v if v <= -9999 => format!("cp -9999{}", suffix(bound | BOUND_LOWER)),
                v if v >= 9999 => format!("cp 9999{}", suffix(bound | BOUND_UPPER)),
                v => format!("cp {}{}", v, suffix(bound)),
            };
            let mut pv = String::new();
            for m in moves.iter().take(max(0, *depth) as usize) {
                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: score },
                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 },
            ]));
        }
    }

    fn queue_best_move(&mut self) {
        let pv = self.tt.extract_pv(&self.position);
        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()),
        });
    }

    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,
            value: report.value,
            searched_nodes: report.searched_nodes,
            duration_millis: duration_millis,
        };

        // Update `self.nps_stats` every 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(report);
        }

        // If principal variations are provided with the report, show them.
        if !report.data.is_empty() {
            self.queue_pv(&report.data);
            self.silent_since = SystemTime::now();
        }

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

/// Serves UCI commands until a "quit" command is received, or an IO
/// error has occurred.
pub fn run_uci<S, T>(name: &'static str, author: &'static str) -> !
    where S: SearchExecutor<ReportData = Vec<Variation>>,
          T: TimeManager<S>
{
    // Obtain the lock.
    {
        let mut engine = ENGINE.lock().unwrap();
        assert!(engine.is_none(), "two engines can not run in parallel");
        *engine = Some(EngineIdentity {
            name: name,
            author: author,
        });
    }

    // Run the engine.
    let result = run_engine::<Engine<S, T>>();

    // Release the lock.
    *ENGINE.lock().unwrap() = None;

    process::exit(match result {
        Ok(_) => 0,
        Err(_) => 1,
    });
}

struct EngineIdentity {
    name: &'static str,
    author: &'static str,
}

lazy_static! {
    static ref ENGINE: Mutex<Option<EngineIdentity>> = Mutex::new(None);
}
<commit_msg>do not release engine's lock.<commit_after>//! Implements a generic chess engine.

use std::process;
use std::marker::PhantomData;
use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use std::time::{SystemTime, Duration};
use std::cmp::{min, max};
use uci::*;
use value::*;
use depth::*;
use hash_table::*;
use search_executor::{SearchParams, SearchReport, SearchExecutor};
use search_node::SearchNode;
use time_manager::{TimeManager, RemainingTime};

struct SearchStatus {
    pub done: bool,
    pub depth: Depth,
    pub value: Value,
    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,
            value: VALUE_UNKNOWN,
            searched_nodes: 0,
            duration_millis: 0,
        }
    }
}

enum PlayWhen<S, T>
    where S: SearchExecutor<ReportData = Vec<Variation>>,
          T: TimeManager<S>
{
    TimeManagement(T), // 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.
    Mate(i16), // Stop when a mate in the given number of moves is found.
    Never(PhantomData<S>), // An infinite search.
}

struct Engine<S, T>
    where S: SearchExecutor<ReportData = Vec<Variation>>,
          T: TimeManager<S>
{
    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<S, T>,
}

impl<S, T> UciEngine for Engine<S, T>
    where S: SearchExecutor<ReportData = Vec<Variation>>,
          T: TimeManager<S>
{
    fn name() -> &'static str {
        ENGINE.lock().unwrap().as_ref().unwrap().name
    }

    fn author() -> &'static str {
        ENGINE.lock().unwrap().as_ref().unwrap().author
    }

    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),
        ];
        options.extend(S::options());
        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<S, T> {
        const START_FEN: &'static str = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 1";
        let tt = Arc::new(S::HashTable::new(tt_size_mb));
        let started_at = SystemTime::now();
        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: started_at,
            status: SearchStatus { done: true, ..Default::default() },
            nps_stats: (0, 0, 0),
            silent_since: started_at,
            is_pondering: false,
            play_when: PlayWhen::Never(PhantomData),
        }
    }

    fn set_option(&mut self, name: &str, value: &str) {
        match name {
            "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);
        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) = S::SearchNode::from_history(fen, moves) {
            self.position = p;
        }
    }

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

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

        // Start a new search.
        let depth = params.depth.map_or(DEPTH_MAX, |x| min(x, DEPTH_MAX as u64) as Depth);
        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(PhantomData)
        } 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(depth)
        } else if params.mate.is_some() {
            PlayWhen::Mate(min(params.mate.unwrap(), (DEPTH_MAX + 1) as u64 / 2) as i16)
        } else {
            PlayWhen::TimeManagement(T::new(&self.position,
                                            RemainingTime {
                                                white_millis: params.wtime.unwrap_or(300_000),
                                                black_millis: params.btime.unwrap_or(300_000),
                                                winc_millis: params.winc.unwrap_or(0),
                                                binc_millis: params.binc.unwrap_or(0),
                                                movestogo: params.movestogo,
                                            }))
        };
        self.searcher.start_search(SearchParams {
            search_id: 0,
            position: self.position.clone(),
            depth: depth,
            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(&self.searcher)
                }
                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::Mate(m) => self.status.done || self.status.value > VALUE_MAX - 2 * m,
                PlayWhen::Never(_) => false,
            } {
                self.stop();
            }
        }

        self.queue.pop_front()
    }

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

impl<S, T> Engine<S, T>
    where S: SearchExecutor<ReportData = Vec<Variation>>,
          T: TimeManager<S>
{
    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) },
        ]));
    }

    fn queue_pv(&mut self, variations: &Vec<Variation>) {
        fn suffix(bound: BoundType) -> &'static str {
            match bound {
                BOUND_UPPER => " upperbound",
                BOUND_LOWER => " lowerbound",
                BOUND_EXACT => "",
                _ => panic!("unexpected bound type"),
            }
        }

        let SearchStatus { ref depth, ref searched_nodes, ref duration_millis, .. } = self.status;
        for (i, &Variation { ref moves, value, bound }) in variations.iter().enumerate() {
            let score = match value {
                v if bound & BOUND_UPPER != 0 && VALUE_MIN < v && v < VALUE_EVAL_MIN => {
                    format!("mate {}", (VALUE_MIN - v - 1) / 2)
                }
                v if bound & BOUND_LOWER != 0 && VALUE_EVAL_MAX < v && v < VALUE_MAX => {
                    format!("mate {}", (VALUE_MAX - v + 1) / 2)
                }
                v if v <= -9999 => format!("cp -9999{}", suffix(bound | BOUND_LOWER)),
                v if v >= 9999 => format!("cp 9999{}", suffix(bound | BOUND_UPPER)),
                v => format!("cp {}{}", v, suffix(bound)),
            };
            let mut pv = String::new();
            for m in moves.iter().take(max(0, *depth) as usize) {
                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: score },
                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 },
            ]));
        }
    }

    fn queue_best_move(&mut self) {
        let pv = self.tt.extract_pv(&self.position);
        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()),
        });
    }

    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,
            value: report.value,
            searched_nodes: report.searched_nodes,
            duration_millis: duration_millis,
        };

        // Update `self.nps_stats` every 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(report);
        }

        // If principal variations are provided with the report, show them.
        if !report.data.is_empty() {
            self.queue_pv(&report.data);
            self.silent_since = SystemTime::now();
        }

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

/// Serves UCI commands until a "quit" command is received, or an IO
/// error has occurred.
pub fn run_uci<S, T>(name: &'static str, author: &'static str) -> !
    where S: SearchExecutor<ReportData = Vec<Variation>>,
          T: TimeManager<S>
{
    // Set engine's identity.
    {
        let mut engine = ENGINE.lock().unwrap();
        assert!(engine.is_none(), "two engines can not run in parallel");
        *engine = Some(EngineIdentity {
            name: name,
            author: author,
        });
    }

    // Run the engine.
    process::exit(match run_engine::<Engine<S, T>>() {
        Ok(_) => 0,
        Err(_) => 1,
    });
}

struct EngineIdentity {
    name: &'static str,
    author: &'static str,
}

lazy_static! {
    static ref ENGINE: Mutex<Option<EngineIdentity>> = Mutex::new(None);
}
<|endoftext|>


<commit_before>use hyper;
use std::convert::From;
use serde_json;

impl From<::json::Error> for HueError {
    fn from(e: ::json::Error) -> HueError {
        HueErrorKind::BridgeError {
                address: e.address,
                description: e.description,
                error: From::from(e.code),
            }
            .into()
    }
}

error_chain! {
    types {
        HueError, HueErrorKind, ResultExt, Result;
    }

    errors {
        /// An error that occured in the bridge
        #[allow(missing_docs)]
        BridgeError {
            address: String,
            description: String,
            error: BridgeError
        } {
            description("bridge error")
            display("Bridge error: {:?} address: {} description: {}", error, address, description)
        }
    }

    foreign_links {
        JsonError(serde_json::Error) #[doc = "Json error"];
        HyperError(hyper::Error)     #[doc = "Hyper error"];
    }
}

macro_rules! error_enum {
    (
        $(#[$meta:meta])*
        pub enum $name:ident{
            $($err:ident = $n:expr),+;
            $other:ident
        }
    ) => (
        $(#[$meta])*
        pub enum $name{
            $($err = $n,)+
            $other
        }
        impl From<u16> for $name{
            fn from(n: u16) -> Self{
                match n {
                    $($n => $name::$err,)+
                    _ => $name::$other
                }
            }
        }
    );
}

error_enum!{
    /// All errors defined in http://www.developers.meethue.com/documentation/error-messages
    #[repr(u16)]
    #[allow(missing_docs)]
    #[derive(Debug, Copy, Clone, PartialEq, Eq)]
    pub enum BridgeError {
        // Generic Errors
        UnauthorizedUser = 1,
        BodyContainsInvalidJson = 2,
        ResourceNotAvailable = 3,
        MethodNotAvailableForResource = 4,
        MissingParametersInBody = 5,
        ParameterNotAvailable = 6,
        InvalidValueForParameter = 7,
        ParameterIsNotModifiable = 8,
        TooManyItemsInList = 11,
        ProtalConnectionRequired = 12,
        InternalError = 901,

        // Command Specific Errors
        LinkButtonNotPressed = 101,
        DHCPCannotBeDisabled = 110,
        InvalidUpdateState = 111,
        DeviceIsSetToOff = 201,
        GroupCouldNotBeCreatedGroupFull = 301,
        DeviceCouldNotBeAddedGroupFull = 302,
        DeviceIsUnreachable = 304,
        UpdateOrDeleteGroupOfThisTypeNotAllowed = 305,
        LightAlreadyUsed = 306,
        SceneCouldNotBeCreated = 401,
        SceneCouldNotBeCreatedBufferFull = 402,
        SceneCouldNotBeRemoved = 403,
        NotAllowedToCreateSensorType = 501,
        SensorListIsFull = 502,
        RuleEngineFull = 601,
        ConditionError = 607,
        ActionError = 608,
        UnableToActivae = 609,
        ScheduleListIsFull = 701,
        ScheduleTimezoneNotValid = 702,
        ScheduleCannotSetTimeAndLocalTime = 703,
        CannotCreateSchedule = 704,
        CannotEnableScheduleTimeInPast = 705,
        CommandError = 706,
        SourceModelInvalid = 801,
        SourceFactoryNew = 802,
        InvalidState = 803;
        Other
    }
}

#[test]
fn bridge_errors() {
    use self::BridgeError::*;

    assert_eq!(BridgeError::from(101), LinkButtonNotPressed);
    assert_eq!(BridgeError::from(0), Other);
    assert_eq!(BridgeError::from(51234), Other);
    assert_eq!(BridgeError::from(4), MethodNotAvailableForResource);
    assert_eq!(SceneCouldNotBeRemoved as u16, 403);
    assert_eq!(InternalError as u16, 901);
}
<commit_msg>Change Display implementation of `BridgeError`<commit_after>use hyper;
use std::convert::From;
use serde_json;

impl From<::json::Error> for HueError {
    fn from(e: ::json::Error) -> HueError {
        HueErrorKind::BridgeError {
                address: e.address,
                description: e.description,
                error: From::from(e.code),
            }
            .into()
    }
}

error_chain! {
    types {
        HueError, HueErrorKind, ResultExt, Result;
    }

    errors {
        /// An error that occured in the bridge
        #[allow(missing_docs)]
        BridgeError {
            address: String,
            description: String,
            error: BridgeError
        } {
            description("bridge error")
            display("Bridge error {:?} on {}: {}", error, address, description)
        }
    }

    foreign_links {
        JsonError(serde_json::Error) #[doc = "Json error"];
        HyperError(hyper::Error)     #[doc = "Hyper error"];
    }
}

macro_rules! error_enum {
    (
        $(#[$meta:meta])*
        pub enum $name:ident{
            $($err:ident = $n:expr),+;
            $other:ident
        }
    ) => (
        $(#[$meta])*
        pub enum $name{
            $($err = $n,)+
            $other
        }
        impl From<u16> for $name{
            fn from(n: u16) -> Self{
                match n {
                    $($n => $name::$err,)+
                    _ => $name::$other
                }
            }
        }
    );
}

error_enum!{
    /// All errors defined in http://www.developers.meethue.com/documentation/error-messages
    #[repr(u16)]
    #[allow(missing_docs)]
    #[derive(Debug, Copy, Clone, PartialEq, Eq)]
    pub enum BridgeError {
        // Generic Errors
        UnauthorizedUser = 1,
        BodyContainsInvalidJson = 2,
        ResourceNotAvailable = 3,
        MethodNotAvailableForResource = 4,
        MissingParametersInBody = 5,
        ParameterNotAvailable = 6,
        InvalidValueForParameter = 7,
        ParameterIsNotModifiable = 8,
        TooManyItemsInList = 11,
        ProtalConnectionRequired = 12,
        InternalError = 901,

        // Command Specific Errors
        LinkButtonNotPressed = 101,
        DHCPCannotBeDisabled = 110,
        InvalidUpdateState = 111,
        DeviceIsSetToOff = 201,
        GroupCouldNotBeCreatedGroupFull = 301,
        DeviceCouldNotBeAddedGroupFull = 302,
        DeviceIsUnreachable = 304,
        UpdateOrDeleteGroupOfThisTypeNotAllowed = 305,
        LightAlreadyUsed = 306,
        SceneCouldNotBeCreated = 401,
        SceneCouldNotBeCreatedBufferFull = 402,
        SceneCouldNotBeRemoved = 403,
        NotAllowedToCreateSensorType = 501,
        SensorListIsFull = 502,
        RuleEngineFull = 601,
        ConditionError = 607,
        ActionError = 608,
        UnableToActivae = 609,
        ScheduleListIsFull = 701,
        ScheduleTimezoneNotValid = 702,
        ScheduleCannotSetTimeAndLocalTime = 703,
        CannotCreateSchedule = 704,
        CannotEnableScheduleTimeInPast = 705,
        CommandError = 706,
        SourceModelInvalid = 801,
        SourceFactoryNew = 802,
        InvalidState = 803;
        Other
    }
}

#[test]
fn bridge_errors() {
    use self::BridgeError::*;

    assert_eq!(BridgeError::from(101), LinkButtonNotPressed);
    assert_eq!(BridgeError::from(0), Other);
    assert_eq!(BridgeError::from(51234), Other);
    assert_eq!(BridgeError::from(4), MethodNotAvailableForResource);
    assert_eq!(SceneCouldNotBeRemoved as u16, 403);
    assert_eq!(InternalError as u16, 901);
}
<|endoftext|>


<commit_before>use crate::{channel::Reply, Error, Promise, PromiseResolver};
use amq_protocol::frame::AMQPFrame;
use log::{log_enabled, trace, Level::Trace};
use parking_lot::Mutex;
use pinky_swear::Cancellable;
use std::{
    collections::{HashMap, VecDeque},
    fmt,
    sync::Arc,
};

pub(crate) struct ExpectedReply(
    pub(crate) Reply,
    pub(crate) Box<dyn Cancellable<Error> + Send>,
);

impl fmt::Debug for ExpectedReply {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_tuple("ExpectedReply").field(&self.0).finish()
    }
}

#[derive(Clone, Default)]
pub(crate) struct Frames {
    inner: Arc<Mutex<Inner>>,
}

impl Frames {
    pub(crate) fn push(
        &self,
        channel_id: u16,
        frame: AMQPFrame,
        resolver: PromiseResolver<()>,
        expected_reply: Option<ExpectedReply>,
    ) {
        self.inner
            .lock()
            .push(channel_id, frame, resolver, expected_reply);
    }

    pub(crate) fn push_frames(&self, frames: Vec<AMQPFrame>) -> Promise<()> {
        self.inner.lock().push_frames(frames)
    }

    pub(crate) fn retry(&self, frame: (AMQPFrame, Option<PromiseResolver<()>>)) {
        self.inner.lock().retry(frame);
    }

    pub(crate) fn pop(&self, flow: bool) -> Option<(AMQPFrame, Option<PromiseResolver<()>>)> {
        self.inner.lock().pop(flow)
    }

    pub(crate) fn next_expected_reply(&self, channel_id: u16) -> Option<Reply> {
        self.inner
            .lock()
            .expected_replies
            .get_mut(&channel_id)
            .and_then(|replies| replies.pop_front())
            .map(|t| t.0)
    }

    pub(crate) fn has_pending(&self) -> bool {
        self.inner.lock().has_pending()
    }

    pub(crate) fn drop_pending(&self, error: Error) {
        self.inner.lock().drop_pending(error);
    }

    pub(crate) fn clear_expected_replies(&self, channel_id: u16, error: Error) {
        self.inner.lock().clear_expected_replies(channel_id, error);
    }
}

struct Inner {
    /* Header frames must follow basic.publish frames directly, otherwise RabbitMQ-server send us an UNEXPECTED_FRAME */
    header_frames: VecDeque<(AMQPFrame, Option<PromiseResolver<()>>)>,
    priority_frames: VecDeque<(AMQPFrame, Option<PromiseResolver<()>>)>,
    frames: VecDeque<(AMQPFrame, Option<PromiseResolver<()>>)>,
    low_prio_frames: VecDeque<(AMQPFrame, Option<PromiseResolver<()>>)>,
    expected_replies: HashMap<u16, VecDeque<ExpectedReply>>,
}

impl Default for Inner {
    fn default() -> Self {
        Self {
            header_frames: VecDeque::default(),
            priority_frames: VecDeque::default(),
            frames: VecDeque::default(),
            low_prio_frames: VecDeque::default(),
            expected_replies: HashMap::default(),
        }
    }
}

impl fmt::Debug for Frames {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut debug = f.debug_struct("Frames");
        if let Some(inner) = self.inner.try_lock() {
            debug.field("expected_replies", &inner.expected_replies);
        }
        debug.finish()
    }
}

impl Inner {
    fn push(
        &mut self,
        channel_id: u16,
        frame: AMQPFrame,
        resolver: PromiseResolver<()>,
        expected_reply: Option<ExpectedReply>,
    ) {
        self.frames.push_back((frame, Some(resolver)));
        if let Some(reply) = expected_reply {
            trace!(
                "channel {} state is now waiting for {:?}",
                channel_id,
                reply
            );
            self.expected_replies
                .entry(channel_id)
                .or_default()
                .push_back(reply);
        }
    }

    fn push_frames(&mut self, mut frames: Vec<AMQPFrame>) -> Promise<()> {
        let (promise, resolver) = Promise::new();
        let last_frame = frames.pop();

        if log_enabled!(Trace) {
            promise.set_marker("Frames".into());
        }

        for frame in frames {
            self.low_prio_frames.push_back((frame, None));
        }
        if let Some(last_frame) = last_frame {
            self.low_prio_frames.push_back((last_frame, Some(resolver)));
        } else {
            resolver.swear(Ok(()));
        }

        promise
    }

    fn pop(&mut self, flow: bool) -> Option<(AMQPFrame, Option<PromiseResolver<()>>)> {
        if let Some(frame) = self
            .header_frames
            .pop_front()
            .or_else(|| self.priority_frames.pop_front())
            .or_else(|| self.frames.pop_front())
        {
            return Some(frame);
        }
        if flow {
            if let Some(frame) = self.low_prio_frames.pop_front() {
                // If the next frame is a header, that means we're a basic.publish
                // Header frame needs to follow directly the basic.publish frame or
                // the AMQP server will close the connection.
                // Push the header into header_frames which is there to handle just that.
                if self
                    .low_prio_frames
                    .front()
                    .map(|(frame, _)| frame.is_header())
                    .unwrap_or(false)
                {
                    // Yes, this will always be Some(), but let's keep our unwrap() count low
                    if let Some(next_frame) = self.low_prio_frames.pop_front() {
                        self.header_frames.push_back(next_frame);
                    }
                }
                return Some(frame);
            }
        }
        None
    }

    fn retry(&mut self, frame: (AMQPFrame, Option<PromiseResolver<()>>)) {
        if frame.0.is_header() {
            self.header_frames.push_front(frame);
        } else {
            self.priority_frames.push_back(frame);
        }
    }

    fn has_pending(&self) -> bool {
        !(self.header_frames.is_empty()
            && self.priority_frames.is_empty()
            && self.frames.is_empty()
            && self.low_prio_frames.is_empty())
    }

    fn drop_pending(&mut self, error: Error) {
        Self::drop_pending_frames(&mut self.header_frames, error.clone());
        Self::drop_pending_frames(&mut self.priority_frames, error.clone());
        Self::drop_pending_frames(&mut self.frames, error.clone());
        Self::drop_pending_frames(&mut self.low_prio_frames, error.clone());
        for (_, replies) in self.expected_replies.drain() {
            Self::cancel_expected_replies(replies, error.clone());
        }
    }

    fn drop_pending_frames(
        frames: &mut VecDeque<(AMQPFrame, Option<PromiseResolver<()>>)>,
        error: Error,
    ) {
        for (_, resolver) in std::mem::take(frames) {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
    }

    fn clear_expected_replies(&mut self, channel_id: u16, error: Error) {
        if let Some(replies) = self.expected_replies.remove(&channel_id) {
            Self::cancel_expected_replies(replies, error);
        }
    }

    fn cancel_expected_replies(replies: VecDeque<ExpectedReply>, error: Error) {
        for ExpectedReply(_, cancel) in replies {
            cancel.cancel(error.clone());
        }
    }
}
<commit_msg>frames: ensure we sent all body frames from a publish before sending another publish<commit_after>use crate::{channel::Reply, Error, Promise, PromiseResolver};
use amq_protocol::frame::AMQPFrame;
use log::{log_enabled, trace, Level::Trace};
use parking_lot::Mutex;
use pinky_swear::Cancellable;
use std::{
    collections::{HashMap, VecDeque},
    fmt,
    sync::Arc,
};

pub(crate) struct ExpectedReply(
    pub(crate) Reply,
    pub(crate) Box<dyn Cancellable<Error> + Send>,
);

impl fmt::Debug for ExpectedReply {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_tuple("ExpectedReply").field(&self.0).finish()
    }
}

#[derive(Clone, Default)]
pub(crate) struct Frames {
    inner: Arc<Mutex<Inner>>,
}

impl Frames {
    pub(crate) fn push(
        &self,
        channel_id: u16,
        frame: AMQPFrame,
        resolver: PromiseResolver<()>,
        expected_reply: Option<ExpectedReply>,
    ) {
        self.inner
            .lock()
            .push(channel_id, frame, resolver, expected_reply);
    }

    pub(crate) fn push_frames(&self, frames: Vec<AMQPFrame>) -> Promise<()> {
        self.inner.lock().push_frames(frames)
    }

    pub(crate) fn retry(&self, frame: (AMQPFrame, Option<PromiseResolver<()>>)) {
        self.inner.lock().retry(frame);
    }

    pub(crate) fn pop(&self, flow: bool) -> Option<(AMQPFrame, Option<PromiseResolver<()>>)> {
        self.inner.lock().pop(flow)
    }

    pub(crate) fn next_expected_reply(&self, channel_id: u16) -> Option<Reply> {
        self.inner
            .lock()
            .expected_replies
            .get_mut(&channel_id)
            .and_then(|replies| replies.pop_front())
            .map(|t| t.0)
    }

    pub(crate) fn has_pending(&self) -> bool {
        self.inner.lock().has_pending()
    }

    pub(crate) fn drop_pending(&self, error: Error) {
        self.inner.lock().drop_pending(error);
    }

    pub(crate) fn clear_expected_replies(&self, channel_id: u16, error: Error) {
        self.inner.lock().clear_expected_replies(channel_id, error);
    }
}

struct Inner {
    /* Header frames must follow basic.publish frames directly, otherwise RabbitMQ-server send us an UNEXPECTED_FRAME */
    header_frames: VecDeque<(AMQPFrame, Option<PromiseResolver<()>>)>,
    /* Body frames must be sent after the header frame but before the next publish */
    body_frames: VecDeque<(AMQPFrame, Option<PromiseResolver<()>>)>,
    priority_frames: VecDeque<(AMQPFrame, Option<PromiseResolver<()>>)>,
    frames: VecDeque<(AMQPFrame, Option<PromiseResolver<()>>)>,
    low_prio_frames: VecDeque<(AMQPFrame, Option<PromiseResolver<()>>)>,
    expected_replies: HashMap<u16, VecDeque<ExpectedReply>>,
}

impl Default for Inner {
    fn default() -> Self {
        Self {
            header_frames: VecDeque::default(),
            body_frames: VecDeque::default(),
            priority_frames: VecDeque::default(),
            frames: VecDeque::default(),
            low_prio_frames: VecDeque::default(),
            expected_replies: HashMap::default(),
        }
    }
}

impl fmt::Debug for Frames {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut debug = f.debug_struct("Frames");
        if let Some(inner) = self.inner.try_lock() {
            debug.field("expected_replies", &inner.expected_replies);
        }
        debug.finish()
    }
}

impl Inner {
    fn push(
        &mut self,
        channel_id: u16,
        frame: AMQPFrame,
        resolver: PromiseResolver<()>,
        expected_reply: Option<ExpectedReply>,
    ) {
        self.frames.push_back((frame, Some(resolver)));
        if let Some(reply) = expected_reply {
            trace!(
                "channel {} state is now waiting for {:?}",
                channel_id,
                reply
            );
            self.expected_replies
                .entry(channel_id)
                .or_default()
                .push_back(reply);
        }
    }

    fn push_frames(&mut self, mut frames: Vec<AMQPFrame>) -> Promise<()> {
        let (promise, resolver) = Promise::new();
        let last_frame = frames.pop();

        if log_enabled!(Trace) {
            promise.set_marker("Frames".into());
        }

        for frame in frames {
            self.low_prio_frames.push_back((frame, None));
        }
        if let Some(last_frame) = last_frame {
            self.low_prio_frames.push_back((last_frame, Some(resolver)));
        } else {
            resolver.swear(Ok(()));
        }

        promise
    }

    fn pop(&mut self, flow: bool) -> Option<(AMQPFrame, Option<PromiseResolver<()>>)> {
        if let Some(frame) = self
            .header_frames
            .pop_front()
            .or_else(|| self.body_frames.pop_front())
            .or_else(|| self.priority_frames.pop_front())
            .or_else(|| self.frames.pop_front())
        {
            return Some(frame);
        }
        if flow {
            if let Some(frame) = self.low_prio_frames.pop_front() {
                // If the next frame is a header, that means we're a basic.publish
                // Header frame needs to follow directly the basic.publish frame or
                // the AMQP server will close the connection.
                // Push the header into header_frames which is there to handle just that.
                if self
                    .low_prio_frames
                    .front()
                    .map(|(frame, _)| frame.is_header())
                    .unwrap_or(false)
                {
                    // Yes, this will always be Some(), but let's keep our unwrap() count low
                    if let Some(next_frame) = self.low_prio_frames.pop_front() {
                        self.header_frames.push_back(next_frame);
                    }
                    while let Some(next_frame) = self.low_prio_frames.pop_front() {
                        match next_frame.0 {
                            AMQPFrame::Body(..) => {
                                self.body_frames.push_back(next_frame);
                            }
                            _ => {
                                self.low_prio_frames.push_front(next_frame);
                                break;
                            }
                        }
                    }
                }
                return Some(frame);
            }
        }
        None
    }

    fn retry(&mut self, frame: (AMQPFrame, Option<PromiseResolver<()>>)) {
        if frame.0.is_header() {
            self.header_frames.push_front(frame);
        } else {
            self.priority_frames.push_back(frame);
        }
    }

    fn has_pending(&self) -> bool {
        !(self.header_frames.is_empty()
            && self.priority_frames.is_empty()
            && self.frames.is_empty()
            && self.low_prio_frames.is_empty())
    }

    fn drop_pending(&mut self, error: Error) {
        Self::drop_pending_frames(&mut self.header_frames, error.clone());
        Self::drop_pending_frames(&mut self.priority_frames, error.clone());
        Self::drop_pending_frames(&mut self.frames, error.clone());
        Self::drop_pending_frames(&mut self.low_prio_frames, error.clone());
        for (_, replies) in self.expected_replies.drain() {
            Self::cancel_expected_replies(replies, error.clone());
        }
    }

    fn drop_pending_frames(
        frames: &mut VecDeque<(AMQPFrame, Option<PromiseResolver<()>>)>,
        error: Error,
    ) {
        for (_, resolver) in std::mem::take(frames) {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
    }

    fn clear_expected_replies(&mut self, channel_id: u16, error: Error) {
        if let Some(replies) = self.expected_replies.remove(&channel_id) {
            Self::cancel_expected_replies(replies, error);
        }
    }

    fn cancel_expected_replies(replies: VecDeque<ExpectedReply>, error: Error) {
        for ExpectedReply(_, cancel) in replies {
            cancel.cancel(error.clone());
        }
    }
}
<|endoftext|>


<commit_before>/// Black-Scholes formula parameters
/// s0 - underlying price (USD per share)
/// x - strike price (USD per share)
/// sigma - volatility (%)
/// r - continuously compounded risk-free interest rate (%)
/// q - continously compounded dividend yield (%)
/// t - time to expiration (% of year)
use std::f64::consts::E;
use std::f64::consts::PI;
use stats::cnd;

pub fn delta_call(s0: f64, x: f64, t: f64, r: f64, q: f64, sigma: f64) -> f64 {
    let d1 = d1(s0, x, t, r, q, sigma);
    let cnd = cnd(d1);
    let e = E.powf(-(q * t));
    return e * cnd;
}

pub fn delta_put(s0: f64, x: f64, t: f64, r: f64, q: f64, sigma: f64) -> f64 {
    let d1 = d1(s0, x, t, r, q, sigma);
    let cnd = cnd(d1);
    let e = E.powf(-(q * t));
    return e * (cnd - 1.0);
}

pub fn gamma(s0: f64, x: f64, t: f64, r: f64, q: f64, sigma: f64) -> f64 {
    let d1 = d1(s0, x, t, r, q, sigma);
    let arg1 = E.powf(-(q * t)) / (s0 * sigma * (t.sqrt()));
    let arg2 = 1.0 / ((2.0 * PI).sqrt());
    let arg3 = E.powf((-d1).powf(2.0)) / 2.0;
    return arg1 * arg2 * arg3;
}

pub fn d1(s0: f64, x: f64, t: f64, r: f64, q: f64, sigma: f64) -> f64 {
    let ln = (s0 / x).ln();
    let t_num = t * (r - q + (sigma.powf(2f64) / 2f64));
    return (ln + t_num) / (sigma * t.sqrt());
}

pub fn d2(s0: f64, x: f64, t: f64, r: f64, q: f64, sigma: f64) -> f64 {
    let d1 = d1(s0, x, t, r, q, sigma);
    return d1 - (t.sqrt() * sigma);
}

#[cfg(test)]
mod tests {

    use greeks::*;

    const UNDERLYING: f64 = 64.68;
    const STRIKE: f64 = 65.00;
    const VOL: f64 = 0.5051;
    const INTEREST_RATE: f64 = 0.0150;
    const DIV_YIELD: f64 = 0.0210;
    const TIME_TO_EXPIRY: f64 = 23.0 / 365.0;

    const E_D1: f64 = -0.1053;
    const E_D2: f64 = 0.02149;
    const E_CALL_DELTA: f64 = 0.5079;
    const E_PUT_DELTA: f64 = -0.4908;
    const E_GAMMA: f64 = 0.0486;

    #[test]
    fn test_d1() {
        let d1 = d1(UNDERLYING,
                    STRIKE,
                    TIME_TO_EXPIRY,
                    INTEREST_RATE,
                    DIV_YIELD,
                    VOL);
        println!("{}", d1);
        let abs = (d1 - E_D1).abs();
        assert!(abs < 0.001);
    }

    #[test]
    fn test_d2() {
        let d2 = d2(UNDERLYING,
                    STRIKE,
                    TIME_TO_EXPIRY,
                    INTEREST_RATE,
                    DIV_YIELD,
                    VOL);
        let abs = (d2 - E_D2).abs();
        assert!(abs < 0.001);
    }

    #[test]
    fn test_delta_call() {
        let call_delta = delta_call(UNDERLYING,
                                    STRIKE,
                                    TIME_TO_EXPIRY,
                                    INTEREST_RATE,
                                    DIV_YIELD,
                                    VOL);
        let abs = (call_delta - E_CALL_DELTA).abs();
        assert!(abs < 0.001);
    }

    #[test]
    fn test_delta_put() {
        let put_delta = delta_put(UNDERLYING,
                                  STRIKE,
                                  TIME_TO_EXPIRY,
                                  INTEREST_RATE,
                                  DIV_YIELD,
                                  VOL);
        let abs = (put_delta - E_PUT_DELTA).abs();
        assert!(abs < 0.001);
    }

    #[test]
    fn test_gamma() {
        let gamma = gamma(UNDERLYING,
                          STRIKE,
                          TIME_TO_EXPIRY,
                          INTEREST_RATE,
                          DIV_YIELD,
                          VOL);
        let abs = (gamma - E_GAMMA).abs();
        println!("{}", gamma);
        assert!(abs < 0.001);
    }
}
<commit_msg>fixed d1, d2 and gamma tests<commit_after>/// Black-Scholes formula parameters
/// s0 - underlying price (USD per share)
/// x - strike price (USD per share)
/// sigma - volatility (%)
/// r - continuously compounded risk-free interest rate (%)
/// q - continously compounded dividend yield (%)
/// t - time to expiration (% of year)
use std::f64::consts::E;
use std::f64::consts::PI;
use stats::cnd;

pub fn delta_call(s0: f64, x: f64, t: f64, r: f64, q: f64, sigma: f64) -> f64 {
    let d1 = d1(s0, x, t, r, q, sigma);
    let cnd = cnd(d1);
    let e = E.powf(-(q * t));
    return e * cnd;
}

pub fn delta_put(s0: f64, x: f64, t: f64, r: f64, q: f64, sigma: f64) -> f64 {
    let d1 = d1(s0, x, t, r, q, sigma);
    let cnd = cnd(d1);
    let e = E.powf(-(q * t));
    return e * (cnd - 1.0);
}

pub fn gamma(s0: f64, x: f64, t: f64, r: f64, q: f64, sigma: f64) -> f64 {
    let d1 = d1(s0, x, t, r, q, sigma);
    let arg1 = E.powf(-(q * t)) / (s0 * sigma * (t.sqrt()));
    let arg2 = 1.0 / ((2.0 * PI).sqrt());
    let arg3 = E.powf((-d1).powf(2.0)) / 2.0;
    return arg1 * arg2 * arg3;
}

pub fn d1(s0: f64, x: f64, t: f64, r: f64, q: f64, sigma: f64) -> f64 {
    let ln = (s0 / x).ln();
    let t_num = t * (r - q + (sigma.powf(2f64) / 2f64));
    return (ln + t_num) / (sigma * t.sqrt());
}

pub fn d2(s0: f64, x: f64, t: f64, r: f64, q: f64, sigma: f64) -> f64 {
    let d1 = d1(s0, x, t, r, q, sigma);
    return d1 - (t.sqrt() * sigma);
}

#[cfg(test)]
mod tests {

    use greeks::*;

    const UNDERLYING: f64 = 64.68;
    const STRIKE: f64 = 65.00;
    const VOL: f64 = 0.5051;
    const INTEREST_RATE: f64 = 0.0150;
    const DIV_YIELD: f64 = 0.0210;
    const TIME_TO_EXPIRY: f64 = 23.0 / 365.0;

    const E_D1: f64 = 0.0214;
    const E_D2: f64 = -0.1053;
    const E_CALL_DELTA: f64 = 0.5079;
    const E_PUT_DELTA: f64 = -0.4908;
    const E_GAMMA: f64 = 0.0243;

    #[test]
    fn test_d1() {
        let d1 = d1(UNDERLYING,
                    STRIKE,
                    TIME_TO_EXPIRY,
                    INTEREST_RATE,
                    DIV_YIELD,
                    VOL);
        println!("{}", d1);
        let abs = (d1 - E_D1).abs();
        assert!(abs < 0.001);
    }

    #[test]
    fn test_d2() {
        let d2 = d2(UNDERLYING,
                    STRIKE,
                    TIME_TO_EXPIRY,
                    INTEREST_RATE,
                    DIV_YIELD,
                    VOL);
        let abs = (d2 - E_D2).abs();
        assert!(abs < 0.001);
    }

    #[test]
    fn test_delta_call() {
        let call_delta = delta_call(UNDERLYING,
                                    STRIKE,
                                    TIME_TO_EXPIRY,
                                    INTEREST_RATE,
                                    DIV_YIELD,
                                    VOL);
        let abs = (call_delta - E_CALL_DELTA).abs();
        assert!(abs < 0.001);
    }

    #[test]
    fn test_delta_put() {
        let put_delta = delta_put(UNDERLYING,
                                  STRIKE,
                                  TIME_TO_EXPIRY,
                                  INTEREST_RATE,
                                  DIV_YIELD,
                                  VOL);
        let abs = (put_delta - E_PUT_DELTA).abs();
        assert!(abs < 0.001);
    }

    #[test]
    fn test_gamma() {
        let gamma = gamma(UNDERLYING,
                          STRIKE,
                          TIME_TO_EXPIRY,
                          INTEREST_RATE,
                          DIV_YIELD,
                          VOL);
        let abs = (gamma - E_GAMMA).abs();
        assert!(abs < 0.001);
    }
}
<|endoftext|>


<commit_before>//! Las headers.

use std::io::Read;

use byteorder::ReadBytesExt;
use byteorder::LittleEndian;
use rustc_serialize::hex::ToHex;

use Error;
use Result;
use io::LasStringExt;
use util::Triplet;

/// Project ID newtype.
#[derive(Clone, Debug, PartialEq)]
pub struct ProjectId(pub [u8; 4], pub [u8; 2], pub [u8; 2], pub [u8; 8]);

impl Default for ProjectId {
    fn default() -> ProjectId {
        ProjectId([0; 4], [0; 2], [0; 2], [0; 8])
    }
}

impl ProjectId {
    /// Returns a hexadecimal string representation of the project ID.
    ///
    /// The format for this string was derived from libLAS's output.
    ///
    /// # Examples
    ///
    /// ```
    /// # use las::header::ProjectId;
    /// let project_id = ProjectId([0; 4], [0; 2], [0; 2], [0; 8]);
    /// assert_eq!("00000000-0000-0000-0000-000000000000", project_id.as_string());
    /// ```
    pub fn as_string(&self) -> String {
        let mut s = self.0.to_hex();
        s.push_str("-");
        s.push_str(&self.1.to_hex());
        s.push_str("-");
        s.push_str(&self.2.to_hex());
        s.push_str("-");
        s.push_str(&self.3[0..2].to_hex());
        s.push_str("-");
        s.push_str(&self.3[2..8].to_hex());
        s
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct Header {
    pub file_signature: [u8; 4],
    pub file_source_id: u16,
    pub global_encoding: u16,
    pub project_id: ProjectId,
    pub version_major: u8,
    pub version_minor: u8,
    pub system_identifier: String,
    pub generating_software: String,
    pub file_creation_day_of_year: u16,
    pub file_creation_year: u16,
    pub header_size: u16,
    pub offset_to_point_data: u32,
    pub number_of_variable_length_records: u32,
    pub point_data_format_id: u8,
    pub point_data_record_length: u16,
    pub number_of_point_records: u32,
    pub number_of_points_by_return: [u32; 5],
    pub scale: Triplet<f64>,
    pub offset: Triplet<f64>,
    pub min: Triplet<f64>,
    pub max: Triplet<f64>,
}

impl Default for Header {
    fn default() -> Header {
        Header {
            file_signature: *b"LASF",
            file_source_id: 0,
            global_encoding: 0,
            project_id: Default::default(),
            version_major: 0,
            version_minor: 0,
            system_identifier: "".to_string(),
            generating_software: "".to_string(),
            file_creation_day_of_year: 0,
            file_creation_year: 0,
            header_size: 0,
            offset_to_point_data: 0,
            number_of_variable_length_records: 0,
            point_data_format_id: 0,
            point_data_record_length: 0,
            number_of_point_records: 0,
            number_of_points_by_return: [0; 5],
            scale: Default::default(),
            offset: Default::default(),
            min: Default::default(),
            max: Default::default(),
        }
    }
}

impl Header {
    /// Creates a new header from a `Read` object.
    ///
    /// # Examples
    ///
    /// ```
    /// # use las::header::Header;
    /// let mut reader = std::fs::File::open("data/1.2_0.las").unwrap();
    /// let header = Header::new(&mut reader);
    /// ```
    pub fn new<R: Read>(reader: &mut R) -> Result<Header> {
        let mut header: Header = Default::default();
        try_read_n!(reader, header.file_signature, 4);
        header.file_source_id = try!(reader.read_u16::<LittleEndian>());
        header.global_encoding = try!(reader.read_u16::<LittleEndian>());
        try_read_n!(reader, header.project_id.0, 4);
        try_read_n!(reader, header.project_id.1, 2);
        try_read_n!(reader, header.project_id.2, 2);
        try_read_n!(reader, header.project_id.3, 8);
        header.version_major = try!(reader.read_u8());
        header.version_minor = try!(reader.read_u8());
        header.system_identifier = try!(reader.read_las_string(32));
        header.generating_software = try!(reader.read_las_string(32));
        header.file_creation_day_of_year = try!(reader.read_u16::<LittleEndian>());
        header.file_creation_year = try!(reader.read_u16::<LittleEndian>());
        header.header_size = try!(reader.read_u16::<LittleEndian>());
        header.offset_to_point_data = try!(reader.read_u32::<LittleEndian>());
        header.number_of_variable_length_records = try!(reader.read_u32::<LittleEndian>());
        header.point_data_format_id = try!(reader.read_u8());
        header.point_data_record_length = try!(reader.read_u16::<LittleEndian>());
        header.number_of_point_records = try!(reader.read_u32::<LittleEndian>());
        header.number_of_points_by_return[0] = try!(reader.read_u32::<LittleEndian>());
        header.number_of_points_by_return[1] = try!(reader.read_u32::<LittleEndian>());
        header.number_of_points_by_return[2] = try!(reader.read_u32::<LittleEndian>());
        header.number_of_points_by_return[3] = try!(reader.read_u32::<LittleEndian>());
        header.number_of_points_by_return[4] = try!(reader.read_u32::<LittleEndian>());
        header.scale.x = try!(reader.read_f64::<LittleEndian>());
        header.scale.y = try!(reader.read_f64::<LittleEndian>());
        header.scale.z = try!(reader.read_f64::<LittleEndian>());
        header.offset.x = try!(reader.read_f64::<LittleEndian>());
        header.offset.y = try!(reader.read_f64::<LittleEndian>());
        header.offset.z = try!(reader.read_f64::<LittleEndian>());
        header.max.x = try!(reader.read_f64::<LittleEndian>());
        header.min.x = try!(reader.read_f64::<LittleEndian>());
        header.max.y = try!(reader.read_f64::<LittleEndian>());
        header.min.y = try!(reader.read_f64::<LittleEndian>());
        header.max.z = try!(reader.read_f64::<LittleEndian>());
        header.min.z = try!(reader.read_f64::<LittleEndian>());
        Ok(header)
    }
}

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

    use std::fs::File;

    #[test]
    fn header() {
        let mut reader = File::open("data/1.2_0.las").unwrap();
        let header = Header::new(&mut reader).unwrap();
        assert_eq!(*b"LASF", header.file_signature);
        assert_eq!(0, header.file_source_id);
        assert_eq!(0, header.global_encoding);
        assert_eq!("b8f18883-1baa-0841-bca3-6bc68e7b062e",
                   header.project_id.as_string());
        assert_eq!(1, header.version_major);
        assert_eq!(2, header.version_minor);
        assert_eq!("libLAS", header.system_identifier);
        assert_eq!("libLAS 1.2", header.generating_software);
        assert_eq!(78, header.file_creation_day_of_year);
        assert_eq!(2008, header.file_creation_year);
        assert_eq!(227, header.header_size);
        assert_eq!(438, header.offset_to_point_data);
        assert_eq!(2, header.number_of_variable_length_records);
        assert_eq!(0, header.point_data_format_id);
        assert_eq!(20, header.point_data_record_length);
        assert_eq!(1, header.number_of_point_records);
        assert_eq!([0, 1, 0, 0, 0], header.number_of_points_by_return);
        assert_eq!(0.01, header.scale.x);
        assert_eq!(0.01, header.scale.y);
        assert_eq!(0.01, header.scale.z);
        assert_eq!(0.0, header.offset.x);
        assert_eq!(0.0, header.offset.y);
        assert_eq!(0.0, header.offset.z);
        assert_eq!(470692.447538, header.min.x);
        assert_eq!(4602888.904642, header.min.y);
        assert_eq!(16.0, header.min.z);
        assert_eq!(470692.447538, header.max.x);
        assert_eq!(4602888.904642, header.max.y);
        assert_eq!(16.0, header.max.z);
    }
}
<commit_msg>Add copy to ProjectId<commit_after>//! Las headers.

use std::io::Read;

use byteorder::ReadBytesExt;
use byteorder::LittleEndian;
use rustc_serialize::hex::ToHex;

use Error;
use Result;
use io::LasStringExt;
use util::Triplet;

/// Project ID newtype.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct ProjectId(pub [u8; 4], pub [u8; 2], pub [u8; 2], pub [u8; 8]);

impl Default for ProjectId {
    fn default() -> ProjectId {
        ProjectId([0; 4], [0; 2], [0; 2], [0; 8])
    }
}

impl ProjectId {
    /// Returns a hexadecimal string representation of the project ID.
    ///
    /// The format for this string was derived from libLAS's output.
    ///
    /// # Examples
    ///
    /// ```
    /// # use las::header::ProjectId;
    /// let project_id = ProjectId([0; 4], [0; 2], [0; 2], [0; 8]);
    /// assert_eq!("00000000-0000-0000-0000-000000000000", project_id.as_string());
    /// ```
    pub fn as_string(&self) -> String {
        let mut s = self.0.to_hex();
        s.push_str("-");
        s.push_str(&self.1.to_hex());
        s.push_str("-");
        s.push_str(&self.2.to_hex());
        s.push_str("-");
        s.push_str(&self.3[0..2].to_hex());
        s.push_str("-");
        s.push_str(&self.3[2..8].to_hex());
        s
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct Header {
    pub file_signature: [u8; 4],
    pub file_source_id: u16,
    pub global_encoding: u16,
    pub project_id: ProjectId,
    pub version_major: u8,
    pub version_minor: u8,
    pub system_identifier: String,
    pub generating_software: String,
    pub file_creation_day_of_year: u16,
    pub file_creation_year: u16,
    pub header_size: u16,
    pub offset_to_point_data: u32,
    pub number_of_variable_length_records: u32,
    pub point_data_format_id: u8,
    pub point_data_record_length: u16,
    pub number_of_point_records: u32,
    pub number_of_points_by_return: [u32; 5],
    pub scale: Triplet<f64>,
    pub offset: Triplet<f64>,
    pub min: Triplet<f64>,
    pub max: Triplet<f64>,
}

impl Default for Header {
    fn default() -> Header {
        Header {
            file_signature: *b"LASF",
            file_source_id: 0,
            global_encoding: 0,
            project_id: Default::default(),
            version_major: 0,
            version_minor: 0,
            system_identifier: "".to_string(),
            generating_software: "".to_string(),
            file_creation_day_of_year: 0,
            file_creation_year: 0,
            header_size: 0,
            offset_to_point_data: 0,
            number_of_variable_length_records: 0,
            point_data_format_id: 0,
            point_data_record_length: 0,
            number_of_point_records: 0,
            number_of_points_by_return: [0; 5],
            scale: Default::default(),
            offset: Default::default(),
            min: Default::default(),
            max: Default::default(),
        }
    }
}

impl Header {
    /// Creates a new header from a `Read` object.
    ///
    /// # Examples
    ///
    /// ```
    /// # use las::header::Header;
    /// let mut reader = std::fs::File::open("data/1.2_0.las").unwrap();
    /// let header = Header::new(&mut reader);
    /// ```
    pub fn new<R: Read>(reader: &mut R) -> Result<Header> {
        let mut header: Header = Default::default();
        try_read_n!(reader, header.file_signature, 4);
        header.file_source_id = try!(reader.read_u16::<LittleEndian>());
        header.global_encoding = try!(reader.read_u16::<LittleEndian>());
        try_read_n!(reader, header.project_id.0, 4);
        try_read_n!(reader, header.project_id.1, 2);
        try_read_n!(reader, header.project_id.2, 2);
        try_read_n!(reader, header.project_id.3, 8);
        header.version_major = try!(reader.read_u8());
        header.version_minor = try!(reader.read_u8());
        header.system_identifier = try!(reader.read_las_string(32));
        header.generating_software = try!(reader.read_las_string(32));
        header.file_creation_day_of_year = try!(reader.read_u16::<LittleEndian>());
        header.file_creation_year = try!(reader.read_u16::<LittleEndian>());
        header.header_size = try!(reader.read_u16::<LittleEndian>());
        header.offset_to_point_data = try!(reader.read_u32::<LittleEndian>());
        header.number_of_variable_length_records = try!(reader.read_u32::<LittleEndian>());
        header.point_data_format_id = try!(reader.read_u8());
        header.point_data_record_length = try!(reader.read_u16::<LittleEndian>());
        header.number_of_point_records = try!(reader.read_u32::<LittleEndian>());
        header.number_of_points_by_return[0] = try!(reader.read_u32::<LittleEndian>());
        header.number_of_points_by_return[1] = try!(reader.read_u32::<LittleEndian>());
        header.number_of_points_by_return[2] = try!(reader.read_u32::<LittleEndian>());
        header.number_of_points_by_return[3] = try!(reader.read_u32::<LittleEndian>());
        header.number_of_points_by_return[4] = try!(reader.read_u32::<LittleEndian>());
        header.scale.x = try!(reader.read_f64::<LittleEndian>());
        header.scale.y = try!(reader.read_f64::<LittleEndian>());
        header.scale.z = try!(reader.read_f64::<LittleEndian>());
        header.offset.x = try!(reader.read_f64::<LittleEndian>());
        header.offset.y = try!(reader.read_f64::<LittleEndian>());
        header.offset.z = try!(reader.read_f64::<LittleEndian>());
        header.max.x = try!(reader.read_f64::<LittleEndian>());
        header.min.x = try!(reader.read_f64::<LittleEndian>());
        header.max.y = try!(reader.read_f64::<LittleEndian>());
        header.min.y = try!(reader.read_f64::<LittleEndian>());
        header.max.z = try!(reader.read_f64::<LittleEndian>());
        header.min.z = try!(reader.read_f64::<LittleEndian>());
        Ok(header)
    }
}

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

    use std::fs::File;

    #[test]
    fn header() {
        let mut reader = File::open("data/1.2_0.las").unwrap();
        let header = Header::new(&mut reader).unwrap();
        assert_eq!(*b"LASF", header.file_signature);
        assert_eq!(0, header.file_source_id);
        assert_eq!(0, header.global_encoding);
        assert_eq!("b8f18883-1baa-0841-bca3-6bc68e7b062e",
                   header.project_id.as_string());
        assert_eq!(1, header.version_major);
        assert_eq!(2, header.version_minor);
        assert_eq!("libLAS", header.system_identifier);
        assert_eq!("libLAS 1.2", header.generating_software);
        assert_eq!(78, header.file_creation_day_of_year);
        assert_eq!(2008, header.file_creation_year);
        assert_eq!(227, header.header_size);
        assert_eq!(438, header.offset_to_point_data);
        assert_eq!(2, header.number_of_variable_length_records);
        assert_eq!(0, header.point_data_format_id);
        assert_eq!(20, header.point_data_record_length);
        assert_eq!(1, header.number_of_point_records);
        assert_eq!([0, 1, 0, 0, 0], header.number_of_points_by_return);
        assert_eq!(0.01, header.scale.x);
        assert_eq!(0.01, header.scale.y);
        assert_eq!(0.01, header.scale.z);
        assert_eq!(0.0, header.offset.x);
        assert_eq!(0.0, header.offset.y);
        assert_eq!(0.0, header.offset.z);
        assert_eq!(470692.447538, header.min.x);
        assert_eq!(4602888.904642, header.min.y);
        assert_eq!(16.0, header.min.z);
        assert_eq!(470692.447538, header.max.x);
        assert_eq!(4602888.904642, header.max.y);
        assert_eq!(16.0, header.max.z);
    }
}
<|endoftext|>


<commit_before>//! The `ztext` module contains encoding functions to encode text in z-ascii characters.
//!

use super::zbytes::Bytes;

pub static ALPHABET: [char; 78] = [
    'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm',
    'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',

    'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M',
    'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z',

    ' ', '\n', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '.',
    ',', '!', '?', '_', '#', '\'','"', '/', '\\','-', ':', '(', ')'];

/// encodes an string to z-characters
/// and returns the length of the used bytes
/// TODO: Only works with lower case letters
///
/// # Examples
///
/// ```ignore
/// let data = Bytes{bytes: Vec::new()};
/// let byteLength = data.encode("hello");
/// ```
pub fn encode(data: &mut Bytes, content: &str, unicode_table: &Vec<u16>) -> u16 {
    let zchars: Vec<u8> = string_to_zchar(content, unicode_table);

    let mut two_bytes: u16 = 0;
    let len = zchars.len();
    for i in 0..len {
        let zasci_id =zchars[i];

        two_bytes |= shift(zasci_id as u16, i as u8);

        if i % 3 == 2 {
            data.write_u16(two_bytes, pos_to_index(i));
            two_bytes = 0;
        }

        // end of string
        if i == len -1 {
            if i % 3 != 2 {
                for j in (i % 3) + 1..3 {
                    two_bytes |= shift(0x05 as u16, j as u8);
                }

                data.write_u16(two_bytes, pos_to_index(i));
            }

            // end bit is written to the first bit of the next to last byte
            let end_index: usize = data.bytes.len() - 2 as usize;
            let mut end_byte: u8 = data.bytes[end_index];

            end_byte |= 0x80;
            data.write_byte(end_byte, end_index as usize);
        }
    }

    data.bytes.len() as u16
}

/// reads the content and converts it to a zasci vector
fn string_to_zchar(content: &str, unicode_table: &Vec<u16>) -> Vec<u8> {
    //let string_bytes = content.to_string().into_bytes();
    let mut zchars: Vec<u8> = Vec::new();

    for character in content.chars() {

        let mut byte: u8 = character as u8;
        let alpha_index = pos_in_alpha(byte as u8);
        if character as u16 <= 126 && alpha_index != -1 {

            if byte == 0x0A {
                // newline
                zchars.push(0x05);
                zchars.push(7);
            } else if byte == 0x20 {
                // space
                zchars.push(0x00);
            } else {
                if alpha_index > 51 {
                    // in A2
                    zchars.push(0x05);
                    zchars.push(alpha_index as u8 % 26 + 6);
                } else if alpha_index < 26 {
                    // in A0
                    zchars.push(alpha_index as u8 % 26 + 6);
                } else {
                    // in A1
                    zchars.push(0x04);
                    zchars.push(alpha_index as u8 % 26 + 6);
                }
            }
        } else {
            // not in alphabet or unicode

            // to change alphabet
            zchars.push(0x05);

            // for special char (10 bit z-ascii)
            zchars.push(0x06);

            //let mut byte: u8;
            if character as u16 <= 126 {
                // not in alphabet, but still ascii
                byte = character as u8;
            } else {
                // unicode
                let unicode_index = pos_in_unicode(character as u16, unicode_table);
                byte = unicode_index as u8 + 155;
            }

            zchars.push(byte >> 5);
            zchars.push(byte & 0x1f);
        }
    }
    zchars
}

/// shifts the z-char in a 2 bytes-array to the right position
/// shift_length has 3 possibilities: 10, 5, 0
/// an z-char use 5 bytes
/// look into two bytes:
/// ____byte one____ ____byte two____
/// 1 2 3 4 5 6  7 8 1 2 3  4 5 6 7 8
///   ^          ^          1
///   1 2 3 4 5  1 2 3 4 5  1 2 3 4 5
///   1. zchar   2. zchar   3. zchar
///   10         5          0
fn shift(zchar: u16, position: u8) -> u16 {
    let shift_length = 10 - (position % 3) * 5;
    zchar << shift_length
}

/// Returns the location of the character of the specified index in the zcode character array
///
/// # Examples
///
fn pos_in_alpha(letter: u8) -> i8 {
    for i in 0..ALPHABET.len() {
        if ALPHABET[i] as u8 == letter {
            return i as i8
        }
    }

    return -1
}

/// returns the position in the unicode-table
pub fn pos_in_unicode(letter: u16, unicode_table: &Vec<u16>) -> i8 {
    for (i, character) in unicode_table.iter().enumerate() {
        if *character == letter {
            return i as i8
        }
    }

    return -1
}

/// position in the vector from the position of an character in the string
/// every 3 zchars encode() writes 2 bytes
/// for example "helloworld"
/// "h": nothing
/// "e": nothing
/// "l": write 2 bytes to position 0   (2 * 2/3))
/// "l": nothing
/// "o": nothing
/// "w": write 2 bytes to position 2   (2 * 5/3))
/// ...
///
/// # Examples
///
fn pos_to_index(position: usize) -> usize {
    2 * (position / 3)
}

#[test]
fn test_pos_in_alpha() {
    assert_eq!(pos_in_alpha('a' as u8), 0);
    assert_eq!(pos_in_alpha('b' as u8), 1);
    assert_eq!(pos_in_alpha('c' as u8), 2);
    assert_eq!(pos_in_alpha('A' as u8), 26);
    assert_eq!(pos_in_alpha('B' as u8), 27);
    assert_eq!(pos_in_alpha('C' as u8), 28);
}

#[test]
fn test_pos_to_index() {
    assert_eq!(pos_to_index(5), 2);
}

#[test]
fn test_shift() {
    assert_eq!(shift(6,2), 6);
    assert_eq!(shift(26,5), 26);
}

#[test]
fn test_string_to_zchar() {
    let mut vec: Vec<u16> = Vec::new();
    assert_eq!(string_to_zchar("i am a string, please test me, no unicode",&vec), vec![14, 0, 6, 18, 0, 6, 0, 24, 25, 23, 14, 19, 12, 5, 19, 0, 21, 17, 10, 6, 24, 10, 0, 25, 10, 24, 25, 0, 18, 10, 5, 19, 0, 19, 20, 0, 26, 19, 14, 8, 20, 9, 10]);
    vec.push('€' as u16);
    assert_eq!(string_to_zchar("nasty char: €",&vec), vec![19, 6, 24, 25, 30, 0, 8, 13, 6, 23, 5, 29, 0, 5, 6, 4, 27]);
}

// #[test]
// fn test_encode() {
//     let mut vec: Vec<u16> = Vec::new();
//     let mut data = Bytes{bytes: Vec::new()};
//     let byteLength = encode(&data,"hello",&vec);
//     assert_eq!(byteLength,5);
//     assert_eq!(data,5);
// }<commit_msg>cleanup<commit_after>//! The `ztext` module contains encoding functions to encode text in z-ascii characters.
//!

use super::zbytes::Bytes;

pub static ALPHABET: [char; 78] = [
    'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm',
    'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',

    'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M',
    'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z',

    ' ', '\n', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '.',
    ',', '!', '?', '_', '#', '\'','"', '/', '\\','-', ':', '(', ')'];

/// encodes an string to z-characters
/// and returns the length of the used bytes
/// TODO: Only works with lower case letters
///
/// # Examples
///
/// ```ignore
/// let data = Bytes{bytes: Vec::new()};
/// let byteLength = data.encode("hello");
/// ```
pub fn encode(data: &mut Bytes, content: &str, unicode_table: &Vec<u16>) -> u16 {
    let zchars: Vec<u8> = string_to_zchar(content, unicode_table);

    let mut two_bytes: u16 = 0;
    let len = zchars.len();
    for i in 0..len {
        let zasci_id =zchars[i];

        two_bytes |= shift(zasci_id as u16, i as u8);

        if i % 3 == 2 {
            data.write_u16(two_bytes, pos_to_index(i));
            two_bytes = 0;
        }

        // end of string
        if i == len -1 {
            if i % 3 != 2 {
                for j in (i % 3) + 1..3 {
                    two_bytes |= shift(0x05 as u16, j as u8);
                }

                data.write_u16(two_bytes, pos_to_index(i));
            }

            // end bit is written to the first bit of the next to last byte
            let end_index: usize = data.bytes.len() - 2 as usize;
            let mut end_byte: u8 = data.bytes[end_index];

            end_byte |= 0x80;
            data.write_byte(end_byte, end_index as usize);
        }
    }

    data.bytes.len() as u16
}

/// reads the content and converts it to a zasci vector
fn string_to_zchar(content: &str, unicode_table: &Vec<u16>) -> Vec<u8> {
    //let string_bytes = content.to_string().into_bytes();
    let mut zchars: Vec<u8> = Vec::new();

    for character in content.chars() {

        let mut byte: u8 = character as u8;
        let alpha_index = pos_in_alpha(byte as u8);
        if character as u16 <= 126 && alpha_index != -1 {

            if byte == 0x0A {
                // newline
                zchars.push(0x05);
                zchars.push(7);
            } else if byte == 0x20 {
                // space
                zchars.push(0x00);
            } else {
                if alpha_index > 51 {
                    // in A2
                    zchars.push(0x05);
                    zchars.push(alpha_index as u8 % 26 + 6);
                } else if alpha_index < 26 {
                    // in A0
                    zchars.push(alpha_index as u8 % 26 + 6);
                } else {
                    // in A1
                    zchars.push(0x04);
                    zchars.push(alpha_index as u8 % 26 + 6);
                }
            }
        } else {
            // not in alphabet or unicode

            // to change alphabet
            zchars.push(0x05);

            // for special char (10 bit z-ascii)
            zchars.push(0x06);

            //let mut byte: u8;
            if character as u16 <= 126 {
                // not in alphabet, but still ascii
                byte = character as u8;
            } else {
                // unicode
                let unicode_index = pos_in_unicode(character as u16, unicode_table);
                byte = unicode_index as u8 + 155;
            }

            zchars.push(byte >> 5);
            zchars.push(byte & 0x1f);
        }
    }
    zchars
}

/// shifts the z-char in a 2 bytes-array to the right position
/// shift_length has 3 possibilities: 10, 5, 0
/// an z-char use 5 bytes
/// look into two bytes:
/// ____byte one____ ____byte two____
/// 1 2 3 4 5 6  7 8 1 2 3  4 5 6 7 8
///   ^          ^          1
///   1 2 3 4 5  1 2 3 4 5  1 2 3 4 5
///   1. zchar   2. zchar   3. zchar
///   10         5          0
fn shift(zchar: u16, position: u8) -> u16 {
    let shift_length = 10 - (position % 3) * 5;
    zchar << shift_length
}

/// Returns the location of the character of the specified index in the zcode character array
///
/// # Examples
///
fn pos_in_alpha(letter: u8) -> i8 {
    for i in 0..ALPHABET.len() {
        if ALPHABET[i] as u8 == letter {
            return i as i8
        }
    }

    return -1
}

/// returns the position in the unicode-table
pub fn pos_in_unicode(letter: u16, unicode_table: &Vec<u16>) -> i8 {
    for (i, character) in unicode_table.iter().enumerate() {
        if *character == letter {
            return i as i8
        }
    }

    return -1
}

/// position in the vector from the position of an character in the string
/// every 3 zchars encode() writes 2 bytes
/// for example "helloworld"
/// "h": nothing
/// "e": nothing
/// "l": write 2 bytes to position 0   (2 * 2/3))
/// "l": nothing
/// "o": nothing
/// "w": write 2 bytes to position 2   (2 * 5/3))
/// ...
///
/// # Examples
///
fn pos_to_index(position: usize) -> usize {
    2 * (position / 3)
}

#[test]
fn test_pos_in_alpha() {
    assert_eq!(pos_in_alpha('a' as u8), 0);
    assert_eq!(pos_in_alpha('b' as u8), 1);
    assert_eq!(pos_in_alpha('c' as u8), 2);
    assert_eq!(pos_in_alpha('A' as u8), 26);
    assert_eq!(pos_in_alpha('B' as u8), 27);
    assert_eq!(pos_in_alpha('C' as u8), 28);
}

#[test]
fn test_pos_to_index() {
    assert_eq!(pos_to_index(5), 2);
}

#[test]
fn test_shift() {
    assert_eq!(shift(6,2), 6);
    assert_eq!(shift(26,5), 26);
}

#[test]
fn test_string_to_zchar() {
    let mut vec: Vec<u16> = Vec::new();
    assert_eq!(string_to_zchar("i am a string, please test me, no unicode",&vec), vec![14, 0, 6, 18, 0, 6, 0, 24, 25, 23, 14, 19, 12, 5, 19, 0, 21, 17, 10, 6, 24, 10, 0, 25, 10, 24, 25, 0, 18, 10, 5, 19, 0, 19, 20, 0, 26, 19, 14, 8, 20, 9, 10]);
    vec.push('€' as u16);
    assert_eq!(string_to_zchar("nasty char: €",&vec), vec![19, 6, 24, 25, 30, 0, 8, 13, 6, 23, 5, 29, 0, 5, 6, 4, 27]);
}<|endoftext|>