<commit_before>use addressable::*;
use memory::*;

enum Operand {
	Implied,							// OPC				// operand implied
	Immediate(u8),						// OPC #$BB			// operand is byte (BB)
	Accumulator,						// OPC A			// operand is AC
	Relative(i8),						// OPC $BB			// branch target is PC + offset (BB), bit 7 signifies negative offset
	Absolute(u16),						// OPC $HHLL		// operand is address $HHLL
	AbsoluteIndexedWithX(u16),			// OPC $HHLL,X		// operand is address $HHLL incremented by X with carry
	AbsoluteIndexedWithY(u16),			// OPC $HHLL,Y		// operand is address $HHLL incremented by Y with carry
	Indirect(u16),						// OPC ($HHLL)		// operand is effective address; effective address is value of address
	ZeroPage(u8),						// OPC $LL			// operand is address $00LL
	ZeroPageIndexedWithX(u8),			// OPC $LL,X		// operand is address $00LL incremented by X, no page transition
	ZeroPageIndexedWithY(u8),			// OPC $LL,Y		// operand is address $00LL incremented by Y, no page transition
	ZeroPageIndexedWithXIndirect(u8),	// OPC ($LL,X)		// operand is effective address; effective address is $00LL incremented by X without carry
	ZeroPageIndirectIndexedWithY(u8),	// OPC ($LL),Y		// operand is effective address incremented by Y with carry; effective address is word at $00LL
}

impl Operand {
	fn addr (&self, cpu: &CPU) -> u16 {
		match * self {
			Implied								=> fail!("mos6510: Implied operand is never targetted to an address"),
			Immediate(_)						=> fail!("mos6510: Immediade operand is never targetted to an address"),
			Accumulator							=> fail!("mos6510: Accumulator operand is never targetted to an address"),
			Relative(offset)					=> cpu.pc + offset as u16,
			Absolute(addr)						=> addr,
			AbsoluteIndexedWithX(addr)			=> addr + cpu.x as u16,
			AbsoluteIndexedWithY(addr)			=> addr + cpu.y as u16,
			Indirect(addr)						=> cpu.get_le(addr),
			ZeroPage(addr)						=> addr as u16,
			ZeroPageIndexedWithX(addr)			=> (addr + cpu.x) as u16,				// no page transition
			ZeroPageIndexedWithY(addr)			=> (addr + cpu.y) as u16,				// no page transition
			ZeroPageIndexedWithXIndirect(addr)	=> cpu.get_le((addr + cpu.x) as u16),	// no page transition
			ZeroPageIndirectIndexedWithY(addr)	=> { let iaddr: u16 = cpu.get_le(addr as u16); iaddr + cpu.y as u16 },
		}
	}

	fn get (&self, cpu: &CPU) -> u8 {
		match *self {
			Implied								=> fail!("mos6510: Implied operand never has a value"),
			Immediate(val)						=> val,
			Accumulator							=> cpu.ac,
			Relative(_target)					=> fail!("mos6510: Relative operand never has a value"),
			op									=> { let addr = op.addr(cpu); cpu.get(addr) },
		}
	}

	fn set (&self, val: u8, cpu: &mut CPU) {
		match *self {
			Implied								=> fail!("mos6510: Implied operand never sets a value"),
			Immediate(_)						=> fail!("mos6510: Immediate operand never sets a value"),
			Accumulator							=> cpu.ac = val,
			Relative(_target)					=> fail!("mos6510: Relative operand never sets a value"),
			op									=> { let addr = op.addr(cpu); cpu.set(addr, val); },
		}
	}
}

pub struct CPU {
	pc: u16,				// program counter
	ac: u8,					// accumulator
	x: u8,					// x register
	y: u8,					// y register
	sr: u8,					// status register (NV-BDIZC: Negative, oVerflow, Break, Decimal, Interrupt, Zero, Carry)
	sp: u8,					// stack pointer
	mem: ~Memory,
}

impl Addressable<u16, u8> for CPU {
	pub fn get (&self, addr: u16) -> u8 {
		self.mem.get(addr as uint)
	}

	pub fn set (&mut self, addr: u16, data: u8) {
		self.mem.set(addr as uint, data);
	}
}
<commit_msg>Added comment about yet missing processor port<commit_after>use addressable::*;
use memory::*;

enum Operand {
	Implied,							// OPC				// operand implied
	Immediate(u8),						// OPC #$BB			// operand is byte (BB)
	Accumulator,						// OPC A			// operand is AC
	Relative(i8),						// OPC $BB			// branch target is PC + offset (BB), bit 7 signifies negative offset
	Absolute(u16),						// OPC $HHLL		// operand is address $HHLL
	AbsoluteIndexedWithX(u16),			// OPC $HHLL,X		// operand is address $HHLL incremented by X with carry
	AbsoluteIndexedWithY(u16),			// OPC $HHLL,Y		// operand is address $HHLL incremented by Y with carry
	Indirect(u16),						// OPC ($HHLL)		// operand is effective address; effective address is value of address
	ZeroPage(u8),						// OPC $LL			// operand is address $00LL
	ZeroPageIndexedWithX(u8),			// OPC $LL,X		// operand is address $00LL incremented by X, no page transition
	ZeroPageIndexedWithY(u8),			// OPC $LL,Y		// operand is address $00LL incremented by Y, no page transition
	ZeroPageIndexedWithXIndirect(u8),	// OPC ($LL,X)		// operand is effective address; effective address is $00LL incremented by X without carry
	ZeroPageIndirectIndexedWithY(u8),	// OPC ($LL),Y		// operand is effective address incremented by Y with carry; effective address is word at $00LL
}

impl Operand {
	fn addr (&self, cpu: &CPU) -> u16 {
		match * self {
			Implied								=> fail!("mos6510: Implied operand is never targetted to an address"),
			Immediate(_)						=> fail!("mos6510: Immediade operand is never targetted to an address"),
			Accumulator							=> fail!("mos6510: Accumulator operand is never targetted to an address"),
			Relative(offset)					=> cpu.pc + offset as u16,
			Absolute(addr)						=> addr,
			AbsoluteIndexedWithX(addr)			=> addr + cpu.x as u16,
			AbsoluteIndexedWithY(addr)			=> addr + cpu.y as u16,
			Indirect(addr)						=> cpu.get_le(addr),
			ZeroPage(addr)						=> addr as u16,
			ZeroPageIndexedWithX(addr)			=> (addr + cpu.x) as u16,				// no page transition
			ZeroPageIndexedWithY(addr)			=> (addr + cpu.y) as u16,				// no page transition
			ZeroPageIndexedWithXIndirect(addr)	=> cpu.get_le((addr + cpu.x) as u16),	// no page transition
			ZeroPageIndirectIndexedWithY(addr)	=> { let iaddr: u16 = cpu.get_le(addr as u16); iaddr + cpu.y as u16 },
		}
	}

	fn get (&self, cpu: &CPU) -> u8 {
		match *self {
			Implied								=> fail!("mos6510: Implied operand never has a value"),
			Immediate(val)						=> val,
			Accumulator							=> cpu.ac,
			Relative(_target)					=> fail!("mos6510: Relative operand never has a value"),
			op									=> { let addr = op.addr(cpu); cpu.get(addr) },
		}
	}

	fn set (&self, val: u8, cpu: &mut CPU) {
		match *self {
			Implied								=> fail!("mos6510: Implied operand never sets a value"),
			Immediate(_)						=> fail!("mos6510: Immediate operand never sets a value"),
			Accumulator							=> cpu.ac = val,
			Relative(_target)					=> fail!("mos6510: Relative operand never sets a value"),
			op									=> { let addr = op.addr(cpu); cpu.set(addr, val); },
		}
	}
}

pub struct CPU {
	pc: u16,				// program counter
	ac: u8,					// accumulator
	x: u8,					// x register
	y: u8,					// y register
	sr: u8,					// status register (NV-BDIZC: Negative, oVerflow, Break, Decimal, Interrupt, Zero, Carry)
	sp: u8,					// stack pointer
	mem: ~Memory,
}

impl Addressable<u16, u8> for CPU {
	pub fn get (&self, addr: u16) -> u8 {
		// TODO: addresses $0000 (data direction) and $0001 (data) are hardwired for the processor I/O port
		self.mem.get(addr as uint)
	}

