<commit_before>use std::thread;
use std::cell::UnsafeCell;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver, RecvError};
use basetypes::*;
use chess_move::*;
use tt::*;
use position::Position;

pub enum Command {
    Search {
        search_id: usize,
        position: Position,
        depth: u8,
        lower_bound: Value,
        upper_bound: Value,
    },
    Stop,
    Exit,
}

pub enum Report {
    Progress {
        search_id: usize,
        searched_nodes: NodeCount,
        depth: u8,
    },
    Done {
        search_id: usize,
        searched_nodes: NodeCount,
        value: Option<Value>,
    },
}

pub fn run_deepening(tt: Arc<TranspositionTable>,
                     commands: Receiver<Command>,
                     reports: Sender<Report>) {
    let (slave_commands_tx, slave_commands_rx) = channel();
    let (slave_reports_tx, slave_reports_rx) = channel();
    let slave = thread::spawn(move || {
        run(tt, slave_commands_rx, slave_reports_tx);
    });
    let mut pending_command = None;
    loop {
        // If there is a pending command, we take it, otherwise we
        // block and wait to receive a new one.
        let command = match pending_command.take() {
            Some(cmd) => cmd,
            None => commands.recv().or::<RecvError>(Ok(Command::Exit)).unwrap(),
        };

        match command {
            Command::Search { search_id, position, depth, lower_bound, upper_bound } => {
                let mut searched_nodes_final = 0;
                let mut value_final = None;
                'depthloop: for n in 1..(depth + 1) {
                    slave_commands_tx.send(Command::Search {
                                         search_id: n as usize,
                                         position: position.clone(),
                                         depth: n,
                                         lower_bound: lower_bound,
                                         upper_bound: upper_bound,
                                     })
                                     .unwrap();
                    loop {
                        match slave_reports_rx.recv().unwrap() {
                            Report::Progress { searched_nodes, .. } => {
                                reports.send(Report::Progress {
                                           search_id: search_id,
                                           searched_nodes: searched_nodes_final + searched_nodes,
                                           depth: n,
                                       })
                                       .ok();
                                if pending_command.is_none() {
                                    pending_command = match commands.try_recv() {
                                        Ok(cmd) => {
                                            slave_commands_tx.send(Command::Stop).unwrap();
                                            Some(cmd)
                                        }
                                        _ => None,
                                    }
                                }
                            }
                            Report::Done { searched_nodes, value, .. } => {
                                searched_nodes_final += searched_nodes;
                                if n == depth {
                                    value_final = value;
                                }
                                if pending_command.is_some() {
                                    break 'depthloop;
                                }
                                break;
                            }
                        }
                    }
                }
                reports.send(Report::Done {
                           search_id: search_id,
                           searched_nodes: searched_nodes_final,
                           value: value_final,
                       })
                       .ok();
            }
            Command::Stop => {
                slave_commands_tx.send(Command::Stop).unwrap();
                continue;
            }
            Command::Exit => {
                slave_commands_tx.send(Command::Exit).unwrap();
                break;
            }
        }
    }
    slave.join().unwrap();
}

pub fn run(tt: Arc<TranspositionTable>, commands: Receiver<Command>, reports: Sender<Report>) {
    thread_local!(
        static MOVE_STACK: UnsafeCell<MoveStack> = UnsafeCell::new(MoveStack::new())
    );
    MOVE_STACK.with(|s| {
        let mut move_stack = unsafe { &mut *s.get() };
        let mut pending_command = None;
        loop {
            // If there is a pending command, we take it, otherwise we
            // block and wait to receive a new one.
            let command = match pending_command.take() {
                Some(cmd) => cmd,
                None => commands.recv().or::<RecvError>(Ok(Command::Exit)).unwrap(),
            };

            match command {
                Command::Search { search_id, mut position, depth, lower_bound, upper_bound } => {
                    let mut reported_nodes = 0;
                    let mut unreported_nodes = 0;
                    let value = search(&tt,
                                       &mut position,
                                       move_stack,
                                       &mut unreported_nodes,
                                       &mut |n| {
                                           reported_nodes += n;
                                           reports.send(Report::Progress {
                                                      search_id: search_id,
                                                      searched_nodes: reported_nodes,
                                                      depth: depth,
                                                  })
                                                  .ok();
                                           match commands.try_recv() {
                                               Ok(x) => {
                                                   pending_command = Some(x);
                                                   Err(TerminatedSearch)
                                               }
                                               _ => Ok(()),
                                           }
                                       },
                                       lower_bound,
                                       upper_bound,
                                       depth)
                                    .ok();
                    reported_nodes += unreported_nodes;
                    reports.send(Report::Progress {
                               search_id: search_id,
                               searched_nodes: reported_nodes,
                               depth: depth,
                           })
                           .ok();
                    reports.send(Report::Done {
                               search_id: search_id,
                               searched_nodes: reported_nodes,
                               value: value,
                           })
                           .ok();
                    move_stack.clear();
                }
                Command::Stop => continue,
                Command::Exit => break,
            }
        }
    })
}

// Represents a terminated search condition.
struct TerminatedSearch;

// Helper function for `run()`. It implements the principal variation
// search algorithm. When returning `Err(TerminatedSearch)`, this
// function may leave un-restored move lists in `moves`.
fn search(tt: &TranspositionTable,
          p: &mut Position,
          moves: &mut MoveStack,
          nc: &mut NodeCount,
          report: &mut FnMut(NodeCount) -> Result<(), TerminatedSearch>,
          mut alpha: Value, // lower bound
          beta: Value, // upper bound
          depth: u8)
          -> Result<Value, TerminatedSearch> {
    assert!(alpha < beta);

    // Consult the transposition table.
    let (hash_move16, eval_value) = if let Some(entry) = tt.probe(p.hash()) {
        if entry.depth() >= depth {
            let value = entry.value();
            let bound = entry.bound();
            if (value >= beta && bound == BOUND_LOWER) ||
               (value <= alpha && bound == BOUND_UPPER) || (bound == BOUND_EXACT) {
                return Ok(value);
            }
        }
        (entry.move16(), entry.eval_value())
    } else {
        (0, p.evaluate_static())
    };

    let mut bound_type = BOUND_UPPER;
    let mut best_move = Move::invalid();

    if depth == 0 {
        // On leaf nodes, do quiescence search.
        let (value, nodes) = p.evaluate_quiescence(alpha, beta, Some(eval_value));
        *nc += nodes;
        if value >= beta {
            alpha = beta;
            bound_type = BOUND_LOWER;
        } else if value > alpha {
            alpha = value;
            bound_type = BOUND_EXACT;
        }
    } else {
        moves.save();

        // TODO: use `eval_value`.

        if hash_move16 != 0 {
            if let Some(m) = p.try_move_digest(hash_move16) {
                moves.push(m);
            }
        }
        let mut generated = false;

        let mut no_moves_yet = true;
        loop {
            // TODO: This is ugly, and probably inefficient.
            let m = match moves.remove_best_move() {
                None if generated => break,
                None => {
                    p.generate_moves(moves);
                    moves.remove_move(hash_move16);
                    generated = true;
                    continue;
                }
                Some(x) => x,
            };

            if p.do_move(m) {
                // From time to time, we report how many nodes had
                // been searched since the last report. This also
                // gives an opportunity for the search to be
                // terminated.
                *nc += 1;
                if *nc > NODE_COUNT_REPORT_INTERVAL {
                    try!(report(*nc));
                    *nc = 0;
                }

                let value = if no_moves_yet {
                    // The first move we analyze with a fully open
                    // window (alpha, beta).
                    -try!(search(tt, p, moves, nc, report, -beta, -alpha, depth - 1))
                } else {
                    // For the next moves we first try to prove that
                    // they are not better than our current best
                    // move. For this purpose we analyze them with a
                    // null window (alpha, alpha + 1). This is faster
                    // than a full window search. Only when we are
                    // certain that the move is better than our
                    // current best move, we do a full-window search.
                    match -try!(search(tt, p, moves, nc, report, -alpha - 1, -alpha, depth - 1)) {
                        x if x <= alpha => x,
                        _ => -try!(search(tt, p, moves, nc, report, -beta, -alpha, depth - 1)),
                    }
                };

                no_moves_yet = false;
                p.undo_move();

                if value >= beta {
                    // This move is too good, so that the opponent
                    // will not allow this line of play to
                    // happen. Therefore we can stop here.
                    alpha = beta;
                    bound_type = BOUND_LOWER;
                    best_move = m;
                    break;
                }
                if value > alpha {
                    // We found ourselves a new best move.
                    alpha = value;
                    bound_type = BOUND_EXACT;
                    best_move = m;
                }
            }
        }
        moves.restore();
        if no_moves_yet {
            // No legal moves -- this is a final position.
            alpha = p.evaluate_final();
            bound_type = BOUND_EXACT;
        }
    }

    tt.store(p.hash(),
             EntryData::new(alpha, bound_type, depth, best_move.digest(), eval_value));
    Ok(alpha)
}

const NODE_COUNT_REPORT_INTERVAL: NodeCount = 10000;

#[cfg(test)]
mod tests {
    use super::search;
    use chess_move::*;
    use tt::*;
    use position::Position;

    #[test]
    fn test_search() {
        let mut p = Position::from_fen("8/8/8/8/3q3k/7n/6PP/2Q2R1K b - - 0 1").ok().unwrap();
        let value = search(&TranspositionTable::new(),
                           &mut p,
                           &mut MoveStack::new(),
                           &mut 0,
                           &mut |_| Ok(()),
                           -30000,
                           30000,
                           2)
                        .ok()
                        .unwrap();
        assert!(value < -300);
        let value = search(&TranspositionTable::new(),
                           &mut p,
                           &mut MoveStack::new(),
                           &mut 0,
                           &mut |_| Ok(()),
                           -30000,
                           30000,
                           4)
                        .ok()
                        .unwrap();
        assert!(value >= 20000);
    }
}
<commit_msg>more beautifications<commit_after>use std::thread;
use std::cell::UnsafeCell;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver, RecvError};
use basetypes::*;
use chess_move::*;
use tt::*;
use position::Position;

pub enum Command {
    Search {
        search_id: usize,
        position: Position,
        depth: u8,
        lower_bound: Value,
        upper_bound: Value,
    },
    Stop,
    Exit,
}

pub enum Report {
    Progress {
        search_id: usize,
        searched_nodes: NodeCount,
        depth: u8,
    },
    Done {
        search_id: usize,
        searched_nodes: NodeCount,
        value: Option<Value>,
    },
}

