<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::MoveStack;
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,
    },
    Done {
        search_id: usize,
        searched_nodes: NodeCount,
        value: Option<Value>,
    },
}

pub fn run_search(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 = pending_command.take()
                                         .unwrap_or(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,
                                                  })
                                                  .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,
                           })
                           .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,
            }
        }
    })
}

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_search(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.
        match pending_command.take().unwrap_or(commands.recv().unwrap()) {
            Command::Search { search_id, position, depth, lower_bound, upper_bound } => {
                let mut searched_nodes_final = 0;
                let mut value_final = None;
                // TODO: May be use MIN_PLY..(depth+1).
                '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,
                                       })
                                       .unwrap();
                                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_none() {
                                    break;
                                } else {
                                    break 'depthloop;
                                }
                            }
                        }
                    }
                }
                reports.send(Report::Done {
                           search_id: search_id,
                           searched_nodes: searched_nodes_final,
                           value: value_final,
                       })
                       .unwrap();
            }
            Command::Stop => {
                slave_commands_tx.send(Command::Stop).unwrap();
                continue;
            }
            Command::Exit => {
                slave_commands_tx.send(Command::Exit).unwrap();
                break;
            }
        }
    }
    slave.join().unwrap();
}

// 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);
    if depth == 0 {
        // On leaf nodes, do quiescence search.
        let (value, nodes) = p.evaluate_quiescence(alpha, beta);
        *nc += nodes;
        Ok(value)
    } else {
        moves.save();
        p.generate_moves(moves);
        // TODO: Consult `tt` here.

        let mut bound_type = BOUND_UPPER;
        let mut move16 = 0;
        let mut no_moves_yet = true;
        while let Some(m) = moves.remove_best_move() {
            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;
                    move16 = m.move16();
                    break;
                }
                if value > alpha {
                    // We found ourselves a new best move.
                    alpha = value;
                    bound_type = BOUND_EXACT;
                    move16 = m.move16();
                }
            }
        }
        moves.restore();
        if no_moves_yet {
            // No legal moves -- this is a final position.
            alpha = p.evaluate_final();
        }
        // TODO: Do not store values with too small `depth`.
        tt.store(p.hash(),
                 EntryData::new(alpha, bound_type, depth, move16, 0));
        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(()),
                           -20000,
                           20000,
                           2)
                        .ok()
                        .unwrap();
        assert!(value < -300);
        let value = search(&TranspositionTable::new(),
                           &mut p,
                           &mut MoveStack::new(),
                           &mut 0,
                           &mut |_| Ok(()),
                           -20000,
                           20000,
                           3)
                        .ok()
                        .unwrap();
        assert!(value >= 20000);
    }
}
<commit_msg>better handling of disconnections<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::MoveStack;
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,
    },
    Done {
        search_id: usize,
        searched_nodes: NodeCount,
        value: Option<Value>,
    },
}

pub fn run_search(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 = pending_command.take()
                                         .unwrap_or(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,
                                                  })
                                                  .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,
                           })
                           .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,
            }
        }
    })
}

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_search(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.
        match pending_command.take().unwrap_or(commands.recv()
                                                       .or::<RecvError>(Ok(Command::Exit))
                                                       .unwrap()) {
            Command::Search { search_id, position, depth, lower_bound, upper_bound } => {
                let mut searched_nodes_final = 0;
                let mut value_final = None;
                // TODO: May be use MIN_PLY..(depth+1).
                '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,
                                       })
                                       .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_none() {
                                    break;
                                } else {
                                    break 'depthloop;
                                }
                            }
                        }
                    }
                }
                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();
}

// 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);
    if depth == 0 {
        // On leaf nodes, do quiescence search.
        let (value, nodes) = p.evaluate_quiescence(alpha, beta);
        *nc += nodes;
        Ok(value)
    } else {
        moves.save();
        p.generate_moves(moves);
        // TODO: Consult `tt` here.

        let mut bound_type = BOUND_UPPER;
        let mut move16 = 0;
        let mut no_moves_yet = true;
        while let Some(m) = moves.remove_best_move() {
            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;
                    move16 = m.move16();
                    break;
                }
                if value > alpha {
                    // We found ourselves a new best move.
                    alpha = value;
                    bound_type = BOUND_EXACT;
                    move16 = m.move16();
                }
            }
        }
        moves.restore();
        if no_moves_yet {
            // No legal moves -- this is a final position.
            alpha = p.evaluate_final();
        }
        // TODO: Do not store values with too small `depth`.
        tt.store(p.hash(),
                 EntryData::new(alpha, bound_type, depth, move16, 0));
        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(()),
                           -20000,
                           20000,
                           2)
                        .ok()
                        .unwrap();
        assert!(value < -300);
        let value = search(&TranspositionTable::new(),
                           &mut p,
                           &mut MoveStack::new(),
                           &mut 0,
                           &mut |_| Ok(()),
                           -20000,
                           20000,
                           3)
                        .ok()
                        .unwrap();
        assert!(value >= 20000);
    }
}
<|endoftext|>


<commit_before>//! Unimplemented.
<commit_msg>Completed with R_ext/Parse.h with ext/parse.rs<commit_after>use super::super::internals::SEXP;
#[repr(C)]
#[derive(Copy, Clone, Debug)]
pub enum ParseStatus {
    PARSE_NULL = 0,
    PARSE_OK,
    PARSE_INCOMPLETE,
    PARSE_ERROR,
    PARSE_EOF,
}

extern {
    pub fn R_ParseVector(arg1: SEXP,
                         arg2: ::libc::c_int,
                         arg3: *mut ParseStatus,
                         arg4: SEXP)
                         -> SEXP;
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

use core::*;

#[derive(Debug, Clone)]
pub enum Error {
    ApplicationErr,
    UnificationErr(Term, Term),
    UnknownVariable(Name),
    MkErr
}

/// A global context for type checking containing the necessary information
/// needed across type checking all definitions.
pub struct TyCtxt {
    types: HashMap<Name, Data>,
    functions: HashMap<Name, Function>,
    globals: HashMap<Name, Term>
}

pub struct LocalCx<'tcx> {
    ty_cx: &'tcx TyCtxt,
    locals: HashMap<Name, Term>
}

impl TyCtxt {
    pub fn empty() -> TyCtxt {
        TyCtxt {
            types: HashMap::new(),
            functions: HashMap::new(),
            globals: HashMap::new(),
        }
    }

    pub fn from_module(module: &Module) -> TyCtxt {
        let mut tycx = TyCtxt::empty();
        for def in &module.defs {
            match def {
                &Definition::Data(ref d) => {
                    tycx.types.insert(d.name.clone(), d.clone());
                    // TODO: add each constructor here
                }
                &Definition::Fn(ref f) => {
                    tycx.functions.insert(f.name.clone(), f.clone());
                    tycx.globals.insert(f.name.clone(), f.ty());
                }
                &Definition::Extern(ref e) => {
                    tycx.globals.insert(e.0.clone(), e.1.clone());
                }
            }
        }

        return tycx;
    }

    pub fn type_check_def(&self, def: &Definition) -> Result<(), Error> {
        match def {
            &Definition::Fn(ref fun) => {
                let lcx = LocalCx::from_cx(self)
                                  .extend(fun.args.clone());
                try!(lcx.type_check_term(&fun.body, &fun.ty));
                Ok(())
            }
            _ => Ok(())
        }
    }
}

impl<'tcx> LocalCx<'tcx> {
    fn from_cx(ty_cx: &'tcx TyCtxt) -> LocalCx<'tcx> {
        LocalCx {
            ty_cx: ty_cx,
            locals: HashMap::new(),
        }
    }

    fn datatype(&self, ty: &Term) -> Option<&Data> {
        match ty {
            &Term::Var(ref n) => {
                self.ty_cx.types.get(n)
            }
            _ => None
        }
    }

    fn ctors(&self, ty: &Term) -> Option<Vec<(Name, Term)>> {
        self.datatype(ty).map(|dt| {
            dt.ctors
        })
    }

    fn extend(mut self, bindings: Vec<(Name, Term)>) -> LocalCx<'tcx> {
        for (x, t) in bindings {
            self.locals.insert(x, t);
        }

        self
    }

    fn lookup(&self, name: &Name) -> Result<&Term, Error> {
        match self.locals.get(name) {
            None => match self.ty_cx.globals.get(name) {
                None => Err(Error::UnknownVariable(name.clone())),
                Some(t) => Ok(t)
            },
            Some(t) => Ok(t)
        }
    }

    pub fn unify(&self, t: &Term, u: &Term) -> Result<Term, Error> {
        if t == u {
            Ok(t.clone())
        } else {
            Err(Error::UnificationErr(t.clone(), u.clone()))
        }
    }

    pub fn type_check_term(&self, term: &Term, ty: &Term) -> Result<Term, Error> {
        let infer_ty = try!(self.type_infer_term(term));
        self.unify(&infer_ty, ty)
        // match term {
        //     &Term::Literal(ref lit) => self.type_check_lit(lit, ty),
        //     &Term::Var(ref n) => self.unify(self.lookup(n), ty),
        //     &Term::Match(..) => panic!(),
        //     &Term::App(ref f, ref g) => {
        //         let g_ty = try!(self.type_infer_term(g));
        //         let f_ty = try!(self.type_infer_term(f));
        //         println!("typechecking app: {:?} {:?}", g_ty, f_ty);
        //         Err(Error::MkErr)
        //     }
        //     &Term::Forall(..) => panic!(),
        //     &Term::Lambda(..) => panic!(),
        //     &Term::Type => self.unify(&Term::Type, ty),
        // }
    }