	pub fn set (&mut self, addr: u16, data: u8) {
		// TODO: addresses $0000 (data direction) and $0001 (data) are hardwired for the processor I/O port
		self.mem.set(addr as uint, data);
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Linus Färnstrand <faern@faern.net>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support for sending and receiving data link layer packets on a fake network managed
//! by in memory FIFO queues. Useful for writing tests.

use {DataLinkReceiver, DataLinkSender, MacAddr, NetworkInterface};

use std::io;
use std::sync::mpsc::{self, Receiver, Sender};
use std::thread;
use std::time;

/// Configuration for the dummy datalink backend. Contains `std::sync::mpsc`
/// channels that are used to communicate with the fake network.
#[derive(Debug)]
pub struct Config {
    receiver: Receiver<io::Result<Box<[u8]>>>,
    inject_handle: Option<Sender<io::Result<Box<[u8]>>>>,

    sender: Sender<Box<[u8]>>,
    read_handle: Option<Receiver<Box<[u8]>>>,
}

impl Config {
    /// Creates a new `Config` with the given channels as the backing network.
    /// When using this constructor `inject_handle` and `read_handle` will return `None`.
    /// Those handles must be kept track of elsewhere.
    ///
    /// The `DataLinkReceiver` created by the dummy backend will read packets from
    /// `receiver`. Both network errors and data can be sent on this channel.
    /// When the `receiver` channel is closed (`Sender` is dropped)
    /// `DataLinkReceiver::next()` will sleep forever, simlating an idle network.
    ///
    /// The `DataLinkSender` created by the dummy backend will send all packets sent
    /// through `build_and_send()` and `send_to()` to the `sender` channel.
    pub fn new(receiver: Receiver<io::Result<Box<[u8]>>>, sender: Sender<Box<[u8]>>) -> Config {
        Config {
            receiver: receiver,
            inject_handle: None,
            sender: sender,
            read_handle: None,
        }
    }

    /// Get the `Sender` handle that can inject packets in the fake network.
    /// Only usable with `Config`s generated from `default()`.
    pub fn inject_handle(&mut self) -> Option<Sender<io::Result<Box<[u8]>>>> {
        self.inject_handle.take()
    }

    /// Get the `Receiver` handle where packets sent to the fake network can be read.
    /// Only usable with `Config`s generated from `default()`.
    pub fn read_handle(&mut self) -> Option<Receiver<Box<[u8]>>> {
        self.read_handle.take()
    }
}

impl<'a> From<&'a super::Config> for Config {
    /// Will not use the `super::Config`. This will simply call `dummy::Config::default()`.
    fn from(_config: &super::Config) -> Config {
        Config::default()
    }
}

impl Default for Config {
    /// Creates a default config with one input and one output channel. The handles used to inject
    /// to and read form the network can be fetched with `inject_handle()` and `read_handle()`.
    fn default() -> Config {
        let (in_tx, in_rx) = mpsc::channel();
        let (out_tx, out_rx) = mpsc::channel();
        Config {
            receiver: in_rx,
            inject_handle: Some(in_tx),
            sender: out_tx,
            read_handle: Some(out_rx),
        }
    }
}

/// Create a data link channel backed by FIFO queues. Useful for debugging and testing.
/// See `Config` for how to inject and read packets on this fake network.
pub fn channel(_: &NetworkInterface, config: Config) -> io::Result<super::Channel> {
    let sender = Box::new(MockDataLinkSender {
        sender: config.sender,
    });
    let receiver = Box::new(MockDataLinkReceiver {
        receiver: config.receiver,
        used_packets: Vec::new(),
    });

    Ok(super::Channel::Ethernet(sender, receiver))
}

struct MockDataLinkSender {
    sender: Sender<Box<[u8]>>,
}

impl DataLinkSender for MockDataLinkSender {
    fn build_and_send(
        &mut self,
        num_packets: usize,
        packet_size: usize,
        func: &mut dyn FnMut(&mut [u8]),
    ) -> Option<io::Result<()>> {
        for _ in 0..num_packets {
            let mut buffer = vec![0; packet_size];
            func(&mut buffer);
            // Send the data to the queue. Don't care if it's closed
            self.sender.send(buffer.into_boxed_slice()).unwrap_or(());
        }
        Some(Ok(()))
    }

    fn send_to(&mut self, packet: &[u8], _dst: Option<NetworkInterface>) -> Option<io::Result<()>> {
        let buffer = packet.to_vec();
        self.sender.send(buffer.into_boxed_slice()).unwrap_or(());
        Some(Ok(()))
    }
}

struct MockDataLinkReceiver {
    receiver: Receiver<io::Result<Box<[u8]>>>,
    used_packets: Vec<Box<[u8]>>,
}

impl DataLinkReceiver for MockDataLinkReceiver {
    fn next(&mut self) -> io::Result<&[u8]> {
        match self.receiver.recv() {
            Ok(result) => {
                // A network event happened. Might be a packet or a simulated error
                match result {
                    Ok(buffer) => {
                        self.used_packets.push(buffer);
                        let buffer_ref = &*self.used_packets[self.used_packets.len() - 1];
                        Ok(buffer_ref)
                    }
                    Err(e) => Err(e),
                }
            }
            Err(_) => {
                // The channel supplying fake packets is broken. The user lost/destroyed their
                // inject_handle. This means there will never be any more packets sent to this
                // dummy network. To simulate an idle network we block and sleep forever here.
                loop {
                    thread::sleep(time::Duration::new(10, 0));
                }
            }
        }
    }
}

/// Get three fake interfaces generated with `dummy_interface(0..3)`.
pub fn interfaces() -> Vec<NetworkInterface> {
    (0..3).map(|i| dummy_interface(i)).collect()
}

/// Generates a fake `NetworkInterface`.
/// The name of the interface will be `ethX` where X is the integer `i`.
/// The index will be `i`.
/// The MAC will be `01:02:03:04:05:i`.
pub fn dummy_interface(i: u8) -> NetworkInterface {
    NetworkInterface {
        name: format!("eth{}", i),
        index: i as u32,
        mac: Some(MacAddr::new(1, 2, 3, 4, 5, i)),
        ips: Vec::new(),
        flags: 0,
    }
}

#[cfg(test)]
mod tests {
    use {DataLinkReceiver, DataLinkSender};

    use std::io;
    use std::sync::mpsc::{self, Receiver, Sender, TryRecvError};
    use std::thread::{sleep, spawn};
    use std::time::Duration;

    #[test]
    fn send_nothing() {
        let (_, read_handle, mut tx, _) = create_net();
        // Check that sending zero packets yields zero packets
        let mut builder = |_: &mut [u8]| {
            panic!("Should not be called");
        };
        tx.build_and_send(0, 20, &mut builder).unwrap().unwrap();
        assert!(read_handle.try_recv().is_err());
    }

    #[test]
    fn send_one_packet() {
        let (_, read_handle, mut tx, _) = create_net();
        // Check that sending one packet yields one packet
        let mut builder = |pkg: &mut [u8]| {
            assert_eq!(pkg.len(), 20);
            pkg[0] = 9;
            pkg[19] = 201;
        };
        tx.build_and_send(1, 20, &mut builder).unwrap().unwrap();
        let pkg = read_handle
            .try_recv()
            .expect("Expected one packet to be sent");
        assert!(read_handle.try_recv().is_err());
        assert_eq!(pkg.len(), 20);
        assert_eq!(pkg[0], 9);
        assert_eq!(pkg[19], 201);
    }

    #[test]
    fn send_multiple_packets() {
        let (_, read_handle, mut tx, _) = create_net();
        // Check that sending multiple packets does the correct thing
        let mut closure_counter = 0;
        let mut builder = |pkg: &mut [u8]| {
            pkg[0] = closure_counter;
            closure_counter += 1;
        };
        tx.build_and_send(3, 20, &mut builder).unwrap().unwrap();
        for i in 0..3 {
            let pkg = read_handle.try_recv().expect("Expected a packet");
            assert_eq!(pkg[0], i);
        }
        assert!(read_handle.try_recv().is_err());
    }

    #[test]
    fn send_to() {
        let (_, read_handle, mut tx, _) = create_net();
        let mut buffer = vec![0; 20];
        buffer[1] = 34;
        buffer[18] = 76;

        tx.send_to(&buffer, None).unwrap().unwrap();
        let pkg = read_handle
            .try_recv()
            .expect("Expected one packet to be sent");
        assert!(read_handle.try_recv().is_err());
        assert_eq!(pkg.len(), 20);
        assert_eq!(pkg[1], 34);
        assert_eq!(pkg[18], 76);
    }

    #[test]
    fn read_nothing() {
        let (_, _, _, mut rx) = create_net();
        let (control_tx, control_rx) = mpsc::channel();
        spawn(move || {
            rx.next().expect("Should not happen 1");
            control_tx.send(()).expect("Should not happen 2");
        });
        sleep(Duration::new(0, 1_000_000));
        match control_rx.try_recv() {
            Ok(_) => panic!("Nothing should have arrived"),
            Err(TryRecvError::Disconnected) => panic!("Thread should not have quit"),
            Err(TryRecvError::Empty) => (),
        }
    }

    #[test]
    fn read_one_pkg() {
        let (inject_handle, _, _, mut rx) = create_net();

        let buffer = vec![0; 20];
        inject_handle.send(Ok(buffer.into_boxed_slice())).unwrap();

        let pkg = rx.next().expect("Expected a packet");
        assert_eq!(pkg.len(), 20);
    }

    #[test]
    fn read_multiple_pkgs() {
        let (inject_handle, _, _, mut rx) = create_net();

        for i in 0..3 {
            let buffer = vec![i; 20];
            inject_handle.send(Ok(buffer.into_boxed_slice())).unwrap();
        }

        {
            let pkg1 = rx.next().expect("Expected a packet");
            assert_eq!(pkg1[0], 0);
        }
        {
            let pkg2 = rx.next().expect("Expected a packet");
            assert_eq!(pkg2[0], 1);
        }
        {
            let pkg3 = rx.next().expect("Expected a packet");
            assert_eq!(pkg3[0], 2);
        }
    }

    fn create_net() -> (
        Sender<io::Result<Box<[u8]>>>,
        Receiver<Box<[u8]>>,
        Box<DataLinkSender>,
        Box<DataLinkReceiver>,
    ) {
        let interface = super::dummy_interface(56);
        let mut config = super::Config::default();
        let inject_handle = config.inject_handle().unwrap();
        let read_handle = config.read_handle().unwrap();

        let channel = super::channel(&interface, config);
        let (tx, rx) = match channel {
            Ok(super::super::Channel::Ethernet(tx, rx)) => (tx, rx),
            _ => panic!("Not a valid channel returned"),
        };
        (inject_handle, read_handle, tx, rx)
    }
}
<commit_msg>Fix warnings in pnet_datalink tests<commit_after>// Copyright (c) 2016 Linus Färnstrand <faern@faern.net>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support for sending and receiving data link layer packets on a fake network managed
//! by in memory FIFO queues. Useful for writing tests.

use {DataLinkReceiver, DataLinkSender, MacAddr, NetworkInterface};

use std::io;
use std::sync::mpsc::{self, Receiver, Sender};
use std::thread;
use std::time;

/// Configuration for the dummy datalink backend. Contains `std::sync::mpsc`
/// channels that are used to communicate with the fake network.
#[derive(Debug)]
pub struct Config {
    receiver: Receiver<io::Result<Box<[u8]>>>,
    inject_handle: Option<Sender<io::Result<Box<[u8]>>>>,

    sender: Sender<Box<[u8]>>,
    read_handle: Option<Receiver<Box<[u8]>>>,
}

impl Config {
    /// Creates a new `Config` with the given channels as the backing network.
    /// When using this constructor `inject_handle` and `read_handle` will return `None`.
    /// Those handles must be kept track of elsewhere.
    ///
    /// The `DataLinkReceiver` created by the dummy backend will read packets from
    /// `receiver`. Both network errors and data can be sent on this channel.
    /// When the `receiver` channel is closed (`Sender` is dropped)
    /// `DataLinkReceiver::next()` will sleep forever, simlating an idle network.
    ///
    /// The `DataLinkSender` created by the dummy backend will send all packets sent
    /// through `build_and_send()` and `send_to()` to the `sender` channel.
    pub fn new(receiver: Receiver<io::Result<Box<[u8]>>>, sender: Sender<Box<[u8]>>) -> Config {
        Config {
            receiver: receiver,
            inject_handle: None,
            sender: sender,
            read_handle: None,
        }
    }

    /// Get the `Sender` handle that can inject packets in the fake network.
    /// Only usable with `Config`s generated from `default()`.
    pub fn inject_handle(&mut self) -> Option<Sender<io::Result<Box<[u8]>>>> {
        self.inject_handle.take()
    }

    /// Get the `Receiver` handle where packets sent to the fake network can be read.
    /// Only usable with `Config`s generated from `default()`.
    pub fn read_handle(&mut self) -> Option<Receiver<Box<[u8]>>> {
        self.read_handle.take()
    }
}

impl<'a> From<&'a super::Config> for Config {
    /// Will not use the `super::Config`. This will simply call `dummy::Config::default()`.
    fn from(_config: &super::Config) -> Config {
        Config::default()
    }
}

impl Default for Config {
    /// Creates a default config with one input and one output channel. The handles used to inject
    /// to and read form the network can be fetched with `inject_handle()` and `read_handle()`.
    fn default() -> Config {
        let (in_tx, in_rx) = mpsc::channel();
        let (out_tx, out_rx) = mpsc::channel();
        Config {
            receiver: in_rx,
            inject_handle: Some(in_tx),
            sender: out_tx,
            read_handle: Some(out_rx),
        }
    }
}

/// Create a data link channel backed by FIFO queues. Useful for debugging and testing.
/// See `Config` for how to inject and read packets on this fake network.
pub fn channel(_: &NetworkInterface, config: Config) -> io::Result<super::Channel> {
    let sender = Box::new(MockDataLinkSender {
        sender: config.sender,
    });
    let receiver = Box::new(MockDataLinkReceiver {
        receiver: config.receiver,
        used_packets: Vec::new(),
    });

    Ok(super::Channel::Ethernet(sender, receiver))
}

struct MockDataLinkSender {
    sender: Sender<Box<[u8]>>,
}

impl DataLinkSender for MockDataLinkSender {
    fn build_and_send(
        &mut self,
        num_packets: usize,
        packet_size: usize,
        func: &mut dyn FnMut(&mut [u8]),
    ) -> Option<io::Result<()>> {
        for _ in 0..num_packets {
            let mut buffer = vec![0; packet_size];
            func(&mut buffer);
            // Send the data to the queue. Don't care if it's closed
            self.sender.send(buffer.into_boxed_slice()).unwrap_or(());
        }
        Some(Ok(()))
    }

    fn send_to(&mut self, packet: &[u8], _dst: Option<NetworkInterface>) -> Option<io::Result<()>> {
        let buffer = packet.to_vec();
        self.sender.send(buffer.into_boxed_slice()).unwrap_or(());
        Some(Ok(()))
    }
}

struct MockDataLinkReceiver {
    receiver: Receiver<io::Result<Box<[u8]>>>,
    used_packets: Vec<Box<[u8]>>,
}

impl DataLinkReceiver for MockDataLinkReceiver {
    fn next(&mut self) -> io::Result<&[u8]> {
        match self.receiver.recv() {
            Ok(result) => {
                // A network event happened. Might be a packet or a simulated error
                match result {
                    Ok(buffer) => {
                        self.used_packets.push(buffer);
                        let buffer_ref = &*self.used_packets[self.used_packets.len() - 1];
                        Ok(buffer_ref)
                    }
                    Err(e) => Err(e),
                }
            }
            Err(_) => {
                // The channel supplying fake packets is broken. The user lost/destroyed their
                // inject_handle. This means there will never be any more packets sent to this
                // dummy network. To simulate an idle network we block and sleep forever here.
                loop {
                    thread::sleep(time::Duration::new(10, 0));
                }
            }
        }
    }
}

/// Get three fake interfaces generated with `dummy_interface(0..3)`.
pub fn interfaces() -> Vec<NetworkInterface> {
    (0..3).map(|i| dummy_interface(i)).collect()
}

/// Generates a fake `NetworkInterface`.
/// The name of the interface will be `ethX` where X is the integer `i`.
/// The index will be `i`.
/// The MAC will be `01:02:03:04:05:i`.
pub fn dummy_interface(i: u8) -> NetworkInterface {
    NetworkInterface {
        name: format!("eth{}", i),
        index: i as u32,
        mac: Some(MacAddr::new(1, 2, 3, 4, 5, i)),
        ips: Vec::new(),
        flags: 0,
    }
}

#[cfg(test)]
mod tests {
    use {DataLinkReceiver, DataLinkSender};

    use std::io;
    use std::sync::mpsc::{self, Receiver, Sender, TryRecvError};
    use std::thread::{sleep, spawn};
    use std::time::Duration;

    #[test]
    fn send_nothing() {
        let (_, read_handle, mut tx, _) = create_net();
        // Check that sending zero packets yields zero packets
        let mut builder = |_: &mut [u8]| {
            panic!("Should not be called");
        };
        tx.build_and_send(0, 20, &mut builder).unwrap().unwrap();
        assert!(read_handle.try_recv().is_err());
    }

    #[test]
    fn send_one_packet() {
        let (_, read_handle, mut tx, _) = create_net();
        // Check that sending one packet yields one packet
        let mut builder = |pkg: &mut [u8]| {
            assert_eq!(pkg.len(), 20);
            pkg[0] = 9;
            pkg[19] = 201;
        };
        tx.build_and_send(1, 20, &mut builder).unwrap().unwrap();
        let pkg = read_handle
            .try_recv()
            .expect("Expected one packet to be sent");
        assert!(read_handle.try_recv().is_err());
        assert_eq!(pkg.len(), 20);
        assert_eq!(pkg[0], 9);
        assert_eq!(pkg[19], 201);
    }

    #[test]
    fn send_multiple_packets() {
        let (_, read_handle, mut tx, _) = create_net();
        // Check that sending multiple packets does the correct thing
        let mut closure_counter = 0;
        let mut builder = |pkg: &mut [u8]| {
            pkg[0] = closure_counter;
            closure_counter += 1;
        };
        tx.build_and_send(3, 20, &mut builder).unwrap().unwrap();
        for i in 0..3 {
            let pkg = read_handle.try_recv().expect("Expected a packet");
            assert_eq!(pkg[0], i);
        }
        assert!(read_handle.try_recv().is_err());
    }

    #[test]
    fn send_to() {
        let (_, read_handle, mut tx, _) = create_net();
        let mut buffer = vec![0; 20];
        buffer[1] = 34;
        buffer[18] = 76;

        tx.send_to(&buffer, None).unwrap().unwrap();
        let pkg = read_handle
            .try_recv()
            .expect("Expected one packet to be sent");
        assert!(read_handle.try_recv().is_err());
        assert_eq!(pkg.len(), 20);
        assert_eq!(pkg[1], 34);
        assert_eq!(pkg[18], 76);
    }

    #[test]
    fn read_nothing() {
        let (_, _, _, mut rx) = create_net();
        let (control_tx, control_rx) = mpsc::channel();
        spawn(move || {
            rx.next().expect("Should not happen 1");
            control_tx.send(()).expect("Should not happen 2");
        });
        sleep(Duration::new(0, 1_000_000));
        match control_rx.try_recv() {
            Ok(_) => panic!("Nothing should have arrived"),
            Err(TryRecvError::Disconnected) => panic!("Thread should not have quit"),
            Err(TryRecvError::Empty) => (),
        }
    }

    #[test]
    fn read_one_pkg() {
        let (inject_handle, _, _, mut rx) = create_net();

        let buffer = vec![0; 20];
        inject_handle.send(Ok(buffer.into_boxed_slice())).unwrap();

        let pkg = rx.next().expect("Expected a packet");
        assert_eq!(pkg.len(), 20);
    }

    #[test]
    fn read_multiple_pkgs() {
        let (inject_handle, _, _, mut rx) = create_net();

        for i in 0..3 {
            let buffer = vec![i; 20];
            inject_handle.send(Ok(buffer.into_boxed_slice())).unwrap();
        }

        {
            let pkg1 = rx.next().expect("Expected a packet");
            assert_eq!(pkg1[0], 0);
        }
        {
            let pkg2 = rx.next().expect("Expected a packet");
            assert_eq!(pkg2[0], 1);
        }
        {
            let pkg3 = rx.next().expect("Expected a packet");
            assert_eq!(pkg3[0], 2);
        }
    }

    fn create_net() -> (
        Sender<io::Result<Box<[u8]>>>,
        Receiver<Box<[u8]>>,
        Box<dyn DataLinkSender>,
        Box<dyn DataLinkReceiver>,
    ) {
        let interface = super::dummy_interface(56);
        let mut config = super::Config::default();
        let inject_handle = config.inject_handle().unwrap();
        let read_handle = config.read_handle().unwrap();

        let channel = super::channel(&interface, config);
        let (tx, rx) = match channel {
            Ok(super::super::Channel::Ethernet(tx, rx)) => (tx, rx),
            _ => panic!("Not a valid channel returned"),
        };
        (inject_handle, read_handle, tx, rx)
    }
}
<|endoftext|>


<commit_before>use csv;
use docopt;

use types::{CliError, Delimiter, InputReader, OutputWriter};
use util;

docopt!(Args, "
Concatenates CSV data by column or by row.

When concatenating by column, the columns will be written in the same order
as the inputs given. The number of rows in the result is always equivalent to
to the minimum number of rows across all given CSV data. (This behavior can
be reversed with the '--pad' flag.)

When concatenating by row, all CSV data must have the same number of columns.
If you need to rearrange the columns or fix the lengths of records, use the
'slice' or 'fixlengths' commands. Also, only the headers of the *first* CSV
data given are used. Headers in subsequent inputs are ignored. (This behavior
can be disabled with --no-headers.)

Usage:
    xcsv cat (rows | columns) [options] [<input>...]

cat options:
    -p, --pad              When concatenating columns, this flag will cause
                           all records to appear. It will pad each row if
                           other CSV data isn't long enough.

Common options:
    -h, --help             Display this message
    -o, --output <file>    Write output to <file> instead of stdout.
    -n, --no-headers       When set, the first row will NOT be interpreted
                           as column names. Note that this has no effect when
                           concatenating columns.
    -d, --delimiter <arg>  The field delimiter for reading CSV data.
                           Must be a single character. [default: ,]
", arg_input: Vec<InputReader>, flag_output: OutputWriter,
   flag_delimiter: Delimiter, flag_out_delimiter: Delimiter,
   flag_length: Option<uint>)

pub fn main() -> Result<(), CliError> {
    let mut args: Args = try!(util::get_args());
    let mut wtr = csv::Encoder::to_writer(args.flag_output.by_ref());
    if args.arg_input.is_empty() {
        args.arg_input.push(ctry!(InputReader::new(None))); // stdin
    }
    ctry!(util::at_most_one_stdin(args.arg_input.as_slice()));

    if args.cmd_rows {
        for (i, inp) in args.arg_input.move_iter().enumerate() {
            let mut rdr = csv_reader!(args, inp);
            if !args.flag_no_headers && i == 0 {
                csv_write_headers!(args, rdr, wtr);
            }
            for r in rdr.iter_bytes() {
                ctry!(wtr.record_bytes(ctry!(r).move_iter()));
            }
        }
    } else if args.cmd_columns {
        let delim = args.flag_delimiter.to_byte();
        let mut rdrs: Vec<csv::Decoder<InputReader>> =
            args.arg_input
                .move_iter().map(|inp| column_reader(inp, delim)).collect();
        // Find the lengths of each record. If a length varies, then an error
        // will occur so we can rely on the first length being the correct one.
        let mut lengths = vec!();
        for rdr in rdrs.mut_iter() {
            lengths.push(ctry!(rdr.headers_bytes()).len());
        }
        'OUTER: loop {
            let mut records: Vec<Vec<csv::ByteString>> = vec!();
            let mut num_done = 0;
            for (rdr, &len) in rdrs.mut_iter().zip(lengths.iter()) {
                match rdr.record_bytes() {
                    Err(ref e) if e.is_eof() => {
                        num_done += 1;
                        if args.flag_pad {
                            // This can probably be optimized by pre-allocating.
                            // It would avoid the intermediate `Vec`.
                            let pad = Vec::from_elem(len, util::empty_field());
                            records.push(pad);
                        } else {
                            break 'OUTER;
                        }
                    }
                    err @ Err(_) => { ctry!(err); }
                    Ok(next) => records.push(next),
                }
            }
            // Only needed when `--pad` is set.
            if num_done >= rdrs.len() {
                break 'OUTER;
            }
            let row = records.as_slice().concat_vec();
            ctry!(wtr.record_bytes(row.move_iter()));
        }
    } else {
        unreachable!();
    }
    ctry!(wtr.flush());
    Ok(())
}

fn column_reader<R: Reader>(rdr: R, delim: u8) -> csv::Decoder<R> {
    // We always set no_headers here because there's no need to distinguish.
    csv::Decoder::from_reader(rdr).separator(delim).no_headers()
}
<commit_msg>Add --help to cat command.<commit_after>use csv;
use docopt;

use types::{CliError, Delimiter, InputReader, OutputWriter};
use util;

docopt!(Args, "
Concatenates CSV data by column or by row.

When concatenating by column, the columns will be written in the same order
as the inputs given. The number of rows in the result is always equivalent to
to the minimum number of rows across all given CSV data. (This behavior can
be reversed with the '--pad' flag.)

When concatenating by row, all CSV data must have the same number of columns.
If you need to rearrange the columns or fix the lengths of records, use the
'slice' or 'fixlengths' commands. Also, only the headers of the *first* CSV
data given are used. Headers in subsequent inputs are ignored. (This behavior
can be disabled with --no-headers.)

Usage:
    xcsv cat (rows | columns) [options] [<input>...]
    xcsv cat --help

cat options:
    -p, --pad              When concatenating columns, this flag will cause
                           all records to appear. It will pad each row if
                           other CSV data isn't long enough.

Common options:
    -h, --help             Display this message
    -o, --output <file>    Write output to <file> instead of stdout.
    -n, --no-headers       When set, the first row will NOT be interpreted
                           as column names. Note that this has no effect when
                           concatenating columns.
    -d, --delimiter <arg>  The field delimiter for reading CSV data.
                           Must be a single character. [default: ,]
", arg_input: Vec<InputReader>, flag_output: OutputWriter,
   flag_delimiter: Delimiter, flag_out_delimiter: Delimiter,
   flag_length: Option<uint>)

pub fn main() -> Result<(), CliError> {
    let mut args: Args = try!(util::get_args());
    let mut wtr = csv::Encoder::to_writer(args.flag_output.by_ref());
    if args.arg_input.is_empty() {
        args.arg_input.push(ctry!(InputReader::new(None))); // stdin
    }
    ctry!(util::at_most_one_stdin(args.arg_input.as_slice()));

    if args.cmd_rows {
        for (i, inp) in args.arg_input.move_iter().enumerate() {
            let mut rdr = csv_reader!(args, inp);
            if !args.flag_no_headers && i == 0 {
                csv_write_headers!(args, rdr, wtr);
            }
            for r in rdr.iter_bytes() {
                ctry!(wtr.record_bytes(ctry!(r).move_iter()));
            }
        }
    } else if args.cmd_columns {
        let delim = args.flag_delimiter.to_byte();
        let mut rdrs: Vec<csv::Decoder<InputReader>> =
            args.arg_input
                .move_iter().map(|inp| column_reader(inp, delim)).collect();
        // Find the lengths of each record. If a length varies, then an error
        // will occur so we can rely on the first length being the correct one.
        let mut lengths = vec!();
        for rdr in rdrs.mut_iter() {
            lengths.push(ctry!(rdr.headers_bytes()).len());
        }
        'OUTER: loop {
            let mut records: Vec<Vec<csv::ByteString>> = vec!();
            let mut num_done = 0;
            for (rdr, &len) in rdrs.mut_iter().zip(lengths.iter()) {
                match rdr.record_bytes() {
                    Err(ref e) if e.is_eof() => {
                        num_done += 1;
                        if args.flag_pad {
                            // This can probably be optimized by pre-allocating.
                            // It would avoid the intermediate `Vec`.
                            let pad = Vec::from_elem(len, util::empty_field());
                            records.push(pad);
                        } else {
                            break 'OUTER;
                        }
                    }
                    err @ Err(_) => { ctry!(err); }
                    Ok(next) => records.push(next),
                }
            }
            // Only needed when `--pad` is set.
            if num_done >= rdrs.len() {
                break 'OUTER;
            }
            let row = records.as_slice().concat_vec();
            ctry!(wtr.record_bytes(row.move_iter()));
        }
    } else {
        unreachable!();
    }
    ctry!(wtr.flush());
    Ok(())
}

