<commit_before>use Configuration;
use deque;
use deque::{Worker, Stealer, Stolen};
use job::JobRef;
use latch::{Latch, LockLatch, SpinLatch};
#[allow(unused_imports)]
use log::Event::*;
use rand;
use std::cell::Cell;
use std::sync::{Arc, Condvar, Mutex, Once, ONCE_INIT};
use std::thread;
use std::collections::VecDeque;
use std::mem;
use util::leak;
use num_cpus;

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

pub struct Registry {
    thread_infos: Vec<ThreadInfo>,
    state: Mutex<RegistryState>,
    work_available: Condvar,
}

struct RegistryState {
    terminate: bool,
    threads_at_work: usize,
    injected_jobs: VecDeque<JobRef>,
}

///////////////////////////////////////////////////////////////////////////
// Initialization

static mut THE_REGISTRY: Option<&'static Registry> = None;
static THE_REGISTRY_SET: Once = ONCE_INIT;

/// Starts the worker threads (if that has not already happened). If
/// initialization has not already occurred, use the default
/// configuration.
pub fn get_registry() -> &'static Registry {
    THE_REGISTRY_SET.call_once(|| unsafe { init_registry(Configuration::new()) });
    unsafe { THE_REGISTRY.unwrap() }
}

/// Starts the worker threads (if that has not already happened) with
/// the given configuration.
pub fn get_registry_with_config(config: Configuration) -> &'static Registry {
    THE_REGISTRY_SET.call_once(|| unsafe { init_registry(config) });
    unsafe { THE_REGISTRY.unwrap() }
}

/// Initializes the global registry with the given configuration.
/// Meant to be called from within the `THE_REGISTRY_SET` once
/// function. Declared `unsafe` because it writes to `THE_REGISTRY` in
/// an unsynchronized fashion.
unsafe fn init_registry(config: Configuration) {
    let registry = leak(Registry::new(config.num_threads()));
    THE_REGISTRY = Some(registry);
}

enum Work {
    None,
    Job(JobRef),
    Terminate,
}

impl Registry {
    pub fn new(num_threads: Option<usize>) -> Arc<Registry> {
        let limit_value = match num_threads {
            Some(value) => value,
            None => num_cpus::get()
        };

        let (workers, stealers) : (Vec<_>, Vec<_>) =
            (0..limit_value).map(|_| deque::new()).unzip();

        let registry = Arc::new(Registry {
            thread_infos: stealers.into_iter().map(|s| ThreadInfo::new(s))
                                          .collect(),
            state: Mutex::new(RegistryState::new()),
            work_available: Condvar::new(),
        });

        for (index, worker) in workers.into_iter().enumerate() {
            let registry = registry.clone();
            thread::spawn(move || unsafe { main_loop(worker, registry, index) });
        }

        registry
    }

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

    /// Waits for the worker threads to get up and running.  This is
    /// meant to be used for benchmarking purposes, primarily, so that
    /// you can get more consistent numbers by having everything
    /// "ready to go".
    pub fn wait_until_primed(&self) {
        for info in &self.thread_infos {
            info.primed.wait();
        }
    }

    ///////////////////////////////////////////////////////////////////////////
    // MAIN LOOP
    //
    // So long as all of the worker threads are hanging out in their
    // top-level loop, there is no work to be done.

    fn start_working(&self, index: usize) {
        log!(StartWorking { index: index });
        let mut state = self.state.lock().unwrap();
        state.threads_at_work += 1;
        self.work_available.notify_all();
    }

    pub unsafe fn inject(&self, injected_jobs: &[JobRef]) {
        log!(InjectJobs { count: injected_jobs.len() });
        let mut state = self.state.lock().unwrap();
        if state.terminate {
            drop(state);
            // We can only reach this point if these conditions are met:
            // 1) This must be the global thread pool
            // 2) The only way to terminate the global thread pool is if a
            //    worker thread panics.
            // 3) A worker thread can only panic if a job panicked and was not
            //    caught by panic::recover.
            panic!("rayon thread pool is contaminated by a previous panic; recovery is only available on nightly compilers");
        }
        state.injected_jobs.extend(injected_jobs);
        self.work_available.notify_all();
    }

    fn wait_for_work(&self, _worker: usize, was_active: bool) -> Work {
        log!(WaitForWork { worker: _worker, was_active: was_active });

        let mut state = self.state.lock().unwrap();

        if was_active {
            state.threads_at_work -= 1;
        }

        loop {
            // Check if we need to terminate.
            if state.terminate {
                return Work::Terminate;
            }

            // Otherwise, if anything was injected from outside,
            // return that.  Note that this gives preference to
            // injected items over stealing from others, which is a
            // bit dubious, but then so is the opposite.
            if let Some(job) = state.injected_jobs.pop_front() {
                state.threads_at_work += 1;
                self.work_available.notify_all();
                return Work::Job(job);
            }

            // If any of the threads are running a job, we should spin
            // up, since they may generate subworkitems.
            if state.threads_at_work > 0 {
                return Work::None;
            }

            state = self.work_available.wait(state).unwrap();
        }
    }

    pub fn terminate(&self) {
        let mut state = self.state.lock().unwrap();
        state.terminate = true;
        for job in state.injected_jobs.drain(..) {
            unsafe {
                (*job.0).abort();
            }
        }
        self.work_available.notify_all();
    }
}

impl RegistryState {
    pub fn new() -> RegistryState {
        RegistryState {
            threads_at_work: 0,
            injected_jobs: VecDeque::new(),
            terminate: false,
        }
    }
}

struct ThreadInfo {
    // latch is set once thread has started and we are entering into
    // the main loop
    primed: LockLatch,
    stealer: Stealer<JobRef>,
}

impl ThreadInfo {
    fn new(stealer: Stealer<JobRef>) -> ThreadInfo {
        ThreadInfo {
            primed: LockLatch::new(),
            stealer: stealer,
        }
    }
}

///////////////////////////////////////////////////////////////////////////
// WorkerThread identifiers

pub struct WorkerThread {
    worker: Worker<JobRef>,
    stealers: Vec<Stealer<JobRef>>,
    registry: Arc<Registry>,
    index: usize
}

// This is a bit sketchy, but basically: the WorkerThread is
// allocated on the stack of the worker on entry and stored into this
// thread local variable. So it will remain valid at least until the
// worker is fully unwound. Using an unsafe pointer avoids the need
// for a RefCell<T> etc.
thread_local! {
    static WORKER_THREAD_STATE: Cell<*const WorkerThread> =
        Cell::new(0 as *const WorkerThread)
}

impl WorkerThread {
    /// Gets the `WorkerThread` index for the current thread; returns
    /// NULL if this is not a worker thread. This pointer is valid
    /// anywhere on the current thread.
    #[inline]
    pub unsafe fn current() -> *const WorkerThread {
        WORKER_THREAD_STATE.with(|t| t.get())
    }

    /// Sets `self` as the worker thread index for the current thread.
    /// This is done during worker thread startup.
    unsafe fn set_current(&self) {
        WORKER_THREAD_STATE.with(|t| {
            assert!(t.get().is_null());
            t.set(self);
        });
    }

    #[inline]
    pub fn index(&self) -> usize {
        self.index
    }

    #[inline]
    pub unsafe fn push(&self, job: JobRef) {
        self.worker.push(job);
    }

    /// Pop `job` from top of stack, returning `false` if it has been
    /// stolen.
    #[inline]
    pub unsafe fn pop(&self) -> bool {
        self.worker.pop().is_some()
    }

    #[cold]
    pub unsafe fn steal_until(&self, latch: &SpinLatch) {
        // If another thread stole our job when we panic, we must halt unwinding
        // until that thread is finished using it.
        struct PanicGuard<'a>(&'a SpinLatch);
        impl<'a> Drop for PanicGuard<'a> {
            fn drop(&mut self) {
                while !self.0.probe() {
                    thread::yield_now();
                }
            }
        }

        let guard = PanicGuard(&latch);
        while !latch.probe() {
            if let Some(job) = steal_work(&self.stealers) {
                (*job.0).execute();
            } else {
                thread::yield_now();
            }
        }
        mem::forget(guard);
    }
}

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

unsafe fn main_loop(worker: Worker<JobRef>, registry: Arc<Registry>, index: usize) {
    let stealers = registry.thread_infos.iter()
        .enumerate().filter(|&(i, _)| i != index)
        .map(|(_, ti)| ti.stealer.clone())
        .collect();

    let worker_thread = WorkerThread {
        worker: worker,
        stealers: stealers,
        registry: registry.clone(),
        index: index,
    };
    worker_thread.set_current();

    // let registry know we are ready to do work
    registry.thread_infos[index].primed.set();

    // If a worker thread panics, bring down the entire thread pool
    struct PanicGuard;
    impl Drop for PanicGuard {
        fn drop(&mut self) {
            unsafe {
                (*WorkerThread::current()).registry.terminate();
            }
        }
    }
    let _guard = PanicGuard;

    let mut was_active = false;
    loop {
        match registry.wait_for_work(index, was_active) {
            Work::Job(injected_job) => {
                (*injected_job.0).execute();
                was_active = true;
                continue;
            }
            Work::Terminate => break,
            Work::None => {},
        }

        if let Some(stolen_job) = steal_work(&worker_thread.stealers) {
            log!(StoleWork { worker: index });
            registry.start_working(index);
            (*stolen_job.0).execute();
            was_active = true;
        } else {
            was_active = false;
        }
    }
}