pub fn run_deepening(tt: Arc<TranspositionTable>,
                     commands: Receiver<Command>,
                     reports: Sender<Report>) {
    let (slave_commands_tx, slave_commands_rx) = channel();
    let (slave_reports_tx, slave_reports_rx) = channel();
    let slave = thread::spawn(move || {
        run(tt, slave_commands_rx, slave_reports_tx);
    });
    let mut pending_command = None;
    loop {
        // If there is a pending command, we take it, otherwise we
        // block and wait to receive a new one.
        let command = match pending_command.take() {
            Some(cmd) => cmd,
            None => commands.recv().or::<RecvError>(Ok(Command::Exit)).unwrap(),
        };

        match command {
            Command::Search { search_id, position, depth, lower_bound, upper_bound } => {
                let mut searched_nodes_final = 0;
                let mut value_final = None;
                'depthloop: for n in 1..(depth + 1) {
                    slave_commands_tx.send(Command::Search {
                                         search_id: n as usize,
                                         position: position.clone(),
                                         depth: n,
                                         lower_bound: lower_bound,
                                         upper_bound: upper_bound,
                                     })
                                     .unwrap();
                    loop {
                        match slave_reports_rx.recv().unwrap() {
                            Report::Progress { searched_nodes, .. } => {
                                reports.send(Report::Progress {
                                           search_id: search_id,
                                           searched_nodes: searched_nodes_final + searched_nodes,
                                           depth: n,
                                       })
                                       .ok();
                                if pending_command.is_none() {
                                    pending_command = match commands.try_recv() {
                                        Ok(cmd) => {
                                            slave_commands_tx.send(Command::Stop).unwrap();
                                            Some(cmd)
                                        }
                                        _ => None,
                                    }
                                }
                            }
                            Report::Done { searched_nodes, value, .. } => {
                                searched_nodes_final += searched_nodes;
                                if n == depth {
                                    value_final = value;
                                }
                                if pending_command.is_some() {
                                    break 'depthloop;
                                }
                                break;
                            }
                        }
                    }
                }
                reports.send(Report::Done {
                           search_id: search_id,
                           searched_nodes: searched_nodes_final,
                           value: value_final,
                       })
                       .ok();
            }
            Command::Stop => {
                slave_commands_tx.send(Command::Stop).unwrap();
                continue;
            }
            Command::Exit => {
                slave_commands_tx.send(Command::Exit).unwrap();
                break;
            }
        }
    }
    slave.join().unwrap();
}

pub fn run(tt: Arc<TranspositionTable>, commands: Receiver<Command>, reports: Sender<Report>) {
    thread_local!(
        static MOVE_STACK: UnsafeCell<MoveStack> = UnsafeCell::new(MoveStack::new())
    );
    MOVE_STACK.with(|s| {
        let mut move_stack = unsafe { &mut *s.get() };
        let mut pending_command = None;
        loop {
            // If there is a pending command, we take it, otherwise we
            // block and wait to receive a new one.
            let command = match pending_command.take() {
                Some(cmd) => cmd,
                None => commands.recv().or::<RecvError>(Ok(Command::Exit)).unwrap(),
            };

            match command {
                Command::Search { search_id, mut position, depth, lower_bound, upper_bound } => {
                    let mut reported_nodes = 0;
                    let mut unreported_nodes = 0;
                    let value = search(&tt,
                                       &mut position,
                                       move_stack,
                                       &mut unreported_nodes,
                                       &mut |n| {
                                           reported_nodes += n;
                                           reports.send(Report::Progress {
                                                      search_id: search_id,
                                                      searched_nodes: reported_nodes,
                                                      depth: depth,
                                                  })
                                                  .ok();
                                           match commands.try_recv() {
                                               Ok(x) => {
                                                   pending_command = Some(x);
                                                   Err(TerminatedSearch)
                                               }
                                               _ => Ok(()),
                                           }
                                       },
                                       lower_bound,
                                       upper_bound,
                                       depth)
                                    .ok();
                    reported_nodes += unreported_nodes;
                    reports.send(Report::Progress {
                               search_id: search_id,
                               searched_nodes: reported_nodes,
                               depth: depth,
                           })
                           .ok();
                    reports.send(Report::Done {
                               search_id: search_id,
                               searched_nodes: reported_nodes,
                               value: value,
                           })
                           .ok();
                    move_stack.clear();
                }
                Command::Stop => continue,
                Command::Exit => break,
            }
        }
    })
}

// Represents a terminated search condition.
struct TerminatedSearch;

// Helper function for `run()`. It implements the principal variation
// search algorithm. When returning `Err(TerminatedSearch)`, this
// function may leave un-restored move lists in `moves`.
fn search(tt: &TranspositionTable,
          p: &mut Position,
          moves: &mut MoveStack,
          nc: &mut NodeCount,
          report: &mut FnMut(NodeCount) -> Result<(), TerminatedSearch>,
          mut alpha: Value, // lower bound
          beta: Value, // upper bound
          depth: u8)
          -> Result<Value, TerminatedSearch> {
    assert!(alpha < beta);

    // Consult the transposition table.
    let (move16, eval_value) = if let Some(entry) = tt.probe(p.hash()) {
        if entry.depth() >= depth {
            let value = entry.value();
            let bound = entry.bound();
            if (value >= beta && bound == BOUND_LOWER) ||
               (value <= alpha && bound == BOUND_UPPER) || (bound == BOUND_EXACT) {
                return Ok(value);
            }
        }
        (entry.move16(), entry.eval_value())
    } else {
        (0, p.evaluate_static())
    };

    // Initial guests for the final result.
    let mut bound_type = BOUND_UPPER;
    let mut best_move = Move::invalid();

    if depth == 0 {
        // On leaf nodes, do quiescence search.
        let (value, nodes) = p.evaluate_quiescence(alpha, beta, Some(eval_value));
        *nc += nodes;
        if value >= beta {
            alpha = beta;
            bound_type = BOUND_LOWER;
        } else if value > alpha {
            alpha = value;
            bound_type = BOUND_EXACT;
        }

    } else {
        // On non-leaf nodes, try moves and make recursive calls.
        moves.save();
        let mut no_moves_yet = true;

        // TODO: This is ugly.
        if move16 != 0 {
            if let Some(m) = p.try_move_digest(move16) {
                moves.push(m);
            }
        }
        let mut generated = false;

        // Try some moves.
        loop {
            // TODO: This is ugly.
            let m = match moves.remove_best_move() {
                None if generated => break,
                None => {
                    p.generate_moves(moves);
                    moves.remove_move(move16);
                    generated = true;
                    continue;
                }
                Some(x) => x,
            };

            if p.do_move(m) {
                // From time to time, we report how many nodes had
                // been searched since the last report. This also
                // gives an opportunity for the search to be
                // terminated.
                *nc += 1;
                if *nc > NODE_COUNT_REPORT_INTERVAL {
                    try!(report(*nc));
                    *nc = 0;
                }

                // Make a recursive call.
                let value = if no_moves_yet {
                    // The first move we analyze with a fully open
                    // window (alpha, beta).
                    -try!(search(tt, p, moves, nc, report, -beta, -alpha, depth - 1))
                } else {
                    // For the next moves we first try to prove that
                    // they are not better than our current best
                    // move. For this purpose we analyze them with a
                    // null window (alpha, alpha + 1). This is faster
                    // than a full window search. Only when we are
                    // certain that the move is better than our
                    // current best move, we do a full-window search.
                    match -try!(search(tt, p, moves, nc, report, -alpha - 1, -alpha, depth - 1)) {
                        x if x <= alpha => x,
                        _ => -try!(search(tt, p, moves, nc, report, -beta, -alpha, depth - 1)),
                    }
                };

                // See how good this move is.
                p.undo_move();
                no_moves_yet = false;
                if value >= beta {
                    // This move is too good, so that the opponent
                    // will not allow this line of play to
                    // happen. Therefore we can stop here.
                    alpha = beta;
                    bound_type = BOUND_LOWER;
                    best_move = m;
                    break;
                }
                if value > alpha {
                    // We found ourselves a new best move.
                    alpha = value;
                    bound_type = BOUND_EXACT;
                    best_move = m;
                }
            }
        }

        // Check if we are in a final position (no legal moves). Then
        // we are done.
        if no_moves_yet {
            // Final positions we can evaluate 100% correctly.
            alpha = p.evaluate_final();
            bound_type = BOUND_EXACT;
        }
        moves.restore();
    }

    // Store the final result in the transposition table and return.
    tt.store(p.hash(),
             EntryData::new(alpha, bound_type, depth, best_move.digest(), eval_value));
    Ok(alpha)
}

const NODE_COUNT_REPORT_INTERVAL: NodeCount = 10000;

#[cfg(test)]
mod tests {
    use super::search;
    use chess_move::*;
    use tt::*;
    use position::Position;

    #[test]
    fn test_search() {
        let mut p = Position::from_fen("8/8/8/8/3q3k/7n/6PP/2Q2R1K b - - 0 1").ok().unwrap();
        let value = search(&TranspositionTable::new(),
                           &mut p,
                           &mut MoveStack::new(),
                           &mut 0,
                           &mut |_| Ok(()),
                           -30000,
                           30000,
                           2)
                        .ok()
                        .unwrap();
        assert!(value < -300);
        let value = search(&TranspositionTable::new(),
                           &mut p,
                           &mut MoveStack::new(),
                           &mut 0,
                           &mut |_| Ok(()),
                           -30000,
                           30000,
                           4)
                        .ok()
                        .unwrap();
        assert!(value >= 20000);
    }
}
<|endoftext|>


<commit_before>//! Implements single-threaded game tree search.

use basetypes::*;
use bitsets::*;
use chess_move::*;
use tt::*;
use position::Position;

/// Represents a terminated search condition.
pub struct TerminatedSearch;

/// Represents a game tree search.
pub struct Search<'a> {
    tt: &'a TranspositionTable,
    position: Position,
    moves: &'a mut MoveStack,
    moves_starting_ply: usize,
    state_stack: Vec<NodeState>,
    reported_nodes: NodeCount,
    unreported_nodes: NodeCount,
    report_function: &'a mut FnMut(NodeCount) -> bool,
}