fn column_reader<R: Reader>(rdr: R, delim: u8) -> csv::Decoder<R> {
    // We always set no_headers here because there's no need to distinguish.
    csv::Decoder::from_reader(rdr).separator(delim).no_headers()
}
<|endoftext|>


<commit_before>extern crate piston;
extern crate graphics;
extern crate sdl2_windows;
extern crate opengl_graphics;

use std::rc::Rc;
use std::cell::RefCell;
use piston::window::WindowSetting;
use piston::event::{ events, RenderArgs, RenderEvent, UpdateArgs,UpdateEvent };

use graphics::{ Context, rectangle, RelativeTransform };
use sdl2_window::Sdl2Window as Window;
use opengl_graphics::{ GlGraphics, OpenGL };

pub struct App {
	gl: GlGraphics,
	rotation: f64
}
impl App {
	fn render(&mut self, args: &RenderArgs) {
		const GREEN: [f32; 4] = [0.0, 1.0, 0.0, 1.0];
		const BLUE: [f32; 4] = [0.0, 0.0, 1.0, 1.0];

		let context = &Context::abs(args.width as f64, args.height as f64);

		let center_context = &context.trans((args.width/2) as f64, (args.height / 2) as f64)
									 .rot_rad(self.rotation)
									 .trans(-25.0, -25.0);
		let square = rectangle::square(0.0, 0.0, 50.0);

		self.gl.draw([0, 0, args.width as i32, args.height as i32], |_, gl| {
			graphics::clear(GREEN, gl);

			graphics::rectangle(BLUE, square, center_context.transform, gl);
		});

	}