    pub fn type_infer_term(&self, term: &Term) -> Result<Term, Error> {
        match term {
            &Term::Literal(ref lit) => match lit {
                &Literal::Int(..) => Ok(panic!()),
                &Literal::Unit => Ok(Term::Var("Unit".to_string())),
            },
            &Term::Var(ref n) => Ok(try!(self.lookup(n)).clone()),
            &Term::Match(ref scrut, ref cases) => {
                let scrut_ty = try!(self.type_infer_term(scrut));
                let ctors = self.ctors(&scrut_ty);

            }
            &Term::App(ref f, ref g) => {
                match try!(self.type_infer_term(f)) {
                    Term::Forall(n, arg_ty, body_ty) => {
                        try!(self.type_check_term(g, &*arg_ty));
                        // try!(self.evaluate())
                        Ok(*body_ty)
                    },
                    _ => Err(Error::ApplicationErr),
                }
            }
            &Term::Forall(..) => panic!(),
            &Term::Lambda(..) => panic!(),
            &Term::Type => Ok(Term::Type),
        }
    }

    pub fn evaluate(&self, term: &Term) -> Term {
        term.clone()
    }
}
<commit_msg>Get pattern matches to CPS phase<commit_after>use std::collections::HashMap;

use core::*;

#[derive(Debug, Clone)]
pub enum Error {
    ApplicationErr,
    UnificationErr(Term, Term),
    UnknownVariable(Name),
    MkErr
}

/// A global context for type checking containing the necessary information
/// needed across type checking all definitions.
pub struct TyCtxt {
    types: HashMap<Name, Data>,
    functions: HashMap<Name, Function>,
    globals: HashMap<Name, Term>
}

#[derive(Clone)]
pub struct LocalCx<'tcx> {
    ty_cx: &'tcx TyCtxt,
    locals: HashMap<Name, Term>
}

impl TyCtxt {
    pub fn empty() -> TyCtxt {
        TyCtxt {
            types: HashMap::new(),
            functions: HashMap::new(),
            globals: HashMap::new(),
        }
    }

    pub fn from_module(module: &Module) -> TyCtxt {
        let mut tycx = TyCtxt::empty();
        for def in &module.defs {
            match def {
                &Definition::Data(ref d) => {
                    tycx.types.insert(d.name.clone(), d.clone());
                    // TODO: add each constructor here
                }
                &Definition::Fn(ref f) => {
                    tycx.functions.insert(f.name.clone(), f.clone());
                    tycx.globals.insert(f.name.clone(), f.ty());
                }
                &Definition::Extern(ref e) => {
                    tycx.globals.insert(e.0.clone(), e.1.clone());
                }
            }
        }

        return tycx;
    }

    pub fn type_check_def(&self, def: &Definition) -> Result<(), Error> {
        match def {
            &Definition::Fn(ref fun) => {
                let lcx = LocalCx::from_cx(self)
                                  .extend(fun.args.clone());
                try!(lcx.type_check_term(&fun.body, &fun.ty));
                Ok(())
            }
            _ => Ok(())
        }
    }
}

impl<'tcx> LocalCx<'tcx> {
    fn from_cx(ty_cx: &'tcx TyCtxt) -> LocalCx<'tcx> {
        LocalCx {
            ty_cx: ty_cx,
            locals: HashMap::new(),
        }
    }

    fn datatype(&self, ty: &Term) -> Option<&Data> {
        match ty {
            &Term::Var(ref n) => {
                self.ty_cx.types.get(n)
            }
            _ => None
        }
    }

    fn ctors(&self, ty: &Term) -> Option<Vec<(Name, Term)>> {
        self.datatype(ty).map(|dt| {
            dt.ctors.clone()
        })
    }

    fn extend(mut self, bindings: Vec<(Name, Term)>) -> LocalCx<'tcx> {
        for (x, t) in bindings {
            self.locals.insert(x, t);
        }

        self
    }

    fn lookup(&self, name: &Name) -> Result<&Term, Error> {
        match self.locals.get(name) {
            None => match self.ty_cx.globals.get(name) {
                None => Err(Error::UnknownVariable(name.clone())),
                Some(t) => Ok(t)
            },
            Some(t) => Ok(t)
        }
    }

    pub fn unify(&self, t: &Term, u: &Term) -> Result<Term, Error> {
        if t == u {
            Ok(t.clone())
        } else {
            Err(Error::UnificationErr(t.clone(), u.clone()))
        }
    }

    pub fn type_check_term(&self, term: &Term, ty: &Term) -> Result<Term, Error> {
        let infer_ty = try!(self.type_infer_term(term));
        self.unify(&infer_ty, ty)
        // match term {
        //     &Term::Literal(ref lit) => self.type_check_lit(lit, ty),
        //     &Term::Var(ref n) => self.unify(self.lookup(n), ty),
        //     &Term::Match(..) => panic!(),
        //     &Term::App(ref f, ref g) => {
        //         let g_ty = try!(self.type_infer_term(g));
        //         let f_ty = try!(self.type_infer_term(f));
        //         println!("typechecking app: {:?} {:?}", g_ty, f_ty);
        //         Err(Error::MkErr)
        //     }
        //     &Term::Forall(..) => panic!(),
        //     &Term::Lambda(..) => panic!(),
        //     &Term::Type => self.unify(&Term::Type, ty),
        // }
    }

    pub fn type_infer_term(&self, term: &Term) -> Result<Term, Error> {
        match term {
            &Term::Literal(ref lit) => match lit {
                &Literal::Int(..) => Ok(panic!()),
                &Literal::Unit => Ok(Term::Var("Unit".to_string())),
            },
            &Term::Var(ref n) => Ok(try!(self.lookup(n)).clone()),
            &Term::Match(ref scrut, ref cases) => {
                let scrut_ty = try!(self.type_infer_term(scrut));
                let ctors: HashMap<_, _> = self.ctors(&scrut_ty).unwrap().into_iter().collect();
                let arm_types = cases.iter().map(|case| {
                    match &case.pattern {
                        &Pattern::Name(ref n) => {
                            let cx = self.clone();
                            let cx = cx.extend(vec![(n.clone(), scrut_ty.clone())]);
                            cx.type_infer_term(&case.rhs)
                        }
                        &Pattern::Constructor(ref n, ref patterns) => {
                            if patterns.len() == 0 {
                                self.type_infer_term(&case.rhs)
                            } else {
                                let ctor = ctors.get(n).unwrap();
                                self.bind_pattern(ctor, patterns, |cx| {
                                    cx.type_infer_term(&case.rhs)
                                })
                            }
                        }
                        &Pattern::Placeholder => {
                            self.type_infer_term(&case.rhs)
                        }
                    }
                }).collect::<Vec<_>>();
                arm_types[0].clone()
            }
            &Term::App(ref f, ref g) => {
                match try!(self.type_infer_term(f)) {
                    Term::Forall(n, arg_ty, body_ty) => {
                        try!(self.type_check_term(g, &*arg_ty));
                        // try!(self.evaluate())
                        Ok(*body_ty)
                    },
                    _ => Err(Error::ApplicationErr),
                }
            }
            &Term::Forall(..) => panic!(),
            &Term::Lambda(..) => panic!(),
            &Term::Type => Ok(Term::Type),
        }
    }

    pub fn bind_pattern<R, F : FnOnce(LocalCx) -> R>(
        &self,
        ty: &Term,
        patterns: &Vec<Pattern>,
        f: F) -> R {
        let mut ty = ty;
        let mut pat_start = 0;

        /* this is mostly definitely not right */
        match ty {
            &Term::Forall(ref n, ref dom, ref codom) => {
                let cx = self.clone();
                f(cx.extend(vec![(n.clone(), *dom.clone())]))
            },
            _ => panic!("invalid pattern binding")
        }
    }

    pub fn evaluate(&self, term: &Term) -> Term {
        term.clone()
    }
}
<|endoftext|>


<commit_before>#[cfg(test)]
pub mod aux {
    use crate::shapes::{GeoKind, Scratchpad};
    use crate::util::SQRT_2;

    // to test if variables are cleared before sum
    pub const NOISE: f64 = 1234.56;

    /// Generates pad Lin2 for tests
    pub fn gen_pad_lin2(l: f64) -> Scratchpad {
        let mut pad = Scratchpad::new(2, GeoKind::Lin2).unwrap();
        pad.set_xx(0, 0, 3.0);
        pad.set_xx(0, 1, 4.0);
        pad.set_xx(1, 0, 3.0 + l / SQRT_2);
        pad.set_xx(1, 1, 4.0 + l / SQRT_2);
        pad
    }

    /// Generates pad Tri3 for tests
    pub fn gen_pad_tri3() -> Scratchpad {
        let mut pad = Scratchpad::new(2, GeoKind::Tri3).unwrap();
        pad.set_xx(0, 0, 3.0);
        pad.set_xx(0, 1, 4.0);
        pad.set_xx(1, 0, 8.0);
        pad.set_xx(1, 1, 4.0);
        pad.set_xx(2, 0, 5.0);
        pad.set_xx(2, 1, 9.0);
        pad
    }

    /// Generates pad Qua4 for tests
    pub fn gen_pad_qua4(xc: f64, yc: f64, a: f64, b: f64) -> Scratchpad {
        let mut pad = Scratchpad::new(2, GeoKind::Qua4).unwrap();
        pad.set_xx(0, 0, xc - a);
        pad.set_xx(0, 1, yc - b);
        pad.set_xx(1, 0, xc + a);
        pad.set_xx(1, 1, yc - b);
        pad.set_xx(2, 0, xc + a);
        pad.set_xx(2, 1, yc + b);
        pad.set_xx(3, 0, xc - a);
        pad.set_xx(3, 1, yc + b);
        pad
    }

    /// Generates pad Qua8 for tests
    pub fn gen_pad_qua8(xc: f64, yc: f64, a: f64, b: f64) -> Scratchpad {
        let mut pad = Scratchpad::new(2, GeoKind::Qua8).unwrap();
        pad.set_xx(0, 0, xc - a);
        pad.set_xx(0, 1, yc - b);
        pad.set_xx(1, 0, xc + a);
        pad.set_xx(1, 1, yc - b);
        pad.set_xx(2, 0, xc + a);
        pad.set_xx(2, 1, yc + b);
        pad.set_xx(3, 0, xc - a);
        pad.set_xx(3, 1, yc + b);
        pad.set_xx(4, 0, xc);
        pad.set_xx(4, 1, yc - b);
        pad.set_xx(5, 0, xc + a);
        pad.set_xx(5, 1, yc);
        pad.set_xx(6, 0, xc);
        pad.set_xx(6, 1, yc + b);
        pad.set_xx(7, 0, xc - a);
        pad.set_xx(7, 1, yc);
        pad
    }