impl<'a> Search<'a> {
    /// Creates a new instance.
    ///
    /// `report_function` should be a function that registers the
    /// search progress. It will be called with the number of searched
    /// positions from the beginning of the search to this moment. The
    /// function should return `true` if the search should be
    /// terminated, otherwise it should return `false`.
    pub fn new(root: Position,
               tt: &'a TranspositionTable,
               move_stack: &'a mut MoveStack,
               report_function: &'a mut FnMut(NodeCount) -> bool)
               -> Search<'a> {
        let moves_starting_ply = move_stack.ply();
        Search {
            tt: tt,
            position: root,
            moves: move_stack,
            moves_starting_ply: moves_starting_ply,
            state_stack: Vec::with_capacity(32),
            reported_nodes: 0,
            unreported_nodes: 0,
            report_function: report_function,
        }
    }

    /// Performs a principal variation search (PVS) and returns a
    /// result.
    ///
    /// The alpha-beta algorithm is a significant enhancement to the
    /// minimax search algorithm that eliminates the need to search
    /// large portions of the game tree applying a branch-and-bound
    /// technique. Remarkably, it does this without any potential of
    /// overlooking a better move. If one already has found a quite
    /// good move and search for alternatives, one refutation is
    /// enough to avoid it. No need to look for even stronger
    /// refutations. The algorithm maintains two values, alpha and
    /// beta. They represent the minimum score that the maximizing
    /// player is assured of and the maximum score that the minimizing
    /// player is assured of respectively. PVS is an enhancement to
    /// the alpha-beta search, based on zero (null) window searches of
    /// none PV-nodes, to prove a move is worse or not than an already
    /// safe score from the principal variation. (The principal
    /// variation (PV) is a sequence of moves that the program
    /// considers best and therefore expects to be played.)
    ///
    /// **Important note**: This method may leave un-restored move
    /// lists in the move stack. Call `reset` if you want the move
    /// stack to be restored to the state it had when the search
    /// instance was created.
    pub fn run(&mut self,
               mut alpha: Value, // lower bound
               beta: Value, // upper bound
               depth: u8)
               -> Result<Value, TerminatedSearch> {
        assert!(alpha < beta);

        if let Some(value) = self.node_begin(alpha, beta, depth) {
            // We already have the position's value.
            alpha = value;

        } else {
            // Initial guests for the final result.
            let mut bound = BOUND_UPPER;
            let mut best_move = Move::invalid();

            // On non-leaf nodes, try moves.
            let mut no_moves_yet = true;
            while let Some(m) = self.do_move() {
                try!(self.report_progress(1));

                // Make a recursive call.
                let value = if no_moves_yet {
                    // The first move we analyze with a fully open window
                    // (alpha, beta). If this happens to be a good move,
                    // it will probably raise `alpha`.
                    no_moves_yet = false;
                    -try!(self.run(-beta, -alpha, depth - 1))
                } else {
                    // For the next moves we first try to prove that they
                    // are not better than our current best move. For this
                    // purpose we analyze them with a null window (alpha,
                    // alpha + 1). This is faster than a full window
                    // search. Only if we are certain that the move is
                    // better than our current best move, we do a
                    // full-window search.
                    match -try!(self.run(-alpha - 1, -alpha, depth - 1)) {
                        x if x <= alpha => x,
                        _ => -try!(self.run(-beta, -alpha, depth - 1)),
                    }
                };
                self.undo_move();

                // See how good this move was.
                if value >= beta {
                    // This move is so good, that the opponent will
                    // probably not allow this line of play to
                    // happen. Therefore we should not lose any more time
                    // on this position.
                    alpha = beta;
                    bound = BOUND_LOWER;
                    best_move = m;
                    break;
                }
                if value > alpha {
                    // We found a new best move.
                    alpha = value;
                    bound = BOUND_EXACT;
                    best_move = m;
                }
            }

            // Check if we are in a final position (no legal moves).
            if no_moves_yet {
                let value = self.position.evaluate_final();
                if value >= beta {
                    alpha = beta;
                    bound = BOUND_LOWER;
                }
                if value > alpha {
                    alpha = value;
                    bound = BOUND_EXACT;
                }
            }

            // Store the result to the TT.
            self.store(alpha, bound, depth, best_move);
        }

        self.node_end();
        Ok(alpha)
    }

    /// Returns the number of searched positions.
    #[inline(always)]
    pub fn node_count(&self) -> NodeCount {
        self.reported_nodes + self.unreported_nodes
    }

    /// Resets the instance to the state it had when it was created.
    #[inline]
    pub fn reset(&mut self) {
        while self.moves.ply() > self.moves_starting_ply {
            self.moves.restore();
        }
        self.state_stack.clear();
        self.reported_nodes = 0;
        self.unreported_nodes = 0;
    }

    // Declares that we are starting to process a new node.
    //
    // Each recursive call to `run` begins with a call to
    // `node_begin`. The returned value (if not `None`) is the value
    // assigned to the node (taken from the TT or, on leaf nodes --
    // calculated by performing quiescence search).
    #[inline]
    fn node_begin(&mut self, alpha: Value, beta: Value, depth: u8) -> Option<Value> {
        // Consult the transposition table.
        let entry = if let Some(e) = self.tt.probe(self.position.hash()) {
            e
        } else {
            EntryData::new(0, BOUND_NONE, 0, 0, self.position.evaluate_static())
        };
        self.state_stack.push(NodeState {
            phase: NodePhase::Pristine,
            entry: entry,
            checkers: BB_UNIVERSAL_SET,
        });

        // Check if the TT entry gives the result.
        if entry.depth() >= depth {
            let value = entry.value();
            let bound = entry.bound();
            if value >= beta && bound & BOUND_LOWER != 0 {
                return Some(beta);
            }
            if value <= alpha && bound & BOUND_UPPER != 0 {
                return Some(alpha);
            }
            if bound == BOUND_EXACT {
                return Some(value);
            };
        };

        // On leaf nodes, do quiescence search.
        if depth == 0 {
            let eval_value = entry.eval_value();
            let (mut value, nodes) = self.position
                                         .evaluate_quiescence(alpha, beta, Some(eval_value));
            self.report_progress(nodes).ok();
            let bound = if value >= beta {
                value = beta;
                BOUND_LOWER
            } else if value <= alpha {
                value = alpha;
                BOUND_UPPER
            } else {
                BOUND_EXACT
            };
            self.tt.store(self.position.hash(),
                          EntryData::new(value, bound, 0, 0, eval_value));
            return Some(value);
        }

        None
    }

    // Declares that we are done processing the current node.
    //
    // Each recursive call to `run` ends with a call to `node_end`.
    #[inline]
    fn node_end(&mut self) {
        if let NodePhase::Pristine = self.state_stack.last().unwrap().phase {
            // For pristine nodes we have not saved the move list
            // yet, so we should not restore it.
        } else {
            self.moves.restore();
        }
        self.state_stack.pop();
    }

    // TODO: Implement do_null_move() method?

    // Plays the next legal move in the current position and returns
    // it.
    //
    // Each call to `do_move` for the same position will play and
    // return a different move. When all legal moves has been played,
    // `None` will be returned. `do_move` will do whatever it can to
    // play the best moves first, and the worst last. It will also try
    // to be efficient, for example it will generate the list of all
    // pseudo-legal moves at the last possible moment.
    #[inline]
    fn do_move(&mut self) -> Option<Move> {
        let state = self.state_stack.last_mut().unwrap();

        // Try the hash move first.
        if let NodePhase::Pristine = state.phase {
            // We save the move list at the last possible moment,
            // because most of the nodes are leafs.
            self.moves.save();

            // We save the checkers bitboard also, because we will
            // need this information later many times, and we do not
            // want to recalculate it needlessly.
            //
            // TODO: Verify if this is the best place to calculate
            // checkers, especially in relation to trying null moves.
            state.checkers = self.position.board().checkers();

            state.phase = NodePhase::TriedHashMove;
            if state.entry.move16() != 0 {
                if let Some(mut m) = self.position.try_move_digest(state.entry.move16()) {
                    if self.position.do_move(m) {
                        m.set_score(MAX_MOVE_SCORE);
                        return Some(m);
                    }
                }
            }
        }

        // After the hash move, we generate all pseudo-legal
        // moves. But we should not forget to remove the already tried
        // hash move from the list.
        if let NodePhase::TriedHashMove = state.phase {
            // TODO: `generate_moves` needs `_checkers` and `_pinned`
            // to do its work. At this time we already have generated
            // them, so we should use them.
            self.position.board()._checkers.set(state.checkers);

            self.position.generate_moves(self.moves);
            if state.entry.move16() != 0 {
                self.moves.remove_move(state.entry.move16());
            }
            state.phase = NodePhase::GeneratedMoves;
        }

        // Spit out the generated moves.
        while let Some(mut m) = self.moves.remove_best_move() {

            // First, the good captures.
            if let NodePhase::GeneratedMoves = state.phase {
                if m.score() == MAX_MOVE_SCORE {
                    if self.position.evaluate_move(m) >= 0 {
                        if self.position.do_move(m) {
                            return Some(m);
                        }
                        continue;
                    }
                    // This is a bad capture -- push it back to the
                    // move stack.
                    m.set_score(MAX_MOVE_SCORE - 1);
                    self.moves.push(m);
                    continue;
                }
                state.phase = NodePhase::TriedGoodCaptures;
            }

            // Second, the bad captures.
            if let NodePhase::TriedGoodCaptures = state.phase {
                if m.score() == MAX_MOVE_SCORE - 1 {
                    if self.position.do_move(m) {
                        return Some(m);
                    }
                    continue;
                }
                state.phase = NodePhase::TriedBadCaptures;
            }

            // Before trying the quiet moves, we should assign proper
            // move scores to them.
            if let NodePhase::TriedBadCaptures = state.phase {
                // TODO: Assign the moves scores here using the killer
                // move heuristics and the history heuristics.

                // We use the score field (2 bits) to properly order
                // quiet movies. Moves which destination square is
                // more advanced into enemy's territory are tried
                // first. The logic is that those moves are riskier,
                // so if such a move loses material this will be
                // detected early and the search tree will be pruned,
                // but if the move does not lose material, chances are
                // that it is a very good move.
                const SCORE_LOOKUP: [[u32; 8]; 2] = [// white
                                                     [0, 1, 2, 3, 4, 5, 6, 7],
                                                     // black
                                                     [7, 6, 5, 4, 3, 2, 1, 0]];
                for m in self.moves.iter_mut() {
                    let rank = rank(m.dest_square());
                    m.set_score(unsafe {
                        *SCORE_LOOKUP.get_unchecked(self.position.board().to_move())
                                     .get_unchecked(rank)
                    });
                }

                state.phase = NodePhase::SortedQuietMoves;
            }

            // Last, quiet moves.
            if self.position.do_move(m) {
                if state.checkers != 0 {
                    // When in check, we set a high move score to all
                    // moves to avoid search depth reductions.
                    m.set_score(MAX_MOVE_SCORE - 1);
                }
                return Some(m);
            }
        }
        None
    }