	fn update(&mut self, args: &UpdateArgs) {
		self.rotation += 2.0 * args.dt;
	}
}

fn main() {
	let ayy = 12;
	let mut lmao = 11;
	println!("so the value of this integer I call ayy is: {} and lmao is: {}", ayy, lmao);

	lmao += 1;
	println!("now the values are: {} and {}", ayy, lmao);

	println!("Oh btw I'm starting the awesomepart now");
	let window = Window::new(
		OpenGL::_3_2,
		WindowSetting::default()
	);
	let window = Rc::new(RefCell::new(window));

	let mut app = App{
		gl: GlGraphics::new(OpenGL::_3_2),
		rotation: 0.0
	};

	for e in events(window) {
		if let Some(r) = e.render_args() {
			app.render(&r);
		}

		if let Some(u) = e.update_args() {
			app.update(&u);
		}
	}
}
<commit_msg>the program compiles on my laptop<commit_after>extern crate piston;
extern crate graphics;
extern crate sdl2_window;
extern crate opengl_graphics;

use std::rc::Rc;
use std::cell::RefCell;
use piston::window::WindowSettings;
use piston::event::{ events, RenderArgs, RenderEvent, UpdateArgs,UpdateEvent };

use graphics::{ Context, rectangle, RelativeTransform };
use sdl2_window::Sdl2Window as Window;
use opengl_graphics::{ GlGraphics, OpenGL };

pub struct App {
	gl: GlGraphics,
	rotation: f64
}
impl App {
	fn render(&mut self, args: &RenderArgs) {
		const DARKGREEN: [f32; 4] = [0.0, 0.5, 0.0, 1.0];
		const BLUE: [f32; 4] = [0.0, 0.0, 1.0, 1.0];

		let context = &Context::abs(args.width as f64, args.height as f64);

		let center_context = &context.trans((args.width/2) as f64, (args.height / 2) as f64)
									 .rot_rad(self.rotation)
									 .trans(-25.0, -25.0);
		let square = rectangle::square(0.0, 0.0, 50.0);

		self.gl.draw([0, 0, args.width as i32, args.height as i32], |_, gl| {
			graphics::clear(DARKGREEN, gl);

			graphics::rectangle(BLUE, square, center_context.transform, gl);
		});

	}

	fn update(&mut self, args: &UpdateArgs) {
		self.rotation += 2.0 * args.dt;
	}
}

fn main() {
	let ayy = 12;
	let mut lmao = 11;
	println!("so the value of this integer I call ayy is: {} and lmao is: {}", ayy, lmao);

	lmao += 1;
	println!("now the values are: {} and {}", ayy, lmao);

	println!("Oh btw I'm starting the awesomepart now");
	let window = Window::new(
		OpenGL::_3_2,
		WindowSettings::default()
	);
	let window = Rc::new(RefCell::new(window));

	let mut app = App{
		gl: GlGraphics::new(OpenGL::_3_2),
		rotation: 0.0
	};

	for e in events(window) {
		if let Some(r) = e.render_args() {
			app.render(&r);
		}

		if let Some(u) = e.update_args() {
			app.update(&u);
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

use std::path::{Path, PathBuf};
use std::io::{Read, Write, Result};

const PROTOS: &'static [&'static str] = &[
    "grr/grr/proto/grr_response_proto/semantic.proto",
    "grr/grr/proto/grr_response_proto/sysinfo.proto",
    "grr/grr/proto/grr_response_proto/knowledge_base.proto",
    "grr/grr/proto/grr_response_proto/jobs.proto",
    "grr/grr/proto/grr_response_proto/timeline.proto",
    "grr/grr/proto/grr_response_proto/anomaly.proto",
    "grr/grr/proto/grr_response_proto/export.proto",
    "grr/grr/proto/grr_response_proto/objects.proto",
    "grr/grr/proto/grr_response_proto/output_plugin.proto",
    "grr/grr/proto/grr_response_proto/flows.proto",
    "grr/grr/proto/grr_response_proto/sysinfo.proto",
    "grr/grr/proto/grr_response_proto/user.proto",
];

const INCLUDES: &'static [&'static str] = &[
    "grr/grr/proto",
];

fn main() {
    // Because GRR proto files are not something that PROST! can accept (because
    // of the missing package declaration), we create a temporary directory with
    // recreated file structure. We apply the patch to each of them and feed the
    // Protocol Buffers compiler with updated files.
    let tempdir = tempfile::tempdir()
        .expect("failed to create temp dir");

    let mut protos = Vec::new();
    let mut includes = Vec::new();

    for path_before in PROTOS {
        let path_after = tempdir.path().join(path_before);

        patch_path(&path_before, &path_after)
            .expect(&format!("failed to patch file '{}'", path_before));

        protos.push(path_after);
    }

    for path_before in INCLUDES {
        let path_after = tempdir.path().join(path_before);
        includes.push(path_after);
    }

    prost_build::compile_protos(&protos, &includes)
        .expect("failed to compile proto files");

    // There is also a problem with one enum generated by PROST!: it's values
    // use name mangling, but it's default value does not. This is likely a bug
    // in PROST! itself, but for now we hack around it by replacing the spurious
    // line in the output file ourselves.
    let outdir: PathBuf = std::env::var("OUT_DIR")
        .expect("no output directory")
        .into();

    let target = outdir.join("grr.rs");

    let grr = std::fs::read_to_string(&target)
        .expect("failed to read the generated output file")
        .replace("TskFsAttrTypeDefault", "Default");

    std::fs::write(&target, grr)
        .expect("failed to write updated output file");
}

/// Patches given file at path `input`, writing patched content at `output`.
///
/// This function takes a path to malformed (i.e. lacking package definition)
/// GRR proto file and converts it to something that PROST! can understand. This
/// workaround has to be used as long as GRR does not fixes its proto files
/// upstream, which might be hard because of compatibility reasons.
fn patch_path<PI, PO>(input: PI, output: PO) -> Result<()>
where
    PI: AsRef<Path>,
    PO: AsRef<Path>,
{
    let mut input = file::open(&input)?;
    let mut output = file::create(&output)?;

    patch_buffer(&mut input, &mut output)
}

/// Patches given `input` buffer, writing patched content to `output`.
///
/// This function takes a buffer with malformed (lacking package definition)
/// GRR proto file and converts it to something that PROST! can understand. This
/// workaround has to be used as long as GRR does not fixes its proto files
/// upstream, which might be hard because of compatibility reasons.
fn patch_buffer<R, W>(input: &mut R, output: &mut W) -> Result<()>
where
    R: Read,
    W: Write,
{
    let mut buffer = String::new();
    input.read_to_string(&mut buffer)?;

    for line in buffer.lines() {
        writeln!(output, "{}", line)?;
    }

    Ok(())
}

mod file {
    use std::fs::File;
    use std::io::Result;
    use std::path::Path;

    /// Opens a file at the specified `path`.
    ///
    /// This function simply delegates to the standard library `File::open` and
    /// exists purely for aesthetic purposes.
    pub fn open<P: AsRef<Path>>(path: P) -> Result<File> {
        File::open(path)
    }

    /// Creates a file at the specified `path` and all the necessary folders
    /// along the way.
    ///
    /// Except for directory creation, this function should behave identically
    /// to the standard library's `File::create`.
    pub fn create<P: AsRef<Path>>(path: P) -> Result<File> {
        if let Some(parent) = path.as_ref().parent() {
            std::fs::create_dir_all(parent)?;
        }

        File::create(path)
    }
}
<commit_msg>Remove workaround for missing proto package names.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

use std::path::PathBuf;

const PROTOS: &'static [&'static str] = &[
    "grr/grr/proto/grr_response_proto/semantic.proto",
    "grr/grr/proto/grr_response_proto/sysinfo.proto",
    "grr/grr/proto/grr_response_proto/knowledge_base.proto",
    "grr/grr/proto/grr_response_proto/jobs.proto",
    "grr/grr/proto/grr_response_proto/timeline.proto",
    "grr/grr/proto/grr_response_proto/anomaly.proto",
    "grr/grr/proto/grr_response_proto/export.proto",
    "grr/grr/proto/grr_response_proto/objects.proto",
    "grr/grr/proto/grr_response_proto/output_plugin.proto",
    "grr/grr/proto/grr_response_proto/flows.proto",
    "grr/grr/proto/grr_response_proto/sysinfo.proto",
    "grr/grr/proto/grr_response_proto/user.proto",
];

const INCLUDES: &'static [&'static str] = &[
    "grr/grr/proto",
];

fn main() {
    prost_build::compile_protos(&PROTOS, &INCLUDES)
        .expect("failed to compile proto files");

    // There is a problem with one enum generated by PROST!: it's values use
    // name mangling, but it's default value does not. This is likely a bug in
    // PROST! itself, but for now we hack around it by replacing the spurious
    // line in the output file ourselves.
    let outdir: PathBuf = std::env::var("OUT_DIR")
        .expect("no output directory")
        .into();

    let target = outdir.join("grr.rs");

    let grr = std::fs::read_to_string(&target)
        .expect("failed to read the generated output file")
        .replace("TskFsAttrTypeDefault", "Default");

    std::fs::write(&target, grr)
        .expect("failed to write updated output file");
}
<|endoftext|>