unsafe fn steal_work(stealers: &[Stealer<JobRef>]) -> Option<JobRef> {
    if stealers.is_empty() { return None }
    let start = rand::random::<usize>() % stealers.len();
    stealers[start..].iter()
        .chain(&stealers[..start])
        .filter_map(|stealer| match stealer.steal() {
            Stolen::Empty => None,
            Stolen::Abort => None, // loop?
            Stolen::Data(v) => Some(v),
        })
        .next()
}
<commit_msg>thread_pool: remove unused registry from WorkerThread<commit_after>use Configuration;
use deque;
use deque::{Worker, Stealer, Stolen};
use job::JobRef;
use latch::{Latch, LockLatch, SpinLatch};
#[allow(unused_imports)]
use log::Event::*;
use rand;
use std::cell::Cell;
use std::sync::{Arc, Condvar, Mutex, Once, ONCE_INIT};
use std::thread;
use std::collections::VecDeque;
use std::mem;
use util::leak;
use num_cpus;

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

pub struct Registry {
    thread_infos: Vec<ThreadInfo>,
    state: Mutex<RegistryState>,
    work_available: Condvar,
}

struct RegistryState {
    terminate: bool,
    threads_at_work: usize,
    injected_jobs: VecDeque<JobRef>,
}

///////////////////////////////////////////////////////////////////////////
// Initialization

static mut THE_REGISTRY: Option<&'static Registry> = None;
static THE_REGISTRY_SET: Once = ONCE_INIT;

/// Starts the worker threads (if that has not already happened). If
/// initialization has not already occurred, use the default
/// configuration.
pub fn get_registry() -> &'static Registry {
    THE_REGISTRY_SET.call_once(|| unsafe { init_registry(Configuration::new()) });
    unsafe { THE_REGISTRY.unwrap() }
}

/// Starts the worker threads (if that has not already happened) with
/// the given configuration.
pub fn get_registry_with_config(config: Configuration) -> &'static Registry {
    THE_REGISTRY_SET.call_once(|| unsafe { init_registry(config) });
    unsafe { THE_REGISTRY.unwrap() }
}

/// Initializes the global registry with the given configuration.
/// Meant to be called from within the `THE_REGISTRY_SET` once
/// function. Declared `unsafe` because it writes to `THE_REGISTRY` in
/// an unsynchronized fashion.
unsafe fn init_registry(config: Configuration) {
    let registry = leak(Registry::new(config.num_threads()));
    THE_REGISTRY = Some(registry);
}

enum Work {
    None,
    Job(JobRef),
    Terminate,
}

impl Registry {
    pub fn new(num_threads: Option<usize>) -> Arc<Registry> {
        let limit_value = match num_threads {
            Some(value) => value,
            None => num_cpus::get()
        };

        let (workers, stealers) : (Vec<_>, Vec<_>) =
            (0..limit_value).map(|_| deque::new()).unzip();

        let registry = Arc::new(Registry {
            thread_infos: stealers.into_iter().map(|s| ThreadInfo::new(s))
                                          .collect(),
            state: Mutex::new(RegistryState::new()),
            work_available: Condvar::new(),
        });

        for (index, worker) in workers.into_iter().enumerate() {
            let registry = registry.clone();
            thread::spawn(move || unsafe { main_loop(worker, registry, index) });
        }

        registry
    }

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

    /// Waits for the worker threads to get up and running.  This is
    /// meant to be used for benchmarking purposes, primarily, so that
    /// you can get more consistent numbers by having everything
    /// "ready to go".
    pub fn wait_until_primed(&self) {
        for info in &self.thread_infos {
            info.primed.wait();
        }
    }

    ///////////////////////////////////////////////////////////////////////////
    // MAIN LOOP
    //
    // So long as all of the worker threads are hanging out in their
    // top-level loop, there is no work to be done.

    fn start_working(&self, index: usize) {
        log!(StartWorking { index: index });
        let mut state = self.state.lock().unwrap();
        state.threads_at_work += 1;
        self.work_available.notify_all();
    }

    pub unsafe fn inject(&self, injected_jobs: &[JobRef]) {
        log!(InjectJobs { count: injected_jobs.len() });
        let mut state = self.state.lock().unwrap();
        if state.terminate {
            drop(state);
            // We can only reach this point if these conditions are met:
            // 1) This must be the global thread pool
            // 2) The only way to terminate the global thread pool is if a
            //    worker thread panics.
            // 3) A worker thread can only panic if a job panicked and was not
            //    caught by panic::recover.
            panic!("rayon thread pool is contaminated by a previous panic; recovery is only available on nightly compilers");
        }
        state.injected_jobs.extend(injected_jobs);
        self.work_available.notify_all();
    }

    fn wait_for_work(&self, _worker: usize, was_active: bool) -> Work {
        log!(WaitForWork { worker: _worker, was_active: was_active });

        let mut state = self.state.lock().unwrap();

        if was_active {
            state.threads_at_work -= 1;
        }

        loop {
            // Check if we need to terminate.
            if state.terminate {
                return Work::Terminate;
            }

            // Otherwise, if anything was injected from outside,
            // return that.  Note that this gives preference to
            // injected items over stealing from others, which is a
            // bit dubious, but then so is the opposite.
            if let Some(job) = state.injected_jobs.pop_front() {
                state.threads_at_work += 1;
                self.work_available.notify_all();
                return Work::Job(job);
            }

            // If any of the threads are running a job, we should spin
            // up, since they may generate subworkitems.
            if state.threads_at_work > 0 {
                return Work::None;
            }

            state = self.work_available.wait(state).unwrap();
        }
    }

    pub fn terminate(&self) {
        let mut state = self.state.lock().unwrap();
        state.terminate = true;
        for job in state.injected_jobs.drain(..) {
            unsafe {
                (*job.0).abort();
            }
        }
        self.work_available.notify_all();
    }
}

impl RegistryState {
    pub fn new() -> RegistryState {
        RegistryState {
            threads_at_work: 0,
            injected_jobs: VecDeque::new(),
            terminate: false,
        }
    }
}

struct ThreadInfo {
    // latch is set once thread has started and we are entering into
    // the main loop
    primed: LockLatch,
    stealer: Stealer<JobRef>,
}

impl ThreadInfo {
    fn new(stealer: Stealer<JobRef>) -> ThreadInfo {
        ThreadInfo {
            primed: LockLatch::new(),
            stealer: stealer,
        }
    }
}

///////////////////////////////////////////////////////////////////////////
// WorkerThread identifiers

pub struct WorkerThread {
    worker: Worker<JobRef>,
    stealers: Vec<Stealer<JobRef>>,
    index: usize
}

// This is a bit sketchy, but basically: the WorkerThread is
// allocated on the stack of the worker on entry and stored into this
// thread local variable. So it will remain valid at least until the
// worker is fully unwound. Using an unsafe pointer avoids the need
// for a RefCell<T> etc.
thread_local! {
    static WORKER_THREAD_STATE: Cell<*const WorkerThread> =
        Cell::new(0 as *const WorkerThread)
}

impl WorkerThread {
    /// Gets the `WorkerThread` index for the current thread; returns
    /// NULL if this is not a worker thread. This pointer is valid
    /// anywhere on the current thread.
    #[inline]
    pub unsafe fn current() -> *const WorkerThread {
        WORKER_THREAD_STATE.with(|t| t.get())
    }

    /// Sets `self` as the worker thread index for the current thread.
    /// This is done during worker thread startup.
    unsafe fn set_current(&self) {
        WORKER_THREAD_STATE.with(|t| {
            assert!(t.get().is_null());
            t.set(self);
        });
    }

    #[inline]
    pub fn index(&self) -> usize {
        self.index
    }

    #[inline]
    pub unsafe fn push(&self, job: JobRef) {
        self.worker.push(job);
    }

    /// Pop `job` from top of stack, returning `false` if it has been
    /// stolen.
    #[inline]
    pub unsafe fn pop(&self) -> bool {
        self.worker.pop().is_some()
    }

    #[cold]
    pub unsafe fn steal_until(&self, latch: &SpinLatch) {
        // If another thread stole our job when we panic, we must halt unwinding
        // until that thread is finished using it.
        struct PanicGuard<'a>(&'a SpinLatch);
        impl<'a> Drop for PanicGuard<'a> {
            fn drop(&mut self) {
                while !self.0.probe() {
                    thread::yield_now();
                }
            }
        }

        let guard = PanicGuard(&latch);
        while !latch.probe() {
            if let Some(job) = steal_work(&self.stealers) {
                (*job.0).execute();
            } else {
                thread::yield_now();
            }
        }
        mem::forget(guard);
    }
}

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

unsafe fn main_loop(worker: Worker<JobRef>, registry: Arc<Registry>, index: usize) {
    let stealers = registry.thread_infos.iter()
        .enumerate().filter(|&(i, _)| i != index)
        .map(|(_, ti)| ti.stealer.clone())
        .collect();

    let worker_thread = WorkerThread {
        worker: worker,
        stealers: stealers,
        index: index,
    };
    worker_thread.set_current();

    // let registry know we are ready to do work
    registry.thread_infos[index].primed.set();

    // If a worker thread panics, bring down the entire thread pool
    struct PanicGuard<'a>(&'a Registry);
    impl<'a> Drop for PanicGuard<'a> {
        fn drop(&mut self) {
            self.0.terminate();
        }
    }
    let _guard = PanicGuard(&registry);

    let mut was_active = false;
    loop {
        match registry.wait_for_work(index, was_active) {
            Work::Job(injected_job) => {
                (*injected_job.0).execute();
                was_active = true;
                continue;
            }
            Work::Terminate => break,
            Work::None => {},
        }

        if let Some(stolen_job) = steal_work(&worker_thread.stealers) {
            log!(StoleWork { worker: index });
            registry.start_working(index);
            (*stolen_job.0).execute();
            was_active = true;
        } else {
            was_active = false;
        }
    }
}