    // Takes the last played move back.
    #[inline]
    fn undo_move(&mut self) {
        self.position.undo_move();
    }

    // Stores updated node information in the transposition table.
    #[inline]
    fn store(&mut self, value: Value, bound: BoundType, depth: u8, best_move: Move) {
        let entry = &self.state_stack.last().unwrap().entry;
        self.tt.store(self.position.hash(),
                      EntryData::new(value, bound, depth, best_move.digest(), entry.eval_value()));
    }

    // Reports search progress.
    //
    // From time to time, we should report how many nodes had been
    // searched since the beginning of the search. This also gives an
    // opportunity for the search to be terminated.
    #[inline]
    fn report_progress(&mut self, new_nodes: NodeCount) -> Result<(), TerminatedSearch> {
        self.unreported_nodes += new_nodes;
        if self.unreported_nodes >= super::NODE_COUNT_REPORT_INTERVAL {
            self.reported_nodes += self.unreported_nodes;
            self.unreported_nodes = 0;
            if (*self.report_function)(self.reported_nodes) {
                return Err(TerminatedSearch);
            }
        }
        Ok(())
    }
}

enum NodePhase {
    Pristine,
    TriedHashMove,
    GeneratedMoves,
    TriedGoodCaptures,
    TriedBadCaptures,
    SortedQuietMoves,
}

struct NodeState {
    phase: NodePhase,
    entry: EntryData,
    checkers: u64,
}

#[cfg(test)]
mod tests {
    use super::Search;
    use chess_move::*;
    use tt::*;
    use position::Position;

    #[test]
    fn test_search() {
        let p = Position::from_fen("8/8/8/8/3q3k/7n/6PP/2Q2R1K b - - 0 1").ok().unwrap();
        let tt = TranspositionTable::new();
        let mut moves = MoveStack::new();
        let mut report = |_| false;
        let mut search = Search::new(p, &tt, &mut moves, &mut report);
        let value = search.run(-30000, 30000, 2)
                          .ok()
                          .unwrap();
        assert!(value < -300);
        search.reset();
        let value = search.run(-30000, 30000, 4)
                          .ok()
                          .unwrap();
        assert!(value >= 20000);
    }
}
<commit_msg>fixed comments<commit_after>//! Implements single-threaded game tree search.

use basetypes::*;
use bitsets::*;
use chess_move::*;
use tt::*;
use position::Position;

/// Represents a terminated search condition.
pub struct TerminatedSearch;

/// Represents a game tree search.
pub struct Search<'a> {
    tt: &'a TranspositionTable,
    position: Position,
    moves: &'a mut MoveStack,
    moves_starting_ply: usize,
    state_stack: Vec<NodeState>,
    reported_nodes: NodeCount,
    unreported_nodes: NodeCount,
    report_function: &'a mut FnMut(NodeCount) -> bool,
}

impl<'a> Search<'a> {
    /// Creates a new instance.
    ///
    /// `report_function` should be a function that registers the
    /// search progress. It will be called with the number of searched
    /// positions from the beginning of the search to this moment. The
    /// function should return `true` if the search should be
    /// terminated, otherwise it should return `false`.
    pub fn new(root: Position,
               tt: &'a TranspositionTable,
               move_stack: &'a mut MoveStack,
               report_function: &'a mut FnMut(NodeCount) -> bool)
               -> Search<'a> {
        let moves_starting_ply = move_stack.ply();
        Search {
            tt: tt,
            position: root,
            moves: move_stack,
            moves_starting_ply: moves_starting_ply,
            state_stack: Vec::with_capacity(32),
            reported_nodes: 0,
            unreported_nodes: 0,
            report_function: report_function,
        }
    }

    /// Performs a principal variation search (PVS) and returns a
    /// result.
    ///
    /// The alpha-beta algorithm is a significant enhancement to the
    /// minimax search algorithm that eliminates the need to search
    /// large portions of the game tree applying a branch-and-bound
    /// technique. Remarkably, it does this without any potential of
    /// overlooking a better move. If one already has found a quite
    /// good move and search for alternatives, one refutation is
    /// enough to avoid it. No need to look for even stronger
    /// refutations. The algorithm maintains two values, alpha and
    /// beta. They represent the minimum score that the maximizing
    /// player is assured of and the maximum score that the minimizing
    /// player is assured of respectively. PVS is an enhancement to
    /// the alpha-beta search, based on zero (null) window searches of
    /// none PV-nodes, to prove a move is worse or not than an already
    /// safe score from the principal variation. (The principal
    /// variation (PV) is a sequence of moves that the program
    /// considers best and therefore expects to be played.)
    ///
    /// **Important note**: This method may leave un-restored move
    /// lists in the move stack. Call `reset` if you want the move
    /// stack to be restored to the state it had when the search
    /// instance was created.
    pub fn run(&mut self,
               mut alpha: Value, // lower bound
               beta: Value, // upper bound
               depth: u8)
               -> Result<Value, TerminatedSearch> {
        assert!(alpha < beta);

        if let Some(value) = self.node_begin(alpha, beta, depth) {
            // We already have the final result.
            alpha = value;

        } else {
            // Initial guests.
            let mut bound = BOUND_UPPER;
            let mut best_move = Move::invalid();

            // Try moves.
            let mut no_moves_yet = true;
            while let Some(m) = self.do_move() {
                try!(self.report_progress(1));

                // Make a recursive call.
                let value = if no_moves_yet {
                    // The first move we analyze with a fully open window
                    // (alpha, beta). If this happens to be a good move,
                    // it will probably raise `alpha`.
                    no_moves_yet = false;
                    -try!(self.run(-beta, -alpha, depth - 1))
                } else {
                    // For the next moves we first try to prove that they
                    // are not better than our current best move. For this
                    // purpose we analyze them with a null window (alpha,
                    // alpha + 1). This is faster than a full window
                    // search. Only if we are certain that the move is
                    // better than our current best move, we do a
                    // full-window search.
                    match -try!(self.run(-alpha - 1, -alpha, depth - 1)) {
                        x if x <= alpha => x,
                        _ => -try!(self.run(-beta, -alpha, depth - 1)),
                    }
                };
                self.undo_move();

                // See how good this move was.
                if value >= beta {
                    // This move is so good, that the opponent will
                    // probably not allow this line of play to
                    // happen. Therefore we should not lose any more time
                    // on this position.
                    alpha = beta;
                    bound = BOUND_LOWER;
                    best_move = m;
                    break;
                }
                if value > alpha {
                    // We found a new best move.
                    alpha = value;
                    bound = BOUND_EXACT;
                    best_move = m;
                }
            }

            // Check if we are in a final position (no legal moves).
            if no_moves_yet {
                let value = self.position.evaluate_final();
                if value >= beta {
                    alpha = beta;
                    bound = BOUND_LOWER;
                }
                if value > alpha {
                    alpha = value;
                    bound = BOUND_EXACT;
                }
            }

            // Store the result to the TT.
            self.store(alpha, bound, depth, best_move);
        }

        self.node_end();
        Ok(alpha)
    }

    /// Returns the number of searched positions.
    #[inline(always)]
    pub fn node_count(&self) -> NodeCount {
        self.reported_nodes + self.unreported_nodes
    }

    /// Resets the instance to the state it had when it was created.
    #[inline]
    pub fn reset(&mut self) {
        while self.moves.ply() > self.moves_starting_ply {
            self.moves.restore();
        }
        self.state_stack.clear();
        self.reported_nodes = 0;
        self.unreported_nodes = 0;
    }

    // Declares that we are starting to process a new node.
    //
    // Each recursive call to `run` begins with a call to
    // `node_begin`. The returned value (if not `None`) is the value
    // assigned to the node (taken from the TT or, on leaf nodes --
    // calculated by performing quiescence search).
    #[inline]
    fn node_begin(&mut self, alpha: Value, beta: Value, depth: u8) -> Option<Value> {
        // Consult the transposition table.
        let entry = if let Some(e) = self.tt.probe(self.position.hash()) {
            e
        } else {
            EntryData::new(0, BOUND_NONE, 0, 0, self.position.evaluate_static())
        };
        self.state_stack.push(NodeState {
            phase: NodePhase::Pristine,
            entry: entry,
            checkers: BB_UNIVERSAL_SET,
        });

        // Check if the TT entry gives the result.
        if entry.depth() >= depth {
            let value = entry.value();
            let bound = entry.bound();
            if value >= beta && bound & BOUND_LOWER != 0 {
                return Some(beta);
            }
            if value <= alpha && bound & BOUND_UPPER != 0 {
                return Some(alpha);
            }
            if bound == BOUND_EXACT {
                return Some(value);
            };
        };

        // On leaf nodes, do quiescence search.
        if depth == 0 {
            let eval_value = entry.eval_value();
            let (mut value, nodes) = self.position
                                         .evaluate_quiescence(alpha, beta, Some(eval_value));
            self.report_progress(nodes).ok();
            let bound = if value >= beta {
                value = beta;
                BOUND_LOWER
            } else if value <= alpha {
                value = alpha;
                BOUND_UPPER
            } else {
                BOUND_EXACT
            };
            self.tt.store(self.position.hash(),
                          EntryData::new(value, bound, 0, 0, eval_value));
            return Some(value);
        }

        None
    }

    // Declares that we are done processing the current node.
    //
    // Each recursive call to `run` ends with a call to `node_end`.
    #[inline]
    fn node_end(&mut self) {
        if let NodePhase::Pristine = self.state_stack.last().unwrap().phase {
            // For pristine nodes we have not saved the move list
            // yet, so we should not restore it.
        } else {
            self.moves.restore();
        }
        self.state_stack.pop();
    }

    // TODO: Implement do_null_move() method?