<commit_before>use std::io;
use std::ops::Deref;
use std::collections::HashMap;
use rmp_serialize::decode::Error;
use rmp_serialize::{Encoder, Decoder};
use rustc_serialize::{Encodable, Decodable};
use sim::{Agent, Simulation};
use redis::{Commands, Client, Connection, PipelineCommands, pipe};
use uuid::Uuid;

fn decode<R: Decodable>(inp: Vec<u8>) -> Result<R, Error> {
    let mut decoder = Decoder::new(&inp[..]);
    Decodable::decode(&mut decoder)
}

fn encode<R: Encodable>(data: R) -> Result<Vec<u8>, io::Error> {
    let mut buf = Vec::<u8>::new();
    match data.encode(&mut Encoder::new(&mut buf)) {
        Ok(_) => Ok(buf),
        Err(e) => Err(io::Error::new(io::ErrorKind::Other, format!("{}", e))),
    }
}

/// An interface to the Redis-backed agent population.
pub struct Population<S: Simulation, C: Commands> {
    conn: C,
    simulation: S,
}

impl<S: Simulation, C: Commands> Population<S, C> {
    pub fn new(simulation: S, conn: C) -> Population<S, C> {
        Population {
            conn: conn,
            simulation: simulation,
        }
    }

    pub fn count(&self) -> usize {
        self.conn.scard::<&str, usize>("population").unwrap()
    }

    pub fn world(&self) -> S::World {
        let data = self.conn.get("world").unwrap();
        decode(data).unwrap()
    }

    pub fn set_world(&self, world: S::World) {
        let data = encode(&world).unwrap();
        let _: () = self.conn.set("world", data).unwrap();
    }

    /// Create a new agent with the specified state, returning the new agent's id.
    pub fn spawn(&self, state: S::State) -> Uuid {
        let id = Uuid::new_v4();
        self.set_agent(id, state.clone());
        let _: () = self.conn.sadd("population", id.to_string()).unwrap();
        self.simulation.setup(Agent {
                                  id: id,
                                  state: state.clone(),
                              },
                              &self);
        id
    }

    /// Get an agent by id.
    pub fn get_agent(&self, id: Uuid) -> Option<Agent<S::State>> {
        let data = self.conn.get(id.to_string()).unwrap();
        Some(Agent {
            id: id,
            state: decode(data).unwrap(),
        })
    }

    /// Set an agent state by id.
    pub fn set_agent(&self, id: Uuid, state: S::State) {
        let data = encode(&state).unwrap();
        let _: () = self.conn.set(id.to_string(), data).unwrap();
    }

    /// Deletes an agent by id.
    pub fn kill(&self, id: Uuid) {
        let _: () = self.conn.del(id.to_string()).unwrap();
        let _: () = self.conn.srem("population", id.to_string()).unwrap();
    }

    /// Lookup agents at a particular index.
    pub fn lookup(&self, index: &str) -> Vec<Agent<S::State>> {
        let ids: Vec<String> = self.conn.smembers(format!("idx:{}", index)).unwrap();
        match ids.len() {
            0 => Vec::new(),
            1 => {
                let id = ids[0].as_ref();
                let state_data: Vec<u8> = self.conn.get(id).unwrap();
                let agent = Agent {
                    id: Uuid::parse_str(id).unwrap(),
                    state: decode(state_data).unwrap(),
                };
                vec![agent]
            }
            _ => {
                let states_data: Vec<Vec<u8>> = self.conn.get(ids.clone()).unwrap();
                ids.iter()
                    .zip(states_data)
                    .map(|(id, state_data)| {
                        Agent {
                            id: Uuid::parse_str(id).unwrap(),
                            state: decode(state_data).unwrap(),
                        }
                    })
                    .collect()
            }
        }
    }

    /// Add an agent (id) to an index.
    pub fn index(&self, index: &str, id: Uuid) {
        let _: () = self.conn.sadd(format!("idx:{}", index), id.to_string()).unwrap();
    }

    /// Reset indices.
    pub fn reset_indices(&self) {
        let keys: Vec<String> = self.conn.keys("idx:*").unwrap();
        if keys.len() > 0 {
            let _: () = self.conn.del(keys).unwrap();
        }
    }

    /// Reset the population.
    pub fn reset(&self) {
        // reset sets
        let _: () = self.conn.del("population").unwrap();
        self.reset_indices();
    }
}

pub struct Manager<S: Simulation, C: Commands> {
    conn: Connection,
    reporters: HashMap<usize, Box<Fn(&Population<S, C>, &Connection) -> () + Send>>,
    pub population: Population<S, C>,
}

impl<S: Simulation, C: Commands> Manager<S, C> {
    pub fn new(addr: &str, conn: C, simulation: S, world: S::World) -> Manager<S, C> {
        let population = Population::new(simulation, conn);
        population.reset();
        population.set_world(world);

        // conn for commanding workers
        let client = Client::open(addr).unwrap();
        let conn = client.get_connection().unwrap();

        let m = Manager {
            population: population,
            reporters: HashMap::new(),
            conn: conn,
        };
        m.reset();
        m
    }

    /// Reset the manager. This unregisters all workers and queues.
    pub fn reset(&self) {
        // start with "idle" phase
        let _: () = self.conn.set("current_phase", "idle").unwrap();

        // reset sets
        let _: () = self.conn.del("workers").unwrap();
        let _: () = self.conn.del("finished").unwrap();
        let _: () = self.conn.del("to_decide").unwrap();
        let _: () = self.conn.del("to_update").unwrap();
    }

    /// Run the simulation for `n_steps`.
    pub fn run(&self, n_steps: usize) -> () {
        let mut steps = 0;

        if self.n_workers() == 0 {
            println!("Waiting for at least one worker...");
            while self.n_workers() == 0 {
            }
            println!("Ok");
        }

        // copy population to the "to_decide" set
        let _: () = self.conn.sunionstore("to_decide", "population").unwrap();

        while steps < n_steps {
            // run any register reporters, if appropriate
            for (interval, reporter) in self.reporters.iter() {
                if steps % interval == 0 {
                    reporter(&self.population, &self.conn);
                }
            }

            let _: () = self.conn.publish("command", "decide").unwrap();
            let _: () = self.conn.set("current_phase", "decide").unwrap();
            self.wait_until_finished();
            let _: () = self.conn.publish("command", "update").unwrap();
            let _: () = self.conn.set("current_phase", "update").unwrap();
            self.wait_until_finished();
            steps += 1;
        }

        let _: () = self.conn.publish("command", "terminate").unwrap();
    }

    /// Register a reporter function to be called every `n_steps`.
    /// It receives a `Population` which can be used to query agents,
    /// compute aggregate statistics, etc, and a Redis connection
    /// that can be used, for example, to send reports via pubsub.
    pub fn register_reporter<F>(&mut self, n_steps: usize, func: F) -> ()
        where F: Fn(&Population<S, C>, &Connection) -> () + Send + 'static
    {
        self.reporters.insert(n_steps, Box::new(func));
    }

    fn wait_until_finished(&self) {
        while self.conn.scard::<&str, usize>("finished").unwrap() != self.n_workers() {
        }
    }

    /// Get the number of workers.
    pub fn n_workers(&self) -> usize {
        self.conn.scard::<&str, usize>("workers").unwrap()
    }
}

pub struct Worker<S: Simulation, C: Commands> {
    id: Uuid,
    addr: String,
    population: Population<S, C>,
    simulation: S,
}

impl<S: Simulation, C: Commands> Worker<S, C> {
    pub fn new(addr: &str, conn: C, simulation: S) -> Worker<S, C> {
        Worker {
            id: Uuid::new_v4(),
            addr: addr.to_owned(),
            population: Population::new(simulation.clone(), conn),
            simulation: simulation,
        }
    }

    pub fn start(&self) {
        let client = Client::open(self.addr.deref()).unwrap();
        let manager = client.get_connection().unwrap();

        // register with the manager
        let _: () = manager.sadd("workers", self.id.to_string()).unwrap();

        // subscribe to the command channel
        let mut pubsub = client.get_pubsub().unwrap();
        pubsub.subscribe("command").unwrap();

        // check what the current phase is
        let phase: String = manager.get("current_phase").unwrap();
        self.process_cmd(phase.as_ref(), &manager);

        loop {
            let msg = pubsub.get_message().unwrap();
            let payload: String = msg.get_payload().unwrap();
            self.process_cmd(payload.as_ref(), &manager);
            if payload == "terminate" {
                break;
            }
        }
    }

    fn process_cmd(&self, cmd: &str, conn: &Connection) {
        match cmd {
            "terminate" => {
                let _: () = conn.srem("workers", self.id.to_string()).unwrap();
            }
            "decide" => {
                self.decide(&self.simulation, &conn);
                let _: () = conn.sadd("finished", self.id.to_string()).unwrap();
            }
            "update" => {
                self.update(&self.simulation, &conn);
                let _: () = conn.sadd("finished", self.id.to_string()).unwrap();
            }
            "idle" => (),
            s => println!("Unrecognized command: {}", s),
        }
    }

    fn decide(&self, simulation: &S, conn: &Connection) {
        let world = self.population.world();
        while let Ok(id) = conn.spop::<&str, String>("to_decide") {
            let id = Uuid::parse_str(&id).unwrap();
            match self.population.get_agent(id) {
                Some(agent) => {
                    let updates = simulation.decide(agent, world.clone(), &self.population);
                    let mut rpipe = pipe();
                    for (id, update) in updates {
                        let data = encode(&update).unwrap();
                        rpipe.lpush(format!("updates:{}", id), data).ignore();
                    }
                    rpipe.sadd("to_update", id.to_string()).ignore();
                    let _: () = rpipe.query(conn).unwrap();
                }
                None => (),
            }
        }
    }

    fn update(&self, simulation: &S, conn: &Connection) {
        while let Ok(id) = conn.spop::<&str, String>("to_update") {
            let updates: Vec<S::Update> = {
                let key = format!("updates:{}", id);
                let updates_data: Vec<Vec<u8>> = conn.lrange(&key, 0, -1)
                    .unwrap();
                let _: () = conn.del(&key).unwrap();
                if updates_data.len() == 0 {
                    Vec::new()
                } else {
                    updates_data.iter().map(|data| decode(data.clone()).unwrap()).collect()
                }
            };
            let id = Uuid::parse_str(&id).unwrap();
            match self.population.get_agent(id) {
                Some(agent) => {
                    let new_state = simulation.update(agent.state.clone(), updates);
                    self.population.set_agent(id, new_state);
                    let _: () = conn.sadd("to_decide", id.to_string()).unwrap();
                }
                None => (),
            }
        }
    }
}
<commit_msg>worker tweak<commit_after>use std::io;
use std::collections::HashMap;
use rmp_serialize::decode::Error;
use rmp_serialize::{Encoder, Decoder};
use rustc_serialize::{Encodable, Decodable};
use sim::{Agent, Simulation};
use redis::{Commands, Client, Connection, PipelineCommands, pipe};
use uuid::Uuid;

fn decode<R: Decodable>(inp: Vec<u8>) -> Result<R, Error> {
    let mut decoder = Decoder::new(&inp[..]);
    Decodable::decode(&mut decoder)
}

fn encode<R: Encodable>(data: R) -> Result<Vec<u8>, io::Error> {
    let mut buf = Vec::<u8>::new();
    match data.encode(&mut Encoder::new(&mut buf)) {
        Ok(_) => Ok(buf),
        Err(e) => Err(io::Error::new(io::ErrorKind::Other, format!("{}", e))),
    }
}