unsafe fn steal_work(stealers: &[Stealer<JobRef>]) -> Option<JobRef> {
    if stealers.is_empty() { return None }
    let start = rand::random::<usize>() % stealers.len();
    stealers[start..].iter()
        .chain(&stealers[..start])
        .filter_map(|stealer| match stealer.steal() {
            Stolen::Empty => None,
            Stolen::Abort => None, // loop?
            Stolen::Data(v) => Some(v),
        })
        .next()
}
<|endoftext|>


<commit_before>//! Process `-Z time-passes` output

use rustc_perf_collector::Pass;

use errors::{Result, ResultExt};

pub fn process_output(name: &str, output: Vec<u8>) -> Result<Vec<Pass>> {
    let output = String::from_utf8(output)
        .chain_err(|| format!("unable to convert output of {} to UTF-8", name))?;
    let mut passes = Vec::new();

    for line in output.lines() {
        if line.starts_with("time: ") {
            let time = &line["time: ".len()..line.find(";").unwrap()];
            let time: f64 = time.parse().chain_err(|| format!("parsed time: {:?}", time))?;

            let mem = &line[line.find("rss: ").unwrap() + 5..line.find("MB").unwrap()];
            let mem: u64 = mem.parse().chain_err(|| format!("parsed memory: {:?}", mem))?;

            let name = line[line.find("MB\t").unwrap() + 3..].to_string();
            passes.push(Pass {
                name: name,
                time: time,
                mem: mem,
            });
        } else if line.starts_with("  time: ") {
            // XXX: sub passes currently cause problems because their names are inconsistent in
            // incremental runs, especially codegen passes
            //let time = &line["  time: ".len()..line.find(";").unwrap()];
            //let time: f64 = time.parse().chain_err(|| format!("parsed time: {:?}", time))?;

            //let mem = &line[line.find("rss: ").unwrap() + 5..line.find("MB").unwrap()];
            //let mem: u64 = mem.parse().chain_err(|| format!("parsed memory: {:?}", mem))?;

            //let name = line[line.find("MB\t").unwrap() + 3..].to_string();
            //sub_passes.insert(name.clone(), Pass {
            //    name: name,
            //    time: time,
            //    mem: mem,
            //    sub_passes: HashMap::new(),
            //});
        } else {
            //info!("unhandled line: {}", line);
        }
    }

    Ok(passes)
}

pub struct PassAverager {
    pub state: Vec<Pass>,
    runs: u64
}

impl PassAverager {
    pub fn new(state: Vec<Pass>) -> Self {
        Self {
            state: state,
            runs: 0,
        }
    }

    pub fn average_with(
        &mut self,
        b: Vec<Pass>,
    ) -> Result<()> {
        self.runs += 1;

        for a in &mut self.state {
            let b = match b.iter().find(|p| p.name == a.name) {
                Some(b) => b,
                None => bail!("expected name {} to exist in both a and b", a.name),
            };
            a.time = a.time + ((b.time - a.time) / (self.runs as f64));
            a.mem = a.mem + ((b.mem - b.mem) / self.runs);
        }

        Ok(())
    }
}
<commit_msg>Fix parsing nightly's `-Z time-passes` output<commit_after>//! Process `-Z time-passes` output

use rustc_perf_collector::Pass;

use errors::{Result, ResultExt};

pub fn process_output(name: &str, output: Vec<u8>) -> Result<Vec<Pass>> {
    let output = String::from_utf8(output)
        .chain_err(|| format!("unable to convert output of {} to UTF-8", name))?;
    let mut passes = Vec::new();

    let mut time_indent = None;
    for line in output.lines() {
        if line.trim().starts_with("time: ") {
            let indent = line.find("time:").unwrap();
            if time_indent.is_some() && Some(indent) != time_indent {
                // XXX: sub passes currently cause problems because their names are inconsistent in
                // incremental runs, especially codegen passes
                //let time = &line["  time: ".len()..line.find(";").unwrap()];
                //let time: f64 = time.parse().chain_err(|| format!("parsed time: {:?}", time))?;

                //let mem = &line[line.find("rss: ").unwrap() + 5..line.find("MB").unwrap()];
                //let mem: u64 = mem.parse().chain_err(|| format!("parsed memory: {:?}", mem))?;

                //let name = line[line.find("MB\t").unwrap() + 3..].to_string();
                //sub_passes.insert(name.clone(), Pass {
                //    name: name,
                //    time: time,
                //    mem: mem,
                //    sub_passes: HashMap::new(),
                //});
                continue
            }
            time_indent = Some(indent);

            let time = &line[indent + "time: ".len()..line.find(";").unwrap()];
            let time: f64 = time.parse().chain_err(|| format!("parsed time: {:?}", time))?;

            let mem = &line[line.find("rss: ").unwrap() + 5..line.find("MB").unwrap()];
            let mem: u64 = mem.parse().chain_err(|| format!("parsed memory: {:?}", mem))?;

            let name = line[line.find("MB\t").unwrap() + 3..].to_string();
            passes.push(Pass {
                name: name,
                time: time,
                mem: mem,
            });
        } else {
            //info!("unhandled line: {}", line);
        }
    }

    Ok(passes)
}

pub struct PassAverager {
    pub state: Vec<Pass>,
    runs: u64
}

impl PassAverager {
    pub fn new(state: Vec<Pass>) -> Self {
        Self {
            state: state,
            runs: 0,
        }
    }