    // Plays the next legal move in the current position and returns
    // it.
    //
    // Each call to `do_move` for the same position will play and
    // return a different move. When all legal moves has been played,
    // `None` will be returned. `do_move` will do whatever it can to
    // play the best moves first, and the worst last. It will also try
    // to be efficient, for example it will generate the list of all
    // pseudo-legal moves at the last possible moment.
    #[inline]
    fn do_move(&mut self) -> Option<Move> {
        let state = self.state_stack.last_mut().unwrap();

        // Try the hash move first.
        if let NodePhase::Pristine = state.phase {
            // We save the move list at the last possible moment,
            // because most of the nodes are leafs.
            self.moves.save();

            // We save the checkers bitboard also, because we will
            // need this information later many times, and we do not
            // want to recalculate it needlessly.
            //
            // TODO: Verify if this is the best place to calculate
            // checkers, especially in relation to trying null moves.
            state.checkers = self.position.board().checkers();

            state.phase = NodePhase::TriedHashMove;
            if state.entry.move16() != 0 {
                if let Some(mut m) = self.position.try_move_digest(state.entry.move16()) {
                    if self.position.do_move(m) {
                        m.set_score(MAX_MOVE_SCORE);
                        return Some(m);
                    }
                }
            }
        }

        // After the hash move, we generate all pseudo-legal
        // moves. But we should not forget to remove the already tried
        // hash move from the list.
        if let NodePhase::TriedHashMove = state.phase {
            // TODO: `generate_moves` needs `_checkers` and `_pinned`
            // to do its work. At this time we already have generated
            // them, so we should use them.
            self.position.board()._checkers.set(state.checkers);

            self.position.generate_moves(self.moves);
            if state.entry.move16() != 0 {
                self.moves.remove_move(state.entry.move16());
            }
            state.phase = NodePhase::GeneratedMoves;
        }

        // Spit out the generated moves.
        while let Some(mut m) = self.moves.remove_best_move() {

            // First, the good captures.
            if let NodePhase::GeneratedMoves = state.phase {
                if m.score() == MAX_MOVE_SCORE {
                    if self.position.evaluate_move(m) >= 0 {
                        if self.position.do_move(m) {
                            return Some(m);
                        }
                        continue;
                    }
                    // This is a bad capture -- push it back to the
                    // move stack.
                    m.set_score(MAX_MOVE_SCORE - 1);
                    self.moves.push(m);
                    continue;
                }
                state.phase = NodePhase::TriedGoodCaptures;
            }

            // Second, the bad captures.
            if let NodePhase::TriedGoodCaptures = state.phase {
                if m.score() == MAX_MOVE_SCORE - 1 {
                    if self.position.do_move(m) {
                        return Some(m);
                    }
                    continue;
                }
                state.phase = NodePhase::TriedBadCaptures;
            }

            // Before trying the quiet moves, we should assign proper
            // move scores to them.
            if let NodePhase::TriedBadCaptures = state.phase {
                // TODO: Assign the moves scores here using the killer
                // move heuristics and the history heuristics.

                // We use the score field (2 bits) to properly order
                // quiet movies. Moves which destination square is
                // more advanced into enemy's territory are tried
                // first. The logic is that those moves are riskier,
                // so if such a move loses material this will be
                // detected early and the search tree will be pruned,
                // but if the move does not lose material, chances are
                // that it is a very good move.
                const SCORE_LOOKUP: [[u32; 8]; 2] = [// white
                                                     [0, 1, 2, 3, 4, 5, 6, 7],
                                                     // black
                                                     [7, 6, 5, 4, 3, 2, 1, 0]];
                for m in self.moves.iter_mut() {
                    let rank = rank(m.dest_square());
                    m.set_score(unsafe {
                        *SCORE_LOOKUP.get_unchecked(self.position.board().to_move())
                                     .get_unchecked(rank)
                    });
                }

                state.phase = NodePhase::SortedQuietMoves;
            }

            // Last, quiet moves.
            if self.position.do_move(m) {
                if state.checkers != 0 {
                    // When in check, we set a high move score to all
                    // moves to avoid search depth reductions.
                    m.set_score(MAX_MOVE_SCORE - 1);
                }
                return Some(m);
            }
        }
        None
    }

    // Takes the last played move back.
    #[inline]
    fn undo_move(&mut self) {
        self.position.undo_move();
    }

    // Stores updated node information in the transposition table.
    #[inline]
    fn store(&mut self, value: Value, bound: BoundType, depth: u8, best_move: Move) {
        let entry = &self.state_stack.last().unwrap().entry;
        self.tt.store(self.position.hash(),
                      EntryData::new(value, bound, depth, best_move.digest(), entry.eval_value()));
    }

    // Reports search progress.
    //
    // From time to time, we should report how many nodes had been
    // searched since the beginning of the search. This also gives an
    // opportunity for the search to be terminated.
    #[inline]
    fn report_progress(&mut self, new_nodes: NodeCount) -> Result<(), TerminatedSearch> {
        self.unreported_nodes += new_nodes;
        if self.unreported_nodes >= super::NODE_COUNT_REPORT_INTERVAL {
            self.reported_nodes += self.unreported_nodes;
            self.unreported_nodes = 0;
            if (*self.report_function)(self.reported_nodes) {
                return Err(TerminatedSearch);
            }
        }
        Ok(())
    }
}

enum NodePhase {
    Pristine,
    TriedHashMove,
    GeneratedMoves,
    TriedGoodCaptures,
    TriedBadCaptures,
    SortedQuietMoves,
}

struct NodeState {
    phase: NodePhase,
    entry: EntryData,
    checkers: u64,
}

#[cfg(test)]
mod tests {
    use super::Search;
    use chess_move::*;
    use tt::*;
    use position::Position;

    #[test]
    fn test_search() {
        let p = Position::from_fen("8/8/8/8/3q3k/7n/6PP/2Q2R1K b - - 0 1").ok().unwrap();
        let tt = TranspositionTable::new();
        let mut moves = MoveStack::new();
        let mut report = |_| false;
        let mut search = Search::new(p, &tt, &mut moves, &mut report);
        let value = search.run(-30000, 30000, 2)
                          .ok()
                          .unwrap();
        assert!(value < -300);
        search.reset();
        let value = search.run(-30000, 30000, 4)
                          .ok()
                          .unwrap();
        assert!(value >= 20000);
    }
}
<|endoftext|>


<commit_before>use ::Key;
use super::rootptr::RootPtr;

pub struct RudyMap<K: Key, V> {
    root: RootPtr<K, V>
}

impl<K: Key, V> RudyMap<K, V> {
    pub fn new() -> RudyMap<K, V> {
        RudyMap {
            root: RootPtr::empty()
        }
    }

    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
        self.root.insert(key, value)
    }

    /*
    fn remove(&mut self, key: K) -> Option<V> {
        self.root.remove(key)
    }
    */

    /*
    fn clear(&mut self) {
        self.root.clear()
    }
    */

    pub fn contains_key(&self, key: K) -> bool {
        // Because this does not visit the pointer returned by get(), there
        // will not be unnecessary cache fills
        self.get(key).is_some()
    }

    pub fn get(&self, key: K) -> Option<&V>{
        self.root.get(key)
    }

    pub fn get_mut(&mut self, key: K) -> Option<&mut V> {
        self.root.get_mut(key)
    }

    pub fn len(&self) -> usize {
        self.root.len()
    }

    /*
    pub fn iter(&self) -> impl Iterator<Item=(K, &V)> {
        Iter::new(&self.root)
    }

    pub fn iter_mut(&self) -> impl Iterator<Item=(K, &mut V)> {
        IterMut::new(&mut self.root)
    }
    */
}
<commit_msg>Add tests to RudyMap for sizes 1 and 2<commit_after>use ::Key;
use super::rootptr::RootPtr;

pub struct RudyMap<K: Key, V> {
    root: RootPtr<K, V>
}

impl<K: Key, V> RudyMap<K, V> {
    pub fn new() -> RudyMap<K, V> {
        RudyMap {
            root: RootPtr::empty()
        }
    }

    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
        self.root.insert(key, value)
    }

    /*
    fn remove(&mut self, key: K) -> Option<V> {
        self.root.remove(key)
    }
    */

    /*
    fn clear(&mut self) {
        self.root.clear()
    }
    */

    pub fn contains_key(&self, key: K) -> bool {
        // Because this does not visit the pointer returned by get(), there
        // will not be unnecessary cache fills
        self.get(key).is_some()
    }

    pub fn get(&self, key: K) -> Option<&V>{
        self.root.get(key)
    }

    pub fn get_mut(&mut self, key: K) -> Option<&mut V> {
        self.root.get_mut(key)
    }

    pub fn len(&self) -> usize {
        self.root.len()
    }

    /*
    pub fn iter(&self) -> impl Iterator<Item=(K, &V)> {
        Iter::new(&self.root)
    }

    pub fn iter_mut(&self) -> impl Iterator<Item=(K, &mut V)> {
        IterMut::new(&mut self.root)
    }
    */
}

#[test]
fn test_insert_get_1() {
    let mut map = RudyMap::<usize, usize>::new();
    map.insert(4usize, 10usize);
    assert_eq!(map.get(4), Some(&10));
}

#[test]
fn test_insert_get_2() {
    let mut map = RudyMap::<usize, usize>::new();
    map.insert(0usize, 10usize);
    map.insert(1usize, 20usize);
    //assert_eq!(map.get(0), Some(&10));
    //assert_eq!(map.get(1), Some(&20));
}
<|endoftext|>


<commit_before>// Copyright 2014-2015 Sam Doshi (sam@metal-fish.co.uk)
// Copyright 2013-2014 Philippe Delrieu (philippe.delrieu@free.fr)
//
// Licensed under either of
//          Apache License, Version 2.0, (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0)
//          MIT license (LICENSE-MIT or http://opensource.org/licenses/MIT).
// This file may not be copied, modified, or distributed except according to those terms.

use std::os::raw::{c_char, c_void, c_int};
use ffi::types::*;

