<commit_before>//! Lua functionality

use hlua;
use hlua::{Lua, LuaError};
use hlua::any::AnyLuaValue;

use std::thread;

use std::fs::{File};
use std::path::Path;
use std::io::Write;

use std::sync::{Mutex, RwLock};
use std::sync::mpsc::{channel, Sender, Receiver};

lazy_static! {
    /// Sends requests to the lua thread
    static ref SENDER: Mutex<Sender<LuaQuery>> = {
        let (tx, rx) = channel::<LuaQuery>();
        Mutex::new(tx)
    };

    /// Receives data back from the lua thread
    /// This should only be accessed by the lua thread itself.
    static ref RECEIVER: Mutex<Receiver<LuaResponse>> = {
        let (tx, rx) = channel::<LuaResponse>();
        Mutex::new(rx)
    };

    /// Whether the lua thread is currently running
    pub static ref RUNNING: RwLock<bool> = RwLock::new(false);
}

/// Messages sent to the lua thread
pub enum LuaQuery {
    /// Halt the lua thread
    Terminate,
    // Restart the lua thread
    Restart,
    /// Execute a string
    Execute(String),
    /// Execute a file
    ExecuteFile(String),
    /// Get a variable
    GetVariable(String),
    /// Set a value
    SetValue {
        name: Box<::std::borrow::Borrow<str> + Sized>,
        val: Box<hlua::Push<&'static mut Lua<'static>> + Sized>
    },
    /// Create a new array
    EmptyArray(String),
    /// Unused send type
    Unused,
}

/// Messages received from lua thread
pub enum LuaResponse {
    /// Lua variable obtained
    Variable(Option<AnyLuaValue>),
    /// Lua error
    Error(hlua::LuaError),
    /// A function is returned
    Function(hlua::functions_read::LuaFunction<String>)
}

unsafe impl Send for LuaQuery { }
unsafe impl Send for LuaResponse { }
unsafe impl Sync for LuaQuery { }
unsafe impl Sync for LuaResponse { }

/// Whether the lua thread is currently available
pub fn thread_running() -> bool {
    *RUNNING.read().unwrap()
}

/// Errors which may arise from attempting
/// to sending a message to the lua thread.
pub enum LuaSendError {
    ThreadClosed,
    Sender
}

/// Attemps to send a LuaQuery to the lua thread.
pub fn try_send(query: LuaQuery) -> Result<(), LuaSendError> {
    if !thread_running() { LuaSendError::ThreadClosed }
    else {
        match SENDER.lock().unwrap().send(query) {
            Ok(_) => Ok(()),
            Err(_) => LuaSendError::Sender
        }
    }
}

/// Sends a value to the lua thread.
///
/// # Panics
/// * If the lua thread is currently not running (check `lua::thread_running()`)
/// * If the sender was unable to send, which should only happen
/// if the lua thread isn't running.
pub fn send(query: LuaQuery) {
    if !thread_running() {
        panic!("lua: Attempted to message the lua thread while it is not active!");
    }
    SENDER.lock().unwrap().send(query).unwrap();
}

/// Initialize the lua thread
pub fn init() {
    trace!("Initializing...");
    let (query_tx, query_rx) = channel::<LuaQuery>();
    let (answer_tx, answer_rx) = channel::<LuaResponse>();
    {
        let mut sender = SENDER.lock().unwrap();
        let mut receiver = RECEIVER.lock().unwrap();

        *sender = query_tx;
        *receiver = answer_rx;
    }

    thread::spawn(move || {
        trace!("thread: Inside thread!");
        let sender = answer_tx;
        let receiver = query_rx;
        let mut lua = Lua::new();
        trace!("thread: Loading libraries...");
        lua.openlibs();
        trace!("thread: Libraries loaded");
        *RUNNING.write().unwrap() = true;
        trace!("thread: Testing hello world...");
        let mut file = File::create("/tmp/init.lua").unwrap();
        file.write(b"print('Hello world!')").unwrap();
        lua.execute_from_reader::<(), File>(File::open("/tmp/init.lua")
                                            .unwrap()).unwrap();
        trace!("thread: Done!");
        trace!("thread: Entering loop...");
        thread_main_loop(sender, receiver, &mut lua);
    });
    trace!("Created thread. Init finished.");
}

fn thread_main_loop(sender: Sender<LuaResponse>, receiver: Receiver<LuaQuery>,
                    lua: &mut Lua) {
    loop {
        let request = receiver.recv();
        match request {
            Err(e) => {
                error!("Lua thread: unable to receive message: {}", e);
                error!("Lua thread: now panicking!");
                *RUNNING.write().unwrap() = false;

                panic!("Lua thread: lost contact with host, exiting!");
            }
            Ok(message) => {
                trace!("Handling a request");
                thread_handle_message(&sender, message, lua);
            }
        }
    }
}

fn thread_handle_message(sender: &Sender<LuaResponse>, request: LuaQuery, lua: &mut Lua) {
    match request {
        LuaQuery::Terminate => {
            trace!("thread: Received terminate signal");
            *RUNNING.write().unwrap() = false;

            info!("thread: Lua thread terminating!");
            return;
        },

        LuaQuery::Restart => {
            trace!("thread: Received restart signal!");
            error!("thread: Lua thread restart not supported!");

            *RUNNING.write().unwrap() = false;

            panic!("Lua thread: Restart not supported!");
        },

        LuaQuery::Execute(code) => {
            trace!("thread: Received request to execute code");
            trace!("thread: Executing {:?}", code);

            match lua.execute::<()>(&code) {
                Err(error) => {
                    warn!("thread: Error executing code: {:?}", error);
                    let response = LuaResponse::Error(error);

                    thread_send(&sender, response);
                }
                Ok(_) => {
                    // This is gonna be really spammy one day
                    trace!("thread: Code executed okay.");
                }
            }
        },

        LuaQuery::ExecuteFile(name) => {
            trace!("thread: Received request to execute file {}", name);
            info!("thread: Executing {}", name);

            let path = Path::new(&name);
            let try_file = File::open(path);

            if let Ok(file) = try_file {
                let result = lua.execute_from_reader::<(), File>(file);
                if let Err(err) = result {
                    warn!("thread: Error executing {}!", name);

                    thread_send(&sender, LuaResponse::Error(err));
                }
                else {
                    trace!("thread: Execution of {} successful.", name);
                }
            }
            else { // Could not open file
                // Unwrap_err is used because we're in the else of let
                let read_error =
                    LuaError::ReadError(try_file.unwrap_err());

                thread_send(&sender, LuaResponse::Error(read_error));
            }
        },

        LuaQuery::GetVariable(varname) => {
            trace!("thread: Received request to get variable {}", varname);
            let var_result = lua.get(varname.as_str());

            match var_result {
                Some(var) => {
                    thread_send(&sender, LuaResponse::Variable(Some(var)));
                }
                None => {
                    warn!("thread: Unable to get variable {}", varname);

                    thread_send(&sender, LuaResponse::Variable(None));
                }
            }
        },

        LuaQuery::SetValue { name: name, val: val } => {

        },

        LuaQuery::EmptyArray(name) => {

        },

        _ => {
            panic!("Unimplemented send type for lua thread!");
        }
    }
}

fn thread_send(sender: &Sender<LuaResponse>, response: LuaResponse) {
    trace!("Called thread_send");
    match sender.send(response) {
        Err(e) => {
            error!("thread: Unable to broadcast response!");
            error!("thread: Shutting down in response to inability \
                    to continue!");
            panic!("Lua thread unable to communicate with main thread, \
                    shutting down!");
        }
        Ok(_) => {}
    }
}
<commit_msg>Fix try_send and build<commit_after>//! Lua functionality

use hlua;
use hlua::{Lua, LuaError};
use hlua::any::AnyLuaValue;

use std::thread;

use std::fs::{File};
use std::path::Path;
use std::io::Write;

use std::sync::{Mutex, RwLock};
use std::sync::mpsc::{channel, Sender, Receiver};

lazy_static! {
    /// Sends requests to the lua thread
    static ref SENDER: Mutex<Sender<LuaQuery>> = {
        let (tx, rx) = channel::<LuaQuery>();
        Mutex::new(tx)
    };

    /// Receives data back from the lua thread
    /// This should only be accessed by the lua thread itself.
    static ref RECEIVER: Mutex<Receiver<LuaResponse>> = {
        let (tx, rx) = channel::<LuaResponse>();
        Mutex::new(rx)
    };

    /// Whether the lua thread is currently running
    pub static ref RUNNING: RwLock<bool> = RwLock::new(false);
}

/// Messages sent to the lua thread
pub enum LuaQuery {
    /// Halt the lua thread
    Terminate,
    // Restart the lua thread
    Restart,
    /// Execute a string
    Execute(String),
    /// Execute a file
    ExecuteFile(String),
    /// Get a variable
    GetVariable(String),
    /// Set a value
    SetValue {
        name: Box<::std::borrow::Borrow<str> + Sized>,
        val: Box<hlua::Push<&'static mut Lua<'static>> + Sized>
    },
    /// Create a new array
    EmptyArray(String),
    /// Unused send type
    Unused,
}

/// Messages received from lua thread
pub enum LuaResponse {
    /// Lua variable obtained
    Variable(Option<AnyLuaValue>),
    /// Lua error
    Error(hlua::LuaError),
    /// A function is returned
    Function(hlua::functions_read::LuaFunction<String>)
}

unsafe impl Send for LuaQuery { }
unsafe impl Send for LuaResponse { }
unsafe impl Sync for LuaQuery { }
unsafe impl Sync for LuaResponse { }

/// Whether the lua thread is currently available
pub fn thread_running() -> bool {
    *RUNNING.read().unwrap()
}

/// Errors which may arise from attempting
/// to sending a message to the lua thread.
pub enum LuaSendError {
    ThreadClosed,
    Sender
}

/// Attemps to send a LuaQuery to the lua thread.
pub fn try_send(query: LuaQuery) -> Result<(), LuaSendError> {
    if !thread_running() { Err(LuaSendError::ThreadClosed) }
    else {
        match SENDER.lock().unwrap().send(query) {
            Ok(_) => Ok(()),
            Err(_) => Err(LuaSendError::Sender)
        }
    }
}

/// Sends a value to the lua thread.
///
/// # Panics
/// * If the lua thread is currently not running (check `lua::thread_running()`)
/// * If the sender was unable to send, which should only happen
/// if the lua thread isn't running.
pub fn send(query: LuaQuery) {
    if !thread_running() {
        panic!("lua: Attempted to message the lua thread while it is not active!");
    }
    SENDER.lock().unwrap().send(query).unwrap();
}