    pub fn average_with(
        &mut self,
        b: Vec<Pass>,
    ) -> Result<()> {
        self.runs += 1;

        for a in &mut self.state {
            let b = match b.iter().find(|p| p.name == a.name) {
                Some(b) => b,
                None => bail!("expected name {} to exist in both a and b", a.name),
            };
            a.time = a.time + ((b.time - a.time) / (self.runs as f64));
            a.mem = a.mem + ((b.mem - b.mem) / self.runs);
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>///
/// An extent represents a series of nodes starting at a specified node
///
pub enum TreeExtent {
    /// Just the initial node
    ThisNode,

    /// The children of this node
    ///
    /// This does not extend beyond the immediate children of the current node.
    Children,

    /// The entire subtree (all children, and their children, and so on)
    ///
    /// Unlike Children, this covers the current node and its entire subtree
    SubTree
}
<commit_msg>Add a 'covers' function to TreeExtent to see if a relative address is covered by an extent<commit_after>use super::address::*;

///
/// An extent represents a series of nodes starting at a specified node
///
pub enum TreeExtent {
    /// Just the initial node
    ThisNode,

    /// The children of this node
    ///
    /// This does not extend beyond the immediate children of the current node.
    Children,

    /// The entire subtree (all children, and their children, and so on)
    ///
    /// Unlike Children, this covers the current node and its entire subtree
    SubTree
}

impl TreeExtent {
    ///
    /// Returns true if this extent will cover the specified address, which is relative to where the extent starts
    ///
    pub fn covers(&self, address: &TreeAddress) -> bool {
        match *self {
            TreeExtent::ThisNode => {
                match *address {
                    TreeAddress::Here   => true,
                    _                   => false
                }
            },

            TreeExtent::Children => {
                match *address {
                    TreeAddress::ChildAtIndex(_, ref child_address) => TreeExtent::ThisNode.covers(child_address),
                    TreeAddress::ChildWithTag(_, ref child_address) => TreeExtent::ThisNode.covers(child_address),
                    _                                               => false
                }
            },

            TreeExtent::SubTree => true
        }
    }
}

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

    #[test]
    fn thisnode_covers_only_here() {
        assert!(TreeExtent::ThisNode.covers(&TreeAddress::Here));
        assert!(!TreeExtent::ThisNode.covers(&(1.to_tree_address())));
    }

    #[test]
    fn children_covers_only_immediate_children() {
        assert!(TreeExtent::Children.covers(&(1.to_tree_address())));
        assert!(TreeExtent::Children.covers(&("tag".to_tree_address())));

        assert!(!TreeExtent::Children.covers(&((1, 2).to_tree_address())));
        assert!(!TreeExtent::Children.covers(&(("tag", "othertag").to_tree_address())));
        assert!(!TreeExtent::Children.covers(&TreeAddress::Here));
    }

    #[test]
    fn subtree_covers_everything() {
        assert!(TreeExtent::SubTree.covers(&(1.to_tree_address())));
        assert!(TreeExtent::SubTree.covers(&("tag".to_tree_address())));

        assert!(TreeExtent::SubTree.covers(&((1, 2).to_tree_address())));
        assert!(TreeExtent::SubTree.covers(&(("tag", "othertag").to_tree_address())));
        assert!(TreeExtent::SubTree.covers(&TreeAddress::Here));
    }
}
<|endoftext|>


<commit_before>
use std::iter::range_step;
use collections::hashmap::HashSet;
use rand::random;
use graphics::*;
use piston::*;
use settings;
use tile::{
    Tile,
    TileStatic,
};

pub struct Board {
    tiles: Vec<Tile>,
}

impl Board {
    pub fn new() -> Board {
        Board {
            tiles: Vec::<Tile>::new(),
        }
    }

    pub fn generate_tile(&mut self) {
        if self.tiles.len() == (settings::TILE_WIDTH * settings::TILE_HEIGHT) as uint {
            return;
        }

        loop {
            let x = (random::<uint>() % settings::TILE_WIDTH as uint) as int;
            let y = (random::<uint>() % settings::TILE_HEIGHT as uint) as int;
            if self.get_tile(x, y).is_none() {
                self.tiles.push(Tile::new(2, x, y));
                break;
            }
        }
    }

    pub fn update(&mut self, dt: f64) {
        for tile in self.tiles.mut_iter() {
            tile.update(dt);
        }

        if self.is_locking() {
            return;
        }
        let mut tiles_need_removed = HashSet::<uint>::new();
        let mut tiles_need_added = Vec::<Tile>::new();
        for i in range(0, self.tiles.len()) {
            let tile1 = self.tiles.get(i);
            if tile1.status != TileStatic {
                continue;
            }
            for j in range(i + 1, self.tiles.len()) {
                let tile2 = self.tiles.get(j);
                if tile2.status != TileStatic
                   || tile1.tile_x != tile2.tile_x
                   || tile1.tile_y != tile2.tile_y {
                    continue;
                }

                tiles_need_removed.insert(i);
                tiles_need_removed.insert(j);
                tiles_need_added.push(Tile::new_combined(tile1.score + tile2.score, tile1.tile_x, tile1.tile_y));
                break;
            }
        }

        if tiles_need_removed.len() > 0 {
            let mut tiles = Vec::<Tile>::new();
            for i in range(0, self.tiles.len()) {
                if !tiles_need_removed.contains(&i) {
                    tiles.push(*self.tiles.get(i));
                }
            }
            self.tiles = tiles.append(tiles_need_added.as_slice());
        }
    }

    pub fn render(&self, c: &Context, gl: &mut Gl) {
        self.render_board(c, gl);
        self.render_tiles(c, gl);
    }

    pub fn merge_from_bottom_to_top(&mut self) {
        self.merge_col(0, settings::TILE_HEIGHT, 1);
    }

    pub fn merge_from_top_to_bottom(&mut self) {
        self.merge_col(settings::TILE_HEIGHT - 1, -1, -1);
    }

    fn merge_col(&mut self, y_start: int, y_end: int, y_step: int) {
        if self.is_locking() {
            return;
        }

        let mut need_generate = false;
        loop {
            // move all tiles to right place
            for col in range(0, settings::TILE_WIDTH) {
                for row in range_step(y_start, y_end, y_step) {
                    match self.get_mut_tile(col, row) {
                        None => {
                            match self.get_mut_next_tile(col, row, 0, y_step) {
                                Some(ref mut tile) => {
                                    println!("move ({}, {}) to ({}, {})",
                                             tile.tile_x, tile.tile_y, col, row);
                                    need_generate = true;
                                    tile.start_moving(col, row);
                                },
                                _ => {},
                            }
                        },
                        _ => {},
                    }
                }
            }

            let mut did_merged = false;
            for col in range(0, settings::TILE_WIDTH) {
                let mut found = false;
                let mut sx = 0;
                let mut sy = 0;
                let mut dx = 0;
                let mut dy = 0;
                for row in range_step(y_start, y_end, y_step) {
                    match self.get_tile(col, row) {
                        Some(ref d_tile) => {
                            match self.get_next_tile(col, row, 0, y_step) {
                                Some(ref s_tile)
                                if d_tile.score == s_tile.score
                                   && self.get_tile_count(d_tile.tile_x, d_tile.tile_y) == 1 => {
                                    found = true;
                                    dx = d_tile.tile_x;
                                    dy = d_tile.tile_y;
                                    sx = s_tile.tile_x;
                                    sy = s_tile.tile_y;
                                    break;
                                },
                                _ => {},
                            }
                        },
                        None => {
                            break;
                        }
                    }
                }
                if found {
                    need_generate = true;
                    did_merged = true;
                    let mut tile = self.get_mut_tile(sx, sy);
                    let tile = tile.get_mut_ref();
                    tile.start_moving(dx, dy);
                    println!("merge ({}, {}) to ({}, {})", sx, sy, dx, dy);
                }
            }

            if !did_merged {
                break;
            }
        }
        if need_generate {
            self.generate_tile();
        }
    }

    pub fn merge_from_left_to_right(&mut self) {
        self.merge_row(settings::TILE_WIDTH - 1, -1, -1);
    }

    pub fn merge_from_right_to_left(&mut self) {
        self.merge_row(0, settings::TILE_WIDTH, 1);
    }

    fn merge_row(&mut self, x_start: int, x_end: int, x_step: int) {
        if self.is_locking() {
            return;
        }

        let mut need_generate = false;
        loop {
            // move all tiles to right place
            for row in range(0, settings::TILE_HEIGHT) {
                for col in range_step(x_start, x_end, x_step) {
                    match self.get_mut_tile(col, row) {
                        None => {
                            match self.get_mut_next_tile(col, row, x_step, 0) {
                                Some(ref mut tile) => {
                                    println!("move ({}, {}) to ({}, {})", tile.tile_x, tile.tile_y, col, row);
                                    need_generate = true;
                                    tile.start_moving(col, row);
                                },
                                _ => {},
                            }
                        },
                        _ => {},
                    }
                }
            }

            // merge
            let mut did_merged = false;
            for row in range(0, settings::TILE_HEIGHT) {
                let mut found = false;
                let mut sx = 0;
                let mut sy = 0;
                let mut dx = 0;
                let mut dy = 0;
                for col in range_step(x_start, x_end, x_step) {
                    match self.get_tile(col, row) {
                        Some(ref d_tile) => {
                            match self.get_next_tile(col, row, x_step, 0) {
                                Some(ref s_tile)
                                if d_tile.score == s_tile.score
                                   && self.get_tile_count(d_tile.tile_x, d_tile.tile_y) == 1 => {
                                    found = true;
                                    dx = d_tile.tile_x;
                                    dy = d_tile.tile_y;
                                    sx = s_tile.tile_x;
                                    sy = s_tile.tile_y;
                                    break;
                                },
                                _ => {},
                            }
                        },
                        None => {
                            break;
                        }
                    }
                }
                if found {
                    need_generate = true;
                    did_merged = true;
                    let mut tile = self.get_mut_tile(sx, sy);
                    let tile = tile.get_mut_ref();
                    tile.start_moving(dx, dy);
                    println!("merge ({}, {}) to ({}, {})", sx, sy, dx, dy);
                }
            }

            if !did_merged {
                break;
            }
        }

        if need_generate {
            self.generate_tile();
        }
    }

    fn is_locking(&self) -> bool {
        for tile in self.tiles.iter() {
            if tile.status != TileStatic {
                return true;
            }
        }
        false
    }

    /// Returns next tile right besides (x, y)
    fn get_next_tile<'a>(&'a self, x: int, y: int, step_x: int, step_y: int) -> Option<&'a Tile> {
        let mut x = x + step_x;
        let mut y = y + step_y;
        while x >= 0 && x < settings::TILE_WIDTH
              && y >= 0 && y < settings::TILE_HEIGHT {
            let tile = self.get_tile(x, y);
            if tile.is_some() {
                return tile;
            }
            x += step_x;
            y += step_y;
        }
        None
    }

    fn get_mut_next_tile<'a>(&'a mut self, x: int, y: int, step_x: int, step_y: int) -> Option<&'a mut Tile> {
        let mut x = x + step_x;
        let mut y = y + step_y;
        let mut found = false;
        while x >= 0 && x < settings::TILE_WIDTH
              && y >= 0 && y < settings::TILE_HEIGHT {
            let tile = self.get_tile(x, y);
            if tile.is_some() {
                found = true;
                break;
            }
            x += step_x;
            y += step_y;
        }

        if found {
            self.get_mut_tile(x, y)
        } else {
            None
        }
    }

    fn get_tile<'a>(&'a self, x: int, y: int) -> Option<&'a Tile> {
        for tile in self.tiles.iter() {
            if tile.tile_x == x && tile.tile_y == y {
                return Some(tile);
            }
        }
        None
    }

    fn get_mut_tile<'a>(&'a mut self, x: int, y: int) -> Option<&'a mut Tile> {
        for tile in self.tiles.mut_iter() {
            if tile.tile_x == x && tile.tile_y == y {
                return Some(tile);
            }
        }
        None
    }

    fn get_tile_count(&self, x: int, y: int) -> int {
        let mut count = 0;
        for tile in self.tiles.iter() {
            if tile.tile_x == x && tile.tile_y == y {
                count += 1;
            }
        }
        count
    }

    fn render_board(&self, c: &Context, gl: &mut Gl) {
        c.view()
         .rect(settings::BOARD_PADDING,
               settings::BOARD_PADDING + settings::BOARD_OFFSET_Y,
               settings::BOARD_SIZE[0],
               settings::BOARD_SIZE[1])
         .rgba(settings::TILE_BACKGROUND_COLOR[0],
               settings::TILE_BACKGROUND_COLOR[1],
               settings::TILE_BACKGROUND_COLOR[2],
               settings::TILE_BACKGROUND_COLOR[3])
         .fill(gl);

        let mut x = settings::BOARD_PADDING + settings::TILE_PADDING;
        let mut y = settings::BOARD_PADDING + settings::BOARD_OFFSET_Y + settings::TILE_PADDING;
        for _ in range(0, settings::TILE_HEIGHT) {
            for _ in range(0, settings::TILE_WIDTH) {
                c.view()
                 .rect(x, y, settings::TILE_SIZE, settings::TILE_SIZE)
                 .rgba(settings::TILES_COLOR[0][0],
                       settings::TILES_COLOR[0][1],
                       settings::TILES_COLOR[0][2],
                       settings::TILES_COLOR[0][3])
                 .fill(gl);

                x += settings::TILE_PADDING + settings::TILE_SIZE;
            }
            x = settings::BOARD_PADDING + settings::TILE_PADDING;
            y += settings::TILE_PADDING + settings::TILE_SIZE;
        }
    }

    fn render_tiles(&self, c: &Context, gl: &mut Gl) {
        for tile in self.tiles.iter() {
            tile.render(c, gl);
        }
    }
}

<commit_msg>Added board score<commit_after>
use std::iter::range_step;
use collections::hashmap::HashSet;
use rand::random;
use graphics::*;
use piston::*;
use settings;
use tile::{
    Tile,
    TileStatic,
};

pub struct Board {
    tiles: Vec<Tile>,
    score: int,
    highest_score: int,
}