#[link(name = "portmidi")]
extern "C" {
    pub fn Pm_Initialize() -> PmError;
    pub fn Pm_Terminate() -> PmError;
    pub fn Pm_HasHostError(stream: *const PortMidiStream) -> c_int;
    pub fn Pm_GetErrorText(errorCode: PmError) -> *const c_char;
    pub fn Pm_GetHostErrorText(msg: *const c_char, len: c_int);
    pub fn Pm_CountDevices() -> c_int;
    pub fn Pm_GetDefaultInputDeviceID() -> PmDeviceId;
    pub fn Pm_GetDefaultOutputDeviceID() -> PmDeviceId;
    pub fn Pm_GetDeviceInfo(id: PmDeviceId) -> *const PmDeviceInfo;
    pub fn Pm_OpenInput(stream: *const *const PortMidiStream,
                        inputDevice: PmDeviceId,
                        inputDriverInfo: Option<*const c_void>,
                        bufferSize: c_int,
                        time_proc: Option<*const c_void>,
                        time_info: Option<*const c_void>)
                        -> PmError;
    pub fn Pm_OpenOutput(stream: *const *const PortMidiStream,
                         outputDevice: PmDeviceId,
                         inputDriverInfo: Option<*const c_void>,
                         bufferSize: c_int,
                         time_proc: Option<*const c_void>,
                         time_info: Option<*const c_void>,
                         latency: c_int)
                         -> PmError;
    pub fn Pm_Read(stream: *const PortMidiStream, buffer: *mut PmEvent, length: c_int) -> c_int;
    pub fn Pm_Abort(stream: *const PortMidiStream) -> PmError;
    pub fn Pm_Close(stream: *const PortMidiStream) -> PmError;
    pub fn Pm_Poll(stream: *const PortMidiStream) -> PmError;
    pub fn Pm_Write(stream: *const PortMidiStream, buffer: *const PmEvent, length: c_int) -> PmError;
    pub fn Pm_WriteShort(stream: *const PortMidiStream,
                         timestamp: PmTimestamp,
                         message: PmMessage)
                         -> PmError;
}
<commit_msg>Revert "Make optional parameters (can be NULL) to `Option`s"<commit_after>// Copyright 2014-2015 Sam Doshi (sam@metal-fish.co.uk)
// Copyright 2013-2014 Philippe Delrieu (philippe.delrieu@free.fr)
//
// Licensed under either of
//          Apache License, Version 2.0, (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0)
//          MIT license (LICENSE-MIT or http://opensource.org/licenses/MIT).
// This file may not be copied, modified, or distributed except according to those terms.

use std::os::raw::{c_char, c_void, c_int};
use ffi::types::*;

#[link(name = "portmidi")]
extern "C" {
    pub fn Pm_Initialize() -> PmError;
    pub fn Pm_Terminate() -> PmError;
    pub fn Pm_HasHostError(stream: *const PortMidiStream) -> c_int;
    pub fn Pm_GetErrorText(errorCode: PmError) -> *const c_char;
    pub fn Pm_GetHostErrorText(msg: *const c_char, len: c_int);
    pub fn Pm_CountDevices() -> c_int;
    pub fn Pm_GetDefaultInputDeviceID() -> PmDeviceId;
    pub fn Pm_GetDefaultOutputDeviceID() -> PmDeviceId;
    pub fn Pm_GetDeviceInfo(id: PmDeviceId) -> *const PmDeviceInfo;
    pub fn Pm_OpenInput(stream: *const *const PortMidiStream,
                        inputDevice: PmDeviceId,
                        inputDriverInfo: *const c_void,
                        bufferSize: i32,
                        time_proc: *const c_void,
                        time_info: *const c_void)
                        -> PmError;
    pub fn Pm_OpenOutput(stream: *const *const PortMidiStream,
                         outputDevice: PmDeviceId,
                         inputDriverInfo: *const c_void,
                         bufferSize: i32,
                         time_proc: *const c_void,
                         time_info: *const c_void,
                         latency: i32)
                         -> PmError;
    pub fn Pm_Read(stream: *const PortMidiStream, buffer: *mut PmEvent, length: c_int) -> c_int;
    pub fn Pm_Abort(stream: *const PortMidiStream) -> PmError;
    pub fn Pm_Close(stream: *const PortMidiStream) -> PmError;
    pub fn Pm_Poll(stream: *const PortMidiStream) -> PmError;
    pub fn Pm_Write(stream: *const PortMidiStream, buffer: *const PmEvent, length: c_int) -> PmError;
    pub fn Pm_WriteShort(stream: *const PortMidiStream,
                         timestamp: PmTimestamp,
                         message: PmMessage)
                         -> PmError;
}
<|endoftext|>


<commit_before>use std::env;
use std::ffi::{CString, OsStr};
use std::fs::{self, File, OpenOptions};
use std::io;
cfg_if! {
    if #[cfg(not(target_os = "wasi"))] {
        use std::os::unix::ffi::OsStrExt;
        use std::os::unix::fs::{MetadataExt, OpenOptionsExt};
    } else {
        use std::os::wasi::ffi::OsStrExt;
        #[cfg(feature = "nightly")]
        use std::os::wasi::fs::MetadataExt;
    }
}
use crate::util;
use std::path::Path;

#[cfg(not(target_os = "redox"))]
use libc::{c_char, c_int, link, rename, unlink};

#[cfg(not(target_os = "redox"))]
#[inline(always)]
pub fn cvt_err(result: c_int) -> io::Result<c_int> {
    if result == -1 {
        Err(io::Error::last_os_error())
    } else {
        Ok(result)
    }
}

#[cfg(target_os = "redox")]
#[inline(always)]
pub fn cvt_err(result: Result<usize, syscall::Error>) -> io::Result<usize> {
    result.map_err(|err| io::Error::from_raw_os_error(err.errno))
}

// Stolen from std.
pub fn cstr(path: &Path) -> io::Result<CString> {
    CString::new(path.as_os_str().as_bytes())
        .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "path contained a null"))
}

pub fn create_named(path: &Path, open_options: &mut OpenOptions) -> io::Result<File> {
    open_options.read(true).write(true).create_new(true);

    #[cfg(not(target_os = "wasi"))]
    {
        open_options.mode(0o600);
    }

    open_options.open(path)
}

fn create_unlinked(path: &Path) -> io::Result<File> {
    let tmp;
    // shadow this to decrease the lifetime. It can't live longer than `tmp`.
    let mut path = path;
    if !path.is_absolute() {
        let cur_dir = env::current_dir()?;
        tmp = cur_dir.join(path);
        path = &tmp;
    }

    let f = create_named(path, &mut OpenOptions::new())?;
    // don't care whether the path has already been unlinked,
    // but perhaps there are some IO error conditions we should send up?
    let _ = fs::remove_file(path);
    Ok(f)
}

#[cfg(target_os = "linux")]
pub fn create(dir: &Path) -> io::Result<File> {
    use libc::{EISDIR, ENOENT, EOPNOTSUPP, O_EXCL, O_TMPFILE};
    OpenOptions::new()
        .read(true)
        .write(true)
        .custom_flags(O_TMPFILE | O_EXCL) // do not mix with `create_new(true)`
        .open(dir)
        .or_else(|e| {
            match e.raw_os_error() {
                // These are the three "not supported" error codes for O_TMPFILE.
                Some(EOPNOTSUPP) | Some(EISDIR) | Some(ENOENT) => create_unix(dir),
                _ => Err(e),
            }
        })
}

#[cfg(not(target_os = "linux"))]
pub fn create(dir: &Path) -> io::Result<File> {
    create_unix(dir)
}

fn create_unix(dir: &Path) -> io::Result<File> {
    util::create_helper(
        dir,
        OsStr::new(".tmp"),
        OsStr::new(""),
        crate::NUM_RAND_CHARS,
        |path| create_unlinked(&path),
    )
}

#[cfg(any(not(target_os = "wasi"), feature = "nightly"))]
pub fn reopen(file: &File, path: &Path) -> io::Result<File> {
    let new_file = OpenOptions::new().read(true).write(true).open(path)?;
    let old_meta = file.metadata()?;
    let new_meta = new_file.metadata()?;
    if old_meta.dev() != new_meta.dev() || old_meta.ino() != new_meta.ino() {
        return Err(io::Error::new(
            io::ErrorKind::NotFound,
            "original tempfile has been replaced",
        ));
    }
    Ok(new_file)
}

#[cfg(all(target_os = "wasi", not(feature = "nightly")))]
pub fn reopen(file: &File, path: &Path) -> io::Result<File> {
    return Err(io::Error::new(
        io::ErrorKind::Other,
        "this operation is supported on WASI only on nightly Rust (with `nightly` feature enabled)",
    ));
}

#[cfg(not(target_os = "redox"))]
pub fn persist(old_path: &Path, new_path: &Path, overwrite: bool) -> io::Result<()> {
    unsafe {
        let old_path = cstr(old_path)?;
        let new_path = cstr(new_path)?;
        if overwrite {
            cvt_err(rename(
                old_path.as_ptr() as *const c_char,
                new_path.as_ptr() as *const c_char,
            ))?;
        } else {
            cvt_err(link(
                old_path.as_ptr() as *const c_char,
                new_path.as_ptr() as *const c_char,
            ))?;
            // Ignore unlink errors. Can we do better?
            // On recent linux, we can use renameat2 to do this atomically.
            let _ = unlink(old_path.as_ptr() as *const c_char);
        }
        Ok(())
    }
}

#[cfg(target_os = "redox")]
pub fn persist(old_path: &Path, new_path: &Path, overwrite: bool) -> io::Result<()> {
    // XXX implement when possible
    Err(io::Error::from_raw_os_error(syscall::ENOSYS))
}

pub fn keep(_: &Path) -> io::Result<()> {
    Ok(())
}
<commit_msg>Fixup for unused vars<commit_after>use std::env;
use std::ffi::{CString, OsStr};
use std::fs::{self, File, OpenOptions};
use std::io;
cfg_if! {
    if #[cfg(not(target_os = "wasi"))] {
        use std::os::unix::ffi::OsStrExt;
        use std::os::unix::fs::{MetadataExt, OpenOptionsExt};
    } else {
        use std::os::wasi::ffi::OsStrExt;
        #[cfg(feature = "nightly")]
        use std::os::wasi::fs::MetadataExt;
    }
}
use crate::util;
use std::path::Path;

#[cfg(not(target_os = "redox"))]
use libc::{c_char, c_int, link, rename, unlink};

#[cfg(not(target_os = "redox"))]
#[inline(always)]
pub fn cvt_err(result: c_int) -> io::Result<c_int> {
    if result == -1 {
        Err(io::Error::last_os_error())
    } else {
        Ok(result)
    }
}

#[cfg(target_os = "redox")]
#[inline(always)]
pub fn cvt_err(result: Result<usize, syscall::Error>) -> io::Result<usize> {
    result.map_err(|err| io::Error::from_raw_os_error(err.errno))
}

// Stolen from std.
pub fn cstr(path: &Path) -> io::Result<CString> {
    CString::new(path.as_os_str().as_bytes())
        .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "path contained a null"))
}

pub fn create_named(path: &Path, open_options: &mut OpenOptions) -> io::Result<File> {
    open_options.read(true).write(true).create_new(true);

    #[cfg(not(target_os = "wasi"))]
    {
        open_options.mode(0o600);
    }

    open_options.open(path)
}