    /// Generates pad Tet4 for tests
    pub fn gen_pad_tet4() -> Scratchpad {
        let mut pad = Scratchpad::new(3, GeoKind::Tet4).unwrap();
        pad.set_xx(0, 0, 2.0);
        pad.set_xx(0, 1, 3.0);
        pad.set_xx(0, 2, 4.0);
        pad.set_xx(1, 0, 6.0);
        pad.set_xx(1, 1, 3.0);
        pad.set_xx(1, 2, 2.0);
        pad.set_xx(2, 0, 2.0);
        pad.set_xx(2, 1, 5.0);
        pad.set_xx(2, 2, 1.0);
        pad.set_xx(3, 0, 4.0);
        pad.set_xx(3, 1, 3.0);
        pad.set_xx(3, 2, 6.0);
        pad
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#[cfg(test)]
mod tests {
    use super::aux;
    use plotpy::Plot;

    #[test]
    fn gen_pad_functions_figs() {
        let pad = aux::gen_pad_lin2(1.0);
        if false {
            let mut plot = Plot::new();
            pad.draw_shape(&mut plot, "blue", true, true).unwrap();
            plot.set_equal_axes(true)
                .grid_and_labels("x", "y")
                .save("/tmp/gemlab/test_gen_pad_functions_lin2.svg")
                .unwrap();
        }

        let pad = aux::gen_pad_tri3();
        if false {
            let mut plot = Plot::new();
            pad.draw_shape(&mut plot, "blue", true, true).unwrap();
            plot.set_equal_axes(true)
                .grid_and_labels("x", "y")
                .save("/tmp/gemlab/test_gen_pad_functions_tri3.svg")
                .unwrap();
        }

        let pad = aux::gen_pad_qua4(2.0, 1.0, 2.0, 1.5);
        if false {
            let mut plot = Plot::new();
            pad.draw_shape(&mut plot, "blue", true, true).unwrap();
            plot.set_equal_axes(true)
                .grid_and_labels("x", "y")
                .save("/tmp/gemlab/test_gen_pad_functions_qua4.svg")
                .unwrap();
        }

        let pad = aux::gen_pad_qua8(2.0, 1.0, 2.0, 1.5);
        if false {
            let mut plot = Plot::new();
            pad.draw_shape(&mut plot, "blue", true, true).unwrap();
            plot.set_equal_axes(true)
                .grid_and_labels("x", "y")
                .save("/tmp/gemlab/test_gen_pad_functions_qua8.svg")
                .unwrap();
        }

        let pad = aux::gen_pad_tet4();
        if false {
            let mut plot = Plot::new();
            pad.draw_shape(&mut plot, "blue", true, true).unwrap();
            plot.set_equal_axes(true)
                .save("/tmp/gemlab/test_gen_pad_functions_tet4.svg")
                .unwrap();
        }
    }
}
<commit_msg>Disable unnecessary tests<commit_after>#[cfg(test)]
pub mod aux {
    use crate::shapes::{GeoKind, Scratchpad};
    use crate::util::SQRT_2;

    // to test if variables are cleared before sum
    pub const NOISE: f64 = 1234.56;

    /// Generates pad Lin2 for tests
    pub fn gen_pad_lin2(l: f64) -> Scratchpad {
        let mut pad = Scratchpad::new(2, GeoKind::Lin2).unwrap();
        pad.set_xx(0, 0, 3.0);
        pad.set_xx(0, 1, 4.0);
        pad.set_xx(1, 0, 3.0 + l / SQRT_2);
        pad.set_xx(1, 1, 4.0 + l / SQRT_2);
        pad
    }

    /// Generates pad Tri3 for tests
    pub fn gen_pad_tri3() -> Scratchpad {
        let mut pad = Scratchpad::new(2, GeoKind::Tri3).unwrap();
        pad.set_xx(0, 0, 3.0);
        pad.set_xx(0, 1, 4.0);
        pad.set_xx(1, 0, 8.0);
        pad.set_xx(1, 1, 4.0);
        pad.set_xx(2, 0, 5.0);
        pad.set_xx(2, 1, 9.0);
        pad
    }

    /// Generates pad Qua4 for tests
    pub fn gen_pad_qua4(xc: f64, yc: f64, a: f64, b: f64) -> Scratchpad {
        let mut pad = Scratchpad::new(2, GeoKind::Qua4).unwrap();
        pad.set_xx(0, 0, xc - a);
        pad.set_xx(0, 1, yc - b);
        pad.set_xx(1, 0, xc + a);
        pad.set_xx(1, 1, yc - b);
        pad.set_xx(2, 0, xc + a);
        pad.set_xx(2, 1, yc + b);
        pad.set_xx(3, 0, xc - a);
        pad.set_xx(3, 1, yc + b);
        pad
    }

    /// Generates pad Qua8 for tests
    pub fn gen_pad_qua8(xc: f64, yc: f64, a: f64, b: f64) -> Scratchpad {
        let mut pad = Scratchpad::new(2, GeoKind::Qua8).unwrap();
        pad.set_xx(0, 0, xc - a);
        pad.set_xx(0, 1, yc - b);
        pad.set_xx(1, 0, xc + a);
        pad.set_xx(1, 1, yc - b);
        pad.set_xx(2, 0, xc + a);
        pad.set_xx(2, 1, yc + b);
        pad.set_xx(3, 0, xc - a);
        pad.set_xx(3, 1, yc + b);
        pad.set_xx(4, 0, xc);
        pad.set_xx(4, 1, yc - b);
        pad.set_xx(5, 0, xc + a);
        pad.set_xx(5, 1, yc);
        pad.set_xx(6, 0, xc);
        pad.set_xx(6, 1, yc + b);
        pad.set_xx(7, 0, xc - a);
        pad.set_xx(7, 1, yc);
        pad
    }

    /// Generates pad Tet4 for tests
    pub fn gen_pad_tet4() -> Scratchpad {
        let mut pad = Scratchpad::new(3, GeoKind::Tet4).unwrap();
        pad.set_xx(0, 0, 2.0);
        pad.set_xx(0, 1, 3.0);
        pad.set_xx(0, 2, 4.0);
        pad.set_xx(1, 0, 6.0);
        pad.set_xx(1, 1, 3.0);
        pad.set_xx(1, 2, 2.0);
        pad.set_xx(2, 0, 2.0);
        pad.set_xx(2, 1, 5.0);
        pad.set_xx(2, 2, 1.0);
        pad.set_xx(3, 0, 4.0);
        pad.set_xx(3, 1, 3.0);
        pad.set_xx(3, 2, 6.0);
        pad
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

/*

//
// NOTE: we don't need the functions below to be activated and making the coverage tool to go crazy.
// We only use these functions to create the figures that are already saved in data/figures.
//
#[cfg(test)]
mod tests {
    use super::aux;
    use plotpy::Plot;

    #[test]
    fn gen_pad_functions_figs() {
        let pad = aux::gen_pad_lin2(1.0);
        if false {
            let mut plot = Plot::new();
            pad.draw_shape(&mut plot, "blue", true, true).unwrap();
            plot.set_equal_axes(true)
                .grid_and_labels("x", "y")
                .save("/tmp/gemlab/test_gen_pad_functions_lin2.svg")
                .unwrap();
        }

        let pad = aux::gen_pad_tri3();
        if false {
            let mut plot = Plot::new();
            pad.draw_shape(&mut plot, "blue", true, true).unwrap();
            plot.set_equal_axes(true)
                .grid_and_labels("x", "y")
                .save("/tmp/gemlab/test_gen_pad_functions_tri3.svg")
                .unwrap();
        }

        let pad = aux::gen_pad_qua4(2.0, 1.0, 2.0, 1.5);
        if false {
            let mut plot = Plot::new();
            pad.draw_shape(&mut plot, "blue", true, true).unwrap();
            plot.set_equal_axes(true)
                .grid_and_labels("x", "y")
                .save("/tmp/gemlab/test_gen_pad_functions_qua4.svg")
                .unwrap();
        }

        let pad = aux::gen_pad_qua8(2.0, 1.0, 2.0, 1.5);
        if false {
            let mut plot = Plot::new();
            pad.draw_shape(&mut plot, "blue", true, true).unwrap();
            plot.set_equal_axes(true)
                .grid_and_labels("x", "y")
                .save("/tmp/gemlab/test_gen_pad_functions_qua8.svg")
                .unwrap();
        }

        let pad = aux::gen_pad_tet4();
        if false {
            let mut plot = Plot::new();
            pad.draw_shape(&mut plot, "blue", true, true).unwrap();
            plot.set_equal_axes(true)
                .save("/tmp/gemlab/test_gen_pad_functions_tet4.svg")
                .unwrap();
        }
    }
}
*/
<|endoftext|>


<commit_before>// Copyright 2012-2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Base64 binary-to-text encoding

use core::iter;
use core::str;
use core::vec;

pub trait ToBase64 {
    fn to_base64(&self) -> ~str;
}

static CHARS: [char, ..64] = [
    'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M',
    'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z',
    'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm',
    'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',
    '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '+', '/'
];

impl<'self> ToBase64 for &'self [u8] {
    /**
     * Turn a vector of `u8` bytes into a base64 string.
     *
     * *Example*:
     *
     * ~~~~
     * extern mod std;
     * use std::base64::ToBase64;
     *
     * fn main () {
     *   let str = [52,32].to_base64();
     *   println(fmt!("%s", str));
     * }
     * ~~~~
     */
    fn to_base64(&self) -> ~str {
        let mut s = ~"";
        unsafe {
            let len = self.len();
            str::reserve(&mut s, ((len + 3u) / 4u) * 3u);

            let mut i = 0u;

            while i < len - (len % 3u) {
                let n = (self[i] as uint) << 16u |
                        (self[i + 1u] as uint) << 8u |
                        (self[i + 2u] as uint);

                // This 24-bit number gets separated into four 6-bit numbers.
                str::push_char(&mut s, CHARS[(n >> 18u) & 63u]);
                str::push_char(&mut s, CHARS[(n >> 12u) & 63u]);
                str::push_char(&mut s, CHARS[(n >> 6u) & 63u]);
                str::push_char(&mut s, CHARS[n & 63u]);

                i += 3u;
            }

            // Heh, would be cool if we knew this was exhaustive
            // (the dream of bounded integer types)
            match len % 3 {
              0 => (),
              1 => {
                let n = (self[i] as uint) << 16u;
                str::push_char(&mut s, CHARS[(n >> 18u) & 63u]);
                str::push_char(&mut s, CHARS[(n >> 12u) & 63u]);
                str::push_char(&mut s, '=');
                str::push_char(&mut s, '=');
              }
              2 => {
                let n = (self[i] as uint) << 16u |
                    (self[i + 1u] as uint) << 8u;
                str::push_char(&mut s, CHARS[(n >> 18u) & 63u]);
                str::push_char(&mut s, CHARS[(n >> 12u) & 63u]);
                str::push_char(&mut s, CHARS[(n >> 6u) & 63u]);
                str::push_char(&mut s, '=');
              }
              _ => fail!(~"Algebra is broken, please alert the math police")
            }
        }
        s
    }
}

impl<'self> ToBase64 for &'self str {
    /**
     * Convert any string (literal, `@`, `&`, or `~`) to base64 encoding.
     *
     *
     * *Example*:
     *
     * ~~~~
     * extern mod std;
     * use std::base64::ToBase64;
     *
     * fn main () {
     *   let str = "Hello, World".to_base64();
     *   println(fmt!("%s",str));
     * }
     * ~~~~
     *
     */
    fn to_base64(&self) -> ~str {
        str::to_bytes(*self).to_base64()
    }
}

pub trait FromBase64 {
    fn from_base64(&self) -> ~[u8];
}

impl FromBase64 for ~[u8] {
    /**
     * Convert base64 `u8` vector into u8 byte values.
     * Every 4 encoded characters is converted into 3 octets, modulo padding.
     *
     * *Example*:
     *
     * ~~~~
     * extern mod std;
     * use std::base64::ToBase64;
     * use std::base64::FromBase64;
     *
     * fn main () {
     *   let str = [52,32].to_base64();
     *   println(fmt!("%s", str));
     *   let bytes = str.from_base64();
     *   println(fmt!("%?",bytes));
     * }
     * ~~~~
     */
    fn from_base64(&self) -> ~[u8] {
        if self.len() % 4u != 0u { fail!(~"invalid base64 length"); }

        let len = self.len();
        let mut padding = 0u;

        if len != 0u {
            if self[len - 1u] == '=' as u8 { padding += 1u; }
            if self[len - 2u] == '=' as u8 { padding += 1u; }
        }

        let mut r = vec::with_capacity((len / 4u) * 3u - padding);

        unsafe {
            let mut i = 0u;
            while i < len {
                let mut n = 0u;

                for iter::repeat(4u) {
                    let ch = self[i] as char;
                    n <<= 6u;

                    if ch >= 'A' && ch <= 'Z' {
                        n |= (ch as uint) - 0x41u;
                    } else if ch >= 'a' && ch <= 'z' {
                        n |= (ch as uint) - 0x47u;
                    } else if ch >= '0' && ch <= '9' {
                        n |= (ch as uint) + 0x04u;
                    } else if ch == '+' {
                        n |= 0x3Eu;
                    } else if ch == '/' {
                        n |= 0x3Fu;
                    } else if ch == '=' {
                        match len - i {
                          1u => {
                            r.push(((n >> 16u) & 0xFFu) as u8);
                            r.push(((n >> 8u ) & 0xFFu) as u8);
                            return copy r;
                          }
                          2u => {
                            r.push(((n >> 10u) & 0xFFu) as u8);
                            return copy r;
                          }
                          _ => fail!(~"invalid base64 padding")
                        }
                    } else {
                        fail!(~"invalid base64 character");
                    }

                    i += 1u;
                };

                r.push(((n >> 16u) & 0xFFu) as u8);
                r.push(((n >> 8u ) & 0xFFu) as u8);
                r.push(((n       ) & 0xFFu) as u8);
            }
        }
        r
    }
}

impl FromBase64 for ~str {
    /**
     * Convert any base64 encoded string (literal, `@`, `&`, or `~`)
     * to the byte values it encodes.
     *
     * You can use the `from_bytes` function in `core::str`
     * to turn a `[u8]` into a string with characters corresponding to those values.
     *
     * *Example*:
     *
     * This converts a string literal to base64 and back.
     *
     * ~~~~
     * extern mod std;
     * use std::base64::ToBase64;
     * use std::base64::FromBase64;
     * use core::str;
     *
     * fn main () {
     *   let hello_str = "Hello, World".to_base64();
     *   println(fmt!("%s",hello_str));
     *   let bytes = hello_str.from_base64();
     *   println(fmt!("%?",bytes));
     *   let result_str = str::from_bytes(bytes);
     *   println(fmt!("%s",result_str));
     * }
     * ~~~~
     */
    fn from_base64(&self) -> ~[u8] {
        str::to_bytes(*self).from_base64()
    }
}

#[cfg(test)]
mod tests {
    use core::str;

    #[test]
    pub fn test_to_base64() {
        assert!((~"").to_base64()       == ~"");
        assert!((~"f").to_base64()      == ~"Zg==");
        assert!((~"fo").to_base64()     == ~"Zm8=");
        assert!((~"foo").to_base64()    == ~"Zm9v");
        assert!((~"foob").to_base64()   == ~"Zm9vYg==");
        assert!((~"fooba").to_base64()  == ~"Zm9vYmE=");
        assert!((~"foobar").to_base64() == ~"Zm9vYmFy");
    }

    #[test]
    pub fn test_from_base64() {
        assert!((~"").from_base64() == str::to_bytes(~""));
        assert!((~"Zg==").from_base64() == str::to_bytes(~"f"));
        assert!((~"Zm8=").from_base64() == str::to_bytes(~"fo"));
        assert!((~"Zm9v").from_base64() == str::to_bytes(~"foo"));
        assert!((~"Zm9vYg==").from_base64() == str::to_bytes(~"foob"));
        assert!((~"Zm9vYmE=").from_base64() == str::to_bytes(~"fooba"));
        assert!((~"Zm9vYmFy").from_base64() == str::to_bytes(~"foobar"));
    }
}
<commit_msg>Change to 4-space indents in code examples<commit_after>// Copyright 2012-2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Base64 binary-to-text encoding

use core::iter;
use core::str;
use core::vec;

pub trait ToBase64 {
    fn to_base64(&self) -> ~str;
}

static CHARS: [char, ..64] = [
    'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M',
    'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z',
    'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm',
    'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',
    '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '+', '/'
];

impl<'self> ToBase64 for &'self [u8] {
    /**
     * Turn a vector of `u8` bytes into a base64 string.
     *
     * *Example*:
     *
     * ~~~~
     * extern mod std;
     * use std::base64::ToBase64;
     *
     * fn main () {
     *     let str = [52,32].to_base64();
     *     println(fmt!("%s", str));
     * }
     * ~~~~
     */
    fn to_base64(&self) -> ~str {
        let mut s = ~"";
        unsafe {
            let len = self.len();
            str::reserve(&mut s, ((len + 3u) / 4u) * 3u);

            let mut i = 0u;

            while i < len - (len % 3u) {
                let n = (self[i] as uint) << 16u |
                        (self[i + 1u] as uint) << 8u |
                        (self[i + 2u] as uint);

                // This 24-bit number gets separated into four 6-bit numbers.
                str::push_char(&mut s, CHARS[(n >> 18u) & 63u]);
                str::push_char(&mut s, CHARS[(n >> 12u) & 63u]);
                str::push_char(&mut s, CHARS[(n >> 6u) & 63u]);
                str::push_char(&mut s, CHARS[n & 63u]);

                i += 3u;
            }

            // Heh, would be cool if we knew this was exhaustive
            // (the dream of bounded integer types)
            match len % 3 {
              0 => (),
              1 => {
                let n = (self[i] as uint) << 16u;
                str::push_char(&mut s, CHARS[(n >> 18u) & 63u]);
                str::push_char(&mut s, CHARS[(n >> 12u) & 63u]);
                str::push_char(&mut s, '=');
                str::push_char(&mut s, '=');
              }
              2 => {
                let n = (self[i] as uint) << 16u |
                    (self[i + 1u] as uint) << 8u;
                str::push_char(&mut s, CHARS[(n >> 18u) & 63u]);
                str::push_char(&mut s, CHARS[(n >> 12u) & 63u]);
                str::push_char(&mut s, CHARS[(n >> 6u) & 63u]);
                str::push_char(&mut s, '=');
              }
              _ => fail!(~"Algebra is broken, please alert the math police")
            }
        }
        s
    }
}

impl<'self> ToBase64 for &'self str {
    /**
     * Convert any string (literal, `@`, `&`, or `~`) to base64 encoding.
     *
     *
     * *Example*:
     *
     * ~~~~
     * extern mod std;
     * use std::base64::ToBase64;
     *
     * fn main () {
     *     let str = "Hello, World".to_base64();
     *     println(fmt!("%s",str));
     * }
     * ~~~~
     *
     */
    fn to_base64(&self) -> ~str {
        str::to_bytes(*self).to_base64()
    }
}

pub trait FromBase64 {
    fn from_base64(&self) -> ~[u8];
}

impl FromBase64 for ~[u8] {
    /**
     * Convert base64 `u8` vector into u8 byte values.
     * Every 4 encoded characters is converted into 3 octets, modulo padding.
     *
     * *Example*:
     *
     * ~~~~
     * extern mod std;
     * use std::base64::ToBase64;
     * use std::base64::FromBase64;
     *
     * fn main () {
     *     let str = [52,32].to_base64();
     *     println(fmt!("%s", str));
     *     let bytes = str.from_base64();
     *     println(fmt!("%?",bytes));
     * }
     * ~~~~
     */
    fn from_base64(&self) -> ~[u8] {
        if self.len() % 4u != 0u { fail!(~"invalid base64 length"); }

        let len = self.len();
        let mut padding = 0u;

        if len != 0u {
            if self[len - 1u] == '=' as u8 { padding += 1u; }
            if self[len - 2u] == '=' as u8 { padding += 1u; }
        }

        let mut r = vec::with_capacity((len / 4u) * 3u - padding);

        unsafe {
            let mut i = 0u;
            while i < len {
                let mut n = 0u;

                for iter::repeat(4u) {
                    let ch = self[i] as char;
                    n <<= 6u;

                    if ch >= 'A' && ch <= 'Z' {
                        n |= (ch as uint) - 0x41u;
                    } else if ch >= 'a' && ch <= 'z' {
                        n |= (ch as uint) - 0x47u;
                    } else if ch >= '0' && ch <= '9' {
                        n |= (ch as uint) + 0x04u;
                    } else if ch == '+' {
                        n |= 0x3Eu;
                    } else if ch == '/' {
                        n |= 0x3Fu;
                    } else if ch == '=' {
                        match len - i {
                          1u => {
                            r.push(((n >> 16u) & 0xFFu) as u8);
                            r.push(((n >> 8u ) & 0xFFu) as u8);
                            return copy r;
                          }
                          2u => {
                            r.push(((n >> 10u) & 0xFFu) as u8);
                            return copy r;
                          }
                          _ => fail!(~"invalid base64 padding")
                        }
                    } else {
                        fail!(~"invalid base64 character");
                    }

                    i += 1u;
                };

                r.push(((n >> 16u) & 0xFFu) as u8);
                r.push(((n >> 8u ) & 0xFFu) as u8);
                r.push(((n       ) & 0xFFu) as u8);
            }
        }
        r
    }
}

impl FromBase64 for ~str {
    /**
     * Convert any base64 encoded string (literal, `@`, `&`, or `~`)
     * to the byte values it encodes.
     *
     * You can use the `from_bytes` function in `core::str`
     * to turn a `[u8]` into a string with characters corresponding to those values.
     *
     * *Example*:
     *
     * This converts a string literal to base64 and back.
     *
     * ~~~~
     * extern mod std;
     * use std::base64::ToBase64;
     * use std::base64::FromBase64;
     * use core::str;
     *
     * fn main () {
     *     let hello_str = "Hello, World".to_base64();
     *     println(fmt!("%s",hello_str));
     *     let bytes = hello_str.from_base64();
     *     println(fmt!("%?",bytes));
     *     let result_str = str::from_bytes(bytes);
     *     println(fmt!("%s",result_str));
     * }
     * ~~~~
     */
    fn from_base64(&self) -> ~[u8] {
        str::to_bytes(*self).from_base64()
    }
}

#[cfg(test)]
mod tests {
    use core::str;

    #[test]
    pub fn test_to_base64() {
        assert!((~"").to_base64()       == ~"");
        assert!((~"f").to_base64()      == ~"Zg==");
        assert!((~"fo").to_base64()     == ~"Zm8=");
        assert!((~"foo").to_base64()    == ~"Zm9v");
        assert!((~"foob").to_base64()   == ~"Zm9vYg==");
        assert!((~"fooba").to_base64()  == ~"Zm9vYmE=");
        assert!((~"foobar").to_base64() == ~"Zm9vYmFy");
    }

    #[test]
    pub fn test_from_base64() {
        assert!((~"").from_base64() == str::to_bytes(~""));
        assert!((~"Zg==").from_base64() == str::to_bytes(~"f"));
        assert!((~"Zm8=").from_base64() == str::to_bytes(~"fo"));
        assert!((~"Zm9v").from_base64() == str::to_bytes(~"foo"));
        assert!((~"Zm9vYg==").from_base64() == str::to_bytes(~"foob"));
        assert!((~"Zm9vYmE=").from_base64() == str::to_bytes(~"fooba"));
        assert!((~"Zm9vYmFy").from_base64() == str::to_bytes(~"foobar"));
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
//
// Original implementation taken from rust-memchr
// Copyright 2015 Andrew Gallant, bluss and Nicolas Koch

/// A safe interface to `memchr`.
///
/// Returns the index corresponding to the first occurrence of `needle` in
/// `haystack`, or `None` if one is not found.
///
/// memchr reduces to super-optimized machine code at around an order of
/// magnitude faster than `haystack.iter().position(|&b| b == needle)`.
/// (See benchmarks.)
///
/// # Example
///
/// This shows how to find the first position of a byte in a byte string.
///
/// ```rust,ignore
/// use memchr::memchr;
///
/// let haystack = b"the quick brown fox";
/// assert_eq!(memchr(b'k', haystack), Some(8));
/// ```
pub fn memchr(needle: u8, haystack: &[u8]) -> Option<usize> {
    // libc memchr
    #[cfg(not(target_os = "windows"))]
    fn memchr_specific(needle: u8, haystack: &[u8]) -> Option<usize> {
        use libc;

        let p = unsafe {
            libc::memchr(
                haystack.as_ptr() as *const libc::c_void,
                needle as libc::c_int,
                haystack.len() as libc::size_t)
        };
        if p.is_null() {
            None
        } else {
            Some(p as usize - (haystack.as_ptr() as usize))
        }
    }

    // use fallback on windows, since it's faster
    #[cfg(target_os = "windows")]
    fn memchr_specific(needle: u8, haystack: &[u8]) -> Option<usize> {
        fallback::memchr(needle, haystack)
    }

    memchr_specific(needle, haystack)
}

/// A safe interface to `memrchr`.
///
/// Returns the index corresponding to the last occurrence of `needle` in
/// `haystack`, or `None` if one is not found.
///
/// # Example
///
/// This shows how to find the last position of a byte in a byte string.
///
/// ```rust,ignore
/// use memchr::memrchr;
///
/// let haystack = b"the quick brown fox";
/// assert_eq!(memrchr(b'o', haystack), Some(17));
/// ```
pub fn memrchr(needle: u8, haystack: &[u8]) -> Option<usize> {

    #[cfg(target_os = "linux")]
    fn memrchr_specific(needle: u8, haystack: &[u8]) -> Option<usize> {
        use libc;

        // GNU's memrchr() will - unlike memchr() - error if haystack is empty.
        if haystack.is_empty() {return None}
        let p = unsafe {
            libc::memrchr(
                haystack.as_ptr() as *const libc::c_void,
                needle as libc::c_int,
                haystack.len() as libc::size_t)
        };
        if p.is_null() {
            None
        } else {
            Some(p as usize - (haystack.as_ptr() as usize))
        }
    }

    #[cfg(not(target_os = "linux"))]
    fn memrchr_specific(needle: u8, haystack: &[u8]) -> Option<usize> {
        fallback::memrchr(needle, haystack)
    }

    memrchr_specific(needle, haystack)
}

#[allow(dead_code)]
mod fallback {
    use cmp;
    use mem;

    const LO_U64: u64 = 0x0101010101010101;
    const HI_U64: u64 = 0x8080808080808080;

    // use truncation
    const LO_USIZE: usize = LO_U64 as usize;
    const HI_USIZE: usize = HI_U64 as usize;

    /// Return `true` if `x` contains any zero byte.
    ///
    /// From *Matters Computational*, J. Arndt
    ///
    /// "The idea is to subtract one from each of the bytes and then look for
    /// bytes where the borrow propagated all the way to the most significant
    /// bit."
    #[inline]
    fn contains_zero_byte(x: usize) -> bool {
        x.wrapping_sub(LO_USIZE) & !x & HI_USIZE != 0
    }

    #[cfg(target_pointer_width = "32")]
    #[inline]
    fn repeat_byte(b: u8) -> usize {
        let mut rep = (b as usize) << 8 | b as usize;
        rep = rep << 16 | rep;
        rep
    }

    #[cfg(target_pointer_width = "64")]
    #[inline]
    fn repeat_byte(b: u8) -> usize {
        let mut rep = (b as usize) << 8 | b as usize;
        rep = rep << 16 | rep;
        rep = rep << 32 | rep;
        rep
    }

    /// Return the first index matching the byte `a` in `text`.
    pub fn memchr(x: u8, text: &[u8]) -> Option<usize> {
        // Scan for a single byte value by reading two `usize` words at a time.
        //
        // Split `text` in three parts
        // - unaligned initial part, before the first word aligned address in text
        // - body, scan by 2 words at a time
        // - the last remaining part, < 2 word size
        let len = text.len();
        let ptr = text.as_ptr();
        let usize_bytes = mem::size_of::<usize>();

        // search up to an aligned boundary
        let align = (ptr as usize) & (usize_bytes- 1);
        let mut offset;
        if align > 0 {
            offset = cmp::min(usize_bytes - align, len);
            if let Some(index) = text[..offset].iter().position(|elt| *elt == x) {
                return Some(index);
            }
        } else {
            offset = 0;
        }

        // search the body of the text
        let repeated_x = repeat_byte(x);

        if len >= 2 * usize_bytes {
            while offset <= len - 2 * usize_bytes {
                unsafe {
                    let u = *(ptr.offset(offset as isize) as *const usize);
                    let v = *(ptr.offset((offset + usize_bytes) as isize) as *const usize);

                    // break if there is a matching byte
                    let zu = contains_zero_byte(u ^ repeated_x);
                    let zv = contains_zero_byte(v ^ repeated_x);
                    if zu || zv {
                        break;
                    }
                }
                offset += usize_bytes * 2;
            }
        }

        // find the byte after the point the body loop stopped
        text[offset..].iter().position(|elt| *elt == x).map(|i| offset + i)
    }

    /// Return the last index matching the byte `a` in `text`.
    pub fn memrchr(x: u8, text: &[u8]) -> Option<usize> {
        // Scan for a single byte value by reading two `usize` words at a time.
        //
        // Split `text` in three parts
        // - unaligned tail, after the last word aligned address in text
        // - body, scan by 2 words at a time
        // - the first remaining bytes, < 2 word size
        let len = text.len();
        let ptr = text.as_ptr();
        let usize_bytes = mem::size_of::<usize>();

        // search to an aligned boundary
        let end_align = (ptr as usize + len) & (usize_bytes - 1);
        let mut offset;
        if end_align > 0 {
            offset = len - cmp::min(end_align, len);
            if let Some(index) = text[offset..].iter().rposition(|elt| *elt == x) {
                return Some(offset + index);
            }
        } else {
            offset = len;
        }

        // search the body of the text
        let repeated_x = repeat_byte(x);

        while offset >= 2 * usize_bytes {
            unsafe {
                let u = *(ptr.offset(offset as isize - 2 * usize_bytes as isize) as *const usize);
                let v = *(ptr.offset(offset as isize - usize_bytes as isize) as *const usize);

                // break if there is a matching byte
                let zu = contains_zero_byte(u ^ repeated_x);
                let zv = contains_zero_byte(v ^ repeated_x);
                if zu || zv {
                    break;
                }
            }
            offset -= 2 * usize_bytes;
        }

        // find the byte before the point the body loop stopped
        text[..offset].iter().rposition(|elt| *elt == x)
    }

    // test fallback implementations on all platforms
    #[test]
    fn matches_one() {
        assert_eq!(Some(0), memchr(b'a', b"a"));
    }

    #[test]
    fn matches_begin() {
        assert_eq!(Some(0), memchr(b'a', b"aaaa"));
    }

    #[test]
    fn matches_end() {
        assert_eq!(Some(4), memchr(b'z', b"aaaaz"));
    }

    #[test]
    fn matches_nul() {
        assert_eq!(Some(4), memchr(b'\x00', b"aaaa\x00"));
    }

    #[test]
    fn matches_past_nul() {
        assert_eq!(Some(5), memchr(b'z', b"aaaa\x00z"));
    }

    #[test]
    fn no_match_empty() {
        assert_eq!(None, memchr(b'a', b""));
    }

    #[test]
    fn no_match() {
        assert_eq!(None, memchr(b'a', b"xyz"));
    }

    #[test]
    fn matches_one_reversed() {
        assert_eq!(Some(0), memrchr(b'a', b"a"));
    }

    #[test]
    fn matches_begin_reversed() {
        assert_eq!(Some(3), memrchr(b'a', b"aaaa"));
    }

    #[test]
    fn matches_end_reversed() {
        assert_eq!(Some(0), memrchr(b'z', b"zaaaa"));
    }

    #[test]
    fn matches_nul_reversed() {
        assert_eq!(Some(4), memrchr(b'\x00', b"aaaa\x00"));
    }

    #[test]
    fn matches_past_nul_reversed() {
        assert_eq!(Some(0), memrchr(b'z', b"z\x00aaaa"));
    }

    #[test]
    fn no_match_empty_reversed() {
        assert_eq!(None, memrchr(b'a', b""));
    }

    #[test]
    fn no_match_reversed() {
        assert_eq!(None, memrchr(b'a', b"xyz"));
    }

    #[test]
    fn each_alignment_reversed() {
        let mut data = [1u8; 64];
        let needle = 2;
        let pos = 40;
        data[pos] = needle;
        for start in 0..16 {
            assert_eq!(Some(pos - start), memrchr(needle, &data[start..]));
        }
    }
}

#[cfg(test)]
mod tests {
    // test the implementations for the current plattform
    use super::{memchr, memrchr};

    #[test]
    fn matches_one() {
        assert_eq!(Some(0), memchr(b'a', b"a"));
    }

    #[test]
    fn matches_begin() {
        assert_eq!(Some(0), memchr(b'a', b"aaaa"));
    }

    #[test]
    fn matches_end() {
        assert_eq!(Some(4), memchr(b'z', b"aaaaz"));
    }

    #[test]
    fn matches_nul() {
        assert_eq!(Some(4), memchr(b'\x00', b"aaaa\x00"));
    }

    #[test]
    fn matches_past_nul() {
        assert_eq!(Some(5), memchr(b'z', b"aaaa\x00z"));
    }

    #[test]
    fn no_match_empty() {
        assert_eq!(None, memchr(b'a', b""));
    }

    #[test]
    fn no_match() {
        assert_eq!(None, memchr(b'a', b"xyz"));
    }

    #[test]
    fn matches_one_reversed() {
        assert_eq!(Some(0), memrchr(b'a', b"a"));
    }

    #[test]
    fn matches_begin_reversed() {
        assert_eq!(Some(3), memrchr(b'a', b"aaaa"));
    }

    #[test]
    fn matches_end_reversed() {
        assert_eq!(Some(0), memrchr(b'z', b"zaaaa"));
    }

    #[test]
    fn matches_nul_reversed() {
        assert_eq!(Some(4), memrchr(b'\x00', b"aaaa\x00"));
    }

    #[test]
    fn matches_past_nul_reversed() {
        assert_eq!(Some(0), memrchr(b'z', b"z\x00aaaa"));
    }

    #[test]
    fn no_match_empty_reversed() {
        assert_eq!(None, memrchr(b'a', b""));
    }

    #[test]
    fn no_match_reversed() {
        assert_eq!(None, memrchr(b'a', b"xyz"));
    }

    #[test]
    fn each_alignment() {
        let mut data = [1u8; 64];
        let needle = 2;
        let pos = 40;
        data[pos] = needle;
        for start in 0..16 {
            assert_eq!(Some(pos - start), memchr(needle, &data[start..]));
        }
    }
}
<commit_msg>memrchr: Use a conditional instead of subtracting a complicated min<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
//
// Original implementation taken from rust-memchr
// Copyright 2015 Andrew Gallant, bluss and Nicolas Koch

/// A safe interface to `memchr`.
///
/// Returns the index corresponding to the first occurrence of `needle` in
/// `haystack`, or `None` if one is not found.
///
/// memchr reduces to super-optimized machine code at around an order of
/// magnitude faster than `haystack.iter().position(|&b| b == needle)`.
/// (See benchmarks.)
///
/// # Example
///
/// This shows how to find the first position of a byte in a byte string.
///
/// ```rust,ignore
/// use memchr::memchr;
///
/// let haystack = b"the quick brown fox";
/// assert_eq!(memchr(b'k', haystack), Some(8));
/// ```
pub fn memchr(needle: u8, haystack: &[u8]) -> Option<usize> {
    // libc memchr
    #[cfg(not(target_os = "windows"))]
    fn memchr_specific(needle: u8, haystack: &[u8]) -> Option<usize> {
        use libc;

        let p = unsafe {
            libc::memchr(
                haystack.as_ptr() as *const libc::c_void,
                needle as libc::c_int,
                haystack.len() as libc::size_t)
        };
        if p.is_null() {
            None
        } else {
            Some(p as usize - (haystack.as_ptr() as usize))
        }
    }

    // use fallback on windows, since it's faster
    #[cfg(target_os = "windows")]
    fn memchr_specific(needle: u8, haystack: &[u8]) -> Option<usize> {
        fallback::memchr(needle, haystack)
    }

    memchr_specific(needle, haystack)
}

/// A safe interface to `memrchr`.
///
/// Returns the index corresponding to the last occurrence of `needle` in
/// `haystack`, or `None` if one is not found.
///
/// # Example
///
/// This shows how to find the last position of a byte in a byte string.
///
/// ```rust,ignore
/// use memchr::memrchr;
///
/// let haystack = b"the quick brown fox";
/// assert_eq!(memrchr(b'o', haystack), Some(17));
/// ```
pub fn memrchr(needle: u8, haystack: &[u8]) -> Option<usize> {

    #[cfg(target_os = "linux")]
    fn memrchr_specific(needle: u8, haystack: &[u8]) -> Option<usize> {
        use libc;

        // GNU's memrchr() will - unlike memchr() - error if haystack is empty.
        if haystack.is_empty() {return None}
        let p = unsafe {
            libc::memrchr(
                haystack.as_ptr() as *const libc::c_void,
                needle as libc::c_int,
                haystack.len() as libc::size_t)
        };
        if p.is_null() {
            None
        } else {
            Some(p as usize - (haystack.as_ptr() as usize))
        }
    }

    #[cfg(not(target_os = "linux"))]
    fn memrchr_specific(needle: u8, haystack: &[u8]) -> Option<usize> {
        fallback::memrchr(needle, haystack)
    }

    memrchr_specific(needle, haystack)
}

#[allow(dead_code)]
mod fallback {
    use cmp;
    use mem;

    const LO_U64: u64 = 0x0101010101010101;
    const HI_U64: u64 = 0x8080808080808080;

    // use truncation
    const LO_USIZE: usize = LO_U64 as usize;
    const HI_USIZE: usize = HI_U64 as usize;

    /// Return `true` if `x` contains any zero byte.
    ///
    /// From *Matters Computational*, J. Arndt
    ///
    /// "The idea is to subtract one from each of the bytes and then look for
    /// bytes where the borrow propagated all the way to the most significant
    /// bit."
    #[inline]
    fn contains_zero_byte(x: usize) -> bool {
        x.wrapping_sub(LO_USIZE) & !x & HI_USIZE != 0
    }

    #[cfg(target_pointer_width = "32")]
    #[inline]
    fn repeat_byte(b: u8) -> usize {
        let mut rep = (b as usize) << 8 | b as usize;
        rep = rep << 16 | rep;
        rep
    }

    #[cfg(target_pointer_width = "64")]
    #[inline]
    fn repeat_byte(b: u8) -> usize {
        let mut rep = (b as usize) << 8 | b as usize;
        rep = rep << 16 | rep;
        rep = rep << 32 | rep;
        rep
    }

    /// Return the first index matching the byte `a` in `text`.
    pub fn memchr(x: u8, text: &[u8]) -> Option<usize> {
        // Scan for a single byte value by reading two `usize` words at a time.
        //
        // Split `text` in three parts
        // - unaligned initial part, before the first word aligned address in text
        // - body, scan by 2 words at a time
        // - the last remaining part, < 2 word size
        let len = text.len();
        let ptr = text.as_ptr();
        let usize_bytes = mem::size_of::<usize>();

        // search up to an aligned boundary
        let align = (ptr as usize) & (usize_bytes- 1);
        let mut offset;
        if align > 0 {
            offset = cmp::min(usize_bytes - align, len);
            if let Some(index) = text[..offset].iter().position(|elt| *elt == x) {
                return Some(index);
            }
        } else {
            offset = 0;
        }

        // search the body of the text
        let repeated_x = repeat_byte(x);

        if len >= 2 * usize_bytes {
            while offset <= len - 2 * usize_bytes {
                unsafe {
                    let u = *(ptr.offset(offset as isize) as *const usize);
                    let v = *(ptr.offset((offset + usize_bytes) as isize) as *const usize);

                    // break if there is a matching byte
                    let zu = contains_zero_byte(u ^ repeated_x);
                    let zv = contains_zero_byte(v ^ repeated_x);
                    if zu || zv {
                        break;
                    }
                }
                offset += usize_bytes * 2;
            }
        }

        // find the byte after the point the body loop stopped
        text[offset..].iter().position(|elt| *elt == x).map(|i| offset + i)
    }

    /// Return the last index matching the byte `a` in `text`.
    pub fn memrchr(x: u8, text: &[u8]) -> Option<usize> {
        // Scan for a single byte value by reading two `usize` words at a time.
        //
        // Split `text` in three parts
        // - unaligned tail, after the last word aligned address in text
        // - body, scan by 2 words at a time
        // - the first remaining bytes, < 2 word size
        let len = text.len();
        let ptr = text.as_ptr();
        let usize_bytes = mem::size_of::<usize>();

        // search to an aligned boundary
        let end_align = (ptr as usize + len) & (usize_bytes - 1);
        let mut offset;
        if end_align > 0 {
            offset = if end_align >= len { 0 } else { len - end_align };
            if let Some(index) = text[offset..].iter().rposition(|elt| *elt == x) {
                return Some(offset + index);
            }
        } else {
            offset = len;
        }

        // search the body of the text
        let repeated_x = repeat_byte(x);

        while offset >= 2 * usize_bytes {
            unsafe {
                let u = *(ptr.offset(offset as isize - 2 * usize_bytes as isize) as *const usize);
                let v = *(ptr.offset(offset as isize - usize_bytes as isize) as *const usize);

                // break if there is a matching byte
                let zu = contains_zero_byte(u ^ repeated_x);
                let zv = contains_zero_byte(v ^ repeated_x);
                if zu || zv {
                    break;
                }
            }
            offset -= 2 * usize_bytes;
        }

        // find the byte before the point the body loop stopped
        text[..offset].iter().rposition(|elt| *elt == x)
    }

    // test fallback implementations on all platforms
    #[test]
    fn matches_one() {
        assert_eq!(Some(0), memchr(b'a', b"a"));
    }

    #[test]
    fn matches_begin() {
        assert_eq!(Some(0), memchr(b'a', b"aaaa"));
    }

    #[test]
    fn matches_end() {
        assert_eq!(Some(4), memchr(b'z', b"aaaaz"));
    }

    #[test]
    fn matches_nul() {
        assert_eq!(Some(4), memchr(b'\x00', b"aaaa\x00"));
    }

    #[test]
    fn matches_past_nul() {
        assert_eq!(Some(5), memchr(b'z', b"aaaa\x00z"));
    }

    #[test]
    fn no_match_empty() {
        assert_eq!(None, memchr(b'a', b""));
    }

    #[test]
    fn no_match() {
        assert_eq!(None, memchr(b'a', b"xyz"));
    }

    #[test]
    fn matches_one_reversed() {
        assert_eq!(Some(0), memrchr(b'a', b"a"));
    }

    #[test]
    fn matches_begin_reversed() {
        assert_eq!(Some(3), memrchr(b'a', b"aaaa"));
    }

    #[test]
    fn matches_end_reversed() {
        assert_eq!(Some(0), memrchr(b'z', b"zaaaa"));
    }

    #[test]
    fn matches_nul_reversed() {
        assert_eq!(Some(4), memrchr(b'\x00', b"aaaa\x00"));
    }

    #[test]
    fn matches_past_nul_reversed() {
        assert_eq!(Some(0), memrchr(b'z', b"z\x00aaaa"));
    }

    #[test]
    fn no_match_empty_reversed() {
        assert_eq!(None, memrchr(b'a', b""));
    }

    #[test]
    fn no_match_reversed() {
        assert_eq!(None, memrchr(b'a', b"xyz"));
    }

    #[test]
    fn each_alignment_reversed() {
        let mut data = [1u8; 64];
        let needle = 2;
        let pos = 40;
        data[pos] = needle;
        for start in 0..16 {
            assert_eq!(Some(pos - start), memrchr(needle, &data[start..]));
        }
    }
}

#[cfg(test)]
mod tests {
    // test the implementations for the current plattform
    use super::{memchr, memrchr};

    #[test]
    fn matches_one() {
        assert_eq!(Some(0), memchr(b'a', b"a"));
    }

    #[test]
    fn matches_begin() {
        assert_eq!(Some(0), memchr(b'a', b"aaaa"));
    }

    #[test]
    fn matches_end() {
        assert_eq!(Some(4), memchr(b'z', b"aaaaz"));
    }

    #[test]
    fn matches_nul() {
        assert_eq!(Some(4), memchr(b'\x00', b"aaaa\x00"));
    }

    #[test]
    fn matches_past_nul() {
        assert_eq!(Some(5), memchr(b'z', b"aaaa\x00z"));
    }

    #[test]
    fn no_match_empty() {
        assert_eq!(None, memchr(b'a', b""));
    }

    #[test]
    fn no_match() {
        assert_eq!(None, memchr(b'a', b"xyz"));
    }

    #[test]
    fn matches_one_reversed() {
        assert_eq!(Some(0), memrchr(b'a', b"a"));
    }

    #[test]
    fn matches_begin_reversed() {
        assert_eq!(Some(3), memrchr(b'a', b"aaaa"));
    }

    #[test]
    fn matches_end_reversed() {
        assert_eq!(Some(0), memrchr(b'z', b"zaaaa"));
    }

    #[test]
    fn matches_nul_reversed() {
        assert_eq!(Some(4), memrchr(b'\x00', b"aaaa\x00"));
    }

    #[test]
    fn matches_past_nul_reversed() {
        assert_eq!(Some(0), memrchr(b'z', b"z\x00aaaa"));
    }

    #[test]
    fn no_match_empty_reversed() {
        assert_eq!(None, memrchr(b'a', b""));
    }

    #[test]
    fn no_match_reversed() {
        assert_eq!(None, memrchr(b'a', b"xyz"));
    }

    #[test]
    fn each_alignment() {
        let mut data = [1u8; 64];
        let needle = 2;
        let pos = 40;
        data[pos] = needle;
        for start in 0..16 {
            assert_eq!(Some(pos - start), memchr(needle, &data[start..]));
        }
    }
}
<|endoftext|>


<commit_before>use database::models::{ToU64, LogChannel};

use chrono::{Utc, Duration};

use diesel::prelude::*;

use serenity::prelude::{Mutex, Mentionable};
use serenity::client::{Context, EventHandler};
use serenity::model::channel::Message;
use serenity::model::event::MessageUpdateEvent;
use serenity::model::guild::Member;
use serenity::model::id::{GuildId, ChannelId, UserId};
use serenity::model::user::User;

use std::collections::HashMap;
use std::sync::atomic::{Ordering, AtomicUsize};

macro_rules! update_message {
  ($message:expr, $update:expr, $($field:ident),+) => {{
    $(
      if let Some(f) = $update.$field {
        $message.$field = f;
      }
    )+
  }}
}

#[derive(Default)]
pub struct Log {
  messages: Mutex<HashMap<UserId, HashMap<GuildId, Vec<Message>>>>,
  count: AtomicUsize
}

impl Log {
  fn get_log_channel<G: Into<GuildId>>(&self, guild: G) -> Option<ChannelId> {
    let guild = guild.into().to_u64();
    let log_channel: Option<LogChannel> = ::bot::with_connection(|c| {
      use database::schema::log_channels::dsl;
      dsl::log_channels.filter(dsl::server_id.eq(guild)).first(c)
    }).ok();
    log_channel.map(|x| ChannelId(*x.channel_id))
  }

  fn prune_messages(&self) {
    let now = Utc::now();
    let one_day = Duration::days(1);
    self.messages
      .lock()
      .values_mut()
      .flat_map(|x| x.values_mut())
      .for_each(|x| x.retain(|m| m.timestamp.signed_duration_since(now) < one_day));
  }
}

impl EventHandler for Log {
  fn guild_member_removal(&self, _: Context, guild: GuildId, user: User, member: Option<Member>) {
    let channel_id = some_or!(self.get_log_channel(guild), return);
    let mention = member.as_ref().map(|x| x.mention()).unwrap_or_else(|| user.mention());
    channel_id.send_message(|m| m.embed(|mut embed| {
      embed = embed
        .author(|a| a
          .name(&user.tag())
          .icon_url(&user.face()));
      if let Some(member) = member {
        embed = embed.field("Display name", member.display_name(), true);
      }
      embed = embed
        .field("Mention", mention, true)
        .field("Action", "Leave server", true)
        .timestamp(&Utc::now())
        .footer(|f| f.text(user.id));
      embed
    })).ok();
  }

  fn guild_member_addition(&self, _: Context, guild: GuildId, member: Member) {
    let channel_id = some_or!(self.get_log_channel(guild), return);
    channel_id.send_message(|m| m.embed(|mut embed| {
      embed = embed
        .author(|a| a
          .name(&member.user.read().tag())
          .icon_url(&member.user.read().face()));
      embed = embed
        .field("Mention", member.mention(), true)
        .field("Action", "Join server", true)
        .timestamp(&Utc::now())
        .footer(|f| f.text(member.user.read().id));
      embed
    })).ok();
  }

  fn message_update(&self, _: Context, update: MessageUpdateEvent) {
    let author = match update.author {
      Some(ref a) => a,
      None => return
    };
    let new_content = match update.content {
      Some(ref c) => c,
      None => return
    };

    let channel = match update.channel_id.get() {
      Ok(c) => c,
      Err(e) => {
        warn!("could not download channel {} for message history: {}", update.channel_id, e);
        return;
      }
    };

    let guild_channel = match channel.guild() {
      Some(g) => g,
      None => return
    };
    let reader = guild_channel.read();

    let channel_id = some_or!(self.get_log_channel(reader.guild_id), return);

    let guild = match reader.guild_id.find() {
      Some(g) => g,
      None => return
    };
    let guild_reader = guild.read();

    let member = match guild_reader.members.get(&author.id).cloned().or_else(|| guild_reader.member(author.id).ok()) {
      Some(m) => m,
      None => return
    };

    let mut messages = self.messages.lock();
    let message = messages
      .get_mut(&author.id)
      .and_then(|x| x.get_mut(&reader.guild_id))
      .and_then(|x| x.iter_mut().find(|m| m.id == update.id));
    let message = match message {
      Some(m) => m,
      None => return
    };

    let original_content = message.content.clone();
    let channel_mention = update.channel_id.mention();
    let message_id = update.id;

    update_message!(
      message,
      update,
      kind, content, tts, pinned, timestamp, author, mention_everyone, mentions, mention_roles,
      attachments
    );
    message.edited_timestamp = update.edited_timestamp;

    channel_id.send_message(|m| m.embed(|mut embed| {
      embed = embed
        .author(|a| a
          .name(&member.user.read().tag())
          .icon_url(&member.user.read().face()));
      embed = embed
        .field("Mention", member.mention(), true)
        .field("Action", "Edited message", true)
        .field("Channel", channel_mention, true)
        .field("Original message", original_content, false)
        .field("New message", new_content, false)
        .timestamp(&Utc::now())
        .footer(|f| f.text(message_id));
      embed
    })).ok();
  }

  fn message(&self, _: Context, message: Message) {
    if self.count.load(Ordering::SeqCst) == 100 {
      self.count.store(0, Ordering::SeqCst);

      self.prune_messages();
    }

    let channel = match message.channel_id.get() {
      Ok(c) => c,
      Err(e) => {
        warn!("could not download channel {} for message history: {}", message.channel_id, e);
        return;
      }
    };

    let guild_channel = match channel.guild() {
      Some(g) => g,
      None => return
    };

    let reader = guild_channel.read();

    self.messages
      .lock()
      .entry(message.author.id)
      .or_default()
      .entry(reader.guild_id)
      .or_default()
      .push(message);

    self.count.fetch_add(1, Ordering::SeqCst);
  }
}
<commit_msg>feat: show message deletion<commit_after>use database::models::{ToU64, LogChannel};

use chrono::{Utc, Duration};

use diesel::prelude::*;

use serenity::prelude::{Mutex, Mentionable};
use serenity::client::{Context, EventHandler};
use serenity::model::channel::Message;
use serenity::model::event::MessageUpdateEvent;
use serenity::model::guild::{Member, Action};
use serenity::model::id::{GuildId, ChannelId, UserId, MessageId};
use serenity::model::user::User;

use std::collections::HashMap;
use std::sync::atomic::{Ordering, AtomicUsize};
use std::thread;

macro_rules! update_message {
  ($message:expr, $update:expr, $($field:ident),+) => {{
    $(
      if let Some(f) = $update.$field {
        $message.$field = f;
      }
    )+
  }}
}

#[derive(Default)]
pub struct Log {
  messages: Mutex<HashMap<UserId, HashMap<GuildId, Vec<Message>>>>,
  count: AtomicUsize
}

impl Log {
  fn get_log_channel<G: Into<GuildId>>(&self, guild: G) -> Option<ChannelId> {
    let guild = guild.into().to_u64();
    let log_channel: Option<LogChannel> = ::bot::with_connection(|c| {
      use database::schema::log_channels::dsl;
      dsl::log_channels.filter(dsl::server_id.eq(guild)).first(c)
    }).ok();
    log_channel.map(|x| ChannelId(*x.channel_id))
  }

  fn prune_messages(&self) {
    let now = Utc::now();
    let one_day = Duration::days(1);
    self.messages
      .lock()
      .values_mut()
      .flat_map(|x| x.values_mut())
      .for_each(|x| x.retain(|m| m.timestamp.signed_duration_since(now) < one_day));
  }
}

impl EventHandler for Log {
  fn guild_member_removal(&self, _: Context, guild: GuildId, user: User, member: Option<Member>) {
    let channel_id = some_or!(self.get_log_channel(guild), return);
    let mention = member.as_ref().map(|x| x.mention()).unwrap_or_else(|| user.mention());
    channel_id.send_message(|m| m.embed(|mut embed| {
      embed = embed
        .author(|a| a
          .name(&user.tag())
          .icon_url(&user.face()));
      if let Some(member) = member {
        embed = embed.field("Display name", member.display_name(), true);
      }
      embed = embed
        .field("Mention", mention, true)
        .field("Action", "Leave server", true)
        .timestamp(&Utc::now())
        .footer(|f| f.text(user.id));
      embed
    })).ok();
  }

  fn guild_member_addition(&self, _: Context, guild: GuildId, member: Member) {
    let channel_id = some_or!(self.get_log_channel(guild), return);
    channel_id.send_message(|m| m.embed(|mut embed| {
      embed = embed
        .author(|a| a
          .name(&member.user.read().tag())
          .icon_url(&member.user.read().face()));
      embed = embed
        .field("Mention", member.mention(), true)
        .field("Action", "Join server", true)
        .timestamp(&Utc::now())
        .footer(|f| f.text(member.user.read().id));
      embed
    })).ok();
  }

  fn message_update(&self, _: Context, update: MessageUpdateEvent) {
    let author = match update.author {
      Some(ref a) => a,
      None => return
    };
    let new_content = match update.content {
      Some(ref c) => c,
      None => return
    };

    let channel = match update.channel_id.get() {
      Ok(c) => c,
      Err(e) => {
        warn!("could not download channel {} for message history: {}", update.channel_id, e);
        return;
      }
    };

    let guild_channel = match channel.guild() {
      Some(g) => g,
      None => return
    };
    let reader = guild_channel.read();

    let channel_id = some_or!(self.get_log_channel(reader.guild_id), return);

    let guild = match reader.guild_id.find() {
      Some(g) => g,
      None => return
    };
    let guild_reader = guild.read();

    let member = match guild_reader.members.get(&author.id).cloned().or_else(|| guild_reader.member(author.id).ok()) {
      Some(m) => m,
      None => return
    };

    let mut messages = self.messages.lock();
    let message = messages
      .get_mut(&author.id)
      .and_then(|x| x.get_mut(&reader.guild_id))
      .and_then(|x| x.iter_mut().find(|m| m.id == update.id));
    let message = match message {
      Some(m) => m,
      None => return
    };

    let original_content = message.content.clone();
    let channel_mention = update.channel_id.mention();
    let message_id = update.id;

    update_message!(
      message,
      update,
      kind, content, tts, pinned, timestamp, author, mention_everyone, mentions, mention_roles,
      attachments
    );
    message.edited_timestamp = update.edited_timestamp;

    channel_id.send_message(|m| m.embed(|mut embed| {
      embed = embed
        .author(|a| a
          .name(&member.user.read().tag())
          .icon_url(&member.user.read().face()));
      embed = embed
        .field("Mention", member.mention(), true)
        .field("Action", "Edited message", true)
        .field("Channel", channel_mention, true)
        .field("Original message", original_content, false)
        .field("New message", new_content, false)
        .timestamp(&Utc::now())
        .footer(|f| f.text(message_id));
      embed
    })).ok();
  }

  fn message_delete(&self, _: Context, channel_id: ChannelId, message_id: MessageId) {
    let channel = match channel_id.get() {
      Ok(c) => c,
      Err(e) => {
        warn!("could not download channel {} for message history: {}", channel_id, e);
        return;
      }
    };

    let guild_channel = match channel.guild() {
      Some(g) => g,
      None => return
    };
    let reader = guild_channel.read();

    let log_channel = some_or!(self.get_log_channel(reader.guild_id), return);

    let guild = match reader.guild_id.find() {
      Some(g) => g,
      None => return
    };

    let guild_reader = guild.read();

    let mut messages = self.messages.lock();
    let message = messages
      .values_mut()
      .filter_map(|x| x.get_mut(&reader.guild_id))
      .next()
      .and_then(|x| x.iter_mut().find(|m| m.id == message_id));
    let message = match message {
      Some(m) => m.clone(),
      None => return
    };

    let deletee = some_or!(guild_reader.members.get(&message.author.id).cloned().or_else(|| guild_reader.member(message.author.id).ok()), return);

    let original_content = message.content;
    let channel_mention = channel_id.mention();

    log_channel.send_message(|m| m.embed(|mut embed| {
      embed = embed
        .author(|a| a
          .name(&deletee.user.read().tag())
          .icon_url(&deletee.user.read().face()));
      embed = embed
        .field("Mention", deletee.mention(), true)
        .field("Action", "Had message deleted", true)
        .field("Channel", channel_mention, true)
        .field("Content", original_content, false)
        .timestamp(&Utc::now())
        .footer(|f| f.text(message_id));
      embed
    })).ok();
  }

  fn message(&self, _: Context, message: Message) {
    if self.count.load(Ordering::SeqCst) == 100 {
      self.count.store(0, Ordering::SeqCst);

      self.prune_messages();
    }

    let channel = match message.channel_id.get() {
      Ok(c) => c,
      Err(e) => {
        warn!("could not download channel {} for message history: {}", message.channel_id, e);
        return;
      }
    };

    let guild_channel = match channel.guild() {
      Some(g) => g,
      None => return
    };

    let reader = guild_channel.read();

    self.messages
      .lock()
      .entry(message.author.id)
      .or_default()
      .entry(reader.guild_id)
      .or_default()
      .push(message);

    self.count.fetch_add(1, Ordering::SeqCst);
  }
}
<|endoftext|>