impl Board {
    pub fn new() -> Board {
        Board {
            tiles: Vec::<Tile>::new(),
            score: 0,
            highest_score: 0,
        }
    }

    pub fn generate_tile(&mut self) {
        if self.tiles.len() == (settings::TILE_WIDTH * settings::TILE_HEIGHT) as uint {
            return;
        }

        loop {
            let x = (random::<uint>() % settings::TILE_WIDTH as uint) as int;
            let y = (random::<uint>() % settings::TILE_HEIGHT as uint) as int;
            if self.get_tile(x, y).is_none() {
                self.tiles.push(Tile::new(2, x, y));
                break;
            }
        }
    }

    pub fn update(&mut self, dt: f64) {
        for tile in self.tiles.mut_iter() {
            tile.update(dt);
        }

        if self.is_locking() {
            return;
        }
        let mut score_to_added = 0;
        let mut tiles_need_removed = HashSet::<uint>::new();
        let mut tiles_need_added = Vec::<Tile>::new();
        for i in range(0, self.tiles.len()) {
            let tile1 = self.tiles.get(i);
            if tile1.status != TileStatic {
                continue;
            }
            for j in range(i + 1, self.tiles.len()) {
                let tile2 = self.tiles.get(j);
                if tile2.status != TileStatic
                   || tile1.tile_x != tile2.tile_x
                   || tile1.tile_y != tile2.tile_y {
                    continue;
                }

                tiles_need_removed.insert(i);
                tiles_need_removed.insert(j);
                tiles_need_added.push(Tile::new_combined(tile1.score + tile2.score, tile1.tile_x, tile1.tile_y));
                score_to_added += tile1.score + tile2.score;
                break;
            }
        }

        if tiles_need_removed.len() > 0 {
            let mut tiles = Vec::<Tile>::new();
            for i in range(0, self.tiles.len()) {
                if !tiles_need_removed.contains(&i) {
                    tiles.push(*self.tiles.get(i));
                }
            }
            self.tiles = tiles.append(tiles_need_added.as_slice());
            self.add_score(score_to_added);
        }
    }

    pub fn render(&self, c: &Context, gl: &mut Gl) {
        self.render_board(c, gl);
        self.render_tiles(c, gl);
    }

    pub fn merge_from_bottom_to_top(&mut self) {
        self.merge_col(0, settings::TILE_HEIGHT, 1);
    }

    pub fn merge_from_top_to_bottom(&mut self) {
        self.merge_col(settings::TILE_HEIGHT - 1, -1, -1);
    }

    fn merge_col(&mut self, y_start: int, y_end: int, y_step: int) {
        if self.is_locking() {
            return;
        }

        let mut need_generate = false;
        loop {
            // move all tiles to right place
            for col in range(0, settings::TILE_WIDTH) {
                for row in range_step(y_start, y_end, y_step) {
                    match self.get_mut_tile(col, row) {
                        None => {
                            match self.get_mut_next_tile(col, row, 0, y_step) {
                                Some(ref mut tile) => {
                                    println!("move ({}, {}) to ({}, {})",
                                             tile.tile_x, tile.tile_y, col, row);
                                    need_generate = true;
                                    tile.start_moving(col, row);
                                },
                                _ => {},
                            }
                        },
                        _ => {},
                    }
                }
            }

            let mut did_merged = false;
            for col in range(0, settings::TILE_WIDTH) {
                let mut found = false;
                let mut sx = 0;
                let mut sy = 0;
                let mut dx = 0;
                let mut dy = 0;
                for row in range_step(y_start, y_end, y_step) {
                    match self.get_tile(col, row) {
                        Some(ref d_tile) => {
                            match self.get_next_tile(col, row, 0, y_step) {
                                Some(ref s_tile)
                                if d_tile.score == s_tile.score
                                   && self.get_tile_count(d_tile.tile_x, d_tile.tile_y) == 1 => {
                                    found = true;
                                    dx = d_tile.tile_x;
                                    dy = d_tile.tile_y;
                                    sx = s_tile.tile_x;
                                    sy = s_tile.tile_y;
                                    break;
                                },
                                _ => {},
                            }
                        },
                        None => {
                            break;
                        }
                    }
                }
                if found {
                    need_generate = true;
                    did_merged = true;
                    let mut tile = self.get_mut_tile(sx, sy);
                    let tile = tile.get_mut_ref();
                    tile.start_moving(dx, dy);
                    println!("merge ({}, {}) to ({}, {})", sx, sy, dx, dy);
                }
            }

            if !did_merged {
                break;
            }
        }
        if need_generate {
            self.generate_tile();
        }
    }

    pub fn merge_from_left_to_right(&mut self) {
        self.merge_row(settings::TILE_WIDTH - 1, -1, -1);
    }

    pub fn merge_from_right_to_left(&mut self) {
        self.merge_row(0, settings::TILE_WIDTH, 1);
    }

    fn merge_row(&mut self, x_start: int, x_end: int, x_step: int) {
        if self.is_locking() {
            return;
        }

        let mut need_generate = false;
        loop {
            // move all tiles to right place
            for row in range(0, settings::TILE_HEIGHT) {
                for col in range_step(x_start, x_end, x_step) {
                    match self.get_mut_tile(col, row) {
                        None => {
                            match self.get_mut_next_tile(col, row, x_step, 0) {
                                Some(ref mut tile) => {
                                    println!("move ({}, {}) to ({}, {})", tile.tile_x, tile.tile_y, col, row);
                                    need_generate = true;
                                    tile.start_moving(col, row);
                                },
                                _ => {},
                            }
                        },
                        _ => {},
                    }
                }
            }

            // merge
            let mut did_merged = false;
            for row in range(0, settings::TILE_HEIGHT) {
                let mut found = false;
                let mut sx = 0;
                let mut sy = 0;
                let mut dx = 0;
                let mut dy = 0;
                for col in range_step(x_start, x_end, x_step) {
                    match self.get_tile(col, row) {
                        Some(ref d_tile) => {
                            match self.get_next_tile(col, row, x_step, 0) {
                                Some(ref s_tile)
                                if d_tile.score == s_tile.score
                                   && self.get_tile_count(d_tile.tile_x, d_tile.tile_y) == 1 => {
                                    found = true;
                                    dx = d_tile.tile_x;
                                    dy = d_tile.tile_y;
                                    sx = s_tile.tile_x;
                                    sy = s_tile.tile_y;
                                    break;
                                },
                                _ => {},
                            }
                        },
                        None => {
                            break;
                        }
                    }
                }
                if found {
                    need_generate = true;
                    did_merged = true;
                    let mut tile = self.get_mut_tile(sx, sy);
                    let tile = tile.get_mut_ref();
                    tile.start_moving(dx, dy);
                    println!("merge ({}, {}) to ({}, {})", sx, sy, dx, dy);
                }
            }

            if !did_merged {
                break;
            }
        }

        if need_generate {
            self.generate_tile();
        }
    }

    fn is_locking(&self) -> bool {
        for tile in self.tiles.iter() {
            if tile.status != TileStatic {
                return true;
            }
        }
        false
    }

    /// Returns next tile right besides (x, y)
    fn get_next_tile<'a>(&'a self, x: int, y: int, step_x: int, step_y: int) -> Option<&'a Tile> {
        let mut x = x + step_x;
        let mut y = y + step_y;
        while x >= 0 && x < settings::TILE_WIDTH
              && y >= 0 && y < settings::TILE_HEIGHT {
            let tile = self.get_tile(x, y);
            if tile.is_some() {
                return tile;
            }
            x += step_x;
            y += step_y;
        }
        None
    }

    fn get_mut_next_tile<'a>(&'a mut self, x: int, y: int, step_x: int, step_y: int) -> Option<&'a mut Tile> {
        let mut x = x + step_x;
        let mut y = y + step_y;
        let mut found = false;
        while x >= 0 && x < settings::TILE_WIDTH
              && y >= 0 && y < settings::TILE_HEIGHT {
            let tile = self.get_tile(x, y);
            if tile.is_some() {
                found = true;
                break;
            }
            x += step_x;
            y += step_y;
        }

        if found {
            self.get_mut_tile(x, y)
        } else {
            None
        }
    }

    fn get_tile<'a>(&'a self, x: int, y: int) -> Option<&'a Tile> {
        for tile in self.tiles.iter() {
            if tile.tile_x == x && tile.tile_y == y {
                return Some(tile);
            }
        }
        None
    }

    fn get_mut_tile<'a>(&'a mut self, x: int, y: int) -> Option<&'a mut Tile> {
        for tile in self.tiles.mut_iter() {
            if tile.tile_x == x && tile.tile_y == y {
                return Some(tile);
            }
        }
        None
    }

    fn get_tile_count(&self, x: int, y: int) -> int {
        let mut count = 0;
        for tile in self.tiles.iter() {
            if tile.tile_x == x && tile.tile_y == y {
                count += 1;
            }
        }
        count
    }

    fn render_board(&self, c: &Context, gl: &mut Gl) {
        c.view()
         .rect(settings::BOARD_PADDING,
               settings::BOARD_PADDING + settings::BOARD_OFFSET_Y,
               settings::BOARD_SIZE[0],
               settings::BOARD_SIZE[1])
         .rgba(settings::TILE_BACKGROUND_COLOR[0],
               settings::TILE_BACKGROUND_COLOR[1],
               settings::TILE_BACKGROUND_COLOR[2],
               settings::TILE_BACKGROUND_COLOR[3])
         .fill(gl);

        let mut x = settings::BOARD_PADDING + settings::TILE_PADDING;
        let mut y = settings::BOARD_PADDING + settings::BOARD_OFFSET_Y + settings::TILE_PADDING;
        for _ in range(0, settings::TILE_HEIGHT) {
            for _ in range(0, settings::TILE_WIDTH) {
                c.view()
                 .rect(x, y, settings::TILE_SIZE, settings::TILE_SIZE)
                 .rgba(settings::TILES_COLOR[0][0],
                       settings::TILES_COLOR[0][1],
                       settings::TILES_COLOR[0][2],
                       settings::TILES_COLOR[0][3])
                 .fill(gl);

                x += settings::TILE_PADDING + settings::TILE_SIZE;
            }
            x = settings::BOARD_PADDING + settings::TILE_PADDING;
            y += settings::TILE_PADDING + settings::TILE_SIZE;
        }
    }

    fn render_tiles(&self, c: &Context, gl: &mut Gl) {
        for tile in self.tiles.iter() {
            tile.render(c, gl);
        }
    }

    fn add_score(&mut self, score: int) {
        self.score += score;
        if self.score > self.highest_score {
            self.highest_score = self.score;
        }
        println!("Score: {}, Highest Score: {}", self.score, self.highest_score);
    }
}