fn create_unlinked(path: &Path) -> io::Result<File> {
    let tmp;
    // shadow this to decrease the lifetime. It can't live longer than `tmp`.
    let mut path = path;
    if !path.is_absolute() {
        let cur_dir = env::current_dir()?;
        tmp = cur_dir.join(path);
        path = &tmp;
    }

    let f = create_named(path, &mut OpenOptions::new())?;
    // don't care whether the path has already been unlinked,
    // but perhaps there are some IO error conditions we should send up?
    let _ = fs::remove_file(path);
    Ok(f)
}

#[cfg(target_os = "linux")]
pub fn create(dir: &Path) -> io::Result<File> {
    use libc::{EISDIR, ENOENT, EOPNOTSUPP, O_EXCL, O_TMPFILE};
    OpenOptions::new()
        .read(true)
        .write(true)
        .custom_flags(O_TMPFILE | O_EXCL) // do not mix with `create_new(true)`
        .open(dir)
        .or_else(|e| {
            match e.raw_os_error() {
                // These are the three "not supported" error codes for O_TMPFILE.
                Some(EOPNOTSUPP) | Some(EISDIR) | Some(ENOENT) => create_unix(dir),
                _ => Err(e),
            }
        })
}

#[cfg(not(target_os = "linux"))]
pub fn create(dir: &Path) -> io::Result<File> {
    create_unix(dir)
}

fn create_unix(dir: &Path) -> io::Result<File> {
    util::create_helper(
        dir,
        OsStr::new(".tmp"),
        OsStr::new(""),
        crate::NUM_RAND_CHARS,
        |path| create_unlinked(&path),
    )
}

#[cfg(any(not(target_os = "wasi"), feature = "nightly"))]
pub fn reopen(file: &File, path: &Path) -> io::Result<File> {
    let new_file = OpenOptions::new().read(true).write(true).open(path)?;
    let old_meta = file.metadata()?;
    let new_meta = new_file.metadata()?;
    if old_meta.dev() != new_meta.dev() || old_meta.ino() != new_meta.ino() {
        return Err(io::Error::new(
            io::ErrorKind::NotFound,
            "original tempfile has been replaced",
        ));
    }
    Ok(new_file)
}

#[cfg(all(target_os = "wasi", not(feature = "nightly")))]
pub fn reopen(_file: &File, _path: &Path) -> io::Result<File> {
    return Err(io::Error::new(
        io::ErrorKind::Other,
        "this operation is supported on WASI only on nightly Rust (with `nightly` feature enabled)",
    ));
}

#[cfg(not(target_os = "redox"))]
pub fn persist(old_path: &Path, new_path: &Path, overwrite: bool) -> io::Result<()> {
    unsafe {
        let old_path = cstr(old_path)?;
        let new_path = cstr(new_path)?;
        if overwrite {
            cvt_err(rename(
                old_path.as_ptr() as *const c_char,
                new_path.as_ptr() as *const c_char,
            ))?;
        } else {
            cvt_err(link(
                old_path.as_ptr() as *const c_char,
                new_path.as_ptr() as *const c_char,
            ))?;
            // Ignore unlink errors. Can we do better?
            // On recent linux, we can use renameat2 to do this atomically.
            let _ = unlink(old_path.as_ptr() as *const c_char);
        }
        Ok(())
    }
}

#[cfg(target_os = "redox")]
pub fn persist(old_path: &Path, new_path: &Path, overwrite: bool) -> io::Result<()> {
    // XXX implement when possible
    Err(io::Error::from_raw_os_error(syscall::ENOSYS))
}

pub fn keep(_: &Path) -> io::Result<()> {
    Ok(())
}
<|endoftext|>


<commit_before>use std::borrow::Cow;
use std::io::Write;

use httparse;

use header::{self, Headers, ContentLength, TransferEncoding};
use http::{MessageHead, RawStatus, Http1Message, ParseResult, ServerMessage, ClientMessage, RequestLine};
use http::h1::{Encoder, Decoder};
use method::Method;
use status::StatusCode;
use version::HttpVersion::{Http10, Http11};

const MAX_HEADERS: usize = 100;
const AVERAGE_HEADER_SIZE: usize = 30; // totally scientific

pub fn parse<T: Http1Message<Incoming=I>, I>(buf: &[u8]) -> ParseResult<I> {
    if buf.len() == 0 {
        return Ok(None);
    }
    trace!("parse({:?})", buf);
    <T as Http1Message>::parse(buf)
}

impl Http1Message for ServerMessage {
    type Incoming = RequestLine;
    type Outgoing = StatusCode;

    fn parse(buf: &[u8]) -> ParseResult<RequestLine> {
        let mut headers = [httparse::EMPTY_HEADER; MAX_HEADERS];
        trace!("Request.parse([Header; {}], [u8; {}])", headers.len(), buf.len());
        let mut req = httparse::Request::new(&mut headers);
        Ok(match try!(req.parse(buf)) {
            httparse::Status::Complete(len) => {
                trace!("Request.parse Complete({})", len);
                Some((MessageHead {
                    version: if req.version.unwrap() == 1 { Http11 } else { Http10 },
                    subject: RequestLine(
                        try!(req.method.unwrap().parse()),
                        try!(req.path.unwrap().parse())
                    ),
                    headers: try!(Headers::from_raw(req.headers))
                }, len))
            },
            httparse::Status::Partial => None
        })
    }

    fn decoder(head: &MessageHead<Self::Incoming>) -> ::Result<Decoder> {
        use ::header;
        if let Some(&header::ContentLength(len)) = head.headers.get() {
            Ok(Decoder::length(len))
        } else if head.headers.has::<header::TransferEncoding>() {
            //TODO: check for Transfer-Encoding: chunked
            Ok(Decoder::chunked())
        } else {
            Ok(Decoder::length(0))
        }
    }

    fn encode(mut head: MessageHead<Self::Outgoing>, dst: &mut Vec<u8>) -> Encoder {
        use ::header;
        trace!("writing head: {:?}", head);

        if !head.headers.has::<header::Date>() {
            head.headers.set(header::Date(header::HttpDate(::time::now_utc())));
        }

        let mut is_chunked = true;
        let mut body = Encoder::chunked();
        if let Some(cl) = head.headers.get::<header::ContentLength>() {
            body = Encoder::length(**cl);
            is_chunked = false
        }

        if is_chunked {
            let encodings = match head.headers.get_mut::<header::TransferEncoding>() {
                Some(&mut header::TransferEncoding(ref mut encodings)) => {
                    if encodings.last() != Some(&header::Encoding::Chunked) {
                        encodings.push(header::Encoding::Chunked);
                    }
                    false
                },
                None => true
            };

            if encodings {
                head.headers.set(header::TransferEncoding(vec![header::Encoding::Chunked]));
            }
        }

        let init_cap = 30 + head.headers.len() * AVERAGE_HEADER_SIZE;
        dst.reserve(init_cap);
        debug!("writing {:#?}", head.headers);
        let _ = write!(dst, "{} {}\r\n{}\r\n", head.version, head.subject, head.headers);

        body
    }
}

impl Http1Message for ClientMessage {
    type Incoming = RawStatus;
    type Outgoing = RequestLine;

    fn parse(buf: &[u8]) -> ParseResult<RawStatus> {
        let mut headers = [httparse::EMPTY_HEADER; MAX_HEADERS];
        trace!("Response.parse([Header; {}], [u8; {}])", headers.len(), buf.len());
        let mut res = httparse::Response::new(&mut headers);
        Ok(match try!(res.parse(buf)) {
            httparse::Status::Complete(len) => {
                trace!("Response.try_parse Complete({})", len);
                let code = res.code.unwrap();
                let reason = match StatusCode::from_u16(code).canonical_reason() {
                    Some(reason) if reason == res.reason.unwrap() => Cow::Borrowed(reason),
                    _ => Cow::Owned(res.reason.unwrap().to_owned())
                };
                Some((MessageHead {
                    version: if res.version.unwrap() == 1 { Http11 } else { Http10 },
                    subject: RawStatus(code, reason),
                    headers: try!(Headers::from_raw(res.headers))
                }, len))
            },
            httparse::Status::Partial => None
        })
    }

    fn decoder(inc: &MessageHead<Self::Incoming>) -> ::Result<Decoder> {
        use ::header;
        // According to https://tools.ietf.org/html/rfc7230#section-3.3.3
        // 1. HEAD reponses, and Status 1xx, 204, and 304 cannot have a body.
        // 2. Status 2xx to a CONNECT cannot have a body.
        //
        // First two steps taken care of before this method.
        //
        // 3. Transfer-Encoding: chunked has a chunked body.
        // 4. If multiple differing Content-Length headers or invalid, close connection.
        // 5. Content-Length header has a sized body.
        // 6. Not Client.
        // 7. Read till EOF.
        if let Some(&header::TransferEncoding(ref codings)) = inc.headers.get() {
            if codings.last() == Some(&header::Encoding::Chunked) {
                Ok(Decoder::chunked())
            } else {
                trace!("not chuncked. read till eof");
                Ok(Decoder::eof())
            }
        } else if let Some(&header::ContentLength(len)) = inc.headers.get() {
            Ok(Decoder::length(len))
        } else if inc.headers.has::<header::ContentLength>() {
            trace!("illegal Content-Length: {:?}", inc.headers.get_raw("Content-Length"));
            Err(::Error::Header)
        } else {
            trace!("neither Transfer-Encoding nor Content-Length");
            Ok(Decoder::eof())
        }
    }

    fn encode(mut head: MessageHead<Self::Outgoing>, dst: &mut Vec<u8>) -> Encoder {
        trace!("writing head: {:?}", head);

        let mut body = Encoder::length(0);
        let expects_no_body = match head.subject.0 {
            Method::Head | Method::Get | Method::Connect => true,
            _ => false
        };
        let mut chunked = false;

        if let Some(con_len) = head.headers.get::<ContentLength>() {
            body = Encoder::length(**con_len);
        } else {
            chunked = !expects_no_body;
        }

        if chunked {
            body = Encoder::chunked();
            let encodings = match head.headers.get_mut::<TransferEncoding>() {
                Some(encodings) => {
                    if !encodings.contains(&header::Encoding::Chunked) {
                        encodings.push(header::Encoding::Chunked);
                    }
                    true
                },
                None => false
            };

            if !encodings {
                head.headers.set(TransferEncoding(vec![header::Encoding::Chunked]));
            }
        }

        let init_cap = 30 + head.headers.len() * AVERAGE_HEADER_SIZE;
        dst.reserve(init_cap);
        debug!("writing {:#?}", head.headers);
        let _ = write!(dst, "{} {}\r\n{}\r\n", head.subject, head.version, head.headers);

        body
    }
}

#[cfg(test)]
mod tests {
    use http;
    use super::{parse};