/// Initialize the lua thread
pub fn init() {
    trace!("Initializing...");
    let (query_tx, query_rx) = channel::<LuaQuery>();
    let (answer_tx, answer_rx) = channel::<LuaResponse>();
    {
        let mut sender = SENDER.lock().unwrap();
        let mut receiver = RECEIVER.lock().unwrap();

        *sender = query_tx;
        *receiver = answer_rx;
    }

    thread::spawn(move || {
        trace!("thread: Inside thread!");
        let sender = answer_tx;
        let receiver = query_rx;
        let mut lua = Lua::new();
        trace!("thread: Loading libraries...");
        lua.openlibs();
        trace!("thread: Libraries loaded");
        *RUNNING.write().unwrap() = true;
        trace!("thread: Testing hello world...");
        let mut file = File::create("/tmp/init.lua").unwrap();
        file.write(b"print('Hello world!')").unwrap();
        lua.execute_from_reader::<(), File>(File::open("/tmp/init.lua")
                                            .unwrap()).unwrap();
        trace!("thread: Done!");
        trace!("thread: Entering loop...");
        thread_main_loop(sender, receiver, &mut lua);
    });
    trace!("Created thread. Init finished.");
}

fn thread_main_loop(sender: Sender<LuaResponse>, receiver: Receiver<LuaQuery>,
                    lua: &mut Lua) {
    loop {
        let request = receiver.recv();
        match request {
            Err(e) => {
                error!("Lua thread: unable to receive message: {}", e);
                error!("Lua thread: now panicking!");
                *RUNNING.write().unwrap() = false;

                panic!("Lua thread: lost contact with host, exiting!");
            }
            Ok(message) => {
                trace!("Handling a request");
                thread_handle_message(&sender, message, lua);
            }
        }
    }
}

fn thread_handle_message(sender: &Sender<LuaResponse>, request: LuaQuery, lua: &mut Lua) {
    match request {
        LuaQuery::Terminate => {
            trace!("thread: Received terminate signal");
            *RUNNING.write().unwrap() = false;

            info!("thread: Lua thread terminating!");
            return;
        },

        LuaQuery::Restart => {
            trace!("thread: Received restart signal!");
            error!("thread: Lua thread restart not supported!");

            *RUNNING.write().unwrap() = false;

            panic!("Lua thread: Restart not supported!");
        },

        LuaQuery::Execute(code) => {
            trace!("thread: Received request to execute code");
            trace!("thread: Executing {:?}", code);

            match lua.execute::<()>(&code) {
                Err(error) => {
                    warn!("thread: Error executing code: {:?}", error);
                    let response = LuaResponse::Error(error);

                    thread_send(&sender, response);
                }
                Ok(_) => {
                    // This is gonna be really spammy one day
                    trace!("thread: Code executed okay.");
                }
            }
        },

        LuaQuery::ExecuteFile(name) => {
            trace!("thread: Received request to execute file {}", name);
            info!("thread: Executing {}", name);

            let path = Path::new(&name);
            let try_file = File::open(path);

            if let Ok(file) = try_file {
                let result = lua.execute_from_reader::<(), File>(file);
                if let Err(err) = result {
                    warn!("thread: Error executing {}!", name);

                    thread_send(&sender, LuaResponse::Error(err));
                }
                else {
                    trace!("thread: Execution of {} successful.", name);
                }
            }
            else { // Could not open file
                // Unwrap_err is used because we're in the else of let
                let read_error =
                    LuaError::ReadError(try_file.unwrap_err());

                thread_send(&sender, LuaResponse::Error(read_error));
            }
        },

        LuaQuery::GetVariable(varname) => {
            trace!("thread: Received request to get variable {}", varname);
            let var_result = lua.get(varname.as_str());

            match var_result {
                Some(var) => {
                    thread_send(&sender, LuaResponse::Variable(Some(var)));
                }
                None => {
                    warn!("thread: Unable to get variable {}", varname);

                    thread_send(&sender, LuaResponse::Variable(None));
                }
            }
        },

        LuaQuery::SetValue { name: name, val: val } => {
        },

        LuaQuery::EmptyArray(name) => {

        },

        _ => {
            panic!("Unimplemented send type for lua thread!");
        }
    }
}

fn thread_send(sender: &Sender<LuaResponse>, response: LuaResponse) {
    trace!("Called thread_send");
    match sender.send(response) {
        Err(e) => {
            error!("thread: Unable to broadcast response!");
            error!("thread: Shutting down in response to inability \
                    to continue!");
            panic!("Lua thread unable to communicate with main thread, \
                    shutting down!");
        }
        Ok(_) => {}
    }
}
<|endoftext|>


<commit_before>use std::vec::Vec;

pub struct Machine {
    memory: Vec<u32>,
    counter: u32,
}

impl Machine {
    pub fn new() -> Machine {
        Machine {
            memory: vec![0; 2097152],
            counter: 0,
        }
    }

    pub fn with_data(data: Vec<u32>) -> Machine {
        Machine {
            memory: data,
            counter: 0,
        }
    }

    pub fn loc(&self) -> u32 {
        self.counter
    }

    pub fn dump(&self) -> &Vec<u32> {
        &self.memory
    }

    pub fn load(&mut self, data: Vec<u32>) {
        self.memory = data;
    }

    pub fn step(&mut self) {
        let counter = self.counter as usize;
        let opcode = self.memory[counter];
        let a = self.memory[counter + 1];
        let b = self.memory[counter + 2];
        let c = self.memory[counter + 3];
        self.exec(opcode, a, b, c);
    }

    pub fn exec(&mut self, opcode: u32, a: u32, b: u32, c: u32) {
        match opcode {
            // PC <- M[A]
            1 => {
                self.counter = self.memory[a as usize];
            },

            // If M[B] = 0, then PC <- M[A]
            2 => {
                if 0 == self.memory[b as usize] {
                  self.counter = self.memory[a as usize];
                }
            },

            // M[A] <- PC
            3 => {
                self.memory[a as usize] = self.counter;
                self.counter += 4;
            },

            // M[A] <- M[B]
            4 => {
                self.memory[a as usize] = self.memory[b as usize];
                self.counter += 4;
            },

            // M[A] <- M[M[B]]
            5 => {
                let b = self.memory[b as usize];
                self.memory[a as usize] = self.memory[b as usize];
                self.counter += 4;
            },

            // M[M[B]] <- M[A]
            6 => {
                let b = self.memory[b as usize];
                self.memory[b as usize] = self.memory[a as usize];
                self.counter += 4;
            },

            // M[A] <- M[B] + M[C]
            7 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b + c;
                self.counter += 4;
            },

            // M[A] <- M[B] - M[C]
            8 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b - c;
                self.counter += 4;
            },