<|endoftext|>


<commit_before>extern crate lz4;

use std::env;
use std::fs;
use std::fs::{File, PathExt, OpenOptions};
use std::io::{Error, ErrorKind, Read, Write, Seek, SeekFrom};
use std::sync::{Arc, Mutex};
use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::hash::{Hasher, SipHasher};
use std::path::{Path, PathBuf};

use super::compiler::OutputInfo;
use super::utils::hash_write_stream;
use super::utils::DEFAULT_BUF_SIZE;
use super::io::binary::*;

const HEADER: &'static [u8] = b"OBCF\x00\x01";
const FOOTER: &'static [u8] = b"END\x00";
const SUFFIX: &'static str = ".lz4";

struct FileHash {
	hash: String,
	size: u64,
	modified: u64,
}

#[derive(Clone)]
pub struct Cache {
	file_hash: Arc<Mutex<HashMap<PathBuf, Arc<Mutex<Option<FileHash>>>>>>,
	cache_dir: PathBuf
}

struct CacheFile {
	path: PathBuf,
	size: u64,
	accessed: u64,
}

impl Cache {
	pub fn new() -> Self {
		let cache_dir = match env::var("OCTOBUILD_CACHE") {
			Ok(value) => Path::new(&value).to_path_buf(),
			Err(_) => env::home_dir().unwrap().join(".octobuild").join("cache")
		};
		Cache {
			file_hash: Arc::new(Mutex::new(HashMap::new())),
			cache_dir: cache_dir
		}
	}

	pub fn run_cached<F: Fn()->Result<OutputInfo, Error>>(&self, params: &str, inputs: &Vec<PathBuf>, outputs: &Vec<PathBuf>, worker: F) -> Result<OutputInfo, Error> {
		let hash = try! (self.generate_hash(params, inputs));
		let path = self.cache_dir.join(&hash[0..2]).join(&(hash[2..].to_string() + SUFFIX));
		// Try to read data from cache.
		match read_cache(&path, outputs) {
			Ok(output) => {return Ok(output)}
			Err(_) => {}
		}
		// Run task and save result to cache.
		let output = try !(worker());
		try !(write_cache(&path, outputs, &output));
		Ok(output)
	}

	pub fn cleanup(&self, max_cache_size: u64) -> Result<(), Error> {
		let mut files: Vec<CacheFile> = Vec::new();
		for item in try! (fs::walk_dir(&self.cache_dir)) {
			let path = try! (item).path();
			if path.to_str().map_or(false, |v| v.ends_with(SUFFIX)) {
				let attr = try! (fs::metadata(&path));
				if attr.is_file() {
					files.push(CacheFile {
						path: path,
						size: attr.len(),
						accessed: attr.modified(),
					});
				}
			}
		}
		files.as_mut_slice().sort_by(|a, b| b.accessed.cmp(&a.accessed));

		let mut cache_size: u64 = 0;
		for item in files.into_iter() {
			cache_size += item.size;
			if cache_size > max_cache_size {
				let _ = try!(fs::remove_file(&item.path));
			}
		}
		Ok(())
	}

	fn generate_hash(&self, params: &str, inputs: &Vec<PathBuf>) -> Result<String, Error> {
		let mut sip_hash = SipHasher::new();
		let hash: &mut Hasher = &mut sip_hash;
		// str
		hash.write(params.as_bytes());
		hash.write_u8(0);
		// inputs
		for input in inputs.iter() {
			let file_hash = try! (self.get_file_hash(input));
			hash.write(file_hash.as_bytes());
		}
		Ok(format!("{:016x}", hash.finish()))
	}

	pub fn get_file_hash(&self, path: &Path) -> Result<String, Error> {
		// Get/create lock for file entry.
		let hash_lock = match self.file_hash.lock() {
			Ok(mut map) => {
				match map.entry(path.to_path_buf()) {
					Entry::Occupied(entry) => entry.get().clone(),
					Entry::Vacant(entry) => entry.insert(Arc::new(Mutex::new(None))).clone()
				}
			}
			Err(e) => {
				return Err(Error::new(ErrorKind::Other, "Mutex error", Some(e.to_string())));
			}
		};
		// Get file hash.
		let result = match hash_lock.lock() { // rust: #22722
			Ok(mut hash_entry) => {
				// Validate entry, if exists.
				match *hash_entry {
					Some(ref value) => {
						let stat = try! (fs::metadata(path));
						if value.size == stat.len() && value.modified == stat.modified() {
							return Ok(value.hash.clone());
						}
					}
					None => {}
				}
				// Calculate hash value.
				let stat = try! (fs::metadata(path));
				let hash = try! (generate_file_hash(path));
				*hash_entry = Some(FileHash {
					hash: hash.clone(),
					size: stat.len(),
					modified: stat.modified(),
				});
				Ok(hash)
			}
			Err(e) => Err(Error::new(ErrorKind::Other, "Mutex error", Some(e.to_string())))
		};
		result
	}
}

fn generate_file_hash(path: &Path) -> Result<String, Error> {
	let mut hash = SipHasher::new();
	let mut file = try! (File::open(path));
	try! (hash_write_stream(&mut hash, &mut file));
	Ok(format!("{:016x}", hash.finish()))
}

fn write_cache(path: &Path, paths: &Vec<PathBuf>, output: &OutputInfo) -> Result<(), Error> {
	if !output.success() {
		return Ok(());
	}
	match path.parent() {
		Some(parent) => try! (fs::create_dir_all(&parent)),
		None => ()
	}
	let mut stream = try! (lz4::Encoder::new(try! (File::create(path)), 1));
	try! (stream.write_all(HEADER));
	try! (write_le_usize(&mut stream, paths.len()));
	let mut buf: [u8; DEFAULT_BUF_SIZE] = [0; DEFAULT_BUF_SIZE];
	for path in paths.iter() {
		let mut file = try! (File::open(path));
		loop {
			let size = try! (file.read(&mut buf));
			if size <= 0 {
				break;
			}
			try! (write_le_usize(&mut stream, size));
			try! (stream.write_all(&buf[0..size]));
		}
		try! (write_le_usize(&mut stream, 0));
	}
	try! (write_output(&mut stream, output));
	try! (stream.write_all(FOOTER));
	match stream.finish() {
		(_, result) => result
	}
}

fn read_cache(path: &Path, paths: &Vec<PathBuf>) -> Result<OutputInfo, Error> {
	let mut file = try! (OpenOptions::new().read(true).write(true).open(Path::new(path)));
	try! (file.write(&[4]));
	try! (file.seek(SeekFrom::Start(0)));
	let mut stream = try! (lz4::Decoder::new (file));
	if try! (read_exact(&mut stream, HEADER.len())) != HEADER {
		return Err(Error::new(ErrorKind::InvalidInput, "Invalid cache file header", Some(path.display().to_string())));
	}
	if try! (read_le_usize(&mut stream)) != paths.len() {
		return Err(Error::new(ErrorKind::InvalidInput, "Unexpected count of packed cached files", Some(path.display().to_string())));
	}
	for path in paths.iter() {
		let mut file = try! (File::create(path));
		loop {
			let size = try! (read_le_usize(&mut stream));
			if size == 0 {break;}
			let block = try! (read_exact(&mut stream, size));
			try! (file.write_all(&block));
		}
	}
	let output = try! (read_output(&mut stream));
	if try! (read_exact(&mut stream, FOOTER.len())) != FOOTER {
		return Err(Error::new(ErrorKind::InvalidInput, "Invalid cache file footer", Some(path.display().to_string())));
	}
	Ok(output)
}

fn write_blob(stream: &mut Write, blob: &[u8]) -> Result<(), Error> {
	try! (write_le_usize(stream, blob.len()));
	try! (stream.write_all(blob));
	Ok(())
}

fn read_blob(stream: &mut Read) -> Result<Vec<u8>, Error> {
	let size = try! (read_le_usize(stream));
	read_exact(stream, size)
}

fn write_output(stream: &mut Write, output: &OutputInfo) -> Result<(), Error> {
	try! (write_blob(stream, &output.stdout));
	try! (write_blob(stream, &output.stderr));
	Ok(())
}

fn read_output(stream: &mut Read) -> Result<OutputInfo, Error> {
	let stdout = try! (read_blob(stream));
	let stderr = try! (read_blob(stream));
	Ok(OutputInfo {
		status: Some(0),
		stdout: stdout,
		stderr: stderr,
	})
}
<commit_msg>Increment cache format version for invalidating old cache<commit_after>extern crate lz4;

use std::env;
use std::fs;
use std::fs::{File, PathExt, OpenOptions};
use std::io::{Error, ErrorKind, Read, Write, Seek, SeekFrom};
use std::sync::{Arc, Mutex};
use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::hash::{Hasher, SipHasher};
use std::path::{Path, PathBuf};