    #[test]
    fn test_parse_request() {
        let raw = b"GET /echo HTTP/1.1\r\nHost: hyper.rs\r\n\r\n";
        parse::<http::ServerMessage, _>(raw).unwrap();
    }

    #[test]
    fn test_parse_raw_status() {
        let raw = b"HTTP/1.1 200 OK\r\n\r\n";
        let (res, _) = parse::<http::ClientMessage, _>(raw).unwrap().unwrap();
        assert_eq!(res.subject.1, "OK");

        let raw = b"HTTP/1.1 200 Howdy\r\n\r\n";
        let (res, _) = parse::<http::ClientMessage, _>(raw).unwrap().unwrap();
        assert_eq!(res.subject.1, "Howdy");
    }

    #[cfg(feature = "nightly")]
    use test::Bencher;

    #[cfg(feature = "nightly")]
    #[bench]
    fn bench_parse_incoming(b: &mut Bencher) {
        let raw = b"GET /echo HTTP/1.1\r\nHost: hyper.rs\r\n\r\n";
        b.iter(|| {
            parse::<http::ServerMessage, _>(raw).unwrap()
        });
    }

}
<commit_msg>perf(http): encoder headers with a faster  fmt::Write<commit_after>use std::borrow::Cow;
use std::fmt::{self, Write};

use httparse;

use header::{self, Headers, ContentLength, TransferEncoding};
use http::{MessageHead, RawStatus, Http1Message, ParseResult, ServerMessage, ClientMessage, RequestLine};
use http::h1::{Encoder, Decoder};
use method::Method;
use status::StatusCode;
use version::HttpVersion::{Http10, Http11};

const MAX_HEADERS: usize = 100;
const AVERAGE_HEADER_SIZE: usize = 30; // totally scientific

pub fn parse<T: Http1Message<Incoming=I>, I>(buf: &[u8]) -> ParseResult<I> {
    if buf.len() == 0 {
        return Ok(None);
    }
    trace!("parse({:?})", buf);
    <T as Http1Message>::parse(buf)
}

impl Http1Message for ServerMessage {
    type Incoming = RequestLine;
    type Outgoing = StatusCode;

    fn parse(buf: &[u8]) -> ParseResult<RequestLine> {
        let mut headers = [httparse::EMPTY_HEADER; MAX_HEADERS];
        trace!("Request.parse([Header; {}], [u8; {}])", headers.len(), buf.len());
        let mut req = httparse::Request::new(&mut headers);
        Ok(match try!(req.parse(buf)) {
            httparse::Status::Complete(len) => {
                trace!("Request.parse Complete({})", len);
                Some((MessageHead {
                    version: if req.version.unwrap() == 1 { Http11 } else { Http10 },
                    subject: RequestLine(
                        try!(req.method.unwrap().parse()),
                        try!(req.path.unwrap().parse())
                    ),
                    headers: try!(Headers::from_raw(req.headers))
                }, len))
            },
            httparse::Status::Partial => None
        })
    }

    fn decoder(head: &MessageHead<Self::Incoming>) -> ::Result<Decoder> {
        use ::header;
        if let Some(&header::ContentLength(len)) = head.headers.get() {
            Ok(Decoder::length(len))
        } else if head.headers.has::<header::TransferEncoding>() {
            //TODO: check for Transfer-Encoding: chunked
            Ok(Decoder::chunked())
        } else {
            Ok(Decoder::length(0))
        }
    }

    fn encode(mut head: MessageHead<Self::Outgoing>, dst: &mut Vec<u8>) -> Encoder {
        use ::header;
        trace!("writing head: {:?}", head);

        if !head.headers.has::<header::Date>() {
            head.headers.set(header::Date(header::HttpDate(::time::now_utc())));
        }

        let mut is_chunked = true;
        let mut body = Encoder::chunked();
        if let Some(cl) = head.headers.get::<header::ContentLength>() {
            body = Encoder::length(**cl);
            is_chunked = false
        }

        if is_chunked {
            let encodings = match head.headers.get_mut::<header::TransferEncoding>() {
                Some(&mut header::TransferEncoding(ref mut encodings)) => {
                    if encodings.last() != Some(&header::Encoding::Chunked) {
                        encodings.push(header::Encoding::Chunked);
                    }
                    false
                },
                None => true
            };

            if encodings {
                head.headers.set(header::TransferEncoding(vec![header::Encoding::Chunked]));
            }
        }

        let init_cap = 30 + head.headers.len() * AVERAGE_HEADER_SIZE;
        dst.reserve(init_cap);
        debug!("writing {:#?}", head.headers);
        if head.version == ::HttpVersion::Http11 && head.subject == ::StatusCode::Ok {
            extend(dst, b"HTTP/1.1 200 OK\r\n");
            let _ = write!(FastWrite(dst), "{}\r\n", head.headers);
        } else {
            let _ = write!(FastWrite(dst), "{} {}\r\n{}\r\n", head.version, head.subject, head.headers);
        }
        body
    }
}

impl Http1Message for ClientMessage {
    type Incoming = RawStatus;
    type Outgoing = RequestLine;

    fn parse(buf: &[u8]) -> ParseResult<RawStatus> {
        let mut headers = [httparse::EMPTY_HEADER; MAX_HEADERS];
        trace!("Response.parse([Header; {}], [u8; {}])", headers.len(), buf.len());
        let mut res = httparse::Response::new(&mut headers);
        Ok(match try!(res.parse(buf)) {
            httparse::Status::Complete(len) => {
                trace!("Response.try_parse Complete({})", len);
                let code = res.code.unwrap();
                let reason = match StatusCode::from_u16(code).canonical_reason() {
                    Some(reason) if reason == res.reason.unwrap() => Cow::Borrowed(reason),
                    _ => Cow::Owned(res.reason.unwrap().to_owned())
                };
                Some((MessageHead {
                    version: if res.version.unwrap() == 1 { Http11 } else { Http10 },
                    subject: RawStatus(code, reason),
                    headers: try!(Headers::from_raw(res.headers))
                }, len))
            },
            httparse::Status::Partial => None
        })
    }

    fn decoder(inc: &MessageHead<Self::Incoming>) -> ::Result<Decoder> {
        use ::header;
        // According to https://tools.ietf.org/html/rfc7230#section-3.3.3
        // 1. HEAD reponses, and Status 1xx, 204, and 304 cannot have a body.
        // 2. Status 2xx to a CONNECT cannot have a body.
        //
        // First two steps taken care of before this method.
        //
        // 3. Transfer-Encoding: chunked has a chunked body.
        // 4. If multiple differing Content-Length headers or invalid, close connection.
        // 5. Content-Length header has a sized body.
        // 6. Not Client.
        // 7. Read till EOF.
        if let Some(&header::TransferEncoding(ref codings)) = inc.headers.get() {
            if codings.last() == Some(&header::Encoding::Chunked) {
                Ok(Decoder::chunked())
            } else {
                trace!("not chuncked. read till eof");
                Ok(Decoder::eof())
            }
        } else if let Some(&header::ContentLength(len)) = inc.headers.get() {
            Ok(Decoder::length(len))
        } else if inc.headers.has::<header::ContentLength>() {
            trace!("illegal Content-Length: {:?}", inc.headers.get_raw("Content-Length"));
            Err(::Error::Header)
        } else {
            trace!("neither Transfer-Encoding nor Content-Length");
            Ok(Decoder::eof())
        }
    }

    fn encode(mut head: MessageHead<Self::Outgoing>, dst: &mut Vec<u8>) -> Encoder {
        trace!("writing head: {:?}", head);

        let mut body = Encoder::length(0);
        let expects_no_body = match head.subject.0 {
            Method::Head | Method::Get | Method::Connect => true,
            _ => false
        };
        let mut chunked = false;

        if let Some(con_len) = head.headers.get::<ContentLength>() {
            body = Encoder::length(**con_len);
        } else {
            chunked = !expects_no_body;
        }

        if chunked {
            body = Encoder::chunked();
            let encodings = match head.headers.get_mut::<TransferEncoding>() {
                Some(encodings) => {
                    if !encodings.contains(&header::Encoding::Chunked) {
                        encodings.push(header::Encoding::Chunked);
                    }
                    true
                },
                None => false
            };

            if !encodings {
                head.headers.set(TransferEncoding(vec![header::Encoding::Chunked]));
            }
        }

        let init_cap = 30 + head.headers.len() * AVERAGE_HEADER_SIZE;
        dst.reserve(init_cap);
        debug!("writing {:#?}", head.headers);
        let _ = write!(FastWrite(dst), "{} {}\r\n{}\r\n", head.subject, head.version, head.headers);

        body
    }
}

struct FastWrite<'a>(&'a mut Vec<u8>);

impl<'a> fmt::Write for FastWrite<'a> {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        extend(self.0, s.as_bytes());
        Ok(())
    }

    fn write_fmt(&mut self, args: fmt::Arguments) -> fmt::Result {
        fmt::write(self, args)
    }
}

fn extend(dst: &mut Vec<u8>, data: &[u8]) {
    use std::ptr;
    dst.reserve(data.len());
    let prev = dst.len();
    unsafe {
        ptr::copy_nonoverlapping(data.as_ptr(),
                                 dst.as_mut_ptr().offset(prev as isize),
                                 data.len());
        dst.set_len(prev + data.len());
    }
}

#[cfg(test)]
mod tests {
    use http;
    use super::{parse};

    #[test]
    fn test_parse_request() {
        let raw = b"GET /echo HTTP/1.1\r\nHost: hyper.rs\r\n\r\n";
        parse::<http::ServerMessage, _>(raw).unwrap();
    }

    #[test]
    fn test_parse_raw_status() {
        let raw = b"HTTP/1.1 200 OK\r\n\r\n";
        let (res, _) = parse::<http::ClientMessage, _>(raw).unwrap().unwrap();
        assert_eq!(res.subject.1, "OK");

        let raw = b"HTTP/1.1 200 Howdy\r\n\r\n";
        let (res, _) = parse::<http::ClientMessage, _>(raw).unwrap().unwrap();
        assert_eq!(res.subject.1, "Howdy");
    }

    #[cfg(feature = "nightly")]
    use test::Bencher;

    #[cfg(feature = "nightly")]
    #[bench]
    fn bench_parse_incoming(b: &mut Bencher) {
        let raw = b"GET /echo HTTP/1.1\r\nHost: hyper.rs\r\n\r\n";
        b.iter(|| {
            parse::<http::ServerMessage, _>(raw).unwrap()
        });
    }

}
<|endoftext|>