/// An interface to the Redis-backed agent population.
pub struct Population<S: Simulation, C: Commands> {
    conn: C,
    simulation: S,
}

impl<S: Simulation, C: Commands> Population<S, C> {
    pub fn new(simulation: S, conn: C) -> Population<S, C> {
        Population {
            conn: conn,
            simulation: simulation,
        }
    }

    pub fn count(&self) -> usize {
        self.conn.scard::<&str, usize>("population").unwrap()
    }

    pub fn world(&self) -> S::World {
        let data = self.conn.get("world").unwrap();
        decode(data).unwrap()
    }

    pub fn set_world(&self, world: S::World) {
        let data = encode(&world).unwrap();
        let _: () = self.conn.set("world", data).unwrap();
    }

    /// Create a new agent with the specified state, returning the new agent's id.
    pub fn spawn(&self, state: S::State) -> Uuid {
        let id = Uuid::new_v4();
        self.set_agent(id, state.clone());
        let _: () = self.conn.sadd("population", id.to_string()).unwrap();
        self.simulation.setup(Agent {
                                  id: id,
                                  state: state.clone(),
                              },
                              &self);
        id
    }

    /// Get an agent by id.
    pub fn get_agent(&self, id: Uuid) -> Option<Agent<S::State>> {
        let data = self.conn.get(id.to_string()).unwrap();
        Some(Agent {
            id: id,
            state: decode(data).unwrap(),
        })
    }

    /// Set an agent state by id.
    pub fn set_agent(&self, id: Uuid, state: S::State) {
        let data = encode(&state).unwrap();
        let _: () = self.conn.set(id.to_string(), data).unwrap();
    }

    /// Deletes an agent by id.
    pub fn kill(&self, id: Uuid) {
        let _: () = self.conn.del(id.to_string()).unwrap();
        let _: () = self.conn.srem("population", id.to_string()).unwrap();
    }

    /// Lookup agents at a particular index.
    pub fn lookup(&self, index: &str) -> Vec<Agent<S::State>> {
        let ids: Vec<String> = self.conn.smembers(format!("idx:{}", index)).unwrap();
        match ids.len() {
            0 => Vec::new(),
            1 => {
                let id = ids[0].as_ref();
                let state_data: Vec<u8> = self.conn.get(id).unwrap();
                let agent = Agent {
                    id: Uuid::parse_str(id).unwrap(),
                    state: decode(state_data).unwrap(),
                };
                vec![agent]
            }
            _ => {
                let states_data: Vec<Vec<u8>> = self.conn.get(ids.clone()).unwrap();
                ids.iter()
                    .zip(states_data)
                    .map(|(id, state_data)| {
                        Agent {
                            id: Uuid::parse_str(id).unwrap(),
                            state: decode(state_data).unwrap(),
                        }
                    })
                    .collect()
            }
        }
    }

    /// Add an agent (id) to an index.
    pub fn index(&self, index: &str, id: Uuid) {
        let _: () = self.conn.sadd(format!("idx:{}", index), id.to_string()).unwrap();
    }

    /// Reset indices.
    pub fn reset_indices(&self) {
        let keys: Vec<String> = self.conn.keys("idx:*").unwrap();
        if keys.len() > 0 {
            let _: () = self.conn.del(keys).unwrap();
        }
    }

    /// Reset the population.
    pub fn reset(&self) {
        // reset sets
        let _: () = self.conn.del("population").unwrap();
        self.reset_indices();
    }
}

pub struct Manager<S: Simulation, C: Commands> {
    conn: Connection,
    reporters: HashMap<usize, Box<Fn(&Population<S, C>, &Connection) -> () + Send>>,
    pub population: Population<S, C>,
}

impl<S: Simulation, C: Commands> Manager<S, C> {
    pub fn new(addr: &str, conn: C, simulation: S, world: S::World) -> Manager<S, C> {
        let population = Population::new(simulation, conn);
        population.reset();
        population.set_world(world);

        // conn for commanding workers
        let client = Client::open(addr).unwrap();
        let conn = client.get_connection().unwrap();

        let m = Manager {
            population: population,
            reporters: HashMap::new(),
            conn: conn,
        };
        m.reset();
        m
    }

    /// Reset the manager. This unregisters all workers and queues.
    pub fn reset(&self) {
        // start with "idle" phase
        let _: () = self.conn.set("current_phase", "idle").unwrap();

        // reset sets
        let _: () = self.conn.del("workers").unwrap();
        let _: () = self.conn.del("finished").unwrap();
        let _: () = self.conn.del("to_decide").unwrap();
        let _: () = self.conn.del("to_update").unwrap();
    }

    /// Run the simulation for `n_steps`.
    pub fn run(&self, n_steps: usize) -> () {
        let mut steps = 0;

        if self.n_workers() == 0 {
            println!("Waiting for at least one worker...");
            while self.n_workers() == 0 {
            }
            println!("Ok");
        }

        // copy population to the "to_decide" set
        let _: () = self.conn.sunionstore("to_decide", "population").unwrap();

        while steps < n_steps {
            // run any register reporters, if appropriate
            for (interval, reporter) in self.reporters.iter() {
                if steps % interval == 0 {
                    reporter(&self.population, &self.conn);
                }
            }

            let _: () = self.conn.publish("command", "decide").unwrap();
            let _: () = self.conn.set("current_phase", "decide").unwrap();
            self.wait_until_finished();
            let _: () = self.conn.publish("command", "update").unwrap();
            let _: () = self.conn.set("current_phase", "update").unwrap();
            self.wait_until_finished();
            steps += 1;
        }

        let _: () = self.conn.publish("command", "terminate").unwrap();
    }

    /// Register a reporter function to be called every `n_steps`.
    /// It receives a `Population` which can be used to query agents,
    /// compute aggregate statistics, etc, and a Redis connection
    /// that can be used, for example, to send reports via pubsub.
    pub fn register_reporter<F>(&mut self, n_steps: usize, func: F) -> ()
        where F: Fn(&Population<S, C>, &Connection) -> () + Send + 'static
    {
        self.reporters.insert(n_steps, Box::new(func));
    }

    fn wait_until_finished(&self) {
        while self.conn.scard::<&str, usize>("finished").unwrap() != self.n_workers() {
        }
    }

    /// Get the number of workers.
    pub fn n_workers(&self) -> usize {
        self.conn.scard::<&str, usize>("workers").unwrap()
    }
}

pub struct Worker<S: Simulation, C: Commands> {
    id: Uuid,
    manager: Client,
    population: Population<S, C>,
    simulation: S,
}

impl<S: Simulation, C: Commands> Worker<S, C> {
    pub fn new(addr: &str, conn: C, simulation: S) -> Worker<S, C> {
        Worker {
            id: Uuid::new_v4(),
            manager: Client::open(addr).unwrap(),
            population: Population::new(simulation.clone(), conn),
            simulation: simulation,
        }
    }

    pub fn start(&self) {
        // register with the manager
        let _: () = self.manager.sadd("workers", self.id.to_string()).unwrap();

        // subscribe to the command channel
        let mut pubsub = self.manager.get_pubsub().unwrap();
        pubsub.subscribe("command").unwrap();

        // check what the current phase is
        let phase: String = self.manager.get("current_phase").unwrap();
        self.process_cmd(phase.as_ref(), &self.manager);

        loop {
            let msg = pubsub.get_message().unwrap();
            let payload: String = msg.get_payload().unwrap();
            self.process_cmd(payload.as_ref(), &self.manager);
            if payload == "terminate" {
                break;
            }
        }
    }

    fn process_cmd(&self, cmd: &str, conn: &Client) {
        match cmd {
            "terminate" => {
                let _: () = conn.srem("workers", self.id.to_string()).unwrap();
            }
            "decide" => {
                self.decide(&self.simulation, &conn);
                let _: () = conn.sadd("finished", self.id.to_string()).unwrap();
            }
            "update" => {
                self.update(&self.simulation, &conn);
                let _: () = conn.sadd("finished", self.id.to_string()).unwrap();
            }
            "idle" => (),
            s => println!("Unrecognized command: {}", s),
        }
    }

    fn decide(&self, simulation: &S, conn: &Client) {
        let world = self.population.world();
        while let Ok(id) = conn.spop::<&str, String>("to_decide") {
            let id = Uuid::parse_str(&id).unwrap();
            match self.population.get_agent(id) {
                Some(agent) => {
                    let updates = simulation.decide(agent, world.clone(), &self.population);
                    let mut rpipe = pipe();
                    for (id, update) in updates {
                        let data = encode(&update).unwrap();
                        rpipe.lpush(format!("updates:{}", id), data).ignore();
                    }
                    rpipe.sadd("to_update", id.to_string()).ignore();
                    let _: () = rpipe.query(conn).unwrap();
                }
                None => (),
            }
        }
    }

    fn update(&self, simulation: &S, conn: &Client) {
        while let Ok(id) = conn.spop::<&str, String>("to_update") {
            let updates: Vec<S::Update> = {
                let key = format!("updates:{}", id);
                let updates_data: Vec<Vec<u8>> = conn.lrange(&key, 0, -1)
                    .unwrap();
                let _: () = conn.del(&key).unwrap();
                if updates_data.len() == 0 {
                    Vec::new()
                } else {
                    updates_data.iter().map(|data| decode(data.clone()).unwrap()).collect()
                }
            };
            let id = Uuid::parse_str(&id).unwrap();
            match self.population.get_agent(id) {
                Some(agent) => {
                    let new_state = simulation.update(agent.state.clone(), updates);
                    self.population.set_agent(id, new_state);
                    let _: () = conn.sadd("to_decide", id.to_string()).unwrap();
                }
                None => (),
            }
        }
    }
}
<|endoftext|>


<commit_before>extern crate "rust-hl-lua" as lua;
extern crate test;
use lua::{Lua, LuaTable};

use std::collections::{HashMap, HashSet};
use test::Bencher;

#[test]
fn write() {
    let mut lua = Lua::new();

    lua.set("a", vec![9, 8, 7]);

    let mut table: LuaTable<Lua> = lua.load("a").unwrap();

    let values: Vec<(i32, i32)> = table.iter().filter_map(|e| e).collect();
    assert_eq!(values, vec!( (1, 9), (2, 8), (3, 7) ));
}

#[test]
fn write_map() {
    let mut lua = Lua::new();

    let mut map = HashMap::new();
    map.insert(5, 8);
    map.insert(13, 21);
    map.insert(34, 55);

    lua.set("a", map.clone());

    let mut table: LuaTable<Lua> = lua.load("a").unwrap();

    let values: HashMap<i32, i32> = table.iter().filter_map(|e| e).collect();
    assert_eq!(values, map);
}

#[test]
fn write_set() {
    let mut lua = Lua::new();

    let mut set = HashSet::new();
    set.insert(5);
    set.insert(8);
    set.insert(13);
    set.insert(21);
    set.insert(34);
    set.insert(55);

    lua.set("a", set.clone());

    let mut table: LuaTable<Lua> = lua.load("a").unwrap();

    let values: HashSet<i32> = table.iter().filter_map(|e| e)
                                           .map(|(elem, set): (i32, bool)| {
        assert!(set);
        elem
    }).collect();

    assert_eq!(values, set);
}

#[bench]
fn new_map(b: &mut Bencher) {
    let mut lua = Lua::new();

    let mut map = HashMap::new();
    map.insert(5, 8);
    map.insert(13, 21);
    map.insert(34, 55);

    b.iter(|| {
        lua.set("a", map.clone());
    })
}

#[bench]
fn new_large_map(b: &mut Bencher) {
    let mut lua = Lua::new();

    let mut map = HashMap::new();
    for i in range(0, 500) {
        map.insert(i, i + 1);
    }

    b.iter(|| {
        lua.set("a", map.clone());
    })
}
<commit_msg>Fix rust_tables tests<commit_after>extern crate "rust-hl-lua" as lua;
extern crate test;
use lua::{Lua, LuaTable};

use std::collections::{HashMap, HashSet};
use test::Bencher;