use super::compiler::OutputInfo;
use super::utils::hash_write_stream;
use super::utils::DEFAULT_BUF_SIZE;
use super::io::binary::*;

const HEADER: &'static [u8] = b"OBCF\x00\x02";
const FOOTER: &'static [u8] = b"END\x00";
const SUFFIX: &'static str = ".lz4";

struct FileHash {
	hash: String,
	size: u64,
	modified: u64,
}

#[derive(Clone)]
pub struct Cache {
	file_hash: Arc<Mutex<HashMap<PathBuf, Arc<Mutex<Option<FileHash>>>>>>,
	cache_dir: PathBuf
}

struct CacheFile {
	path: PathBuf,
	size: u64,
	accessed: u64,
}

impl Cache {
	pub fn new() -> Self {
		let cache_dir = match env::var("OCTOBUILD_CACHE") {
			Ok(value) => Path::new(&value).to_path_buf(),
			Err(_) => env::home_dir().unwrap().join(".octobuild").join("cache")
		};
		Cache {
			file_hash: Arc::new(Mutex::new(HashMap::new())),
			cache_dir: cache_dir
		}
	}

	pub fn run_cached<F: Fn()->Result<OutputInfo, Error>>(&self, params: &str, inputs: &Vec<PathBuf>, outputs: &Vec<PathBuf>, worker: F) -> Result<OutputInfo, Error> {
		let hash = try! (self.generate_hash(params, inputs));
		let path = self.cache_dir.join(&hash[0..2]).join(&(hash[2..].to_string() + SUFFIX));
		// Try to read data from cache.
		match read_cache(&path, outputs) {
			Ok(output) => {return Ok(output)}
			Err(_) => {}
		}
		// Run task and save result to cache.
		let output = try !(worker());
		try !(write_cache(&path, outputs, &output));
		Ok(output)
	}

	pub fn cleanup(&self, max_cache_size: u64) -> Result<(), Error> {
		let mut files: Vec<CacheFile> = Vec::new();
		for item in try! (fs::walk_dir(&self.cache_dir)) {
			let path = try! (item).path();
			if path.to_str().map_or(false, |v| v.ends_with(SUFFIX)) {
				let attr = try! (fs::metadata(&path));
				if attr.is_file() {
					files.push(CacheFile {
						path: path,
						size: attr.len(),
						accessed: attr.modified(),
					});
				}
			}
		}
		files.as_mut_slice().sort_by(|a, b| b.accessed.cmp(&a.accessed));

		let mut cache_size: u64 = 0;
		for item in files.into_iter() {
			cache_size += item.size;
			if cache_size > max_cache_size {
				let _ = try!(fs::remove_file(&item.path));
			}
		}
		Ok(())
	}

	fn generate_hash(&self, params: &str, inputs: &Vec<PathBuf>) -> Result<String, Error> {
		let mut sip_hash = SipHasher::new();
		let hash: &mut Hasher = &mut sip_hash;
		// str
		hash.write(params.as_bytes());
		hash.write_u8(0);
		// inputs
		for input in inputs.iter() {
			let file_hash = try! (self.get_file_hash(input));
			hash.write(file_hash.as_bytes());
		}
		Ok(format!("{:016x}", hash.finish()))
	}

	pub fn get_file_hash(&self, path: &Path) -> Result<String, Error> {
		// Get/create lock for file entry.
		let hash_lock = match self.file_hash.lock() {
			Ok(mut map) => {
				match map.entry(path.to_path_buf()) {
					Entry::Occupied(entry) => entry.get().clone(),
					Entry::Vacant(entry) => entry.insert(Arc::new(Mutex::new(None))).clone()
				}
			}
			Err(e) => {
				return Err(Error::new(ErrorKind::Other, "Mutex error", Some(e.to_string())));
			}
		};
		// Get file hash.
		let result = match hash_lock.lock() { // rust: #22722
			Ok(mut hash_entry) => {
				// Validate entry, if exists.
				match *hash_entry {
					Some(ref value) => {
						let stat = try! (fs::metadata(path));
						if value.size == stat.len() && value.modified == stat.modified() {
							return Ok(value.hash.clone());
						}
					}
					None => {}
				}
				// Calculate hash value.
				let stat = try! (fs::metadata(path));
				let hash = try! (generate_file_hash(path));
				*hash_entry = Some(FileHash {
					hash: hash.clone(),
					size: stat.len(),
					modified: stat.modified(),
				});
				Ok(hash)
			}
			Err(e) => Err(Error::new(ErrorKind::Other, "Mutex error", Some(e.to_string())))
		};
		result
	}
}

fn generate_file_hash(path: &Path) -> Result<String, Error> {
	let mut hash = SipHasher::new();
	let mut file = try! (File::open(path));
	try! (hash_write_stream(&mut hash, &mut file));
	Ok(format!("{:016x}", hash.finish()))
}

fn write_cache(path: &Path, paths: &Vec<PathBuf>, output: &OutputInfo) -> Result<(), Error> {
	if !output.success() {
		return Ok(());
	}
	match path.parent() {
		Some(parent) => try! (fs::create_dir_all(&parent)),
		None => ()
	}
	let mut stream = try! (lz4::Encoder::new(try! (File::create(path)), 1));
	try! (stream.write_all(HEADER));
	try! (write_le_usize(&mut stream, paths.len()));
	let mut buf: [u8; DEFAULT_BUF_SIZE] = [0; DEFAULT_BUF_SIZE];
	for path in paths.iter() {
		let mut file = try! (File::open(path));
		loop {
			let size = try! (file.read(&mut buf));
			if size <= 0 {
				break;
			}
			try! (write_le_usize(&mut stream, size));
			try! (stream.write_all(&buf[0..size]));
		}
		try! (write_le_usize(&mut stream, 0));
	}
	try! (write_output(&mut stream, output));
	try! (stream.write_all(FOOTER));
	match stream.finish() {
		(_, result) => result
	}
}

fn read_cache(path: &Path, paths: &Vec<PathBuf>) -> Result<OutputInfo, Error> {
	let mut file = try! (OpenOptions::new().read(true).write(true).open(Path::new(path)));
	try! (file.write(&[4]));
	try! (file.seek(SeekFrom::Start(0)));
	let mut stream = try! (lz4::Decoder::new (file));
	if try! (read_exact(&mut stream, HEADER.len())) != HEADER {
		return Err(Error::new(ErrorKind::InvalidInput, "Invalid cache file header", Some(path.display().to_string())));
	}
	if try! (read_le_usize(&mut stream)) != paths.len() {
		return Err(Error::new(ErrorKind::InvalidInput, "Unexpected count of packed cached files", Some(path.display().to_string())));
	}
	for path in paths.iter() {
		let mut file = try! (File::create(path));
		loop {
			let size = try! (read_le_usize(&mut stream));
			if size == 0 {break;}
			let block = try! (read_exact(&mut stream, size));
			try! (file.write_all(&block));
		}
	}
	let output = try! (read_output(&mut stream));
	if try! (read_exact(&mut stream, FOOTER.len())) != FOOTER {
		return Err(Error::new(ErrorKind::InvalidInput, "Invalid cache file footer", Some(path.display().to_string())));
	}
	Ok(output)
}

fn write_blob(stream: &mut Write, blob: &[u8]) -> Result<(), Error> {
	try! (write_le_usize(stream, blob.len()));
	try! (stream.write_all(blob));
	Ok(())
}

fn read_blob(stream: &mut Read) -> Result<Vec<u8>, Error> {
	let size = try! (read_le_usize(stream));
	read_exact(stream, size)
}

fn write_output(stream: &mut Write, output: &OutputInfo) -> Result<(), Error> {
	try! (write_blob(stream, &output.stdout));
	try! (write_blob(stream, &output.stderr));
	Ok(())
}

fn read_output(stream: &mut Read) -> Result<OutputInfo, Error> {
	let stdout = try! (read_blob(stream));
	let stderr = try! (read_blob(stream));
	Ok(OutputInfo {
		status: Some(0),
		stdout: stdout,
		stderr: stderr,
	})
}
<|endoftext|>


<commit_before>use std::cmp::min;
use std::io;
use std::str;

use bytes::{Bytes, BytesMut};
use bytes::{IntoBuf, Buf, BufMut, BigEndian};
use tokio_io::codec::{Encoder, Decoder};
use tokio_proto::streaming::pipeline::Frame;

use constants::*;
use packet::{PacketMagic, PTYPES};

#[derive(Debug)]
pub struct PacketHeader {
    pub magic: PacketMagic,
    pub ptype: u32,
    pub psize: u32,
}

pub struct PacketCodec {
    pub data_todo: Option<usize>,
}

type PacketItem = Frame<PacketHeader, BytesMut, io::Error>;

impl PacketHeader {
    pub fn admin_decode(buf: &mut BytesMut) -> Result<Option<PacketItem>, io::Error> {
        let newline = buf[..].iter().position(|b| *b == b'\n');
        if let Some(n) = newline {
            let line = buf.split_to(n);
            buf.split_to(1); // drop the newline itself
            let data_str = match str::from_utf8(&line[..]) {
                Ok(s) => s,
                Err(_) => return Err(io::Error::new(io::ErrorKind::Other, "invalid string")),
            };
            info!("admin command data: {:?}", data_str);
            let command = match data_str.trim() {
                "version" => ADMIN_VERSION,
                "status" => ADMIN_STATUS,
                _ => ADMIN_UNKNOWN,
            };
            return Ok(Some(Frame::Message {
                message: PacketHeader {
                    magic: PacketMagic::TEXT,
                    ptype: command,
                    psize: 0,
                },
                body: false,
            }));
        }
        Ok(None) // Wait for more data
    }