            // M[A] <- M[B] x M[C]
            9 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b * c;
                self.counter += 4;
            },

            // M[A] <- M[B] / M[C]
            10 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b / c;
                self.counter += 4;
            },

            // M[A] <- M[B] modulo M[C]
            11 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b % c;
                self.counter += 4;
            },

            // If M[B] < M[C], then M[A] <- 1, else M[A] <- 0
            12 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                if b < c {
                  self.memory[a as usize] = 1;
                } else {
                  self.memory[a as usize] = 0;
                }
                self.counter += 4;
            },

            // Unknown opcode
            _ => {
            },
        }
    }
}

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

    #[test]
    fn it_loads_data() {
        let mut m = Machine::new();
        assert_eq!(Some(&0), m.dump().first());

        m.load(vec![1]);
        assert_eq!(Some(&1), m.dump().first());
    }

    #[test]
    fn it_runs_opcode_1() {
        // PC <- M[A]
        let mut m = Machine::with_data(vec![1, 4, 0, 0, 2]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(2, m.loc());
    }

    #[test]
    fn it_runs_opcode_2_then_branch() {
        // If M[B] = 0, then PC <- M[A]
        let mut m = Machine::with_data(vec![2, 4, 5, 0, 2, 0]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(2, m.loc());
    }

    #[test]
    fn it_runs_opcode_2_else_branch() {
        // If M[B] = 0, then PC <- M[A]
        let mut m = Machine::with_data(vec![2, 4, 5, 0, 2, 1]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(0, m.loc());
    }

    #[test]
    fn it_runs_opcode_3() {
        // M[A] <- PC
        let mut m = Machine::with_data(vec![3, 4, 0, 0, 1]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![3, 4, 0, 0, 0], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_4() {
        // M[A] <- M[B]
        let mut m = Machine::with_data(vec![4, 4, 5, 0, 6, 7]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![4, 4, 5, 0, 7, 7], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_5() {
        // M[A] <- M[M[B]]
        let mut m = Machine::with_data(vec![5, 4, 5, 0, 6, 7, 0, 8]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![5, 4, 5, 0, 8, 7, 0, 8], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_6() {
        // M[M[B]] <- M[A]
        let mut m = Machine::with_data(vec![6, 4, 5, 0, 8, 6, 7]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![6, 4, 5, 0, 8, 6, 8], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_7() {
        // M[A] <- M[B] + M[C]
        let mut m = Machine::with_data(vec![7, 4, 2, 3, 0]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![7, 4, 2, 3, 5], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_8() {
        let mut m = Machine::with_data(vec![0, 3, 2]);

        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(8, 0, 1, 2);
        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_9() {
        let mut m = Machine::with_data(vec![1, 2, 3]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(9, 0, 1, 2);
        assert_eq!(Some(&6), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_10() {
        let mut m = Machine::with_data(vec![1, 6, 2]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(10, 0, 1, 2);
        assert_eq!(Some(&3), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_11() {
        let mut m = Machine::with_data(vec![1, 11, 3]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(11, 0, 1, 2);
        assert_eq!(Some(&2), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_12_then_branch() {
        let mut m = Machine::with_data(vec![2, 3, 4]);

        assert_eq!(Some(&2), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(12, 0, 1, 2);
        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_12_else_branch() {
        let mut m = Machine::with_data(vec![2, 4, 3]);

        assert_eq!(Some(&2), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(12, 0, 1, 2);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }
}
<commit_msg>Replace exec with step for opcode 8.<commit_after>use std::vec::Vec;

pub struct Machine {
    memory: Vec<u32>,
    counter: u32,
}

impl Machine {
    pub fn new() -> Machine {
        Machine {
            memory: vec![0; 2097152],
            counter: 0,
        }
    }

    pub fn with_data(data: Vec<u32>) -> Machine {
        Machine {
            memory: data,
            counter: 0,
        }
    }

    pub fn loc(&self) -> u32 {
        self.counter
    }

    pub fn dump(&self) -> &Vec<u32> {
        &self.memory
    }

    pub fn load(&mut self, data: Vec<u32>) {
        self.memory = data;
    }

    pub fn step(&mut self) {
        let counter = self.counter as usize;
        let opcode = self.memory[counter];
        let a = self.memory[counter + 1];
        let b = self.memory[counter + 2];
        let c = self.memory[counter + 3];
        self.exec(opcode, a, b, c);
    }

    pub fn exec(&mut self, opcode: u32, a: u32, b: u32, c: u32) {
        match opcode {
            // PC <- M[A]
            1 => {
                self.counter = self.memory[a as usize];
            },

            // If M[B] = 0, then PC <- M[A]
            2 => {
                if 0 == self.memory[b as usize] {
                  self.counter = self.memory[a as usize];
                }
            },

            // M[A] <- PC
            3 => {
                self.memory[a as usize] = self.counter;
                self.counter += 4;
            },

            // M[A] <- M[B]
            4 => {
                self.memory[a as usize] = self.memory[b as usize];
                self.counter += 4;
            },

            // M[A] <- M[M[B]]
            5 => {
                let b = self.memory[b as usize];
                self.memory[a as usize] = self.memory[b as usize];
                self.counter += 4;
            },

            // M[M[B]] <- M[A]
            6 => {
                let b = self.memory[b as usize];
                self.memory[b as usize] = self.memory[a as usize];
                self.counter += 4;
            },

            // M[A] <- M[B] + M[C]
            7 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b + c;
                self.counter += 4;
            },

            // M[A] <- M[B] - M[C]
            8 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b - c;
                self.counter += 4;
            },

            // M[A] <- M[B] x M[C]
            9 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b * c;
                self.counter += 4;
            },

            // M[A] <- M[B] / M[C]
            10 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b / c;
                self.counter += 4;
            },

            // M[A] <- M[B] modulo M[C]
            11 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b % c;
                self.counter += 4;
            },

            // If M[B] < M[C], then M[A] <- 1, else M[A] <- 0
            12 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                if b < c {
                  self.memory[a as usize] = 1;
                } else {
                  self.memory[a as usize] = 0;
                }
                self.counter += 4;
            },

            // Unknown opcode
            _ => {
            },
        }
    }
}

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

    #[test]
    fn it_loads_data() {
        let mut m = Machine::new();
        assert_eq!(Some(&0), m.dump().first());

        m.load(vec![1]);
        assert_eq!(Some(&1), m.dump().first());
    }

    #[test]
    fn it_runs_opcode_1() {
        // PC <- M[A]
        let mut m = Machine::with_data(vec![1, 4, 0, 0, 2]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(2, m.loc());
    }

    #[test]
    fn it_runs_opcode_2_then_branch() {
        // If M[B] = 0, then PC <- M[A]
        let mut m = Machine::with_data(vec![2, 4, 5, 0, 2, 0]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(2, m.loc());
    }

    #[test]
    fn it_runs_opcode_2_else_branch() {
        // If M[B] = 0, then PC <- M[A]
        let mut m = Machine::with_data(vec![2, 4, 5, 0, 2, 1]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(0, m.loc());
    }

    #[test]
    fn it_runs_opcode_3() {
        // M[A] <- PC
        let mut m = Machine::with_data(vec![3, 4, 0, 0, 1]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![3, 4, 0, 0, 0], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_4() {
        // M[A] <- M[B]
        let mut m = Machine::with_data(vec![4, 4, 5, 0, 6, 7]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![4, 4, 5, 0, 7, 7], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_5() {
        // M[A] <- M[M[B]]
        let mut m = Machine::with_data(vec![5, 4, 5, 0, 6, 7, 0, 8]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![5, 4, 5, 0, 8, 7, 0, 8], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_6() {
        // M[M[B]] <- M[A]
        let mut m = Machine::with_data(vec![6, 4, 5, 0, 8, 6, 7]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![6, 4, 5, 0, 8, 6, 8], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_7() {
        // M[A] <- M[B] + M[C]
        let mut m = Machine::with_data(vec![7, 4, 2, 3, 0]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![7, 4, 2, 3, 5], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_8() {
        // M[A] <- M[B] - M[C]
        let mut m = Machine::with_data(vec![8, 4, 5, 6, 0, 3, 2]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![8, 4, 5, 6, 1, 3, 2], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_9() {
        let mut m = Machine::with_data(vec![1, 2, 3]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(9, 0, 1, 2);
        assert_eq!(Some(&6), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_10() {
        let mut m = Machine::with_data(vec![1, 6, 2]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(10, 0, 1, 2);
        assert_eq!(Some(&3), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_11() {
        let mut m = Machine::with_data(vec![1, 11, 3]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(11, 0, 1, 2);
        assert_eq!(Some(&2), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_12_then_branch() {
        let mut m = Machine::with_data(vec![2, 3, 4]);

        assert_eq!(Some(&2), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(12, 0, 1, 2);
        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_12_else_branch() {
        let mut m = Machine::with_data(vec![2, 4, 3]);

        assert_eq!(Some(&2), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(12, 0, 1, 2);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }
}
<|endoftext|>


<commit_before>use std::io::{Read, BufRead, Error as IoError, ErrorKind};
use std::string::{FromUtf8Error};

use num::{Zero};
use num::bigint::{BigInt, ToBigInt, Sign};
use byteorder::{ReadBytesExt, LittleEndian, BigEndian, Error as ByteorderError};
use from_ascii::{FromAscii, ParseIntError, ParseFloatError};

use string::{unescape, Error as UnescapeError};
use value::{Value};

use opcodes::*;

quick_error! {
    #[derive(Debug)]
    pub enum Error {
        EmptyMarker
        StackTooSmall
        EmptyStack
        InvalidValueOnStack
        InvalidGetValue
        InvalidPutValue
        NotImplemented

        Read(err: ByteorderError) {
            from()
        }
        Io(err: IoError) {
            from()
        }
        UnknownOpcode(opcode: u8) {}

        InvalidInt {
            from(ParseIntError)
        }
        InvalidLong
        InvalidFloat {
            from(ParseFloatError)
        }

        InvalidString {
            from(FromUtf8Error)
        }
        UnescapeError(err: UnescapeError) {
            from()
        }

        InvalidProto(proto: u8)
        NegativeLength
    }
}

#[derive(Debug, PartialEq)]
pub enum BooleanOrInt {
    Boolean(bool),
    Int(i64),
}

struct Memo {
    small_map: Vec<Value>,
}

impl Memo {
    fn new() -> Self {
        Memo {
            small_map: Vec::with_capacity(8),
        }
    }

    #[inline]
    fn insert(&mut self, key: usize, value: Value) -> Option<()> {
        let len = self.small_map.len();
        if len != key {
            return None
        }
        self.small_map.push(value);
        Some(())
    }

    #[inline]
    pub fn get(&self, key: usize) -> Option<&Value> {
        self.small_map.get(key)
    }
}

fn read_exact<R>(rd: &mut R, mut buf: &mut [u8]) -> Result<(), IoError> where R: Read {
    while !buf.is_empty() {
        match rd.read(buf) {
            Ok(0) => break,
            Ok(n) => { let tmp = buf; buf = &mut tmp[n..]; }
            Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
            Err(e) => return Err(e),
        }
    }
    if !buf.is_empty() {
        Err(IoError::new(ErrorKind::Other,
                       "failed to fill whole buffer"))
    } else {
        Ok(())
    }
}

fn read_until_newline<R>(rd: &mut R) -> Result<Vec<u8>, Error> where R: Read + BufRead {
    let mut buf = Vec::new();
    try!(rd.read_until('\n' as u8, &mut buf));

    // Skip last symbol — \n
    match buf.split_last() {
        Some((&b'\n', init)) => Ok(init.to_vec()),
        _ => return Err(Error::InvalidString),
    }
}

fn read_decimal_int<R>(rd: &mut R) -> Result<BooleanOrInt, Error> where R: Read + BufRead {
    let s = try!(read_until_newline(rd));
    let val = match &s[..] {
        b"00" => BooleanOrInt::Boolean(false),
        b"01" => BooleanOrInt::Boolean(true),
        _ => BooleanOrInt::Int(try!(i64::from_ascii(&s)))
    };
    Ok(val)
}

fn read_decimal_long<R>(rd: &mut R) -> Result<BigInt, Error> where R: Read + BufRead {
    let s = try!(read_until_newline(rd));
    let init = match s.split_last() {
        None => return Err(Error::InvalidString),
        Some((&b'L', init)) => init,
        Some(_) => &s[..],
    };

    match BigInt::parse_bytes(&init, 10) {
        Some(i) => Ok(i),
        None => Err(Error::InvalidLong)
    }
}

fn read_long<R>(rd: &mut R, length: usize) -> Result<BigInt, Error> where R: Read + BufRead {
    let mut buf = vec![0; length];
    try!(read_exact(rd, buf.as_mut()));

    let mut n = BigInt::from_bytes_le(Sign::Plus, &buf);

    let last = match buf.last_mut() {
        None => return Err(Error::InvalidLong),
        Some(last) => last,
    };

    if *last > 127 {
        n = n - (1.to_bigint().unwrap() << (length * 8))
    }

    Ok(n)
}

pub struct Machine {
    stack: Vec<Value>,
    memo: Memo,
    marker: Option<usize>,
}

impl Machine {
    pub fn new() -> Self {
        Machine {
            stack: Vec::new(),
            memo: Memo::new(),
            marker: None,
        }
    }

    fn split_off(&mut self) -> Result<Vec<Value>, Error> {
        let at = match self.marker {
            None => return Err(Error::EmptyMarker),
            Some(marker) => marker,
        };

        if at > self.stack.len() {
            return Err(Error::StackTooSmall);
        }

        Ok(self.stack.split_off(at))
    }

    pub fn pop(&mut self) -> Result<Value, Error> {
        match self.stack.pop() {
            None => return Err(Error::EmptyStack),
            Some(value) => Ok(value),
        }
    }

    fn handle_get(&mut self, i: usize) -> Result<(), Error> {
        let value = match self.memo.get(i) {
            None => return Err(Error::InvalidGetValue),
            Some(ref v) => (*v).clone(),
        };
        self.stack.push(value);
        Ok(())
    }

    fn handle_put(&mut self, i: usize) -> Result<(), Error> {
        let value = match self.stack.last() {
            None => return Err(Error::EmptyStack),
            Some(ref v) => (*v).clone(),
        };
        self.memo.insert(i, value);
        Ok(())
    }

    pub fn execute<R>(&mut self, rd: &mut R) -> Result<bool, Error> where R: Read + BufRead {
        macro_rules! ensure_not_negative {
            ($n: expr) => ({
                if $n < Zero::zero() {
                    return Err(Error::NegativeLength)
                }
            })
        }

        let marker = try!(rd.read_u8());
        match marker {
            PROTO => {
                let version = try!(rd.read_u8());
                if version < 2 {
                    return Err(Error::InvalidProto(version))
                }
            },
            STOP => return Ok(true),

            INT => {
                self.stack.push(match try!(read_decimal_int(rd)) {
                    BooleanOrInt::Boolean(v) => Value::Bool(v),
                    BooleanOrInt::Int(v) => Value::Long(BigInt::from(v)),
                })
            },
            BININT => self.stack.push(Value::Int(try!(rd.read_i32::<LittleEndian>()) as usize)),
            BININT1 => self.stack.push(Value::Int(try!(rd.read_u8()) as usize)),
            BININT2 => self.stack.push(Value::Int(try!(rd.read_u16::<LittleEndian>()) as usize)),
            LONG => self.stack.push(Value::Long(BigInt::from(try!(read_decimal_long(rd))))),
            LONG1 => {
                let length = try!(rd.read_u8());
                self.stack.push(Value::Long(BigInt::from(try!(read_long(rd, length as usize)))))
            }
            LONG4 => {
                let length = try!(rd.read_i32::<LittleEndian>());
                self.stack.push(Value::Long(BigInt::from(try!(read_long(rd, length as usize)))))
            }

            STRING => self.stack.push(Value::String(try!(unescape(&try!(read_until_newline(rd)), false)))),
            BINSTRING => {
                let length = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(length);

                let mut buf = vec![0; length as usize];
                try!(read_exact(rd, &mut buf));
                self.stack.push(Value::String(buf))
            },
            SHORT_BINSTRING => {
                let length = try!(rd.read_u8());
                let mut buf = vec![0; length as usize];
                try!(read_exact(rd, &mut buf));
                self.stack.push(Value::String(buf))
            },

            NONE => self.stack.push(Value::None),
            NEWTRUE => self.stack.push(Value::Bool(true)),
            NEWFALSE => self.stack.push(Value::Bool(false)),

            UNICODE => {
                let buf = try!(unescape(&try!(read_until_newline(rd)), true));
                self.stack.push(Value::Unicode(try!(String::from_utf8(buf))))
            },
            BINUNICODE => {
                let length = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(length);
                let mut buf = vec![0; length as usize];
                try!(read_exact(rd, buf.as_mut()));
                self.stack.push(Value::Unicode(try!(String::from_utf8(buf))))
            },

            FLOAT => {
                let s = try!(read_until_newline(rd));
                self.stack.push(Value::Float(try!(f64::from_ascii(&s))))
            },
            BINFLOAT => {
                self.stack.push(Value::Float(try!(rd.read_f64::<BigEndian>())))
            },

            EMPTY_LIST => {
                self.stack.push(Value::List(Vec::new()))
            },
            APPEND => {
                let v = try!(self.pop());
                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::List(ref mut list)) => list.push(v),
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },
            APPENDS => {
                let values = try!(self.split_off());
                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::List(ref mut list)) => {
                        list.extend(values);
                    },
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },
            LIST => {
                let values = try!(self.split_off());
                self.stack.push(Value::List(values));
            },

            EMPTY_TUPLE => self.stack.push(Value::Tuple(Vec::new())),
            TUPLE => {
                let values = try!(self.split_off());
                self.stack.push(Value::Tuple(values));
            },
            TUPLE1 => {
                let v1 = try!(self.pop());
                self.stack.push(Value::Tuple(vec![v1]))
            },
            TUPLE2 => {
                let v1 = try!(self.pop());
                let v2 = try!(self.pop());
                self.stack.push(Value::Tuple(vec![v1, v2]))
            },
            TUPLE3 => {
                let v1 = try!(self.pop());
                let v2 = try!(self.pop());
                let v3 = try!(self.pop());
                self.stack.push(Value::Tuple(vec![v1, v2, v3]))
            }

            EMPTY_DICT => self.stack.push(Value::Dict(Vec::new())),
            DICT => {
                let mut values = try!(self.split_off());
                let mut dict = Vec::new();

                for i in 0 .. values.len() / 2 { // TODO: Check panic
                    let key = values.remove(2 * i);
                    let value = values.remove(2 * i + 1);
                    dict.push((key, value));
                }
                self.stack.push(Value::Dict(dict));
            },
            SETITEM => {
                let value = try!(self.pop());
                let key = try!(self.pop());
                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::Dict(ref mut dict)) => dict.push((key, value)),
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },
            SETITEMS => {
                let mut values = try!(self.split_off());

                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::Dict(ref mut dict)) => {
                        for i in 0 .. values.len() / 2 { // TODO: Check panic
                            let key = values.remove(2 * i);
                            let value = values.remove(2 * i + 1);
                            dict.push((key, value));
                        }
                    },
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },

            POP => {
                try!(self.pop());
            },
            DUP => {
                let value = match self.stack.last() {
                    None => return Err(Error::EmptyStack),
                    Some(ref v) => (*v).clone(),
                };
                self.stack.push(value)
            },
            MARK => {
                self.marker = Some(self.stack.len())
            },
            POP_MARK => {
                try!(self.split_off());
            },

            GET => {
                let n = match try!(read_decimal_int(rd)) {
                    BooleanOrInt::Int(n) => n,
                    BooleanOrInt::Boolean(false) => 0,
                    BooleanOrInt::Boolean(true) => 1,
                };
                ensure_not_negative!(n);
                try!(self.handle_get(n as usize))
            }
            BINGET => {
                try!(self.handle_get(try!(rd.read_u8()) as usize))
            }
            LONG_BINGET => {
                let n = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(n);
                try!(self.handle_get(n as usize))
            }

            PUT => {
                let n = match try!(read_decimal_int(rd)) {
                    BooleanOrInt::Int(n) => n,
                    BooleanOrInt::Boolean(false) => 0,
                    BooleanOrInt::Boolean(true) => 1,
                };
                ensure_not_negative!(n);
                try!(self.handle_put(n as usize))
            }
            BINPUT => {
                try!(self.handle_put(try!(rd.read_u8()) as usize))
            }
            LONG_BINPUT => {
                let n = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(n);
                try!(self.handle_put(n as usize))
            }

            c => return Err(Error::UnknownOpcode(c)),
        }
        Ok(false)
    }
}
<commit_msg>Inline memo<commit_after>use std::io::{Read, BufRead, Error as IoError, ErrorKind};
use std::string::{FromUtf8Error};

use num::{Zero};
use num::bigint::{BigInt, ToBigInt, Sign};
use byteorder::{ReadBytesExt, LittleEndian, BigEndian, Error as ByteorderError};
use from_ascii::{FromAscii, ParseIntError, ParseFloatError};

use string::{unescape, Error as UnescapeError};
use value::{Value};

use opcodes::*;

quick_error! {
    #[derive(Debug)]
    pub enum Error {
        EmptyMarker
        StackTooSmall
        EmptyStack
        InvalidValueOnStack
        InvalidGetValue
        InvalidPutValue
        NotImplemented

        Read(err: ByteorderError) {
            from()
        }
        Io(err: IoError) {
            from()
        }
        UnknownOpcode(opcode: u8) {}

        InvalidInt {
            from(ParseIntError)
        }
        InvalidLong
        InvalidFloat {
            from(ParseFloatError)
        }

        InvalidString {
            from(FromUtf8Error)
        }
        UnescapeError(err: UnescapeError) {
            from()
        }

        InvalidProto(proto: u8)
        NegativeLength
    }
}

#[derive(Debug, PartialEq)]
pub enum BooleanOrInt {
    Boolean(bool),
    Int(i64),
}

fn read_exact<R>(rd: &mut R, mut buf: &mut [u8]) -> Result<(), IoError> where R: Read {
    while !buf.is_empty() {
        match rd.read(buf) {
            Ok(0) => break,
            Ok(n) => { let tmp = buf; buf = &mut tmp[n..]; }
            Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
            Err(e) => return Err(e),
        }
    }
    if !buf.is_empty() {
        Err(IoError::new(ErrorKind::Other,
                       "failed to fill whole buffer"))
    } else {
        Ok(())
    }
}

fn read_until_newline<R>(rd: &mut R) -> Result<Vec<u8>, Error> where R: Read + BufRead {
    let mut buf = Vec::new();
    try!(rd.read_until('\n' as u8, &mut buf));

    // Skip last symbol — \n
    match buf.split_last() {
        Some((&b'\n', init)) => Ok(init.to_vec()),
        _ => return Err(Error::InvalidString),
    }
}

fn read_decimal_int<R>(rd: &mut R) -> Result<BooleanOrInt, Error> where R: Read + BufRead {
    let s = try!(read_until_newline(rd));
    let val = match &s[..] {
        b"00" => BooleanOrInt::Boolean(false),
        b"01" => BooleanOrInt::Boolean(true),
        _ => BooleanOrInt::Int(try!(i64::from_ascii(&s)))
    };
    Ok(val)
}

fn read_decimal_long<R>(rd: &mut R) -> Result<BigInt, Error> where R: Read + BufRead {
    let s = try!(read_until_newline(rd));
    let init = match s.split_last() {
        None => return Err(Error::InvalidString),
        Some((&b'L', init)) => init,
        Some(_) => &s[..],
    };

    match BigInt::parse_bytes(&init, 10) {
        Some(i) => Ok(i),
        None => Err(Error::InvalidLong)
    }
}

fn read_long<R>(rd: &mut R, length: usize) -> Result<BigInt, Error> where R: Read + BufRead {
    let mut buf = vec![0; length];
    try!(read_exact(rd, buf.as_mut()));

    let mut n = BigInt::from_bytes_le(Sign::Plus, &buf);

    let last = match buf.last_mut() {
        None => return Err(Error::InvalidLong),
        Some(last) => last,
    };

    if *last > 127 {
        n = n - (1.to_bigint().unwrap() << (length * 8))
    }

    Ok(n)
}

pub struct Machine {
    stack: Vec<Value>,
    memo: Vec<Value>,
    marker: Option<usize>,
}

impl Machine {
    pub fn new() -> Self {
        Machine {
            stack: Vec::new(),
            memo: Vec::with_capacity(8),
            marker: None,
        }
    }

    fn split_off(&mut self) -> Result<Vec<Value>, Error> {
        let at = match self.marker {
            None => return Err(Error::EmptyMarker),
            Some(marker) => marker,
        };

        if at > self.stack.len() {
            return Err(Error::StackTooSmall);
        }

        Ok(self.stack.split_off(at))
    }

    pub fn pop(&mut self) -> Result<Value, Error> {
        match self.stack.pop() {
            None => return Err(Error::EmptyStack),
            Some(value) => Ok(value),
        }
    }

    fn handle_get(&mut self, i: usize) -> Result<(), Error> {
        let value = match self.memo.get(i) {
            None => return Err(Error::InvalidGetValue),
            Some(ref v) => (*v).clone(),
        };
        self.stack.push(value);
        Ok(())
    }

    fn handle_put(&mut self, i: usize) -> Result<(), Error> {
        let value = match self.stack.last() {
            None => return Err(Error::EmptyStack),
            Some(ref v) => (*v).clone(),
        };
        let len = self.memo.len();
        if len != i {
            return Err(Error::InvalidPutValue)
        }
        self.memo.push(value);
        Ok(())
    }

    pub fn execute<R>(&mut self, rd: &mut R) -> Result<bool, Error> where R: Read + BufRead {
        macro_rules! ensure_not_negative {
            ($n: expr) => ({
                if $n < Zero::zero() {
                    return Err(Error::NegativeLength)
                }
            })
        }

        let marker = try!(rd.read_u8());
        match marker {
            PROTO => {
                let version = try!(rd.read_u8());
                if version < 2 {
                    return Err(Error::InvalidProto(version))
                }
            },
            STOP => return Ok(true),

            INT => {
                self.stack.push(match try!(read_decimal_int(rd)) {
                    BooleanOrInt::Boolean(v) => Value::Bool(v),
                    BooleanOrInt::Int(v) => Value::Long(BigInt::from(v)),
                })
            },
            BININT => self.stack.push(Value::Int(try!(rd.read_i32::<LittleEndian>()) as usize)),
            BININT1 => self.stack.push(Value::Int(try!(rd.read_u8()) as usize)),
            BININT2 => self.stack.push(Value::Int(try!(rd.read_u16::<LittleEndian>()) as usize)),
            LONG => self.stack.push(Value::Long(BigInt::from(try!(read_decimal_long(rd))))),
            LONG1 => {
                let length = try!(rd.read_u8());
                self.stack.push(Value::Long(BigInt::from(try!(read_long(rd, length as usize)))))
            }
            LONG4 => {
                let length = try!(rd.read_i32::<LittleEndian>());
                self.stack.push(Value::Long(BigInt::from(try!(read_long(rd, length as usize)))))
            }

            STRING => self.stack.push(Value::String(try!(unescape(&try!(read_until_newline(rd)), false)))),
            BINSTRING => {
                let length = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(length);

                let mut buf = vec![0; length as usize];
                try!(read_exact(rd, &mut buf));
                self.stack.push(Value::String(buf))
            },
            SHORT_BINSTRING => {
                let length = try!(rd.read_u8());
                let mut buf = vec![0; length as usize];
                try!(read_exact(rd, &mut buf));
                self.stack.push(Value::String(buf))
            },

            NONE => self.stack.push(Value::None),
            NEWTRUE => self.stack.push(Value::Bool(true)),
            NEWFALSE => self.stack.push(Value::Bool(false)),

            UNICODE => {
                let buf = try!(unescape(&try!(read_until_newline(rd)), true));
                self.stack.push(Value::Unicode(try!(String::from_utf8(buf))))
            },
            BINUNICODE => {
                let length = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(length);
                let mut buf = vec![0; length as usize];
                try!(read_exact(rd, buf.as_mut()));
                self.stack.push(Value::Unicode(try!(String::from_utf8(buf))))
            },

            FLOAT => {
                let s = try!(read_until_newline(rd));
                self.stack.push(Value::Float(try!(f64::from_ascii(&s))))
            },
            BINFLOAT => {
                self.stack.push(Value::Float(try!(rd.read_f64::<BigEndian>())))
            },

            EMPTY_LIST => {
                self.stack.push(Value::List(Vec::new()))
            },
            APPEND => {
                let v = try!(self.pop());
                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::List(ref mut list)) => list.push(v),
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },
            APPENDS => {
                let values = try!(self.split_off());
                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::List(ref mut list)) => {
                        list.extend(values);
                    },
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },
            LIST => {
                let values = try!(self.split_off());
                self.stack.push(Value::List(values));
            },

            EMPTY_TUPLE => self.stack.push(Value::Tuple(Vec::new())),
            TUPLE => {
                let values = try!(self.split_off());
                self.stack.push(Value::Tuple(values));
            },
            TUPLE1 => {
                let v1 = try!(self.pop());
                self.stack.push(Value::Tuple(vec![v1]))
            },
            TUPLE2 => {
                let v1 = try!(self.pop());
                let v2 = try!(self.pop());
                self.stack.push(Value::Tuple(vec![v1, v2]))
            },
            TUPLE3 => {
                let v1 = try!(self.pop());
                let v2 = try!(self.pop());
                let v3 = try!(self.pop());
                self.stack.push(Value::Tuple(vec![v1, v2, v3]))
            }

            EMPTY_DICT => self.stack.push(Value::Dict(Vec::new())),
            DICT => {
                let mut values = try!(self.split_off());
                let mut dict = Vec::new();

                for i in 0 .. values.len() / 2 { // TODO: Check panic
                    let key = values.remove(2 * i);
                    let value = values.remove(2 * i + 1);
                    dict.push((key, value));
                }
                self.stack.push(Value::Dict(dict));
            },
            SETITEM => {
                let value = try!(self.pop());
                let key = try!(self.pop());
                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::Dict(ref mut dict)) => dict.push((key, value)),
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },
            SETITEMS => {
                let mut values = try!(self.split_off());

                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::Dict(ref mut dict)) => {
                        for i in 0 .. values.len() / 2 { // TODO: Check panic
                            let key = values.remove(2 * i);
                            let value = values.remove(2 * i + 1);
                            dict.push((key, value));
                        }
                    },
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },

            POP => {
                try!(self.pop());
            },
            DUP => {
                let value = match self.stack.last() {
                    None => return Err(Error::EmptyStack),
                    Some(ref v) => (*v).clone(),
                };
                self.stack.push(value)
            },
            MARK => {
                self.marker = Some(self.stack.len())
            },
            POP_MARK => {
                try!(self.split_off());
            },

            GET => {
                let n = match try!(read_decimal_int(rd)) {
                    BooleanOrInt::Int(n) => n,
                    BooleanOrInt::Boolean(false) => 0,
                    BooleanOrInt::Boolean(true) => 1,
                };
                ensure_not_negative!(n);
                try!(self.handle_get(n as usize))
            }
            BINGET => {
                try!(self.handle_get(try!(rd.read_u8()) as usize))
            }
            LONG_BINGET => {
                let n = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(n);
                try!(self.handle_get(n as usize))
            }

            PUT => {
                let n = match try!(read_decimal_int(rd)) {
                    BooleanOrInt::Int(n) => n,
                    BooleanOrInt::Boolean(false) => 0,
                    BooleanOrInt::Boolean(true) => 1,
                };
                ensure_not_negative!(n);
                try!(self.handle_put(n as usize))
            }
            BINPUT => {
                try!(self.handle_put(try!(rd.read_u8()) as usize))
            }
            LONG_BINPUT => {
                let n = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(n);
                try!(self.handle_put(n as usize))
            }

            c => return Err(Error::UnknownOpcode(c)),
        }
        Ok(false)
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 Deyan Ginev. See the LICENSE
// file at the top-level directory of this distribution.
//
// Licensed under the MIT license <LICENSE-MIT or http://opensource.org/licenses/MIT>.
// This file may not be copied, modified, or distributed
// except according to those terms.
extern crate zmq;
extern crate tempfile;

use zmq::{Error, SNDMORE};
use backend::{Backend, DEFAULT_DB_ADDRESS};
use data::{Task, TaskReport, TaskStatus, Service};

use std::thread;
use std::sync::Arc;
use std::sync::Mutex;

use std::ops::Deref;
use std::collections::HashMap;

use std::path::Path;
use std::fs::File;
// use tempfile::TempFile;
use std::io::{Write};
use std::io::Read;

pub struct TaskManager {
  pub source_port : usize,
  pub result_port : usize,
  pub queue_size : usize,
  pub message_size : usize,
  pub backend_address : String
}
pub struct Server {
  pub port : usize,
  pub queue_size : usize,
  pub message_size : usize,
  pub backend : Backend,
  pub backend_address : String
}

impl Default for TaskManager {
  fn default() -> TaskManager {
    TaskManager {
        source_port : 5555,
        result_port : 5555,
        queue_size : 100,
        message_size : 100000,
        backend_address : DEFAULT_DB_ADDRESS.clone().to_string()
    } } }

impl TaskManager {
  pub fn start<'manager>(&'manager self) -> Result<(), Error> {
    // We'll use some local memoization shared between source and sink:
    let services: HashMap<String, Option<Service>> = HashMap::new();
    let progress_queue: HashMap<i64, Task> = HashMap::new();
    let done_queue: Vec<TaskReport> = Vec::new();

    let services_arc = Arc::new(Mutex::new(services));
    let progress_queue_arc = Arc::new(Mutex::new(progress_queue));
    let done_queue_arc = Arc::new(Mutex::new(done_queue));

    // First prepare the source ventilator
    let source_port = self.source_port.clone();
    let source_queue_size = self.queue_size.clone();
    let source_message_size = self.message_size.clone();
    let source_backend_address = self.backend_address.clone();

    let vent_services_arc = services_arc.clone();
    let vent_progress_queue_arc = progress_queue_arc.clone();
    let vent_thread = thread::spawn(move || {
      let sources = Server {
        port : source_port,
        queue_size : source_queue_size,
        message_size : source_message_size,
        backend : Backend::from_address(&source_backend_address),
        backend_address : source_backend_address.clone()
      };
      sources.start_ventilator(vent_services_arc, vent_progress_queue_arc).unwrap();
    });

    // Next prepare the finalize thread which will persist finished jobs to the DB
    let finalize_backend_address = self.backend_address.clone();
    let finalize_done_queue_arc = done_queue_arc.clone();
    let finalize_thread = thread::spawn(move || {
      let finalize_backend = Backend::from_address(&finalize_backend_address);
      // Persist every 1 second, if there is something to record
      loop {
        match Server::mark_done_arc(&finalize_backend, &finalize_done_queue_arc) {
          true => {
            true;
          } // we did some work, on to the next iteration
          false => { // If we have no reports to process, sleep for a second and recheck
            thread::sleep_ms(1000);
          }
        }
      };
    });

    // Now prepare the results sink
    let result_port = self.result_port.clone();
    let result_queue_size = self.queue_size.clone();
    let result_message_size = self.message_size.clone();
    let result_backend_address = self.backend_address.clone();

    let sink_services_arc = services_arc.clone();
    let sink_progress_queue_arc = progress_queue_arc.clone();

    let sink_done_queue_arc = done_queue_arc.clone();
    let sink_thread = thread::spawn(move || {
      let results = Server {
        port : result_port,
        queue_size : result_queue_size,
        message_size : result_message_size,
        backend : Backend::from_address(&result_backend_address),
        backend_address: result_backend_address.clone()
      };
      results.start_sink(sink_services_arc, sink_progress_queue_arc, sink_done_queue_arc).unwrap();
    });

    vent_thread.join().unwrap();
    sink_thread.join().unwrap();
    finalize_thread.join().unwrap();
    Ok(())
  }
}

impl Server {
  pub fn start_ventilator(&self,
      services_arc : Arc<Mutex<HashMap<String, Option<Service>>>>,
      progress_queue_arc : Arc<Mutex<HashMap<i64, Task>>>)
      -> Result <(),Error> {
    // We have a Ventilator-exclusive "queues" stack for tasks to be dispatched
    let mut queues : HashMap<String, Vec<Task>> = HashMap::new();
    // Assuming this is the only And tidy up the postgres tasks:
    self.backend.clear_limbo_tasks().unwrap();
    // Ok, let's bind to a port and start broadcasting
    let mut context = zmq::Context::new();
    let mut ventilator = context.socket(zmq::ROUTER).unwrap();
    let port_str = self.port.to_string();
    let address = "tcp://*:".to_string() + &port_str;
    assert!(ventilator.bind(&address).is_ok());
    let mut source_job_count = 0;

    loop {
      let mut msg = zmq::Message::new().unwrap();
      let mut identity = zmq::Message::new().unwrap();
      ventilator.recv(&mut identity, 0).unwrap();
      ventilator.recv(&mut msg, 0).unwrap();
      let service_name = msg.as_str().unwrap().to_string();
      println!("Task requested for service: {}", service_name.clone());
      source_job_count += 1;

      let mut dispatched_task : Option<Task> = None;
      match self.get_sync_service_record(&services_arc, service_name.clone()) {
        None => {},
        Some(service) => {
          if !queues.contains_key(&service_name) {
            queues.insert(service_name.clone(), Vec::new());
          }
          let mut task_queue : &mut Vec<Task> = queues.get_mut(&service_name).unwrap();
          if task_queue.is_empty() {
            task_queue.extend(self.backend.fetch_tasks(&service, self.queue_size).unwrap()); }
          match task_queue.pop() {
            Some(current_task) => {
              let taskid = current_task.id.unwrap();
              let serviceid = current_task.serviceid;

              dispatched_task = Some(current_task.clone());

              ventilator.send_msg(identity, SNDMORE).unwrap();
              ventilator.send_str(&taskid.to_string(), SNDMORE).unwrap();
              if serviceid == 1 { // No payload needed for init
                ventilator.send(&[],0).unwrap(); }
              else {
                // Regular services fetch the task payload and transfer it to the worker
                let file_opt = service.prepare_input_stream(current_task.clone());
                if file_opt.is_ok() {
                  let mut file = file_opt.unwrap();
                  let mut total_outgoing = 0;
                  loop {
                    // Stream input data via zmq
                    let mut data = vec![0; self.message_size];
                    let size = file.read(&mut data).unwrap();
                    total_outgoing += size;
                    data.truncate(size);

                    if size < self.message_size {
                      // If exhausted, send the last frame
                      ventilator.send(&data,0).unwrap();
                      // And terminate
                      break;
                    } else {
                      // If more to go, send the frame and indicate there's more to come
                      ventilator.send(&data,SNDMORE).unwrap();
                    }
                  }
                  println!("Source job {}, message size: {}", source_job_count, total_outgoing);
                } else {
                  // TODO: smart handling of failures
                  ventilator.send(&[],0).unwrap();
                }
              }
            },
            None => {}
          };
        }
      };
      // Record that a task has been dispatched in the progress queue
      if dispatched_task.is_some() {
        Server::push_progress_task(&progress_queue_arc, dispatched_task.unwrap());
      }
    }
  }

  pub fn start_sink(&self,
      services_arc : Arc<Mutex<HashMap<String, Option<Service>>>>,
      progress_queue_arc : Arc<Mutex<HashMap<i64, Task>>>,
      done_queue_arc : Arc<Mutex<Vec<TaskReport>>>)
      -> Result <(),Error> {

    // Ok, let's bind to a port and start broadcasting
    let mut context = zmq::Context::new();
    let mut sink = context.socket(zmq::PULL).unwrap();
    let port_str = self.port.to_string();
    let address = "tcp://*:".to_string() + &port_str;
    assert!(sink.bind(&address).is_ok());

    let mut sink_job_count = 0;

    loop {
      let mut recv_msg = zmq::Message::new().unwrap();
      let mut taskid_msg = zmq::Message::new().unwrap();
      let mut service_msg = zmq::Message::new().unwrap();

      sink.recv(&mut service_msg, 0).unwrap();
      let service_name = service_msg.as_str().unwrap();

      sink.recv(&mut taskid_msg, 0).unwrap();
      let taskid_str = taskid_msg.as_str().unwrap();
      let taskid = taskid_str.parse::<i64>().unwrap();
      // We have a job, count it
      sink_job_count += 1;
      let mut total_incoming = 0;
      println!("Incoming sink job {:?} for Service: {:?}, taskid: {:?}", sink_job_count, service_name, taskid_str);

      match Server::pop_progress_task(&progress_queue_arc, taskid) {
        None => {} // TODO: No such task, what to do?
        Some(task) => {

          // println!("{:?}", task);
          let service_option = Server::get_service_record(&services_arc, service_name.to_string());
          match service_option.clone() {
            None => {}, // TODO: Handle errors
            Some(service) => {
              let serviceid = service.id.unwrap();
              // println!("Service: {:?}", serviceid);
              if serviceid == task.serviceid {
                // println!("Task and Service match up.");
                if serviceid == 1 { // No payload needed for init
                  sink.recv(&mut recv_msg, 0).unwrap();
                  let done_report = TaskReport {
                    task : task.clone(),
                    status : TaskStatus::NoProblem,
                    messages : Vec::new()
                  };
                  Server::push_done_queue(&done_queue_arc, done_report);
                }
                else {
                  // Receive the rest of the input in the correct file
                  let recv_dir = Path::new(&task.entry).parent().unwrap();
                  // "_" + &service.version.round().to_string() +
                  let recv_pathname = recv_dir.to_str().unwrap().to_string() + "/" + &service.name + ".zip";
                  let recv_path = Path::new(&recv_pathname);
                  // println!("Will write to {:?}", recv_path);
                  let mut file = File::create(recv_path).unwrap();
                  loop {
                    sink.recv(&mut recv_msg, 0).unwrap();

                    total_incoming += file.write(recv_msg.deref()).unwrap();
                    if !sink.get_rcvmore().unwrap() {
                      break;
                    }
                  }
                  // Then mark the task done. This can be in a new thread later on
                  let done_report = task.generate_report(recv_path);
                  Server::push_done_queue(&done_queue_arc, done_report);
                }

              }
              else {
                // Otherwise just discard the rest of the message
                loop {
                  sink.recv(&mut recv_msg, 0).unwrap();
                  if !sink.get_rcvmore().unwrap() {
                    break;
                  }
                }
              }
            }
          };
        }
      }
      println!("Sink job {}, message size: {}", sink_job_count, total_incoming);

      // let mut file = File::create("/tmp/cortex_sink_".to_string() + &sink_job_count.to_string()).unwrap();
      // file.write_all(&payload).unwrap();
    }
  }

  fn get_sync_service_record(&self, services_arc : &Arc<Mutex<HashMap<String, Option<Service>>>>, service_name : String) -> Option<Service> {
    let mut services = services_arc.lock().unwrap();
    let service_record = services.entry(service_name.clone()).or_insert(
      Service::from_name(&self.backend.connection, service_name.clone()).unwrap()).clone();
    service_record
  }

  fn get_service_record(services_arc : &Arc<Mutex<HashMap<String, Option<Service>>>>, service_name : String) -> Option<Service> {
    let services = services_arc.lock().unwrap();
    let service_record = services.get(&service_name);
    match service_record {
      None => None, // TODO: Handle errors
      Some(service_option) => service_option.clone()
    }
  }

  pub fn mark_done_arc(backend : &Backend, reports_arc: &Arc<Mutex<Vec<TaskReport>>>) -> bool {
    let mut reports = reports_arc.lock().unwrap();
    if reports.len() > 0 {
      backend.mark_done(&reports).unwrap(); // TODO: error handling if DB fails
      reports.clear();
      true
    } else {
      false
    }
  }
  pub fn push_done_queue(reports_arc : &Arc<Mutex<Vec<TaskReport>>>, report : TaskReport) {
    let mut reports = reports_arc.lock().unwrap();
    reports.push(report)
  }

  fn pop_progress_task(progress_queue_arc : &Arc<Mutex<HashMap<i64, Task>>>, taskid: i64) -> Option<Task> {
    let mut progress_queue = progress_queue_arc.lock().unwrap();
    progress_queue.remove(&taskid)
  }

  fn push_progress_task(progress_queue_arc : &Arc<Mutex<HashMap<i64, Task>>>, progress_task: Task) {
    let mut progress_queue = progress_queue_arc.lock().unwrap();
    progress_queue.insert(progress_task.id.unwrap(), progress_task);
  }
}<commit_msg>release reports mutex as soon as we've finished passing them to the backend<commit_after>// Copyright 2015 Deyan Ginev. See the LICENSE
// file at the top-level directory of this distribution.
//
// Licensed under the MIT license <LICENSE-MIT or http://opensource.org/licenses/MIT>.
// This file may not be copied, modified, or distributed
// except according to those terms.
extern crate zmq;
extern crate tempfile;

use zmq::{Error, SNDMORE};
use backend::{Backend, DEFAULT_DB_ADDRESS};
use data::{Task, TaskReport, TaskStatus, Service};

use std::thread;
use std::sync::Arc;
use std::sync::Mutex;

use std::ops::Deref;
use std::collections::HashMap;

use std::path::Path;
use std::fs::File;
// use tempfile::TempFile;
use std::io::{Write};
use std::io::Read;

pub struct TaskManager {
  pub source_port : usize,
  pub result_port : usize,
  pub queue_size : usize,
  pub message_size : usize,
  pub backend_address : String
}
pub struct Server {
  pub port : usize,
  pub queue_size : usize,
  pub message_size : usize,
  pub backend : Backend,
  pub backend_address : String
}

impl Default for TaskManager {
  fn default() -> TaskManager {
    TaskManager {
        source_port : 5555,
        result_port : 5555,
        queue_size : 100,
        message_size : 100000,
        backend_address : DEFAULT_DB_ADDRESS.clone().to_string()
    } } }

impl TaskManager {
  pub fn start<'manager>(&'manager self) -> Result<(), Error> {
    // We'll use some local memoization shared between source and sink:
    let services: HashMap<String, Option<Service>> = HashMap::new();
    let progress_queue: HashMap<i64, Task> = HashMap::new();
    let done_queue: Vec<TaskReport> = Vec::new();

    let services_arc = Arc::new(Mutex::new(services));
    let progress_queue_arc = Arc::new(Mutex::new(progress_queue));
    let done_queue_arc = Arc::new(Mutex::new(done_queue));

    // First prepare the source ventilator
    let source_port = self.source_port.clone();
    let source_queue_size = self.queue_size.clone();
    let source_message_size = self.message_size.clone();
    let source_backend_address = self.backend_address.clone();

    let vent_services_arc = services_arc.clone();
    let vent_progress_queue_arc = progress_queue_arc.clone();
    let vent_thread = thread::spawn(move || {
      let sources = Server {
        port : source_port,
        queue_size : source_queue_size,
        message_size : source_message_size,
        backend : Backend::from_address(&source_backend_address),
        backend_address : source_backend_address.clone()
      };
      sources.start_ventilator(vent_services_arc, vent_progress_queue_arc).unwrap();
    });

    // Next prepare the finalize thread which will persist finished jobs to the DB
    let finalize_backend_address = self.backend_address.clone();
    let finalize_done_queue_arc = done_queue_arc.clone();
    let finalize_thread = thread::spawn(move || {
      let finalize_backend = Backend::from_address(&finalize_backend_address);
      // Persist every 1 second, if there is something to record
      loop {
        match Server::mark_done_arc(&finalize_backend, &finalize_done_queue_arc) {
          true => {
            true;
          } // we did some work, on to the next iteration
          false => { // If we have no reports to process, sleep for a second and recheck
            thread::sleep_ms(1000);
          }
        }
      };
    });

    // Now prepare the results sink
    let result_port = self.result_port.clone();
    let result_queue_size = self.queue_size.clone();
    let result_message_size = self.message_size.clone();
    let result_backend_address = self.backend_address.clone();

    let sink_services_arc = services_arc.clone();
    let sink_progress_queue_arc = progress_queue_arc.clone();

    let sink_done_queue_arc = done_queue_arc.clone();
    let sink_thread = thread::spawn(move || {
      let results = Server {
        port : result_port,
        queue_size : result_queue_size,
        message_size : result_message_size,
        backend : Backend::from_address(&result_backend_address),
        backend_address: result_backend_address.clone()
      };
      results.start_sink(sink_services_arc, sink_progress_queue_arc, sink_done_queue_arc).unwrap();
    });

    vent_thread.join().unwrap();
    sink_thread.join().unwrap();
    finalize_thread.join().unwrap();
    Ok(())
  }
}

impl Server {
  pub fn start_ventilator(&self,
      services_arc : Arc<Mutex<HashMap<String, Option<Service>>>>,
      progress_queue_arc : Arc<Mutex<HashMap<i64, Task>>>)
      -> Result <(),Error> {
    // We have a Ventilator-exclusive "queues" stack for tasks to be dispatched
    let mut queues : HashMap<String, Vec<Task>> = HashMap::new();
    // Assuming this is the only And tidy up the postgres tasks:
    self.backend.clear_limbo_tasks().unwrap();
    // Ok, let's bind to a port and start broadcasting
    let mut context = zmq::Context::new();
    let mut ventilator = context.socket(zmq::ROUTER).unwrap();
    let port_str = self.port.to_string();
    let address = "tcp://*:".to_string() + &port_str;
    assert!(ventilator.bind(&address).is_ok());
    let mut source_job_count = 0;

    loop {
      let mut msg = zmq::Message::new().unwrap();
      let mut identity = zmq::Message::new().unwrap();
      ventilator.recv(&mut identity, 0).unwrap();
      ventilator.recv(&mut msg, 0).unwrap();
      let service_name = msg.as_str().unwrap().to_string();
      println!("Task requested for service: {}", service_name.clone());
      source_job_count += 1;

      let mut dispatched_task : Option<Task> = None;
      match self.get_sync_service_record(&services_arc, service_name.clone()) {
        None => {},
        Some(service) => {
          if !queues.contains_key(&service_name) {
            queues.insert(service_name.clone(), Vec::new());
          }
          let mut task_queue : &mut Vec<Task> = queues.get_mut(&service_name).unwrap();
          if task_queue.is_empty() {
            task_queue.extend(self.backend.fetch_tasks(&service, self.queue_size).unwrap()); }
          match task_queue.pop() {
            Some(current_task) => {
              let taskid = current_task.id.unwrap();
              let serviceid = current_task.serviceid;

              dispatched_task = Some(current_task.clone());

              ventilator.send_msg(identity, SNDMORE).unwrap();
              ventilator.send_str(&taskid.to_string(), SNDMORE).unwrap();
              if serviceid == 1 { // No payload needed for init
                ventilator.send(&[],0).unwrap(); }
              else {
                // Regular services fetch the task payload and transfer it to the worker
                let file_opt = service.prepare_input_stream(current_task.clone());
                if file_opt.is_ok() {
                  let mut file = file_opt.unwrap();
                  let mut total_outgoing = 0;
                  loop {
                    // Stream input data via zmq
                    let mut data = vec![0; self.message_size];
                    let size = file.read(&mut data).unwrap();
                    total_outgoing += size;
                    data.truncate(size);

                    if size < self.message_size {
                      // If exhausted, send the last frame
                      ventilator.send(&data,0).unwrap();
                      // And terminate
                      break;
                    } else {
                      // If more to go, send the frame and indicate there's more to come
                      ventilator.send(&data,SNDMORE).unwrap();
                    }
                  }
                  println!("Source job {}, message size: {}", source_job_count, total_outgoing);
                } else {
                  // TODO: smart handling of failures
                  ventilator.send(&[],0).unwrap();
                }
              }
            },
            None => {}
          };
        }
      };
      // Record that a task has been dispatched in the progress queue
      if dispatched_task.is_some() {
        Server::push_progress_task(&progress_queue_arc, dispatched_task.unwrap());
      }
    }
  }

  pub fn start_sink(&self,
      services_arc : Arc<Mutex<HashMap<String, Option<Service>>>>,
      progress_queue_arc : Arc<Mutex<HashMap<i64, Task>>>,
      done_queue_arc : Arc<Mutex<Vec<TaskReport>>>)
      -> Result <(),Error> {

    // Ok, let's bind to a port and start broadcasting
    let mut context = zmq::Context::new();
    let mut sink = context.socket(zmq::PULL).unwrap();
    let port_str = self.port.to_string();
    let address = "tcp://*:".to_string() + &port_str;
    assert!(sink.bind(&address).is_ok());

    let mut sink_job_count = 0;

    loop {
      let mut recv_msg = zmq::Message::new().unwrap();
      let mut taskid_msg = zmq::Message::new().unwrap();
      let mut service_msg = zmq::Message::new().unwrap();

      sink.recv(&mut service_msg, 0).unwrap();
      let service_name = service_msg.as_str().unwrap();

      sink.recv(&mut taskid_msg, 0).unwrap();
      let taskid_str = taskid_msg.as_str().unwrap();
      let taskid = taskid_str.parse::<i64>().unwrap();
      // We have a job, count it
      sink_job_count += 1;
      let mut total_incoming = 0;
      println!("Incoming sink job {:?} for Service: {:?}, taskid: {:?}", sink_job_count, service_name, taskid_str);

      match Server::pop_progress_task(&progress_queue_arc, taskid) {
        None => {} // TODO: No such task, what to do?
        Some(task) => {

          // println!("{:?}", task);
          let service_option = Server::get_service_record(&services_arc, service_name.to_string());
          match service_option.clone() {
            None => {}, // TODO: Handle errors
            Some(service) => {
              let serviceid = service.id.unwrap();
              // println!("Service: {:?}", serviceid);
              if serviceid == task.serviceid {
                // println!("Task and Service match up.");
                if serviceid == 1 { // No payload needed for init
                  sink.recv(&mut recv_msg, 0).unwrap();
                  let done_report = TaskReport {
                    task : task.clone(),
                    status : TaskStatus::NoProblem,
                    messages : Vec::new()
                  };
                  Server::push_done_queue(&done_queue_arc, done_report);
                }
                else {
                  // Receive the rest of the input in the correct file
                  let recv_dir = Path::new(&task.entry).parent().unwrap();
                  // "_" + &service.version.round().to_string() +
                  let recv_pathname = recv_dir.to_str().unwrap().to_string() + "/" + &service.name + ".zip";
                  let recv_path = Path::new(&recv_pathname);
                  // println!("Will write to {:?}", recv_path);
                  let mut file = File::create(recv_path).unwrap();
                  loop {
                    sink.recv(&mut recv_msg, 0).unwrap();

                    total_incoming += file.write(recv_msg.deref()).unwrap();
                    if !sink.get_rcvmore().unwrap() {
                      break;
                    }
                  }
                  // Then mark the task done. This can be in a new thread later on
                  let done_report = task.generate_report(recv_path);
                  Server::push_done_queue(&done_queue_arc, done_report);
                }

              }
              else {
                // Otherwise just discard the rest of the message
                loop {
                  sink.recv(&mut recv_msg, 0).unwrap();
                  if !sink.get_rcvmore().unwrap() {
                    break;
                  }
                }
              }
            }
          };
        }
      }
      println!("Sink job {}, message size: {}", sink_job_count, total_incoming);

      // let mut file = File::create("/tmp/cortex_sink_".to_string() + &sink_job_count.to_string()).unwrap();
      // file.write_all(&payload).unwrap();
    }
  }

  fn get_sync_service_record(&self, services_arc : &Arc<Mutex<HashMap<String, Option<Service>>>>, service_name : String) -> Option<Service> {
    let mut services = services_arc.lock().unwrap();
    let service_record = services.entry(service_name.clone()).or_insert(
      Service::from_name(&self.backend.connection, service_name.clone()).unwrap()).clone();
    service_record
  }

  fn get_service_record(services_arc : &Arc<Mutex<HashMap<String, Option<Service>>>>, service_name : String) -> Option<Service> {
    let services = services_arc.lock().unwrap();
    let service_record = services.get(&service_name);
    match service_record {
      None => None, // TODO: Handle errors
      Some(service_option) => service_option.clone()
    }
  }

  pub fn mark_done_arc(backend : &Backend, reports_arc: &Arc<Mutex<Vec<TaskReport>>>) -> bool {
    let reports = Server::fetch_shared_vec(reports_arc);
    if reports.len() > 0 {
      backend.mark_done(&reports).unwrap(); // TODO: error handling if DB fails
      true
    } else {
      false
    }
  }
  pub fn push_done_queue(reports_arc : &Arc<Mutex<Vec<TaskReport>>>, report : TaskReport) {
    let mut reports = reports_arc.lock().unwrap();
    reports.push(report)
  }

  fn pop_progress_task(progress_queue_arc : &Arc<Mutex<HashMap<i64, Task>>>, taskid: i64) -> Option<Task> {
    let mut progress_queue = progress_queue_arc.lock().unwrap();
    progress_queue.remove(&taskid)
  }

  fn push_progress_task(progress_queue_arc : &Arc<Mutex<HashMap<i64, Task>>>, progress_task: Task) {
    let mut progress_queue = progress_queue_arc.lock().unwrap();
    progress_queue.insert(progress_task.id.unwrap(), progress_task);
  }

  fn fetch_shared_vec<T: Clone>(vec_arc: &Arc<Mutex<Vec<T>>>) -> Vec<T> {
    let mut vec_mutex_guard = vec_arc.lock().unwrap();
    let fetched_vec : Vec<T> = (*vec_mutex_guard).clone();
    vec_mutex_guard.clear();
    fetched_vec
  }
}<|endoftext|>


<commit_before>//! Functions for creating and modifying messages to send across the message bus.
use std::mem::transmute;
use std::collections::HashMap;

use dbus_serialize::types::{Path,Variant,Value,BasicValue,Signature};

use marshal::{Marshal,pad_to_multiple};

#[derive(Debug,Default)]
pub struct MessageType(pub u8);
pub const MESSAGE_TYPE_INVALID : MessageType        = MessageType(0);
pub const MESSAGE_TYPE_METHOD_CALL : MessageType    = MessageType(1);
pub const MESSAGE_TYPE_METHOD_RETURN : MessageType  = MessageType(2);
pub const MESSAGE_TYPE_ERROR : MessageType          = MessageType(3);
pub const MESSAGE_TYPE_SIGNAL : MessageType         = MessageType(4);

#[derive(Copy,Clone)]
pub enum HeaderFieldName {
    Invalid = 0,
    Path = 1,
    Interface = 2,
    Member = 3,
    ErrorName = 4,
    ReplySerial = 5,
    Destination = 6,
    Sender = 7,
    Signature = 8
}

struct HeaderField (
    HeaderFieldName,
    Variant
);

impl Marshal for HeaderField {
    fn dbus_encode(&self, buf: &mut Vec<u8>) -> usize {
        pad_to_multiple(buf, 8);
        let start_len = buf.len();
        let code = self.0 as u8;
        code.dbus_encode(buf);
        self.1.dbus_encode(buf);
        buf.len() - start_len
    }
    fn get_type(&self) -> String {
        "(yv)".to_string()
    }
}

fn encode_header (msg_type: MessageType, serial: u32) -> Vec<u8> {
    let mut buf = Vec::new();

    let endian = 'l' as u8;
    endian.dbus_encode(&mut buf);
    msg_type.0.dbus_encode(&mut buf);
    (0 as u8).dbus_encode(&mut buf);
    (1 as u8).dbus_encode(&mut buf);
    (0 as u32).dbus_encode(&mut buf);
    serial.dbus_encode(&mut buf);
    buf
}

/// Container for the header of a D-Bus message.  Arguments can be added with add_arg()
pub struct MessageBuf(pub Vec<u8>);

/// Create a MessageBuf for a D-Bus method call.  Once a MessageBuf object is created, arguments
/// can be added to the object with MessageBuf.add_arg
pub fn create_method_call (dest: &str, path: &str, iface: &str, method: &str) -> MessageBuf {
    let mut msg = encode_header(MESSAGE_TYPE_METHOD_CALL, 0);
    let mut headers : Vec<HeaderField> = Vec::new();
    let mut v = Variant::new(Value::BasicValue(BasicValue::String(dest.to_string())), "s");
    headers.push(HeaderField(HeaderFieldName::Destination, v));
    v = Variant::new(Value::BasicValue(BasicValue::ObjectPath(Path(path.to_string()))), "o");
    headers.push(HeaderField(HeaderFieldName::Path, v));
    v = Variant::new(Value::BasicValue(BasicValue::String(iface.to_string())), "s");
    headers.push(HeaderField(HeaderFieldName::Interface, v));
    v = Variant::new(Value::BasicValue(BasicValue::String(method.to_string())), "s");
    headers.push(HeaderField(HeaderFieldName::Member, v));
    headers.dbus_encode(&mut msg);
    pad_to_multiple(&mut msg, 8);

    // Store the length of the header so we can easily compute body length later
    let len = msg.len() as u32;
    let mut lenbuf = Vec::new();
    len.dbus_encode(&mut lenbuf);
    set_length(&mut msg, &lenbuf);
    MessageBuf(msg)
}

impl MessageBuf {
    /// Add the given argument to the MessageBuf.  Accepts anything that implements the Marshal
    /// trait, which is most basic types, as well as the general-purpose
    /// dbus_serialize::types::Value enum.
    ///
    /// Note that these calls can be chained together to add multiple arguments, see the example
    ///
    /// # Examples
    /// ```
    /// dbus_bytestream::message::create_method_call("foo", "/bar", "baz", "bloop")
    ///     .add_arg(&1)
    ///     .add_arg(&"string");
    /// ```
    pub fn add_arg(mut self, arg: &Marshal) -> MessageBuf {
        arg.dbus_encode(&mut self.0);
        self
    }
}

const LEN_OFFSET : usize = 4;

pub fn set_length (msg: &mut [u8], buf: &[u8]) {
    for i in 0..buf.len() {
        msg[i+LEN_OFFSET] = buf[i];
    }
}

pub fn get_length (msg: &[u8]) -> u32 {
    assert!(msg.len() >= LEN_OFFSET+4);
    let mut lenbuf = [0; 4];
    for i in 0..4 {
        lenbuf[i] = msg[i+LEN_OFFSET];
    }
    let len : u32 = unsafe { transmute(lenbuf) };
    // len is already LE, so this is a no-op except on BE systems
    len.to_le()
}

/// Represents a received message from the message bus
#[derive(Debug,Default)]
pub struct Message {
    pub big_endian: bool,
    pub message_type: MessageType,
    pub flags: u8,
    pub version: u8,
    pub serial: u32,
    pub headers: HashMap<u8,Value>,
    pub body: Vec<Value>
}

#[test]
fn test_header () {
    let buf = encode_header(MESSAGE_TYPE_METHOD_CALL, 12);
    println!("{:?}", buf);
}

#[test]
fn test_msg () {
    create_method_call("foo", "bar", "baz", "floob")
        .add_arg(&1)
        .add_arg(&2);
}
<commit_msg>message.rs: add create_method_return()<commit_after>//! Functions for creating and modifying messages to send across the message bus.
use std::mem::transmute;
use std::collections::HashMap;

use dbus_serialize::types::{Path,Variant,Value,BasicValue,Signature};

use marshal::{Marshal,pad_to_multiple};

#[derive(Debug,Default)]
pub struct MessageType(pub u8);
pub const MESSAGE_TYPE_INVALID : MessageType        = MessageType(0);
pub const MESSAGE_TYPE_METHOD_CALL : MessageType    = MessageType(1);
pub const MESSAGE_TYPE_METHOD_RETURN : MessageType  = MessageType(2);
pub const MESSAGE_TYPE_ERROR : MessageType          = MessageType(3);
pub const MESSAGE_TYPE_SIGNAL : MessageType         = MessageType(4);

#[derive(Copy,Clone)]
pub enum HeaderFieldName {
    Invalid = 0,
    Path = 1,
    Interface = 2,
    Member = 3,
    ErrorName = 4,
    ReplySerial = 5,
    Destination = 6,
    Sender = 7,
    Signature = 8
}

struct HeaderField (
    HeaderFieldName,
    Variant
);

impl Marshal for HeaderField {
    fn dbus_encode(&self, buf: &mut Vec<u8>) -> usize {
        pad_to_multiple(buf, 8);
        let start_len = buf.len();
        let code = self.0 as u8;
        code.dbus_encode(buf);
        self.1.dbus_encode(buf);
        buf.len() - start_len
    }
    fn get_type(&self) -> String {
        "(yv)".to_string()
    }
}

fn encode_header (msg_type: MessageType, serial: u32) -> Vec<u8> {
    let mut buf = Vec::new();

    let endian = 'l' as u8;
    endian.dbus_encode(&mut buf);
    msg_type.0.dbus_encode(&mut buf);
    (0 as u8).dbus_encode(&mut buf);
    (1 as u8).dbus_encode(&mut buf);
    (0 as u32).dbus_encode(&mut buf);
    serial.dbus_encode(&mut buf);
    buf
}

/// Container for the header of a D-Bus message.  Arguments can be added with add_arg()
pub struct MessageBuf(pub Vec<u8>);

/// Create a MessageBuf for a D-Bus method call.  Once a MessageBuf object is created, arguments
/// can be added to the object with MessageBuf.add_arg
pub fn create_method_call (dest: &str, path: &str, iface: &str, method: &str) -> MessageBuf {
    let mut msg = encode_header(MESSAGE_TYPE_METHOD_CALL, 0);
    let mut headers : Vec<HeaderField> = Vec::new();
    let mut v = Variant::new(Value::BasicValue(BasicValue::String(dest.to_string())), "s");
    headers.push(HeaderField(HeaderFieldName::Destination, v));
    v = Variant::new(Value::BasicValue(BasicValue::ObjectPath(Path(path.to_string()))), "o");
    headers.push(HeaderField(HeaderFieldName::Path, v));
    v = Variant::new(Value::BasicValue(BasicValue::String(iface.to_string())), "s");
    headers.push(HeaderField(HeaderFieldName::Interface, v));
    v = Variant::new(Value::BasicValue(BasicValue::String(method.to_string())), "s");
    headers.push(HeaderField(HeaderFieldName::Member, v));
    headers.dbus_encode(&mut msg);
    pad_to_multiple(&mut msg, 8);

    // Store the length of the header so we can easily compute body length later
    let len = msg.len() as u32;
    let mut lenbuf = Vec::new();
    len.dbus_encode(&mut lenbuf);
    set_length(&mut msg, &lenbuf);
    MessageBuf(msg)
}

/// Create a MessageBuf for a D-Bus method return.  Once created, return values can be added to the
/// object with MessageBuf.add_arg
pub fn create_method_return(reply_serial: u32) -> MessageBuf {
    let mut msg = encode_header(MESSAGE_TYPE_METHOD_RETURN, 0);
    let mut headers : Vec<HeaderField> = Vec::new();
    let v = Variant::new(Value::from(reply_serial), "u");
    headers.push(HeaderField(HeaderFieldName::ReplySerial, v));
    pad_to_multiple(&mut msg, 8);

    // Store the length of the header so we can easily compute body length later
    let len = msg.len() as u32;
    let mut lenbuf = Vec::new();
    len.dbus_encode(&mut lenbuf);
    set_length(&mut msg, &lenbuf);
    MessageBuf(msg)
}

impl MessageBuf {
    /// Add the given argument to the MessageBuf.  Accepts anything that implements the Marshal
    /// trait, which is most basic types, as well as the general-purpose
    /// dbus_serialize::types::Value enum.
    ///
    /// Note that these calls can be chained together to add multiple arguments, see the example
    ///
    /// # Examples
    /// ```
    /// dbus_bytestream::message::create_method_call("foo", "/bar", "baz", "bloop")
    ///     .add_arg(&1)
    ///     .add_arg(&"string");
    /// ```
    pub fn add_arg(mut self, arg: &Marshal) -> MessageBuf {
        arg.dbus_encode(&mut self.0);
        self
    }
}

const LEN_OFFSET : usize = 4;

pub fn set_length (msg: &mut [u8], buf: &[u8]) {
    for i in 0..buf.len() {
        msg[i+LEN_OFFSET] = buf[i];
    }
}

pub fn get_length (msg: &[u8]) -> u32 {
    assert!(msg.len() >= LEN_OFFSET+4);
    let mut lenbuf = [0; 4];
    for i in 0..4 {
        lenbuf[i] = msg[i+LEN_OFFSET];
    }
    let len : u32 = unsafe { transmute(lenbuf) };
    // len is already LE, so this is a no-op except on BE systems
    len.to_le()
}

/// Represents a received message from the message bus
#[derive(Debug,Default)]
pub struct Message {
    pub big_endian: bool,
    pub message_type: MessageType,
    pub flags: u8,
    pub version: u8,
    pub serial: u32,
    pub headers: HashMap<u8,Value>,
    pub body: Vec<Value>
}

#[test]
fn test_header () {
    let buf = encode_header(MESSAGE_TYPE_METHOD_CALL, 12);
    println!("{:?}", buf);
}

#[test]
fn test_msg () {
    create_method_call("foo", "bar", "baz", "floob")
        .add_arg(&1)
        .add_arg(&2);
}
<|endoftext|>