#[test]
fn write() {
    let mut lua = Lua::new();

    lua.set("a", vec![9, 8, 7]);

    let mut table: LuaTable<_> = lua.get("a").unwrap();

    let values: Vec<(i32, i32)> = table.iter().filter_map(|e| e).collect();
    assert_eq!(values, vec!( (1, 9), (2, 8), (3, 7) ));
}

#[test]
fn write_map() {
    let mut lua = Lua::new();

    let mut map = HashMap::new();
    map.insert(5, 8);
    map.insert(13, 21);
    map.insert(34, 55);

    lua.set("a", map.clone());

    let mut table: LuaTable<_> = lua.get("a").unwrap();

    let values: HashMap<i32, i32> = table.iter().filter_map(|e| e).collect();
    assert_eq!(values, map);
}

#[test]
fn write_set() {
    let mut lua = Lua::new();

    let mut set = HashSet::new();
    set.insert(5);
    set.insert(8);
    set.insert(13);
    set.insert(21);
    set.insert(34);
    set.insert(55);

    lua.set("a", set.clone());

    let mut table: LuaTable<_> = lua.get("a").unwrap();

    let values: HashSet<i32> = table.iter().filter_map(|e| e)
                                           .map(|(elem, set): (i32, bool)| {
        assert!(set);
        elem
    }).collect();

    assert_eq!(values, set);
}

#[bench]
fn new_map(b: &mut Bencher) {
    let mut lua = Lua::new();

    let mut map = HashMap::new();
    map.insert(5, 8);
    map.insert(13, 21);
    map.insert(34, 55);

    b.iter(|| {
        lua.set("a", map.clone());
    })
}

#[bench]
fn new_large_map(b: &mut Bencher) {
    let mut lua = Lua::new();

    let mut map = HashMap::new();
    for i in range(0, 500) {
        map.insert(i, i + 1);
    }

    b.iter(|| {
        lua.set("a", map.clone());
    })
}
<|endoftext|>


<commit_before>use learning::toolkit::kernel::{Kernel, SquaredExp};
use learning::SupModel;
use linalg::matrix::Matrix;

pub trait MeanFunc {
    fn func(&self, x: Matrix<f64>) -> Matrix<f64>;
}

/// Constant mean function
pub struct ConstMean {
    a: f64,
}

impl Default for ConstMean {
    fn default() -> ConstMean {
        ConstMean { a: 0f64 }
    }
}

impl MeanFunc for ConstMean {
    fn func(&self, x: Matrix<f64>) -> Matrix<f64> {
        Matrix::ones(x.rows(), x.cols()) * self.a
    }
}

/// Gaussian Process struct
///
/// Gaussian process with generic kernel and deterministic mean function.
/// Can be used for gaussian process regression with noise.
/// Currently does not support classification.
pub struct GaussianProcess<T: Kernel, U: MeanFunc> {
    ker: T,
    mean: U,
    noise: f64,
    train_data: Option<Matrix<f64>>,
    train_output: Option<Matrix<f64>>,
    train_mat: Option<Matrix<f64>>,
}

impl Default for GaussianProcess<SquaredExp, ConstMean> {
    /// Construct a default Gaussian Process
    ///
    /// # Examples
    ///
    /// ```
    /// use rusty_machine::learning::gp;
    ///
    /// let gp = gp::GaussianProcess::default();
    /// ```
    fn default() -> GaussianProcess<SquaredExp, ConstMean> {
        GaussianProcess {
            ker: SquaredExp::default(),
            mean: ConstMean::default(),
            noise: 1f64,
            train_data: None,
            train_output: None,
            train_mat: None,
        }
    }
}

impl<T: Kernel, U: MeanFunc> GaussianProcess<T, U> {
    /// Construct a new Gaussian Process.
    ///
    /// # Examples
    ///
    /// ```
    /// use rusty_machine::learning::gp;
    /// use rusty_machine::learning::toolkit::kernel;
    ///
    /// let ker = kernel::SquaredExp::default();
    /// let mean = gp::ConstMean::default();
    /// let gaussp = gp::GaussianProcess::new(ker, mean, 1e-3f64);
    /// ```
    pub fn new(ker: T, mean: U, noise: f64) -> GaussianProcess<T, U> {
        GaussianProcess {
            ker: ker,
            mean: mean,
            noise: noise,
            train_data: None,
            train_output: None,
            train_mat: None,
        }
    }

    /// Construct a kernel matrix
    fn ker_mat(&self, m1: &Matrix<f64>, m2: &Matrix<f64>) -> Matrix<f64> {
        assert_eq!(m1.cols(), m2.cols());
        let dim1 = m1.rows();
        let dim2 = m2.rows();

        let mut ker_data = Vec::with_capacity(dim1 * dim2);

        for i in 0..dim1 {
            for j in 0..dim2 {
                ker_data.push(self.ker.kernel(&m1.data[i * m1.cols()..(i + 1) * m1.cols()],
                                              &m2.data[j * m2.cols()..(j + 1) * m2.cols()]));
            }
        }

        Matrix::new(dim1, dim2, ker_data)
    }
}

impl<T: Kernel, U: MeanFunc> SupModel<Matrix<f64>, Matrix<f64>> for GaussianProcess<T, U> {
    /// Predict output from data.
    fn predict(&self, data: &Matrix<f64>) -> Matrix<f64> {
        let mean = self.mean.func(data.clone());

        // Messy referencing for succint syntax
        if let (&Some(ref t_data), &Some(ref t_mat), &Some(ref t_out)) = (&self.train_data,
                                                                          &self.train_mat,
                                                                          &self.train_output) {
            let test_mat = self.ker_mat(data, t_data) * t_mat *
                           (t_out - self.mean.func(t_data.clone()));
            return mean + test_mat;
        }

        panic!("The model has not been trained.");
    }

    /// Train the model using data and outputs.
    fn train(&mut self, data: &Matrix<f64>, value: &Matrix<f64>) {
        let noise_mat = Matrix::identity(data.rows()) * self.noise;

        let ker_mat = self.ker_mat(data, data);

        let train_mat = (ker_mat + noise_mat).inverse();

        self.train_mat = Some(train_mat);
        self.train_data = Some(data.clone());
        self.train_output = Some(value.clone());
    }
}

impl<T: Kernel, U: MeanFunc> GaussianProcess<T, U> {
    /// Compute the posterior distribution
    ///
    /// Requires the model to be trained first.
    /// _Note that this is a messy compromise as GPs do not
    /// fit the SupModel trait as is currently implemented._
    pub fn get_posterior(&self, data: &Matrix<f64>) -> (Matrix<f64>, Matrix<f64>) {
        let mean = self.mean.func(data.clone());

        // Messy referencing for succint syntax
        if let (&Some(ref t_data), &Some(ref t_mat), &Some(ref t_out)) = (&self.train_data,
                                                                          &self.train_mat,
                                                                          &self.train_output) {
            let test_mat = self.ker_mat(data, t_data) * t_mat;
            let post_mean = mean + &test_mat * (t_out - self.mean.func(t_data.clone()));

            let post_cov = self.ker_mat(data, data) - test_mat * self.ker_mat(t_data, data);

            return (post_mean, post_cov);
        }

        panic!("The model has not been trained.");
    }
}
<commit_msg>Minor changes to access<commit_after>use learning::toolkit::kernel::{Kernel, SquaredExp};
use learning::SupModel;
use linalg::matrix::Matrix;

pub trait MeanFunc {
    fn func(&self, x: Matrix<f64>) -> Matrix<f64>;
}

/// Constant mean function
pub struct ConstMean {
    a: f64,
}

impl Default for ConstMean {
    fn default() -> ConstMean {
        ConstMean { a: 0f64 }
    }
}

impl MeanFunc for ConstMean {
    fn func(&self, x: Matrix<f64>) -> Matrix<f64> {
        Matrix::zeros(x.rows(), x.cols()) + self.a
    }
}

/// Gaussian Process struct
///
/// Gaussian process with generic kernel and deterministic mean function.
/// Can be used for gaussian process regression with noise.
/// Currently does not support classification.
pub struct GaussianProcess<T: Kernel, U: MeanFunc> {
    ker: T,
    mean: U,
    pub noise: f64,
    train_data: Option<Matrix<f64>>,
    train_output: Option<Matrix<f64>>,
    train_mat: Option<Matrix<f64>>,
}

impl Default for GaussianProcess<SquaredExp, ConstMean> {
    /// Construct a default Gaussian Process
    ///
    /// # Examples
    ///
    /// ```
    /// use rusty_machine::learning::gp;
    ///
    /// let gp = gp::GaussianProcess::default();
    /// ```
    fn default() -> GaussianProcess<SquaredExp, ConstMean> {
        GaussianProcess {
            ker: SquaredExp::default(),
            mean: ConstMean::default(),
            noise: 1f64,
            train_data: None,
            train_output: None,
            train_mat: None,
        }
    }
}

impl<T: Kernel, U: MeanFunc> GaussianProcess<T, U> {
    /// Construct a new Gaussian Process.
    ///
    /// # Examples
    ///
    /// ```
    /// use rusty_machine::learning::gp;
    /// use rusty_machine::learning::toolkit::kernel;
    ///
    /// let ker = kernel::SquaredExp::default();
    /// let mean = gp::ConstMean::default();
    /// let gaussp = gp::GaussianProcess::new(ker, mean, 1e-3f64);
    /// ```
    pub fn new(ker: T, mean: U, noise: f64) -> GaussianProcess<T, U> {
        GaussianProcess {
            ker: ker,
            mean: mean,
            noise: noise,
            train_data: None,
            train_output: None,
            train_mat: None,
        }
    }

    /// Construct a kernel matrix
    fn ker_mat(&self, m1: &Matrix<f64>, m2: &Matrix<f64>) -> Matrix<f64> {
        assert_eq!(m1.cols(), m2.cols());
        let dim1 = m1.rows();
        let dim2 = m2.rows();

        let mut ker_data = Vec::with_capacity(dim1 * dim2);

        for i in 0..dim1 {
            for j in 0..dim2 {
                ker_data.push(self.ker.kernel(&m1.data[i * m1.cols()..(i + 1) * m1.cols()],
                                              &m2.data[j * m2.cols()..(j + 1) * m2.cols()]));
            }
        }

        Matrix::new(dim1, dim2, ker_data)
    }
}

impl<T: Kernel, U: MeanFunc> SupModel<Matrix<f64>, Matrix<f64>> for GaussianProcess<T, U> {
    /// Predict output from data.
    fn predict(&self, data: &Matrix<f64>) -> Matrix<f64> {
        let mean = self.mean.func(data.clone());

        // Messy referencing for succint syntax
        if let (&Some(ref t_data), &Some(ref t_mat), &Some(ref t_out)) = (&self.train_data,
                                                                          &self.train_mat,
                                                                          &self.train_output) {
            let test_mat = self.ker_mat(data, t_data) * t_mat *
                           (t_out - self.mean.func(t_data.clone()));
            return mean + test_mat;
        }

        panic!("The model has not been trained.");
    }

    /// Train the model using data and outputs.
    fn train(&mut self, data: &Matrix<f64>, value: &Matrix<f64>) {
        let noise_mat = Matrix::identity(data.rows()) * self.noise;

        let ker_mat = self.ker_mat(data, data);

        let train_mat = (ker_mat + noise_mat).inverse();

        self.train_mat = Some(train_mat);
        self.train_data = Some(data.clone());
        self.train_output = Some(value.clone());
    }
}

impl<T: Kernel, U: MeanFunc> GaussianProcess<T, U> {
    /// Compute the posterior distribution
    ///
    /// Requires the model to be trained first.
    /// _Note that this is a messy compromise as GPs do not
    /// fit the SupModel trait as is currently implemented._
    pub fn get_posterior(&self, data: &Matrix<f64>) -> (Matrix<f64>, Matrix<f64>) {
        let mean = self.mean.func(data.clone());

        // Messy referencing for succint syntax
        if let (&Some(ref t_data), &Some(ref t_mat), &Some(ref t_out)) = (&self.train_data,
                                                                          &self.train_mat,
                                                                          &self.train_output) {
            let test_mat = self.ker_mat(data, t_data) * t_mat;
            let post_mean = mean + &test_mat * (t_out - self.mean.func(t_data.clone()));

            let post_cov = self.ker_mat(data, data) - test_mat * self.ker_mat(t_data, data);

            return (post_mean, post_cov);
        }

        panic!("The model has not been trained.");
    }
}<|endoftext|>