    pub fn decode(buf: &mut BytesMut) -> Result<Option<PacketItem>, io::Error> {
        debug!("Decoding {:?}", buf);
        // Peek at first 4
        // Is this a req/res
        if buf.len() < 4 {
            return Ok(None);
        }
        let mut magic_buf: [u8; 4] = [0; 4];
        magic_buf.clone_from_slice(&buf[0..4]);
        let magic = match magic_buf {
            REQ => PacketMagic::REQ,
            RES => PacketMagic::RES,
            // TEXT/ADMIN protocol
            _ => PacketMagic::TEXT,
        };
        debug!("Magic is {:?}", magic);
        if magic == PacketMagic::TEXT {
            debug!("admin protocol detected");
            return PacketHeader::admin_decode(buf);
        }
        if buf.len() < 12 {
            return Ok(None);
        }
        buf.split_to(4);
        // Now get the type
        let ptype = buf.split_to(4).into_buf().get_u32::<BigEndian>();
        debug!("We got a {}", &PTYPES[ptype as usize].name);
        // Now the length
        let psize = buf.split_to(4).into_buf().get_u32::<BigEndian>();
        debug!("Data section is {} bytes", psize);
        Ok(Some(Frame::Message {
            message: PacketHeader {
                magic: magic,
                ptype: ptype,
                psize: psize,
            },
            body: true, // TODO: false for 0 psize?
        }))
    }

    pub fn to_bytes(&self) -> Bytes {
        let magic = match self.magic {
            PacketMagic::UNKNOWN => panic!("Unknown packet magic cannot be sent"),
            PacketMagic::REQ => REQ,
            PacketMagic::RES => RES,
            PacketMagic::TEXT => {
                return Bytes::from_static(b"");
            }
        };
        let mut buf = BytesMut::with_capacity(12);
        buf.extend(magic.iter());
        buf.put_u32::<BigEndian>(self.ptype);
        buf.put_u32::<BigEndian>(self.psize);
        buf.freeze()
    }

    pub fn new_text_res(body: &BytesMut) -> PacketHeader {
        PacketHeader {
            magic: PacketMagic::TEXT,
            ptype: ADMIN_RESPONSE,
            psize: body.len() as u32,
        }
    }

    pub fn noop() -> PacketHeader {
        PacketHeader {
            magic: PacketMagic::TEXT,
            ptype: ADMIN_RESPONSE,
            psize: 0,
        }
    }
}

impl Decoder for PacketCodec {
    type Item = Frame<PacketHeader, BytesMut, io::Error>;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, io::Error> {
        match self.data_todo {
            None => {
                match PacketHeader::decode(buf)? {
                    Some(Frame::Message { message, body }) => {
                        self.data_todo = Some(message.psize as usize);
                        Ok(Some(Frame::Message {
                            message: message,
                            body: body,
                        }))
                    }
                    Some(_) => panic!("Expecting Frame::Message, got something else"),
                    None => Ok(None),
                }
            }
            Some(0) => {
                self.data_todo = None;
                Ok(Some(Frame::Body { chunk: None }))
            }
            Some(data_todo) => {
                let chunk_size = min(buf.len(), data_todo);
                self.data_todo = Some(data_todo - chunk_size);
                Ok(Some(Frame::Body { chunk: Some(buf.split_to(chunk_size)) }))
            }
        }
    }
}

impl Encoder for PacketCodec {
    type Item = Frame<PacketHeader, BytesMut, io::Error>;
    type Error = io::Error;

    fn encode(&mut self, msg: Self::Item, buf: &mut BytesMut) -> io::Result<()> {
        debug!("Encoding {:?}", msg);
        match msg {
            Frame::Message { message, body } => {
                if body {
                    debug!("body follows")
                }
                buf.extend(message.to_bytes())
            }
            Frame::Body { chunk } => {
                match chunk {
                    Some(chunk) => buf.extend_from_slice(&chunk[..]),
                    None => {}
                }
            }
            Frame::Error { error } => return Err(error),
        }
        Ok(())
    }
}
<commit_msg>Support error frames<commit_after>use std::cmp::min;
use std::io;
use std::str;

use bytes::{Bytes, BytesMut};
use bytes::{IntoBuf, Buf, BufMut, BigEndian};
use tokio_io::codec::{Encoder, Decoder};
use tokio_proto::streaming::pipeline::Frame;

use constants::*;
use packet::{PacketMagic, PTYPES};

#[derive(Debug)]
pub struct PacketHeader {
    pub magic: PacketMagic,
    pub ptype: u32,
    pub psize: u32,
}

pub struct PacketCodec {
    pub data_todo: Option<usize>,
}

type PacketItem = Frame<PacketHeader, BytesMut, io::Error>;

impl PacketHeader {
    pub fn admin_decode(buf: &mut BytesMut) -> Result<Option<PacketItem>, io::Error> {
        let newline = buf[..].iter().position(|b| *b == b'\n');
        if let Some(n) = newline {
            let line = buf.split_to(n);
            buf.split_to(1); // drop the newline itself
            let data_str = match str::from_utf8(&line[..]) {
                Ok(s) => s,
                Err(_) => return Err(io::Error::new(io::ErrorKind::Other, "invalid string")),
            };
            info!("admin command data: {:?}", data_str);
            let command = match data_str.trim() {
                "version" => ADMIN_VERSION,
                "status" => ADMIN_STATUS,
                _ => ADMIN_UNKNOWN,
            };
            return Ok(Some(Frame::Message {
                message: PacketHeader {
                    magic: PacketMagic::TEXT,
                    ptype: command,
                    psize: 0,
                },
                body: false,
            }));
        }
        Ok(None) // Wait for more data
    }

    pub fn decode(buf: &mut BytesMut) -> Result<Option<PacketItem>, io::Error> {
        debug!("Decoding {:?}", buf);
        // Peek at first 4
        // Is this a req/res
        if buf.len() < 4 {
            return Ok(None);
        }
        let mut magic_buf: [u8; 4] = [0; 4];
        magic_buf.clone_from_slice(&buf[0..4]);
        let magic = match magic_buf {
            REQ => PacketMagic::REQ,
            RES => PacketMagic::RES,
            // TEXT/ADMIN protocol
            _ => PacketMagic::TEXT,
        };
        debug!("Magic is {:?}", magic);
        if magic == PacketMagic::TEXT {
            debug!("admin protocol detected");
            return PacketHeader::admin_decode(buf);
        }
        if buf.len() < 12 {
            return Ok(None);
        }
        buf.split_to(4);
        // Now get the type
        let ptype = buf.split_to(4).into_buf().get_u32::<BigEndian>();
        debug!("We got a {}", &PTYPES[ptype as usize].name);
        // Now the length
        let psize = buf.split_to(4).into_buf().get_u32::<BigEndian>();
        debug!("Data section is {} bytes", psize);
        Ok(Some(Frame::Message {
            message: PacketHeader {
                magic: magic,
                ptype: ptype,
                psize: psize,
            },
            body: true, // TODO: false for 0 psize?
        }))
    }

    pub fn to_bytes(&self) -> Bytes {
        let magic = match self.magic {
            PacketMagic::UNKNOWN => panic!("Unknown packet magic cannot be sent"),
            PacketMagic::REQ => REQ,
            PacketMagic::RES => RES,
            PacketMagic::TEXT => {
                return Bytes::from_static(b"");
            }
        };
        let mut buf = BytesMut::with_capacity(12);
        buf.extend(magic.iter());
        buf.put_u32::<BigEndian>(self.ptype);
        buf.put_u32::<BigEndian>(self.psize);
        buf.freeze()
    }

    pub fn new_text_res(body: &BytesMut) -> PacketHeader {
        PacketHeader {
            magic: PacketMagic::TEXT,
            ptype: ADMIN_RESPONSE,
            psize: body.len() as u32,
        }
    }

    pub fn noop() -> PacketHeader {
        PacketHeader {
            magic: PacketMagic::TEXT,
            ptype: ADMIN_RESPONSE,
            psize: 0,
        }
    }
}

impl Decoder for PacketCodec {
    type Item = Frame<PacketHeader, BytesMut, io::Error>;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, io::Error> {
        match self.data_todo {
            None => {
                match PacketHeader::decode(buf)? {
                    Some(Frame::Message { message, body }) => {
                        self.data_todo = Some(message.psize as usize);
                        Ok(Some(Frame::Message {
                            message: message,
                            body: body,
                        }))
                    }
                    Some(_) => panic!("Expecting Frame::Message, got something else"),
                    None => Ok(None),
                }
            }
            Some(0) => {
                self.data_todo = None;
                Ok(Some(Frame::Body { chunk: None }))
            }
            Some(data_todo) => {
                let chunk_size = min(buf.len(), data_todo);
                self.data_todo = Some(data_todo - chunk_size);
                Ok(Some(Frame::Body { chunk: Some(buf.split_to(chunk_size)) }))
            }
        }
    }
}

impl Encoder for PacketCodec {
    type Item = Frame<PacketHeader, BytesMut, io::Error>;
    type Error = io::Error;

    fn encode(&mut self, msg: Self::Item, buf: &mut BytesMut) -> io::Result<()> {
        debug!("Encoding {:?}", msg);
        match msg {
            Frame::Message { message, body } => {
                if body {
                    debug!("body follows")
                }
                buf.extend(message.to_bytes())
            }
            Frame::Body { chunk } => {
                match chunk {
                    Some(chunk) => buf.extend_from_slice(&chunk[..]),
                    None => {}
                }
            }
            Frame::Error { error } => {
                error!("Sending error frame. {}", error);
                buf.extend("ERR UNKNOWN_COMMAND Unknown+server+command".bytes())
            }
        }
        Ok(())
    }
}
<|endoftext|>