<commit_before>#![feature(core)]

use std::env;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::ops::Deref;
use std::path::Path;
use std::str::FromStr;
use std::sync::{Arc, Mutex};
use std::thread;

fn load_matrix(file_path: &Path) -> Vec<Vec<u32>> {

    let mat_file = File::open(file_path).unwrap();

    let mat_reader = BufReader::new(mat_file);

    let mut mat_line_iterator = mat_reader.lines();

    let mut mat_content: Vec<Vec<u32>> = Vec::new();

    loop {
        match mat_line_iterator.next() {
            Some(result) => {
                                let mat_line = result.unwrap();
                                let slice_iter = mat_line.as_slice().trim_matches(',').split(',')
                                    .map(|x| {
                                                FromStr::from_str(x).unwrap()
                                            });
                                mat_content.push(slice_iter.collect());
                            },

            None => break,
        }
    }

    mat_content
}

fn single_thread_mul(m1: &Vec<Vec<u32>>, m2: &Vec<Vec<u32>>) -> Vec<Vec<u32>> {

    let mut result: Vec<Vec<u32>> = Vec::new();
    let mul_length = m2.len();

    for y in 0..m1.len() {

        let mut result_row: Vec<u32> = Vec::new();
        let mut entry_value: u32 = 0;

        for x in 0..m2[0].len() {

            for i in 0..mul_length {
                entry_value += m1[y][i] * m2[i][x];
            }

            result_row.push(entry_value);
        }

        result.push(result_row);
    }

    result
}

fn multi_thread_mul(m1: Vec<Vec<u32>>, m2: Vec<Vec<u32>>, thread_num: usize) -> Vec<Vec<u32>> {

    let rows_to_process = m1.len();
    let rows_per_thread = rows_to_process / thread_num;

    let shared_m1 = Arc::new(m1);
    let shared_m2 = Arc::new(m2);

    // is there a better way to initialize a vector?
    let mut result: Vec<Vec<u32>> = Vec::with_capacity(rows_to_process);
    result.resize(rows_to_process, Vec::new());

    let mul_length = shared_m2.len();

    let shared_result = Arc::new(Mutex::new(result));

    let mut threads = vec![];
    for thread_idx in 0..thread_num {

        let child_m1 = shared_m1.clone();
        let child_m2 = shared_m2.clone();
        let child_result = shared_result.clone();

        let join_handle = thread::spawn(move || {

            let start_row_idx = thread_idx * rows_per_thread;
            let end_row_idx =
                if start_row_idx + rows_per_thread < rows_to_process - 1 {
                    start_row_idx + rows_per_thread
                }
                else {
                    rows_to_process
                };

            println!("thread #{}: start={}, end={}", thread_idx, start_row_idx, end_row_idx);

            for y in start_row_idx..end_row_idx {

                let mut result_row: Vec<u32> = Vec::new();
                let mut entry_value: u32 = 0;

                for x in 0..child_m2[0].len() {

                    for i in 0..mul_length {
                        entry_value += child_m1[y][i] * child_m2[i][x];
                    }

                    result_row.push(entry_value);
                }

                let mut child_result = child_result.lock().unwrap();
                child_result[y] = result_row;
            }

            println!("thread #{} finished", thread_idx);

        });

        threads.push(join_handle);
    }

    for t in threads {
        let _ = t.join();
    }

    // could it be more straightforward? put this mutex guard and the threads in a separate block?
    let guard = shared_result.lock().unwrap();
    guard.deref().clone()
}

fn main() {

    let args: Vec<String> = env::args().collect();
    if args.len() < 2 {
        println!("incorrect arguments");
        return;
    }

    let m1: Vec<Vec<u32>> = load_matrix(Path::new(&args[1]));
    let m2: Vec<Vec<u32>> = load_matrix(Path::new(&args[2]));

    assert!(m1[0].len() == m2.len());

    //let result: Vec<Vec<u32>> = single_thread_mul(&m1, &m2);
/*
    for y in 0..result.len() {
        for x in 0..result[y].len() {
            print!("{},", result[y][x]);
        }
        print!("\n");
    }

*/
    let result: Vec<Vec<u32>> = multi_thread_mul(m1, m2, 100);
/*
    for y in 0..result.len() {
        for x in 0..result[y].len() {
            print!("{},", result[y][x]);
        }
        print!("\n");
    }
*/
}
<commit_msg>For now use a more straightforward approach for vector initialization<commit_after>#![feature(core)]

use std::env;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::ops::Deref;
use std::path::Path;
use std::str::FromStr;
use std::sync::{Arc, Mutex};
use std::thread;

fn load_matrix(file_path: &Path) -> Vec<Vec<u32>> {

    let mat_file = File::open(file_path).unwrap();

    let mat_reader = BufReader::new(mat_file);

    let mut mat_line_iterator = mat_reader.lines();

    let mut mat_content: Vec<Vec<u32>> = Vec::new();

    loop {
        match mat_line_iterator.next() {
            Some(result) => {
                                let mat_line = result.unwrap();
                                let slice_iter = mat_line.as_slice().trim_matches(',').split(',')
                                    .map(|x| {
                                                FromStr::from_str(x).unwrap()
                                            });
                                mat_content.push(slice_iter.collect());
                            },

            None => break,
        }
    }

    mat_content
}

fn single_thread_mul(m1: &Vec<Vec<u32>>, m2: &Vec<Vec<u32>>) -> Vec<Vec<u32>> {

    let mut result: Vec<Vec<u32>> = Vec::new();
    let mul_length = m2.len();

    for y in 0..m1.len() {

        let mut result_row: Vec<u32> = Vec::new();
        let mut entry_value: u32 = 0;

        for x in 0..m2[0].len() {

            for i in 0..mul_length {
                entry_value += m1[y][i] * m2[i][x];
            }

            result_row.push(entry_value);
        }

        result.push(result_row);
    }

    result
}

fn multi_thread_mul(m1: Vec<Vec<u32>>, m2: Vec<Vec<u32>>, thread_num: usize) -> Vec<Vec<u32>> {

    let rows_to_process = m1.len();
    let rows_per_thread = rows_to_process / thread_num;

    let shared_m1 = Arc::new(m1);
    let shared_m2 = Arc::new(m2);

    // is there a better way to initialize a vector?
    let mut result: Vec<Vec<u32>> = Vec::with_capacity(rows_to_process);
    //result.resize(rows_to_process, Vec::new());
    for _ in 0..rows_to_process {
        result.push(Vec::new());
    }

    let mul_length = shared_m2.len();

    let shared_result = Arc::new(Mutex::new(result));

    let mut threads = vec![];
    for thread_idx in 0..thread_num {

        let child_m1 = shared_m1.clone();
        let child_m2 = shared_m2.clone();
        let child_result = shared_result.clone();

        let join_handle = thread::spawn(move || {

            let start_row_idx = thread_idx * rows_per_thread;
            let end_row_idx =
                if start_row_idx + rows_per_thread < rows_to_process - 1 {
                    start_row_idx + rows_per_thread
                }
                else {
                    rows_to_process
                };

            println!("thread #{}: start={}, end={}", thread_idx, start_row_idx, end_row_idx);

            for y in start_row_idx..end_row_idx {

                let mut result_row: Vec<u32> = Vec::new();
                let mut entry_value: u32 = 0;

                for x in 0..child_m2[0].len() {

                    for i in 0..mul_length {
                        entry_value += child_m1[y][i] * child_m2[i][x];
                    }

                    result_row.push(entry_value);
                }

                let mut child_result = child_result.lock().unwrap();
                child_result[y] = result_row;
            }

            println!("thread #{} finished", thread_idx);

        });

        threads.push(join_handle);
    }

    for t in threads {
        let _ = t.join();
    }

    // could it be more straightforward? put this mutex guard and the threads in a separate block?
    let guard = shared_result.lock().unwrap();
    guard.deref().clone()
}

fn main() {

    let args: Vec<String> = env::args().collect();
    if args.len() < 2 {
        println!("incorrect arguments");
        return;
    }

    let m1: Vec<Vec<u32>> = load_matrix(Path::new(&args[1]));
    let m2: Vec<Vec<u32>> = load_matrix(Path::new(&args[2]));

    assert!(m1[0].len() == m2.len());

    //let result: Vec<Vec<u32>> = single_thread_mul(&m1, &m2);
/*
    for y in 0..result.len() {
        for x in 0..result[y].len() {
            print!("{},", result[y][x]);
        }
        print!("\n");
    }

*/
    let result: Vec<Vec<u32>> = multi_thread_mul(m1, m2, 100);
/*
    for y in 0..result.len() {
        for x in 0..result[y].len() {
            print!("{},", result[y][x]);
        }
        print!("\n");
    }
*/
}
<|endoftext|>


<commit_before>use super::gl_lib as gl;
use super::buffer::{ArrayBufferBinder, ElementArrayBufferBinder};

pub struct Context {
    pub array_buffer: ArrayBufferBinder,
    pub element_array_buffer: ElementArrayBufferBinder
}

impl Context {
    pub unsafe fn current_context() -> Self {
        Context {
            array_buffer: ArrayBufferBinder,
            element_array_buffer: ElementArrayBufferBinder
        }
    }

    pub fn clear_color(&mut self, color: super::Color) {
        unsafe {
            gl::ClearColor(color.r, color.g, color.b, color.a);
        }
    }

    pub fn clear(&mut self, buffers: super::BufferBits) {
        unsafe {
            gl::Clear(buffers.bits())
        }
    }

    pub fn enable_vertex_attrib_array(&self, attrib: super::ProgramAttrib) {
        unsafe {
            gl::EnableVertexAttribArray(attrib.gl_index);
        }
    }
}

#[macro_export]
macro_rules! bind_array_buffer {
    ($gl:expr, $buffer:expr) => {
        $gl.array_buffer.bind($buffer)
    }
}

#[macro_export]
macro_rules! bind_element_array_buffer {
    ($gl:expr, $buffer:expr) => {
        $gl.element_array_buffer.bind($buffer)
    }
}
<commit_msg>Add a `ProgramBinder` field to `Context`<commit_after>use super::gl_lib as gl;
use super::buffer::{ArrayBufferBinder, ElementArrayBufferBinder};
use super::program::{ProgramBinder};

pub struct Context {
    pub array_buffer: ArrayBufferBinder,
    pub element_array_buffer: ElementArrayBufferBinder,
    pub program: ProgramBinder
}

impl Context {
    pub unsafe fn current_context() -> Self {
        Context {
            array_buffer: ArrayBufferBinder,
            element_array_buffer: ElementArrayBufferBinder,
            program: ProgramBinder
        }
    }

    pub fn clear_color(&mut self, color: super::Color) {
        unsafe {
            gl::ClearColor(color.r, color.g, color.b, color.a);
        }
    }

    pub fn clear(&mut self, buffers: super::BufferBits) {
        unsafe {
            gl::Clear(buffers.bits())
        }
    }

    pub fn enable_vertex_attrib_array(&self, attrib: super::ProgramAttrib) {
        unsafe {
            gl::EnableVertexAttribArray(attrib.gl_index);
        }
    }
}

#[macro_export]
macro_rules! bind_array_buffer {
    ($gl:expr, $buffer:expr) => {
        $gl.array_buffer.bind($buffer)
    }
}

#[macro_export]
macro_rules! bind_element_array_buffer {
    ($gl:expr, $buffer:expr) => {
        $gl.element_array_buffer.bind($buffer)
    }
}
<|endoftext|>