<commit_before>extern crate bytes;
extern crate env_logger;
extern crate mio;
#[macro_use]
extern crate log;
#[macro_use]
extern crate clap;

use clap::{App, Arg};
use mio::tcp::*;
use mio::util::Slab;
use mio::{TryRead, TryWrite};
use std::collections::VecDeque;
use std::ops;

#[derive(Debug)]
enum MiChatCommand {
    Broadcast(String),
    NewConnection(TcpStream),
}

struct MiChat {
    connections: Slab<EventHandler>,
}

trait RuleHandler {
    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet);
    fn register(&self, event_loop: &mut mio::EventLoop<MiChat>, token: mio::Token);
    fn process_rules(
        &mut self,
        event_loop: &mut mio::EventLoop<MiChat>,
        to_parent: &mut VecDeque<MiChatCommand>,
    );
    fn is_closed(&self) -> bool;
}

impl MiChat {
    fn new() -> MiChat {
        MiChat {
            connections: Slab::new(1024),
        }
    }

    fn listen(&mut self, event_loop: &mut mio::EventLoop<MiChat>, listener: TcpListener) {
        let l = EventHandler::Listener(Listener::new(listener));
        let token = self.connections.insert(l).expect("insert listener");
        &self.connections[token].register(event_loop, token);
    }

    fn process_action(&mut self, msg: MiChatCommand, event_loop: &mut mio::EventLoop<MiChat>) {
        trace!("{:p}; got {:?}", self, msg);
        match msg {
            MiChatCommand::Broadcast(s) => self.process_broadcast(s),
            MiChatCommand::NewConnection(socket) => self.process_new_connection(event_loop, socket),
        }
    }

    fn process_broadcast(&mut self, s: String) {
        for c in self.connections.iter_mut() {
            if let &mut EventHandler::Conn(ref mut c) = c {
                c.enqueue(&s);
            }
        }
    }

    fn process_new_connection(
        &mut self,
        event_loop: &mut mio::EventLoop<MiChat>,
        socket: mio::tcp::TcpStream,
    ) {
        let token = self
            .connections
            .insert_with(|token| EventHandler::Conn(Connection::new(socket, token)))
            .expect("token insert");
        &self.connections[token].register(event_loop, token);
    }
}

#[derive(Debug)]
struct Connection {
    socket: TcpStream,
    sock_status: mio::EventSet,
    token: mio::Token,
    read_buf: Vec<u8>,
    read_eof: bool,
    failed: bool,
    write_buf: Vec<u8>,
}

impl RuleHandler for Connection {
    fn register(&self, event_loop: &mut mio::EventLoop<MiChat>, token: mio::Token) {
        event_loop
            .register_opt(
                &self.socket,
                token,
                mio::EventSet::readable(),
                mio::PollOpt::edge() | mio::PollOpt::oneshot(),
            ).expect("event loop register");
    }

    // Event updates arrive here
    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet) {
        self.sock_status.insert(events);
        info!(
            "Connection::handle_event: {:?}; this time: {:?}; now: {:?}",
            self.socket.peer_addr(),
            events,
            self.sock_status
        );
    }

    // actions are processed here on down.

    fn process_rules(
        &mut self,
        event_loop: &mut mio::EventLoop<MiChat>,
        to_parent: &mut VecDeque<MiChatCommand>,
    ) {
        if self.sock_status.is_readable() {
            self.read();
            self.sock_status.remove(mio::EventSet::readable());
        }

        self.process_buffer(to_parent);

        if self.sock_status.is_writable() {
            self.write();
            self.sock_status.remove(mio::EventSet::writable());
        }

        if !self.is_closed() {
            self.reregister(event_loop)
        }
    }

    fn is_closed(&self) -> bool {
        self.failed || (self.read_eof && self.write_buf.is_empty())
    }
}

impl Connection {
    fn new(socket: TcpStream, token: mio::Token) -> Connection {
        Connection {
            socket: socket,
            sock_status: mio::EventSet::none(),
            token: token,
            read_buf: Vec::with_capacity(1024),
            write_buf: Vec::new(),
            read_eof: false,
            failed: false,
        }
    }
    fn process_buffer(&mut self, to_parent: &mut VecDeque<MiChatCommand>) {
        let mut prev = 0;
        info!(
            "{:?}: Read buffer: {:?}",
            self.socket.peer_addr(),
            self.read_buf
        );
        for n in self.read_buf.iter().enumerate().filter_map(|(i, e)| {
            if *e == '\n' as u8 {
                Some(i)
            } else {
                None
            }
        }) {
            info!(
                "{:?}: Pos: {:?}; chunk: {:?}",
                self.socket.peer_addr(),
                n,
                &self.read_buf[prev..n]
            );
            let s = String::from_utf8_lossy(&self.read_buf[prev..n]).to_string();
            let cmd = MiChatCommand::Broadcast(s);
            info!("Send! {:?}", cmd);
            to_parent.push_back(cmd);
            prev = n + 1;
        }
        let remainder = self.read_buf[prev..].to_vec();
        info!(
            "{:?}: read Remainder: {}",
            self.socket.peer_addr(),
            remainder.len()
        );
        self.read_buf = remainder;
    }

    fn read(&mut self) {
        let mut abuf = Vec::new();
        match self.socket.try_read_buf(&mut abuf) {
            Ok(Some(0)) => {
                info!("{:?}: EOF!", self.socket.peer_addr());
                self.read_eof = true
            }
            Ok(Some(n)) => {
                info!("{:?}: Read {}bytes", self.socket.peer_addr(), n);
                self.read_buf.extend(abuf);
            }
            Ok(None) => {
                info!("{:?}: Noop!", self.socket.peer_addr());
            }
            Err(e) => {
                error!("got an error trying to read; err={:?}", e);
                self.failed = true;
            }
        }
    }

    fn write(&mut self) {
        match self.socket.try_write(&mut self.write_buf) {
            Ok(Some(n)) => {
                info!(
                    "{:?}: Wrote {} of {} in buffer",
                    self.socket.peer_addr(),
                    n,
                    self.write_buf.len()
                );
                self.write_buf = self.write_buf[n..].to_vec();
                info!(
                    "{:?}: Now {:?}b",
                    self.socket.peer_addr(),
                    self.write_buf.len()
                );
            }
            Ok(None) => {
                info!("Write unready");
            }
            Err(e) => {
                error!("got an error trying to write; err={:?}", e);
                self.failed = true;
            }
        }
    }

    fn reregister(&mut self, event_loop: &mut mio::EventLoop<MiChat>) {
        let mut flags = mio::EventSet::readable();
        if !self.write_buf.is_empty() {
            flags.insert(mio::EventSet::writable());
        }
        info!("Registering {:?} with {:?}", self, flags);

        event_loop
            .reregister(&self.socket, self.token, flags, mio::PollOpt::oneshot())
            .expect("EventLoop#reregister")
    }

    fn enqueue(&mut self, s: &str) {
        self.write_buf.extend(s.as_bytes());
        self.write_buf.push('\n' as u8);
    }
}

#[derive(Debug)]
struct Listener {
    listener: TcpListener,
    sock_status: mio::EventSet,
}

impl Listener {
    fn new(listener: TcpListener) -> Listener {
        Listener {
            listener: listener,
            sock_status: mio::EventSet::none(),
        }
    }
}

impl RuleHandler for Listener {
    fn register(&self, event_loop: &mut mio::EventLoop<MiChat>, token: mio::Token) {
        event_loop
            .register(&self.listener, token)
            .expect("Register listener");
    }

    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet) {
        assert!(events.is_readable());
        self.sock_status.insert(events);
        info!(
            "Listener::handle_event: {:?}; this time: {:?}; now: {:?}",
            self.listener.local_addr(),
            events,
            self.sock_status
        );
    }
    fn process_rules(
        &mut self,
        event_loop: &mut mio::EventLoop<MiChat>,
        to_parent: &mut VecDeque<MiChatCommand>,
    ) {
        if self.sock_status.is_readable() {
            info!("the listener socket is ready to accept a connection");
            match self.listener.accept() {
                Ok(Some(socket)) => {
                    let cmd = MiChatCommand::NewConnection(socket);
                    to_parent.push_back(cmd);
                }
                Ok(None) => {
                    info!("the listener socket wasn't actually ready");
                }
                Err(e) => {
                    info!("listener.accept() errored: {}", e);
                    event_loop.shutdown();
                }
            }
            self.sock_status.remove(mio::EventSet::readable());
        }
    }

    fn is_closed(&self) -> bool {
        false
    }
}

#[derive(Debug)]
enum EventHandler {
    Listener(Listener),
    Conn(Connection),
}

impl ops::Deref for EventHandler {
    type Target = RuleHandler;
    fn deref(&self) -> &Self::Target {
        match self {
            &EventHandler::Conn(ref conn) => conn,
            &EventHandler::Listener(ref listener) => listener,
        }
    }
}
impl ops::DerefMut for EventHandler {
    fn deref_mut(&mut self) -> &mut Self::Target {
        match self {
            &mut EventHandler::Conn(ref mut conn) => conn,
            &mut EventHandler::Listener(ref mut listener) => listener,
        }
    }
}

impl mio::Handler for MiChat {
    type Timeout = ();
    type Message = ();
    fn ready(
        &mut self,
        event_loop: &mut mio::EventLoop<Self>,
        token: mio::Token,
        events: mio::EventSet,
    ) {
        info!("{:?}: {:?}", token, events);
        self.connections[token].handle_event(event_loop, events);
        if self.connections[token].is_closed() {
            info!("Removing; {:?}", token);
            self.connections.remove(token);
        }

        let mut parent_actions = VecDeque::new();
        loop {
            for conn in self.connections.iter_mut() {
                conn.process_rules(event_loop, &mut parent_actions);
            }
            // Anything left to process?
            if parent_actions.is_empty() {
                break;
            }

            for action in parent_actions.drain(..) {
                self.process_action(action, event_loop);
            }
        }
    }
}

fn main() {
    env_logger::init();
    let matches = App::new("michat")
        .arg(
            Arg::with_name("bind")
                .short("l")
                .takes_value(true)
                .required(true),
        ).get_matches();

    let address = value_t_or_exit!(matches.value_of("bind"), std::net::SocketAddr);
    let listener = TcpListener::bind(&address).expect("bind");

    let mut event_loop = mio::EventLoop::new().expect("Create event loop");
    let mut service = MiChat::new();

    service.listen(&mut event_loop, listener);
    info!("running michat listener at: {:?}", address);
    event_loop.run(&mut service).expect("Run loop");
}
<commit_msg>impl Trait turns out to be nice.<commit_after>extern crate bytes;
extern crate env_logger;
extern crate mio;
#[macro_use]
extern crate log;
#[macro_use]
extern crate clap;

use clap::{App, Arg};
use mio::tcp::*;
use mio::util::Slab;
use mio::{TryRead, TryWrite};
use std::collections::VecDeque;
use std::ops;

#[derive(Debug)]
enum MiChatCommand {
    Broadcast(String),
    NewConnection(TcpStream),
}

struct MiChat {
    connections: Slab<EventHandler>,
}

trait RuleHandler {
    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet);
    fn register(&self, event_loop: &mut mio::EventLoop<MiChat>, token: mio::Token);
    fn process_rules(
        &mut self,
        event_loop: &mut mio::EventLoop<MiChat>,
        to_parent: &mut VecDeque<MiChatCommand>,
    );
    fn is_closed(&self) -> bool;
}

impl MiChat {
    fn new() -> MiChat {
        MiChat {
            connections: Slab::new(1024),
        }
    }

    fn listen(&mut self, event_loop: &mut mio::EventLoop<MiChat>, listener: TcpListener) {
        let l = EventHandler::Listener(Listener::new(listener));
        let token = self.connections.insert(l).expect("insert listener");
        &self.connections[token].register(event_loop, token);
    }

    fn process_action(&mut self, msg: MiChatCommand, event_loop: &mut mio::EventLoop<MiChat>) {
        trace!("{:p}; got {:?}", self, msg);
        match msg {
            MiChatCommand::Broadcast(s) => self.process_broadcast(s),
            MiChatCommand::NewConnection(socket) => self.process_new_connection(event_loop, socket),
        }
    }

    fn process_broadcast(&mut self, s: String) {
        for c in self.client_connections_mut() {
            c.enqueue(&s);
        }
    }

    fn process_new_connection(
        &mut self,
        event_loop: &mut mio::EventLoop<MiChat>,
        socket: mio::tcp::TcpStream,
    ) {
        let token = self
            .connections
            .insert_with(|token| EventHandler::Conn(Connection::new(socket, token)))
            .expect("token insert");
        &self.connections[token].register(event_loop, token);
    }

    fn client_connections_mut<'a>(&'a mut self) -> impl Iterator<Item = &'a mut Connection> + 'a {
        self.connections.iter_mut().filter_map(|it| match it {
            &mut EventHandler::Conn(ref mut c) => Some(c),
            _ => None,
        })
    }
}

#[derive(Debug)]
struct Connection {
    socket: TcpStream,
    sock_status: mio::EventSet,
    token: mio::Token,
    read_buf: Vec<u8>,
    read_eof: bool,
    failed: bool,
    write_buf: Vec<u8>,
}

impl RuleHandler for Connection {
    fn register(&self, event_loop: &mut mio::EventLoop<MiChat>, token: mio::Token) {
        event_loop
            .register_opt(
                &self.socket,
                token,
                mio::EventSet::readable(),
                mio::PollOpt::edge() | mio::PollOpt::oneshot(),
            ).expect("event loop register");
    }

    // Event updates arrive here
    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet) {
        self.sock_status.insert(events);
        info!(
            "Connection::handle_event: {:?}; this time: {:?}; now: {:?}",
            self.socket.peer_addr(),
            events,
            self.sock_status
        );
    }

    // actions are processed here on down.

    fn process_rules(
        &mut self,
        event_loop: &mut mio::EventLoop<MiChat>,
        to_parent: &mut VecDeque<MiChatCommand>,
    ) {
        if self.sock_status.is_readable() {
            self.read();
            self.sock_status.remove(mio::EventSet::readable());
        }

        self.process_buffer(to_parent);

        if self.sock_status.is_writable() {
            self.write();
            self.sock_status.remove(mio::EventSet::writable());
        }

        if !self.is_closed() {
            self.reregister(event_loop)
        }
    }

    fn is_closed(&self) -> bool {
        self.failed || (self.read_eof && self.write_buf.is_empty())
    }
}

impl Connection {
    fn new(socket: TcpStream, token: mio::Token) -> Connection {
        Connection {
            socket: socket,
            sock_status: mio::EventSet::none(),
            token: token,
            read_buf: Vec::with_capacity(1024),
            write_buf: Vec::new(),
            read_eof: false,
            failed: false,
        }
    }
    fn process_buffer(&mut self, to_parent: &mut VecDeque<MiChatCommand>) {
        let mut prev = 0;
        info!(
            "{:?}: Read buffer: {:?}",
            self.socket.peer_addr(),
            self.read_buf
        );
        for n in self.read_buf.iter().enumerate().filter_map(|(i, e)| {
            if *e == '\n' as u8 {
                Some(i)
            } else {
                None
            }
        }) {
            info!(
                "{:?}: Pos: {:?}; chunk: {:?}",
                self.socket.peer_addr(),
                n,
                &self.read_buf[prev..n]
            );
            let s = String::from_utf8_lossy(&self.read_buf[prev..n]).to_string();
            let cmd = MiChatCommand::Broadcast(s);
            info!("Send! {:?}", cmd);
            to_parent.push_back(cmd);
            prev = n + 1;
        }
        let remainder = self.read_buf[prev..].to_vec();
        info!(
            "{:?}: read Remainder: {}",
            self.socket.peer_addr(),
            remainder.len()
        );
        self.read_buf = remainder;
    }

    fn read(&mut self) {
        let mut abuf = Vec::new();
        match self.socket.try_read_buf(&mut abuf) {
            Ok(Some(0)) => {
                info!("{:?}: EOF!", self.socket.peer_addr());
                self.read_eof = true
            }
            Ok(Some(n)) => {
                info!("{:?}: Read {}bytes", self.socket.peer_addr(), n);
                self.read_buf.extend(abuf);
            }
            Ok(None) => {
                info!("{:?}: Noop!", self.socket.peer_addr());
            }
            Err(e) => {
                error!("got an error trying to read; err={:?}", e);
                self.failed = true;
            }
        }
    }

    fn write(&mut self) {
        match self.socket.try_write(&mut self.write_buf) {
            Ok(Some(n)) => {
                info!(
                    "{:?}: Wrote {} of {} in buffer",
                    self.socket.peer_addr(),
                    n,
                    self.write_buf.len()
                );
                self.write_buf = self.write_buf[n..].to_vec();
                info!(
                    "{:?}: Now {:?}b",
                    self.socket.peer_addr(),
                    self.write_buf.len()
                );
            }
            Ok(None) => {
                info!("Write unready");
            }
            Err(e) => {
                error!("got an error trying to write; err={:?}", e);
                self.failed = true;
            }
        }
    }

    fn reregister(&mut self, event_loop: &mut mio::EventLoop<MiChat>) {
        let mut flags = mio::EventSet::readable();
        if !self.write_buf.is_empty() {
            flags.insert(mio::EventSet::writable());
        }
        info!("Registering {:?} with {:?}", self, flags);

        event_loop
            .reregister(&self.socket, self.token, flags, mio::PollOpt::oneshot())
            .expect("EventLoop#reregister")
    }

    fn enqueue(&mut self, s: &str) {
        self.write_buf.extend(s.as_bytes());
        self.write_buf.push('\n' as u8);
    }
}

#[derive(Debug)]
struct Listener {
    listener: TcpListener,
    sock_status: mio::EventSet,
}

impl Listener {
    fn new(listener: TcpListener) -> Listener {
        Listener {
            listener: listener,
            sock_status: mio::EventSet::none(),
        }
    }
}

impl RuleHandler for Listener {
    fn register(&self, event_loop: &mut mio::EventLoop<MiChat>, token: mio::Token) {
        event_loop
            .register(&self.listener, token)
            .expect("Register listener");
    }

    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet) {
        assert!(events.is_readable());
        self.sock_status.insert(events);
        info!(
            "Listener::handle_event: {:?}; this time: {:?}; now: {:?}",
            self.listener.local_addr(),
            events,
            self.sock_status
        );
    }
    fn process_rules(
        &mut self,
        event_loop: &mut mio::EventLoop<MiChat>,
        to_parent: &mut VecDeque<MiChatCommand>,
    ) {
        if self.sock_status.is_readable() {
            info!("the listener socket is ready to accept a connection");
            match self.listener.accept() {
                Ok(Some(socket)) => {
                    let cmd = MiChatCommand::NewConnection(socket);
                    to_parent.push_back(cmd);
                }
                Ok(None) => {
                    info!("the listener socket wasn't actually ready");
                }
                Err(e) => {
                    info!("listener.accept() errored: {}", e);
                    event_loop.shutdown();
                }
            }
            self.sock_status.remove(mio::EventSet::readable());
        }
    }

    fn is_closed(&self) -> bool {
        false
    }
}

#[derive(Debug)]
enum EventHandler {
    Listener(Listener),
    Conn(Connection),
}

impl ops::Deref for EventHandler {
    type Target = RuleHandler;
    fn deref(&self) -> &Self::Target {
        match self {
            &EventHandler::Conn(ref conn) => conn,
            &EventHandler::Listener(ref listener) => listener,
        }
    }
}
impl ops::DerefMut for EventHandler {
    fn deref_mut(&mut self) -> &mut Self::Target {
        match self {
            &mut EventHandler::Conn(ref mut conn) => conn,
            &mut EventHandler::Listener(ref mut listener) => listener,
        }
    }
}

impl mio::Handler for MiChat {
    type Timeout = ();
    type Message = ();
    fn ready(
        &mut self,
        event_loop: &mut mio::EventLoop<Self>,
        token: mio::Token,
        events: mio::EventSet,
    ) {
        info!("{:?}: {:?}", token, events);
        self.connections[token].handle_event(event_loop, events);
        if self.connections[token].is_closed() {
            info!("Removing; {:?}", token);
            self.connections.remove(token);
        }

        let mut parent_actions = VecDeque::new();
        loop {
            for conn in self.connections.iter_mut() {
                conn.process_rules(event_loop, &mut parent_actions);
            }
            // Anything left to process?
            if parent_actions.is_empty() {
                break;
            }

            for action in parent_actions.drain(..) {
                self.process_action(action, event_loop);
            }
        }
    }
}

fn main() {
    env_logger::init();
    let matches = App::new("michat")
        .arg(
            Arg::with_name("bind")
                .short("l")
                .takes_value(true)
                .required(true),
        ).get_matches();

    let address = value_t_or_exit!(matches.value_of("bind"), std::net::SocketAddr);
    let listener = TcpListener::bind(&address).expect("bind");

    let mut event_loop = mio::EventLoop::new().expect("Create event loop");
    let mut service = MiChat::new();

    service.listen(&mut event_loop, listener);
    info!("running michat listener at: {:?}", address);
    event_loop.run(&mut service).expect("Run loop");
}
<|endoftext|>


<commit_before>#![feature(duration)]
#![feature(duration_span)]

extern crate image;
extern crate rand;
extern crate docopt;
extern crate rustc_serialize as rustc_serialize;
extern crate threadpool;
extern crate num_cpus;
extern crate tray_rust;

use std::vec::Vec;
use std::sync::{Arc};
use std::sync::mpsc::{self, Sender, Receiver};
use std::path::Path;
use std::time::Duration;

use threadpool::ThreadPool;
use rand::StdRng;
use docopt::Docopt;

use tray_rust::film;
use tray_rust::geometry::Geometry;
use tray_rust::sampler::{self, Sampler};
use tray_rust::scene;
use tray_rust::integrator::Integrator;

static WIDTH: usize = 800;
static HEIGHT: usize = 600;
static USAGE: &'static str = "
Usage: tray_rust [options]

Options:
  -o <file>     Specify the output file to save the image. Supported formats are
                PNG, JPG and PPM. Default is 'out.png'.
  -n <number>   Specify the number of threads to use for rendering. Defaults to the number of cores
                on the system.
  -h, --help    Show this message.
";

#[derive(RustcDecodable, Debug)]
struct Args {
    flag_o: Option<String>,
    flag_n: Option<usize>,
}

// TODO: AREA LIGHTS NOTES:
// - They should be intersectable so they're like wierd instances.
// - What if point lights were also instances that just always returned None
//   for intersection tests?
// - Then each thread would build a list of &Instance that refer to the emissive
//   instances in the scene and we'd pass this around for the various lighting calculations.
//   It's a bit ugly but maybe cleaner than what was done previously in tray
//   It breaks the awkward handling of mixing point and area lights with each other which
//   is a plus. I guess we can also put point lights in the BVH and their bounds would
//   just be the point they're located at.

/// Threads are each sent a sender end of the channel that is
/// read from by the render target thread which then saves the
/// values recieved to the render target
fn thread_work(tx: Sender<(f32, f32, film::Colorf)>, queue: Arc<sampler::BlockQueue>,
               scene: Arc<scene::Scene>) {
    let mut sampler = sampler::LowDiscrepancy::new(queue.block_dim(), 4);
    let mut samples = Vec::with_capacity(sampler.max_spp());
    let mut sample_pos = Vec::with_capacity(sampler.max_spp());
    let mut rng = match StdRng::new() {
        Ok(r) => r,
        Err(e) => { println!("Failed to get StdRng, {}", e); return }
    };
    // Grab a block from the queue and start working on it, submitting samples
    // to the render target thread after each pixel
    for b in queue.iter() {
        sampler.select_block(b);
        while sampler.has_samples() {
            // Get samples for a pixel and render them
            sampler.get_samples(&mut sample_pos, &mut rng);
            for s in sample_pos.iter() {
                let mut ray = scene.camera.generate_ray(s);
                if let Some(hit) = scene.intersect(&mut ray) {
                    let c = scene.integrator.illumination(&*scene, &ray, &hit, &mut sampler, &mut rng).clamp();
                    samples.push((s.0, s.1, c));
                }
            }
            for s in samples.iter() {
                if let Err(e) = tx.send(*s) {
                    println!("Worker thread exiting with send error {:?}", e);
                    return;
                }
            }
            samples.clear();
        }
    }
}

/// Spawn `n` worker threads to render the scene in parallel. Returns the receive end
/// of the channel where the threads will write their samples so that the receiver
/// can write these samples to the render target
fn spawn_workers(pool: &ThreadPool, n: usize, scene: Arc<scene::Scene>) -> Receiver<(f32, f32, film::Colorf)> {
    let (tx, rx) = mpsc::channel();
    let block_queue = Arc::new(sampler::BlockQueue::new((WIDTH as u32, HEIGHT as u32), (8, 8)));
    for _ in 0..n {
        let q = block_queue.clone();
        let t = tx.clone();
        let s = scene.clone();
        pool.execute(move || {
            thread_work(t, q, s);
        });
    }
    rx
}

/// Render the scene in parallel to the render target
fn render_parallel(rt: &mut film::RenderTarget, n: usize){
    let scene = Arc::new(scene::Scene::new(WIDTH, HEIGHT));
    let pool = ThreadPool::new(n);
    let rx = spawn_workers(&pool, n, scene);
    for m in rx.iter() {
        rt.write(m.0, m.1, &m.2);
    }
}

fn main() {
    let args: Args = Docopt::new(USAGE).and_then(|d| d.decode()).unwrap_or_else(|e| e.exit());

    let mut rt = film::RenderTarget::new(WIDTH, HEIGHT);
    let n = match args.flag_n {
        Some(n) => n,
        None => num_cpus::get(),
    };
    println!("Rendering using {} threads", n);
    let d = Duration::span(|| render_parallel(&mut rt, n));
    println!("Rendering took {}", d);
    let img = rt.get_render();
    let out_file = match args.flag_o {
        Some(f) => f,
        None => "out.png".to_string(),
    };
    match image::save_buffer(&Path::new(&out_file), &img[..], WIDTH as u32, HEIGHT as u32, image::RGB(8)) {
        Ok(_) => {},
        Err(e) => println!("Error saving image, {}", e),
    };
}

<commit_msg>Some better thoughts on area light work<commit_after>#![feature(duration)]
#![feature(duration_span)]

extern crate image;
extern crate rand;
extern crate docopt;
extern crate rustc_serialize as rustc_serialize;
extern crate threadpool;
extern crate num_cpus;
extern crate tray_rust;

use std::vec::Vec;
use std::sync::{Arc};
use std::sync::mpsc::{self, Sender, Receiver};
use std::path::Path;
use std::time::Duration;

use threadpool::ThreadPool;
use rand::StdRng;
use docopt::Docopt;

use tray_rust::film;
use tray_rust::geometry::Geometry;
use tray_rust::sampler::{self, Sampler};
use tray_rust::scene;
use tray_rust::integrator::Integrator;

static WIDTH: usize = 800;
static HEIGHT: usize = 600;
static USAGE: &'static str = "
Usage: tray_rust [options]

Options:
  -o <file>     Specify the output file to save the image. Supported formats are
                PNG, JPG and PPM. Default is 'out.png'.
  -n <number>   Specify the number of threads to use for rendering. Defaults to the number of cores
                on the system.
  -h, --help    Show this message.
";

#[derive(RustcDecodable, Debug)]
struct Args {
    flag_o: Option<String>,
    flag_n: Option<usize>,
}

// TODO: AREA LIGHTS NOTES:
// - They should be intersectable so they're like wierd instances.
// - What if point lights were also instances that just always returned None
//   for intersection tests?
// - Then each thread would build a list of &Instance that refer to the emissive
//   instances in the scene and we'd pass this around for the various lighting calculations.
//   It's a bit ugly but maybe cleaner than what was done previously in tray
//   It breaks the awkward handling of mixing point and area lights with each other which
//   is a plus. I guess we can also put point lights in the BVH and their bounds would
//   just be the point they're located at.
// - If we make them their own instance type we could do something like have the light
//   be the "geometry" for this LightInstance type which would implment the Instance trait
//   then there'd be like light::Sphere which would use the sphere's intersection code
//   and geometry::Sphere but would then have the surface sampling implementations.
//   So we get some re-use of geometry types as well. This would then make
//   our Instance calls go through virtual calls (eg. they become more like geometry)

/// Threads are each sent a sender end of the channel that is
/// read from by the render target thread which then saves the
/// values recieved to the render target
fn thread_work(tx: Sender<(f32, f32, film::Colorf)>, queue: Arc<sampler::BlockQueue>,
               scene: Arc<scene::Scene>) {
    let mut sampler = sampler::LowDiscrepancy::new(queue.block_dim(), 4);
    let mut samples = Vec::with_capacity(sampler.max_spp());
    let mut sample_pos = Vec::with_capacity(sampler.max_spp());
    let mut rng = match StdRng::new() {
        Ok(r) => r,
        Err(e) => { println!("Failed to get StdRng, {}", e); return }
    };
    // Grab a block from the queue and start working on it, submitting samples
    // to the render target thread after each pixel
    for b in queue.iter() {
        sampler.select_block(b);
        while sampler.has_samples() {
            // Get samples for a pixel and render them
            sampler.get_samples(&mut sample_pos, &mut rng);
            for s in sample_pos.iter() {
                let mut ray = scene.camera.generate_ray(s);
                if let Some(hit) = scene.intersect(&mut ray) {
                    let c = scene.integrator.illumination(&*scene, &ray, &hit, &mut sampler, &mut rng).clamp();
                    samples.push((s.0, s.1, c));
                }
            }
            for s in samples.iter() {
                if let Err(e) = tx.send(*s) {
                    println!("Worker thread exiting with send error {:?}", e);
                    return;
                }
            }
            samples.clear();
        }
    }
}

/// Spawn `n` worker threads to render the scene in parallel. Returns the receive end
/// of the channel where the threads will write their samples so that the receiver
/// can write these samples to the render target
fn spawn_workers(pool: &ThreadPool, n: usize, scene: Arc<scene::Scene>) -> Receiver<(f32, f32, film::Colorf)> {
    let (tx, rx) = mpsc::channel();
    let block_queue = Arc::new(sampler::BlockQueue::new((WIDTH as u32, HEIGHT as u32), (8, 8)));
    for _ in 0..n {
        let q = block_queue.clone();
        let t = tx.clone();
        let s = scene.clone();
        pool.execute(move || {
            thread_work(t, q, s);
        });
    }
    rx
}

/// Render the scene in parallel to the render target
fn render_parallel(rt: &mut film::RenderTarget, n: usize){
    let scene = Arc::new(scene::Scene::new(WIDTH, HEIGHT));
    let pool = ThreadPool::new(n);
    let rx = spawn_workers(&pool, n, scene);
    for m in rx.iter() {
        rt.write(m.0, m.1, &m.2);
    }
}

fn main() {
    let args: Args = Docopt::new(USAGE).and_then(|d| d.decode()).unwrap_or_else(|e| e.exit());

    let mut rt = film::RenderTarget::new(WIDTH, HEIGHT);
    let n = match args.flag_n {
        Some(n) => n,
        None => num_cpus::get(),
    };
    println!("Rendering using {} threads", n);
    let d = Duration::span(|| render_parallel(&mut rt, n));
    println!("Rendering took {}", d);
    let img = rt.get_render();
    let out_file = match args.flag_o {
        Some(f) => f,
        None => "out.png".to_string(),
    };
    match image::save_buffer(&Path::new(&out_file), &img[..], WIDTH as u32, HEIGHT as u32, image::RGB(8)) {
        Ok(_) => {},
        Err(e) => println!("Error saving image, {}", e),
    };
}

<|endoftext|>


<commit_before>mod engine {
    extern crate sdl2;

    pub struct Engine {
        event_pump: sdl2::EventPump,
        renderer: sdl2::render::Renderer<'static>,
    }

    impl Engine {
        pub fn run(&mut self) {
            loop {
                for event in self.event_pump.poll_iter() {
                    match event {
                        sdl2::event::Event::Quit { .. } |
                        sdl2::event::Event::KeyDown { .. } => return,
                        _ => (),
                    }
                }
                self.renderer.set_draw_color(sdl2::pixels::Color::RGB(0, 0, 0));
                self.renderer.clear();
                self.renderer.present();
            }
        }
    }

    pub fn init() -> Result<Engine, String> {
        let sdl = sdl2::init()?;
        let video = sdl.video()?;
        let event_pump = sdl.event_pump()?;
        let window = video.window("Snake", 640, 480)
            .build()
            .or_else(|e| Err(format!("{}", e)))?;
        let renderer = window.renderer()
            .present_vsync()
            .build()
            .or_else(|e| Err(format!("{}", e)))?;

        Ok(Engine {
               event_pump: event_pump,
               renderer: renderer,
           })
    }
}

fn main() {
    match engine::init() {
        Ok(mut engine) => engine.run(),
        Err(s) => println!("Engine initialization failed: {}", s),
    }
}
<commit_msg>Draw a playing field<commit_after>mod game {
    extern crate ndarray;

    type TileArray = ndarray::Array2<Tile>;

    pub struct GameState {
        tiles: TileArray,
    }

    impl GameState {
        pub fn new() -> Self {
            let width = 80;
            let height = 50;
            // Arrays are indexed in (row (y), column (x)) order
            let mut tiles = ndarray::Array::from_elem((height, width), Tile::Empty);

            // Build level wall
            for x in 0..width {
                tiles[[0, x]] = Tile::Wall;
                tiles[[height - 1, x]] = Tile::Wall;
            }
            for y in 0..height {
                tiles[[y, 0]] = Tile::Wall;
                tiles[[y, width - 1]] = Tile::Wall;
            }

            GameState { tiles: tiles }
        }

        pub fn tiles(&self) -> &TileArray {
            &self.tiles
        }
    }

    #[derive(Copy, Clone, Debug)]
    pub enum Tile {
        Empty,
        Wall,
    }
}

mod engine {
    extern crate sdl2;

    use game;

    pub struct Engine {
        event_pump: sdl2::EventPump,
        renderer: sdl2::render::Renderer<'static>,
        game_state: game::GameState,
    }

    impl Engine {
        pub fn run(&mut self) -> Result<(), String> {
            loop {
                for event in self.event_pump.poll_iter() {
                    match event {
                        sdl2::event::Event::Quit { .. } |
                        sdl2::event::Event::KeyDown { .. } => return Ok(()),
                        _ => (),
                    }
                }
                self.render()?;
            }
        }

        fn render(&mut self) -> Result<(), String> {
            self.renderer.set_draw_color(sdl2::pixels::Color::RGB(0, 0, 0));
            self.renderer.clear();
            let scale: u32 = 8;
            for ((y, x), &tile) in self.game_state.tiles().indexed_iter() {
                let color = match tile {
                    game::Tile::Empty => sdl2::pixels::Color::RGB(255, 0, 0),
                    game::Tile::Wall => sdl2::pixels::Color::RGB(0, 255, 0),
                };
                self.renderer.set_draw_color(color);
                self.renderer
                    .fill_rect(sdl2::rect::Rect::new(x as i32 * scale as i32,
                                                     y as i32 * scale as i32,
                                                     scale,
                                                     scale))?;
            }
            self.renderer.present();

            Ok(())
        }
    }

    pub fn init() -> Result<Engine, String> {
        let sdl = sdl2::init()?;
        let video = sdl.video()?;
        let event_pump = sdl.event_pump()?;
        let window = video.window("Snake", 640, 480)
            .build()
            .or_else(|e| Err(format!("{}", e)))?;
        let renderer = window.renderer()
            .present_vsync()
            .build()
            .or_else(|e| Err(format!("{}", e)))?;
        let game_state = game::GameState::new();

        Ok(Engine {
               event_pump: event_pump,
               renderer: renderer,
               game_state: game_state,
           })
    }
}

fn main() {
    match engine::init() {
        Ok(mut engine) => {
            if let Err(s) = engine.run() {
                println!("Runtime error: {}", s)
            }
        }
        Err(s) => println!("Engine initialization failed: {}", s),
    }
}
<|endoftext|>


<commit_before>#![cfg_attr(test, allow(dead_code))]
#![feature(collections, env, libc, old_io, os)]
#![cfg_attr(not(windows), feature(old_path, fs, io, std_misc, core))]

extern crate libc;
#[cfg(not(windows))] extern crate collections;

mod args;
mod matching;
mod screen;
#[cfg(not(windows))] mod ansi;

use std::env;
use std::old_io;
use std::cmp::min;
use screen::Screen;
use screen::Key;
use screen::Key::*;
use self::SearchState::*;

fn main() {
    let args = match args::parse_args() {
        Some(args) => args,
        None => {
            env::set_exit_status(1);
            return;
        },
    };

    if args.help { return; }

    let choices = read_choices();
    let initial_search = args.initial_search.clone();
    let choices = choices.iter().map(|x| &x[..]).collect::<Vec<&str>>();
    if args.use_first {
        let matches = matching::compute_matches(&choices, &initial_search);
        println!("{}", matches.get(0).unwrap_or(&""));
        return;
    } else {
        event_loop(&choices, &initial_search);
    }
}

fn event_loop(choices: &[&str], initial_search: &str) {
    let mut search = Search::new(choices, initial_search.to_string());
    let mut screen = Screen::open_screen();

    loop {
        search.recompute_matches();

        match search.state {
            InProgress => draw_screen(&mut screen, &search),
            _ => break,
        }

        let keys = screen.get_buffered_keys();
        for key in keys.iter() {
            handle_key(&mut search, key, screen.visible_choices);
        }
    }

    screen.move_cursor_to_bottom();
    println!("{}", search.get_selection());
}

fn handle_key(search: &mut Search, key: &Key, visible_choices: u16) {
    match *key {
        Char(x) => search.append(x),
        Backspace => search.backspace(),
        Control('h') => search.backspace(),
        Control('w') => search.delete_word(),
        Control('u') => search.clear_query(),
        Control('c') => search.cancel(),
        Control('g') => search.cancel(),
        Control('n') => search.down(visible_choices),
        Control('p') => search.up(visible_choices),
        Enter => search.done(),
        _ => {}
    }
}

struct Search<'a> {
    choices: &'a [&'a str],
    query: String,
    matches: Vec<&'a str>,
    stale: bool,
    index: usize,
    state: SearchState,
}

#[derive(PartialEq, Eq)]
enum SearchState {
    InProgress,
    Done,
    Canceled,
}

impl<'a> Search<'a> {
    fn new(choices: &'a [&'a str], initial_search: String) -> Search<'a> {
        let matches = choices.to_vec();
        Search {
            choices: choices,
            query: initial_search,
            matches: matches,
            stale: true,
            index: 0,
            state: InProgress,
        }
    }

    fn up(&mut self, visible_choices: u16) {
        let limit = min(visible_choices as usize - 1, self.matches.len() - 1);
        self.index = if self.index == 0 { limit } else { self.index - 1 };
    }

    fn down(&mut self, visible_choices: u16) {
        let limit = min(visible_choices as usize - 1, self.matches.len() - 1);
        self.index = if self.index == limit { 0 } else { self.index + 1 }
    }

    fn backspace(&mut self) {
        self.query.pop();
        self.stale = true;
        self.index = 0;
        self.matches = self.choices.to_vec();
    }

    fn delete_word(&mut self) {
        self.stale = true;
        delete_last_word(&mut self.query);
        self.matches = self.choices.to_vec();
    }

    fn append(&mut self, c: char) {
        self.query.push(c);
        self.stale = true;
        self.index = 0;
    }

    fn clear_query(&mut self) {
        self.query.clear();
        self.matches = self.choices.to_vec();
    }

    fn recompute_matches(&mut self) {
        if self.stale {
            self.matches = matching::compute_matches(&self.matches, &self.query);
            self.stale = false;
        }
    }

    fn get_selection(&mut self) -> String {
        if self.state == Canceled {
            "".to_string()
        } else {
            self.recompute_matches();
            self.matches.get(self.index).unwrap_or(&"").to_string()
        }
    }

    fn cancel(&mut self) {
        self.state = Canceled
    }

    fn done(&mut self) {
        self.state = Done
    }
}

fn draw_screen(screen: &mut Screen, search: &Search) {
    screen.hide_cursor();
    screen.blank_screen();
    screen.write(&format!("> {} ({}/{} choices)\n", search.query, search.matches.len(), search.choices.len()));

    print_matches(screen, &search.matches, &search.query, search.index);

    let start = screen.start_line;
    screen.move_cursor(start, 2 + search.query.chars().count() as u16);
    screen.show_cursor();
}

fn print_matches(screen: &mut Screen, matches: &[&str], query: &str, index: usize) {
    let mut i = 1;
    for choice in matches.iter() {
        let indices = matching::visual_score(choice, query);
        let max_width = screen.width;
        print_match(choice, &indices, max_width, &mut |s, highlight| {
            if i == index + 1 {
                if highlight { screen.write_red_inverted(s); } else { screen.write_inverted(s); }
            } else {
                if highlight { screen.write_red(s); } else { screen.write(s); }
            }
        });
        if i >= screen.visible_choices as usize {
            return;
        } else {
            screen.write("\n");
        }
        i += 1;
    }
}

fn print_match(choice: &str, indices: &[usize], max_width: u16, writer: &mut FnMut(&str, bool)) {
    #[cfg(windows)] const MARGIN: u16 = 1;
    #[cfg(not(windows))] const MARGIN: u16 = 0;
    let max_width = max_width - MARGIN;
    let chars_in_choice = choice.chars().collect::<Vec<char>>().len();
    let chars_to_draw = min(chars_in_choice, max_width as usize);
    let mut last_idx = 0;
    for &idx in indices {
        let idx = min(idx, chars_to_draw);
        if last_idx >= chars_to_draw {
            return;
        }
        writer(&choice.slice_chars(last_idx, idx), false);
        if idx == chars_to_draw { return }
        writer(&choice.slice_chars(idx, idx + 1), true);
        last_idx = idx + 1;
    }
    writer(&choice.slice_chars(last_idx, chars_to_draw), false);
}

fn read_choices() -> Vec<String> {
    let mut stdin = old_io::stdio::stdin();
    let mut lines = Vec::new();

    while let Ok(mut s) = stdin.read_line() {
        trim(&mut s);
        lines.push(s);
    }

    lines
}

pub fn trim(s: &mut String) {
    while let Some(x) = s.pop() {
        if x != '\n' {
            s.push(x);
            return;
        }
    }
}

fn delete_last_word(s: &mut String) {
    let mut deleted_something = false;
    while let Some(x) = s.pop() {
        if x == ' ' {
            if deleted_something {
                s.push(x);
                return;
            }
        } else {
            deleted_something = true;
        }
    }
}

#[test]
fn trim_test() {
    fn should_become(before: &str, after: &str) {
        let mut x = before.to_string();
        trim(&mut x);
        assert_eq!(after.to_string(), x);
    }
    should_become("", "");
    should_become("\n", "");
    should_become("\n\n", "");
    should_become("asdf", "asdf");
    should_become("asdf\n", "asdf");
    should_become("asdf\nasdf\n", "asdf\nasdf");
}

#[test]
fn delete_word_test() {
    fn should_become(before: &str, after: &str) {
        let mut x = before.to_string();
        delete_last_word(&mut x);
        assert_eq!(after.to_string(), x);
    }
    should_become("", "");
    should_become("a", "");
    should_become("asdf", "");
    should_become("asdf asdf asdf", "asdf asdf ");
    should_become("asdf asdf asdf ", "asdf asdf ");
    should_become("asdf asdf asdf  ", "asdf asdf ");
}
<commit_msg>Remove unnecessary collect() call<commit_after>#![cfg_attr(test, allow(dead_code))]
#![feature(collections, env, libc, old_io, os)]
#![cfg_attr(not(windows), feature(old_path, fs, io, std_misc, core))]

extern crate libc;
#[cfg(not(windows))] extern crate collections;

mod args;
mod matching;
mod screen;
#[cfg(not(windows))] mod ansi;

use std::env;
use std::old_io;
use std::cmp::min;
use screen::Screen;
use screen::Key;
use screen::Key::*;
use self::SearchState::*;

fn main() {
    let args = match args::parse_args() {
        Some(args) => args,
        None => {
            env::set_exit_status(1);
            return;
        },
    };

    if args.help { return; }

    let choices = read_choices();
    let initial_search = args.initial_search.clone();
    let choices = choices.iter().map(|x| &x[..]).collect::<Vec<&str>>();
    if args.use_first {
        let matches = matching::compute_matches(&choices, &initial_search);
        println!("{}", matches.get(0).unwrap_or(&""));
        return;
    } else {
        event_loop(&choices, &initial_search);
    }
}

fn event_loop(choices: &[&str], initial_search: &str) {
    let mut search = Search::new(choices, initial_search.to_string());
    let mut screen = Screen::open_screen();

    loop {
        search.recompute_matches();

        match search.state {
            InProgress => draw_screen(&mut screen, &search),
            _ => break,
        }

        let keys = screen.get_buffered_keys();
        for key in keys.iter() {
            handle_key(&mut search, key, screen.visible_choices);
        }
    }

    screen.move_cursor_to_bottom();
    println!("{}", search.get_selection());
}

fn handle_key(search: &mut Search, key: &Key, visible_choices: u16) {
    match *key {
        Char(x) => search.append(x),
        Backspace => search.backspace(),
        Control('h') => search.backspace(),
        Control('w') => search.delete_word(),
        Control('u') => search.clear_query(),
        Control('c') => search.cancel(),
        Control('g') => search.cancel(),
        Control('n') => search.down(visible_choices),
        Control('p') => search.up(visible_choices),
        Enter => search.done(),
        _ => {}
    }
}

struct Search<'a> {
    choices: &'a [&'a str],
    query: String,
    matches: Vec<&'a str>,
    stale: bool,
    index: usize,
    state: SearchState,
}

#[derive(PartialEq, Eq)]
enum SearchState {
    InProgress,
    Done,
    Canceled,
}

impl<'a> Search<'a> {
    fn new(choices: &'a [&'a str], initial_search: String) -> Search<'a> {
        let matches = choices.to_vec();
        Search {
            choices: choices,
            query: initial_search,
            matches: matches,
            stale: true,
            index: 0,
            state: InProgress,
        }
    }

    fn up(&mut self, visible_choices: u16) {
        let limit = min(visible_choices as usize - 1, self.matches.len() - 1);
        self.index = if self.index == 0 { limit } else { self.index - 1 };
    }

    fn down(&mut self, visible_choices: u16) {
        let limit = min(visible_choices as usize - 1, self.matches.len() - 1);
        self.index = if self.index == limit { 0 } else { self.index + 1 }
    }

    fn backspace(&mut self) {
        self.query.pop();
        self.stale = true;
        self.index = 0;
        self.matches = self.choices.to_vec();
    }

    fn delete_word(&mut self) {
        self.stale = true;
        delete_last_word(&mut self.query);
        self.matches = self.choices.to_vec();
    }

    fn append(&mut self, c: char) {
        self.query.push(c);
        self.stale = true;
        self.index = 0;
    }

    fn clear_query(&mut self) {
        self.query.clear();
        self.matches = self.choices.to_vec();
    }

    fn recompute_matches(&mut self) {
        if self.stale {
            self.matches = matching::compute_matches(&self.matches, &self.query);
            self.stale = false;
        }
    }

    fn get_selection(&mut self) -> String {
        if self.state == Canceled {
            "".to_string()
        } else {
            self.recompute_matches();
            self.matches.get(self.index).unwrap_or(&"").to_string()
        }
    }

    fn cancel(&mut self) {
        self.state = Canceled
    }

    fn done(&mut self) {
        self.state = Done
    }
}

fn draw_screen(screen: &mut Screen, search: &Search) {
    screen.hide_cursor();
    screen.blank_screen();
    screen.write(&format!("> {} ({}/{} choices)\n", search.query, search.matches.len(), search.choices.len()));

    print_matches(screen, &search.matches, &search.query, search.index);

    let start = screen.start_line;
    screen.move_cursor(start, 2 + search.query.chars().count() as u16);
    screen.show_cursor();
}

fn print_matches(screen: &mut Screen, matches: &[&str], query: &str, index: usize) {
    let mut i = 1;
    for choice in matches.iter() {
        let indices = matching::visual_score(choice, query);
        let max_width = screen.width;
        print_match(choice, &indices, max_width, &mut |s, highlight| {
            if i == index + 1 {
                if highlight { screen.write_red_inverted(s); } else { screen.write_inverted(s); }
            } else {
                if highlight { screen.write_red(s); } else { screen.write(s); }
            }
        });
        if i >= screen.visible_choices as usize {
            return;
        } else {
            screen.write("\n");
        }
        i += 1;
    }
}

fn print_match(choice: &str, indices: &[usize], max_width: u16, writer: &mut FnMut(&str, bool)) {
    #[cfg(windows)] const MARGIN: u16 = 1;
    #[cfg(not(windows))] const MARGIN: u16 = 0;
    let max_width = max_width - MARGIN;
    let chars_in_choice = choice.chars().count();
    let chars_to_draw = min(chars_in_choice, max_width as usize);
    let mut last_idx = 0;
    for &idx in indices {
        let idx = min(idx, chars_to_draw);
        if last_idx >= chars_to_draw {
            return;
        }
        writer(&choice.slice_chars(last_idx, idx), false);
        if idx == chars_to_draw { return }
        writer(&choice.slice_chars(idx, idx + 1), true);
        last_idx = idx + 1;
    }
    writer(&choice.slice_chars(last_idx, chars_to_draw), false);
}

fn read_choices() -> Vec<String> {
    let mut stdin = old_io::stdio::stdin();
    let mut lines = Vec::new();

    while let Ok(mut s) = stdin.read_line() {
        trim(&mut s);
        lines.push(s);
    }

    lines
}

pub fn trim(s: &mut String) {
    while let Some(x) = s.pop() {
        if x != '\n' {
            s.push(x);
            return;
        }
    }
}

fn delete_last_word(s: &mut String) {
    let mut deleted_something = false;
    while let Some(x) = s.pop() {
        if x == ' ' {
            if deleted_something {
                s.push(x);
                return;
            }
        } else {
            deleted_something = true;
        }
    }
}

#[test]
fn trim_test() {
    fn should_become(before: &str, after: &str) {
        let mut x = before.to_string();
        trim(&mut x);
        assert_eq!(after.to_string(), x);
    }
    should_become("", "");
    should_become("\n", "");
    should_become("\n\n", "");
    should_become("asdf", "asdf");
    should_become("asdf\n", "asdf");
    should_become("asdf\nasdf\n", "asdf\nasdf");
}

#[test]
fn delete_word_test() {
    fn should_become(before: &str, after: &str) {
        let mut x = before.to_string();
        delete_last_word(&mut x);
        assert_eq!(after.to_string(), x);
    }
    should_become("", "");
    should_become("a", "");
    should_become("asdf", "");
    should_become("asdf asdf asdf", "asdf asdf ");
    should_become("asdf asdf asdf ", "asdf asdf ");
    should_become("asdf asdf asdf  ", "asdf asdf ");
}
<|endoftext|>


<commit_before>#![feature(path_ext)]
#![feature(collections)]

#[macro_use]
extern crate log;
extern crate git2;
extern crate iron;
extern crate toml;

use std::env;
use std::io::Read;
use std::fs::{self, PathExt, File};
use std::path::PathBuf;
use std::error::Error;

use git2::Repository;

use iron::prelude::*;

/// Contains various logger implementations for Grist.
///
/// TODO: log to a file in addition to the console
/// TODO: log to /var/log/
pub mod loggers;

/// Contains Iron server functions.
pub mod servers;

/// Contains code related to setting Grist's configuration.
///
/// Currently, configurations are set from a TOML configuration
/// file stored in `~/.gristconfig.toml`. Eventually, functionality
/// will be added to also allow configurations to be set from the
/// command-line, falling back to `~/.gristconfig.toml` and finally
/// to the defaults.
///
/// TODO: add command-line argument parsing
/// TODO: allow log file to be set from configs
pub mod config;

fn main() {

    // try to parse the config file at ~/.gristconfig.toml; otherwise, use defaults
    let configs = File::open("~/.gristconfig.toml")
        .and_then(|mut f: File| {
            let mut s = String::new();
            try!(f.read_to_string(&mut s));
            Ok(s) })
        .map_err(|err| String::from_str(err.description()) )
        .and_then(|s: String| config::parse_toml(s.as_ref()))
        .unwrap_or(config::DEFAULTS.clone());

    // set log level from configs
    if configs.debug {
        let _ = log::set_logger(|max_log_level| {
            max_log_level.set(log::LogLevelFilter::Debug);
            Box::new(loggers::DebugLogger)
        });
    } else {
        let _ = log::set_logger(|max_log_level| {
            max_log_level.set(log::LogLevelFilter::Info);
            Box::new(loggers::DefaultLogger)
        });
    };

    let roots = configs.roots.unwrap_or(
        // get the current working dir, assume that we are in a valid directory
        vec![env::current_dir().unwrap()]
    );

    info!("Starting Grist in {:?}", roots);

    let repos: Vec<Repository> = roots.iter().flat_map(|ref root: &PathBuf| {
        fs::read_dir(root) // walk the root dir
            .unwrap() // we are assuming the directory is valid (for now)
            .filter_map( |entry| match entry {
                Ok(ref dir) if dir.path().is_dir() => { // found a dir
                    debug!("Found directory: {:?}", dir.path());
                    match Repository::init(dir.path()) { // attempt to open dir as git repo
                        Ok(repo) => { debug!("Found repository: {:?}", dir.path()); Some(repo) },
                        Err(why) => { warn!("Failed to open {:?}: {}", dir.path(), why); None }
                    }
                },
                Err(why) => { warn!("Could not read entry: {}", why); None },
                Ok(ref e)=> { debug!("{:?} is not a directory.", e.path()); None }
            })
    }).collect();

    info!("Discovered {} repositories.", repos.len());

    // serve sample hello world page for now
    Iron::new(servers::hello_world).http("localhost:3000").unwrap();

}
<commit_msg>Set port from config file<commit_after>#![feature(path_ext)]
#![feature(collections)]

#[macro_use]
extern crate log;
extern crate git2;
extern crate iron;
extern crate toml;

use std::env;
use std::io::Read;
use std::fs::{self, PathExt, File};
use std::path::PathBuf;
use std::error::Error;

use git2::Repository;

use iron::prelude::*;

/// Contains various logger implementations for Grist.
///
/// TODO: log to a file in addition to the console
/// TODO: log to /var/log/
pub mod loggers;

/// Contains Iron server functions.
pub mod servers;

/// Contains code related to setting Grist's configuration.
///
/// Currently, configurations are set from a TOML configuration
/// file stored in `~/.gristconfig.toml`. Eventually, functionality
/// will be added to also allow configurations to be set from the
/// command-line, falling back to `~/.gristconfig.toml` and finally
/// to the defaults.
///
/// TODO: add command-line argument parsing
/// TODO: allow log file to be set from configs
pub mod config;

fn main() {

    // try to parse the config file at ~/.gristconfig.toml; otherwise, use defaults
    let configs = File::open("~/.gristconfig.toml")
        .and_then(|mut f: File| {
            let mut s = String::new();
            try!(f.read_to_string(&mut s));
            Ok(s) })
        .map_err(|err| String::from_str(err.description()) )
        .and_then(|s: String| config::parse_toml(s.as_ref()))
        .unwrap_or(config::DEFAULTS.clone());

    // set log level from configs
    if configs.debug {
        let _ = log::set_logger(|max_log_level| {
            max_log_level.set(log::LogLevelFilter::Debug);
            Box::new(loggers::DebugLogger)
        });
    } else {
        let _ = log::set_logger(|max_log_level| {
            max_log_level.set(log::LogLevelFilter::Info);
            Box::new(loggers::DefaultLogger)
        });
    };

    let roots = configs.roots.unwrap_or(
        // get the current working dir, assume that we are in a valid directory
        vec![env::current_dir().unwrap()]
    );

    info!("Searching {:?} for repositories", roots);

    let repos: Vec<Repository> = roots.iter().flat_map(|ref root: &PathBuf| {
        fs::read_dir(root) // walk the root dir
            .unwrap() // we are assuming the directory is valid (for now)
            .filter_map( |entry| match entry {
                Ok(ref dir) if dir.path().is_dir() => { // found a dir
                    debug!("Found directory: {:?}", dir.path());
                    match Repository::init(dir.path()) { // attempt to open dir as git repo
                        Ok(repo) => { debug!("Found repository: {:?}", dir.path()); Some(repo) },
                        Err(why) => { warn!("Failed to open {:?}: {}", dir.path(), why); None }
                    }
                },
                Err(why) => { warn!("Could not read entry: {}", why); None },
                Ok(ref e)=> { debug!("{:?} is not a directory.", e.path()); None }
            })
    }).collect();

    info!("Discovered {} repositories.", repos.len());

    // serve sample hello world page for now
    info!("Starting Grist on port {}", configs.port);
    Iron::new(servers::hello_world)
        .http(AsRef::<str>::as_ref(&format!("localhost:{}", configs.port)))
        .unwrap();
}
<|endoftext|>


<commit_before>extern crate alpm;
#[macro_use]
extern crate clap;
extern crate curl;
extern crate env_logger;
extern crate itertools;
#[macro_use]
extern crate log;
extern crate serde_json;
extern crate term;

use clap::App;
use curl::easy::Easy;
use itertools::Itertools;
use serde_json::Value;
use std::cmp::Ordering;
use std::collections::btree_map::Entry::{Occupied, Vacant};
use std::collections::BTreeMap;
use std::default::Default;
use std::process::exit;
use std::str;

mod avg;
mod enums;

const WEBSITE: &str = "https://security.archlinux.org";
const ROOT_DIR: &str = "/";
const DB_PATH: &str = "/var/lib/pacman/";

struct Options {
    format: Option<String>,
    quiet: u64,
    upgradable_only: bool,
}

impl Default for Options {
    fn default() -> Options {
        Options {
            format: None,
            quiet: 0,
            upgradable_only: false,
        }
    }
}

fn main() {
    env_logger::init();

    let yaml = load_yaml!("cli.yml");
    let args = App::from_yaml(yaml).get_matches();

    let options = Options {
        format: {
            match args.value_of("format") {
                Some(f) => Some(f.to_string()),
                None => None,
            }
        },
        quiet: args.occurrences_of("quiet"),
        upgradable_only: args.is_present("upgradable"),
    };

    let mut avgs = String::new();
    {
        info!("Downloading AVGs...");
        let avgs_url = format!("{}/issues/all.json", WEBSITE);

        let mut easy = Easy::new();
        easy.fail_on_error(true)
            .expect("curl::Easy::fail_on_error failed");
        easy.url(&avgs_url).expect("curl::Easy::url failed");
        let mut transfer = easy.transfer();
        transfer
            .write_function(|data| {
                avgs.push_str(str::from_utf8(data).expect("str conversion failed"));
                Ok(data.len())
            })
            .expect("write_function failed");
        match transfer.perform() {
            Ok(_) => {}
            Err(_) => {
                println!(
                    "Cannot fetch data from {}, please check your network connection!",
                    WEBSITE
                );
                exit(1)
            }
        };
    }

    let pacman = match args.value_of("dbpath") {
        Some(path) => {
            alpm::Alpm::new(ROOT_DIR, path).expect("alpm::Alpm::new with custom dbpath failed")
        }
        None => alpm::Alpm::new(ROOT_DIR, DB_PATH).expect("alpm::Alpm::new failed"),
    };
    let db = pacman.localdb();

    let mut cves: BTreeMap<String, Vec<_>> = BTreeMap::new();
    {
        let json: Value = serde_json::from_str(&avgs).expect("serde_json::from_str failed");

        for avg in json.as_array().expect("Value::as_array failed") {
            let packages = avg["packages"]
                .as_array()
                .expect("Value::as_array failed")
                .iter()
                .map(|s| s.as_str().expect("Value::as_str failed").to_string())
                .collect::<Vec<_>>();

            if !package_is_installed(&db, &packages) {
                continue;
            }

            let info = to_avg(avg);

            if info.status != enums::Status::NotAffected {
                for p in packages {
                    match cves.entry(p) {
                        Occupied(c) => c.into_mut(),
                        Vacant(c) => c.insert(vec![]),
                    }
                    .push(info.clone());
                }
            }
        }
    }

    let mut affected_avgs: BTreeMap<String, Vec<_>> = BTreeMap::new();
    for (pkg, avgs) in cves {
        for avg in &avgs {
            if system_is_affected(&db, &pkg, avg) {
                match affected_avgs.entry(pkg.clone()) {
                    Occupied(c) => c.into_mut(),
                    Vacant(c) => c.insert(vec![]),
                }
                .push(avg.clone());
            }
        }
    }

    let merged = merge_avgs(&affected_avgs);
    print_avgs(&options, &merged);
}

/// Converts a JSON to an `avg::AVG`
fn to_avg(data: &Value) -> avg::AVG {
    avg::AVG {
        issues: data["issues"]
            .as_array()
            .expect("Value::as_array failed")
            .iter()
            .map(|s| s.as_str().expect("Value::as_str failed").to_string())
            .collect(),
        fixed: match data["fixed"].as_str() {
            Some(s) => Some(s.to_string()),
            None => None,
        },
        severity: data["severity"]
            .as_str()
            .expect("Value::as_str failed")
            .to_string()
            .parse::<enums::Severity>()
            .expect("parse::<Severity> failed"),
        status: data["status"]
            .as_str()
            .expect("Value::as_str failed")
            .to_string()
            .parse::<enums::Status>()
            .expect("parse::<Status> failed"),
    }
}

#[test]
fn test_to_avg() {
    let json: Value = serde_json::from_str(
        "{\"issues\": [\"CVE-1\", \"CVE-2\"], \"fixed\": \"1.0\", \
         \"severity\": \"High\", \"status\": \"Not affected\"}",
    )
    .expect("serde_json::from_str failed");

    let avg1 = to_avg(&json);
    assert_eq!(2, avg1.issues.len());
    assert_eq!(Some("1.0".to_string()), avg1.fixed);
    assert_eq!(enums::Severity::High, avg1.severity);
    assert_eq!(enums::Status::NotAffected, avg1.status);

    let json: Value = serde_json::from_str(
        "{\"issues\": [\"CVE-1\"], \"fixed\": null, \
         \"severity\": \"Low\", \"status\": \"Vulnerable\"}",
    )
    .expect("serde_json::from_str failed");

    let avg2 = to_avg(&json);
    assert_eq!(1, avg2.issues.len());
    assert_eq!(None, avg2.fixed);
    assert_eq!(enums::Severity::Low, avg2.severity);
    assert_eq!(enums::Status::Vulnerable, avg2.status);
}

/// Given a package and an `avg::AVG`, returns true if the system is affected
fn system_is_affected(db: &alpm::Db, pkg: &str, avg: &avg::AVG) -> bool {
    match db.pkg(pkg.clone()) {
        Ok(v) => {
            info!(
                "Found installed version {} for package {}",
                v.version(),
                pkg
            );
            match avg.fixed {
                Some(ref version) => {
                    info!("Comparing with fixed version {}", version);
                    let cmp = alpm::vercmp(v.version().to_string(), version.clone());
                    if let Ordering::Less = cmp {
                        return true;
                    }
                }
                None => return true,
            };
        }
        Err(_) => debug!("Package {} not installed", pkg),
    }

    false
}

#[test]
fn test_system_is_affected() {
    let pacman = alpm::Alpm::new(ROOT_DIR, DB_PATH).expect("Alpm::new failed");
    let db = pacman.localdb();

    let avg1 = avg::AVG {
        issues: vec!["CVE-1".to_string(), "CVE-2".to_string()],
        fixed: Some("1.0.0".to_string()),
        severity: enums::Severity::Unknown,
        status: enums::Status::Unknown,
    };

    assert_eq!(false, system_is_affected(&db, &"pacman".to_string(), &avg1));

    let avg2 = avg::AVG {
        issues: vec!["CVE-1".to_string(), "CVE-2".to_string()],
        fixed: Some("7.0.0".to_string()),
        severity: enums::Severity::Unknown,
        status: enums::Status::Unknown,
    };

    assert!(system_is_affected(&db, &"pacman".to_string(), &avg2));
}

/// Given a list of package names, returns true when at least one is installed
fn package_is_installed(db: &alpm::Db, packages: &[String]) -> bool {
    for pkg in packages {
        match db.pkg(pkg.as_str()) {
            Ok(_) => {
                info!("Package {} is installed", pkg);
                return true;
            }
            Err(_) => debug!("Package {} not installed", pkg),
        }
    }
    false
}

#[test]
fn test_package_is_installed() {
    let pacman = alpm::Alpm::new(ROOT_DIR, DB_PATH).expect("Alpm::new failed");
    let db = pacman.localdb();

    let packages = vec!["pacman".to_string(), "pac".to_string()];
    assert!(package_is_installed(&db, &packages));

    let packages = vec!["pac".to_string()];
    assert_eq!(false, package_is_installed(&db, &packages));
}

/// Merge a list of `avg::AVG` into a single `avg::AVG` using major version as version
fn merge_avgs(cves: &BTreeMap<String, Vec<avg::AVG>>) -> BTreeMap<String, avg::AVG> {
    let mut avgs: BTreeMap<String, avg::AVG> = BTreeMap::new();
    for (pkg, list) in cves.iter() {
        let mut avg_issues = vec![];
        let mut avg_fixed: Option<String> = None;
        let mut avg_severity = enums::Severity::Unknown;
        let mut avg_status = enums::Status::Unknown;

        for a in list.iter() {
            avg_issues.append(&mut a.issues.clone());

            match avg_fixed.clone() {
                Some(ref version) => {
                    if let Some(ref v) = a.fixed {
                        let cmp = alpm::vercmp(version.to_string(), v.to_string());
                        if let Ordering::Greater = cmp {
                            avg_fixed = a.fixed.clone();
                        }
                    }
                }
                None => avg_fixed = a.fixed.clone(),
            }

            if a.severity > avg_severity {
                avg_severity = a.severity;
            }

            if a.status > avg_status {
                avg_status = a.status;
            }
        }

        let avg = avg::AVG {
            issues: avg_issues,
            fixed: avg_fixed,
            severity: avg_severity,
            status: avg_status,
        };
        avgs.insert(pkg.to_string(), avg);
    }

    avgs
}

#[test]
fn test_merge_avgs() {
    let mut avgs: BTreeMap<String, Vec<_>> = BTreeMap::new();

    let avg1 = avg::AVG {
        issues: vec!["CVE-1".to_string(), "CVE-2".to_string()],
        fixed: Some("1.0.0".to_string()),
        severity: enums::Severity::Unknown,
        status: enums::Status::Fixed,
    };

    let avg2 = avg::AVG {
        issues: vec!["CVE-4".to_string(), "CVE-10".to_string()],
        fixed: Some("0.9.8".to_string()),
        severity: enums::Severity::High,
        status: enums::Status::Testing,
    };

    assert!(enums::Severity::Critical > enums::Severity::High);

    avgs.insert("package".to_string(), vec![avg1.clone(), avg2.clone()]);

    avgs.insert("package2".to_string(), vec![avg1, avg2]);

    let merged = merge_avgs(&avgs);

    assert_eq!(2, merged.len());
    assert_eq!(
        4,
        merged
            .get(&"package".to_string())
            .expect("'package' key not found")
            .issues
            .len()
    );
    assert_eq!(
        enums::Severity::High,
        merged
            .get(&"package".to_string())
            .expect("'package' key not found")
            .severity
    );
    assert_eq!(
        enums::Status::Testing,
        merged
            .get(&"package".to_string())
            .expect("'package' key not found")
            .status
    );
}

/// Print a list of `avg::AVG`
fn print_avgs(options: &Options, avgs: &BTreeMap<String, avg::AVG>) {
    let mut t = term::stdout().expect("term::stdout failed");

    for (pkg, avg) in avgs {
        t.fg(avg.severity.to_color()).expect("term::fg failed");

        match avg.fixed {
            Some(ref v) if avg.status != enums::Status::Vulnerable => {
                if options.quiet >= 2 {
                    writeln!(t, "{}", pkg).expect("term::writeln failed");
                } else if options.quiet == 1 {
                    writeln!(t, "{}>={}", pkg, v).expect("term::writeln failed");
                } else {
                    t.fg(term::color::RED).expect("term::color::RED failed");
                    match options.format {
                        Some(ref f) => writeln!(
                            t,
                            "{}",
                            f.replace("%n", pkg.as_str())
                                .replace("%c", avg.issues.iter().join(",").as_str(),)
                        )
                        .expect("term::writeln failed"),
                        None => {
                            let msg = format!(
                                "Package {} is affected by {}. {}!",
                                pkg,
                                avg.issues.join(", "),
                                avg.severity
                            );

                            if avg.status == enums::Status::Testing {
                                writeln!(t, "{} Update to {} from the testing repos!", msg, v)
                                    .expect("term::writeln failed")
                            } else if avg.status == enums::Status::Fixed {
                                writeln!(t, "{} Update to {}!", msg, v)
                                    .expect("term::writeln failed")
                            } else {
                                writeln!(t, "{}", msg).expect("term::writeln failed")
                            }
                        }
                    }
                }
            }
            _ => {
                if !options.upgradable_only {
                    if options.quiet > 0 {
                        println!("{}", pkg);
                    } else {
                        match options.format {
                            Some(ref f) => println!(
                                "{}",
                                f.replace("%n", pkg.as_str())
                                    .replace("%c", avg.issues.iter().join(",").as_str(),)
                            ),
                            None => {
                                println!(
                                    "Package {} is affected by {}. {}!",
                                    pkg,
                                    avg.issues.join(", "),
                                    avg.severity
                                );
                            }
                        }
                    }
                }
            }
        }
    }
}
<commit_msg>Pull Request: Added bold text to package and color to CVE only.<commit_after>extern crate alpm;
#[macro_use]
extern crate clap;
extern crate curl;
extern crate env_logger;
extern crate itertools;
#[macro_use]
extern crate log;
extern crate serde_json;
extern crate term;

use clap::App;
use curl::easy::Easy;
use itertools::Itertools;
use serde_json::Value;
use std::cmp::Ordering;
use std::collections::btree_map::Entry::{Occupied, Vacant};
use std::collections::BTreeMap;
use std::default::Default;
use std::process::exit;
use std::str;

mod avg;
mod enums;

const WEBSITE: &str = "https://security.archlinux.org";
const ROOT_DIR: &str = "/";
const DB_PATH: &str = "/var/lib/pacman/";

struct Options {
    format: Option<String>,
    quiet: u64,
    upgradable_only: bool,
}

impl Default for Options {
    fn default() -> Options {
        Options {
            format: None,
            quiet: 0,
            upgradable_only: false,
        }
    }
}

fn main() {
    env_logger::init();

    let yaml = load_yaml!("cli.yml");
    let args = App::from_yaml(yaml).get_matches();

    let options = Options {
        format: {
            match args.value_of("format") {
                Some(f) => Some(f.to_string()),
                None => None,
            }
        },
        quiet: args.occurrences_of("quiet"),
        upgradable_only: args.is_present("upgradable"),
    };

    let mut avgs = String::new();
    {
        info!("Downloading AVGs...");
        let avgs_url = format!("{}/issues/all.json", WEBSITE);

        let mut easy = Easy::new();
        easy.fail_on_error(true)
            .expect("curl::Easy::fail_on_error failed");
        easy.url(&avgs_url).expect("curl::Easy::url failed");
        let mut transfer = easy.transfer();
        transfer
            .write_function(|data| {
                avgs.push_str(str::from_utf8(data).expect("str conversion failed"));
                Ok(data.len())
            })
            .expect("write_function failed");
        match transfer.perform() {
            Ok(_) => {}
            Err(_) => {
                println!(
                    "Cannot fetch data from {}, please check your network connection!",
                    WEBSITE
                );
                exit(1)
            }
        };
    }

    let pacman = match args.value_of("dbpath") {
        Some(path) => {
            alpm::Alpm::new(ROOT_DIR, path).expect("alpm::Alpm::new with custom dbpath failed")
        }
        None => alpm::Alpm::new(ROOT_DIR, DB_PATH).expect("alpm::Alpm::new failed"),
    };
    let db = pacman.localdb();

    let mut cves: BTreeMap<String, Vec<_>> = BTreeMap::new();
    {
        let json: Value = serde_json::from_str(&avgs).expect("serde_json::from_str failed");

        for avg in json.as_array().expect("Value::as_array failed") {
            let packages = avg["packages"]
                .as_array()
                .expect("Value::as_array failed")
                .iter()
                .map(|s| s.as_str().expect("Value::as_str failed").to_string())
                .collect::<Vec<_>>();

            if !package_is_installed(&db, &packages) {
                continue;
            }

            let info = to_avg(avg);

            if info.status != enums::Status::NotAffected {
                for p in packages {
                    match cves.entry(p) {
                        Occupied(c) => c.into_mut(),
                        Vacant(c) => c.insert(vec![]),
                    }
                    .push(info.clone());
                }
            }
        }
    }

    let mut affected_avgs: BTreeMap<String, Vec<_>> = BTreeMap::new();
    for (pkg, avgs) in cves {
        for avg in &avgs {
            if system_is_affected(&db, &pkg, avg) {
                match affected_avgs.entry(pkg.clone()) {
                    Occupied(c) => c.into_mut(),
                    Vacant(c) => c.insert(vec![]),
                }
                .push(avg.clone());
            }
        }
    }

    let merged = merge_avgs(&affected_avgs);
    print_avgs(&options, &merged);
}

/// Converts a JSON to an `avg::AVG`
fn to_avg(data: &Value) -> avg::AVG {
    avg::AVG {
        issues: data["issues"]
            .as_array()
            .expect("Value::as_array failed")
            .iter()
            .map(|s| s.as_str().expect("Value::as_str failed").to_string())
            .collect(),
        fixed: match data["fixed"].as_str() {
            Some(s) => Some(s.to_string()),
            None => None,
        },
        severity: data["severity"]
            .as_str()
            .expect("Value::as_str failed")
            .to_string()
            .parse::<enums::Severity>()
            .expect("parse::<Severity> failed"),
        status: data["status"]
            .as_str()
            .expect("Value::as_str failed")
            .to_string()
            .parse::<enums::Status>()
            .expect("parse::<Status> failed"),
    }
}

#[test]
fn test_to_avg() {
    let json: Value = serde_json::from_str(
        "{\"issues\": [\"CVE-1\", \"CVE-2\"], \"fixed\": \"1.0\", \
         \"severity\": \"High\", \"status\": \"Not affected\"}",
    )
    .expect("serde_json::from_str failed");

    let avg1 = to_avg(&json);
    assert_eq!(2, avg1.issues.len());
    assert_eq!(Some("1.0".to_string()), avg1.fixed);
    assert_eq!(enums::Severity::High, avg1.severity);
    assert_eq!(enums::Status::NotAffected, avg1.status);

    let json: Value = serde_json::from_str(
        "{\"issues\": [\"CVE-1\"], \"fixed\": null, \
         \"severity\": \"Low\", \"status\": \"Vulnerable\"}",
    )
    .expect("serde_json::from_str failed");

    let avg2 = to_avg(&json);
    assert_eq!(1, avg2.issues.len());
    assert_eq!(None, avg2.fixed);
    assert_eq!(enums::Severity::Low, avg2.severity);
    assert_eq!(enums::Status::Vulnerable, avg2.status);
}

/// Given a package and an `avg::AVG`, returns true if the system is affected
fn system_is_affected(db: &alpm::Db, pkg: &str, avg: &avg::AVG) -> bool {
    match db.pkg(pkg.clone()) {
        Ok(v) => {
            info!(
                "Found installed version {} for package {}",
                v.version(),
                pkg
            );
            match avg.fixed {
                Some(ref version) => {
                    info!("Comparing with fixed version {}", version);
                    let cmp = alpm::vercmp(v.version().to_string(), version.clone());
                    if let Ordering::Less = cmp {
                        return true;
                    }
                }
                None => return true,
            };
        }
        Err(_) => debug!("Package {} not installed", pkg),
    }

    false
}

#[test]
fn test_system_is_affected() {
    let pacman = alpm::Alpm::new(ROOT_DIR, DB_PATH).expect("Alpm::new failed");
    let db = pacman.localdb();

    let avg1 = avg::AVG {
        issues: vec!["CVE-1".to_string(), "CVE-2".to_string()],
        fixed: Some("1.0.0".to_string()),
        severity: enums::Severity::Unknown,
        status: enums::Status::Unknown,
    };

    assert_eq!(false, system_is_affected(&db, &"pacman".to_string(), &avg1));

    let avg2 = avg::AVG {
        issues: vec!["CVE-1".to_string(), "CVE-2".to_string()],
        fixed: Some("7.0.0".to_string()),
        severity: enums::Severity::Unknown,
        status: enums::Status::Unknown,
    };

    assert!(system_is_affected(&db, &"pacman".to_string(), &avg2));
}

/// Given a list of package names, returns true when at least one is installed
fn package_is_installed(db: &alpm::Db, packages: &[String]) -> bool {
    for pkg in packages {
        match db.pkg(pkg.as_str()) {
            Ok(_) => {
                info!("Package {} is installed", pkg);
                return true;
            }
            Err(_) => debug!("Package {} not installed", pkg),
        }
    }
    false
}

#[test]
fn test_package_is_installed() {
    let pacman = alpm::Alpm::new(ROOT_DIR, DB_PATH).expect("Alpm::new failed");
    let db = pacman.localdb();

    let packages = vec!["pacman".to_string(), "pac".to_string()];
    assert!(package_is_installed(&db, &packages));

    let packages = vec!["pac".to_string()];
    assert_eq!(false, package_is_installed(&db, &packages));
}

/// Merge a list of `avg::AVG` into a single `avg::AVG` using major version as version
fn merge_avgs(cves: &BTreeMap<String, Vec<avg::AVG>>) -> BTreeMap<String, avg::AVG> {
    let mut avgs: BTreeMap<String, avg::AVG> = BTreeMap::new();
    for (pkg, list) in cves.iter() {
        let mut avg_issues = vec![];
        let mut avg_fixed: Option<String> = None;
        let mut avg_severity = enums::Severity::Unknown;
        let mut avg_status = enums::Status::Unknown;

        for a in list.iter() {
            avg_issues.append(&mut a.issues.clone());

            match avg_fixed.clone() {
                Some(ref version) => {
                    if let Some(ref v) = a.fixed {
                        let cmp = alpm::vercmp(version.to_string(), v.to_string());
                        if let Ordering::Greater = cmp {
                            avg_fixed = a.fixed.clone();
                        }
                    }
                }
                None => avg_fixed = a.fixed.clone(),
            }

            if a.severity > avg_severity {
                avg_severity = a.severity;
            }

            if a.status > avg_status {
                avg_status = a.status;
            }
        }

        let avg = avg::AVG {
            issues: avg_issues,
            fixed: avg_fixed,
            severity: avg_severity,
            status: avg_status,
        };
        avgs.insert(pkg.to_string(), avg);
    }

    avgs
}

#[test]
fn test_merge_avgs() {
    let mut avgs: BTreeMap<String, Vec<_>> = BTreeMap::new();

    let avg1 = avg::AVG {
        issues: vec!["CVE-1".to_string(), "CVE-2".to_string()],
        fixed: Some("1.0.0".to_string()),
        severity: enums::Severity::Unknown,
        status: enums::Status::Fixed,
    };

    let avg2 = avg::AVG {
        issues: vec!["CVE-4".to_string(), "CVE-10".to_string()],
        fixed: Some("0.9.8".to_string()),
        severity: enums::Severity::High,
        status: enums::Status::Testing,
    };

    assert!(enums::Severity::Critical > enums::Severity::High);

    avgs.insert("package".to_string(), vec![avg1.clone(), avg2.clone()]);

    avgs.insert("package2".to_string(), vec![avg1, avg2]);

    let merged = merge_avgs(&avgs);

    assert_eq!(2, merged.len());
    assert_eq!(
        4,
        merged
            .get(&"package".to_string())
            .expect("'package' key not found")
            .issues
            .len()
    );
    assert_eq!(
        enums::Severity::High,
        merged
            .get(&"package".to_string())
            .expect("'package' key not found")
            .severity
    );
    assert_eq!(
        enums::Status::Testing,
        merged
            .get(&"package".to_string())
            .expect("'package' key not found")
            .status
    );
}

/// Print a list of `avg::AVG`
fn print_avgs(options: &Options, avgs: &BTreeMap<String, avg::AVG>) {
    let mut t = term::stdout().expect("term::stdout failed");
    for (pkg, avg) in avgs {
        match avg.fixed {
            Some(ref v) if avg.status != enums::Status::Vulnerable => {
                // Quiet option
                if options.quiet >= 2 {
                    t.fg(avg.severity.to_color()).expect("term::fg failed");
                    writeln!(t, "{}", pkg).expect("term::writeln failed");
                    t.reset().expect("term::stdout failed");
                } else if options.quiet == 1 {
                    t.fg(avg.severity.to_color()).expect("term::fg failed");
                    writeln!(t, "{}>={}", pkg, v).expect("term::writeln failed");
                    t.reset().expect("term::stdout failed");
                } else {
                    match options.format {
                        Some(ref f) => {
                            t.fg(term::color::RED).expect("term::color::RED failed");
                            writeln!(
                                t,
                                "{}",
                                f.replace("%n", pkg.as_str())
                                    .replace("%c", avg.issues.iter().join(",").as_str(),)
                            )
                            .expect("term::writeln failed");
                            t.reset().expect("term::stdout failed");
                        }
                        None => {
                            print_avg_colored(&mut t, pkg, avg);
                            // Colored update
                            if avg.status == enums::Status::Testing {
                                // Print: Update to {} from the testing repos!"
                                write!(t, " Update to").expect("term::write failed");
                                t.attr(term::Attr::Bold).expect("term::attr failed");
                                t.fg(term::color::GREEN).expect("term::fg failed");
                                write!(t, " {}", v).expect("term::write failed");
                                t.reset().expect("term::stdout failed");
                                writeln!(t, " from the testing repos!")
                                    .expect("term::writeln failed");
                            } else if avg.status == enums::Status::Fixed {
                                // Print: Update to {}!
                                write!(t, " Update to").expect("term::write failed");
                                t.attr(term::Attr::Bold).expect("term::attr failed");
                                t.fg(term::color::GREEN).expect("term::fg failed");
                                write!(t, " {}", v).expect("term::write failed");
                                t.reset().expect("term::stdout failed");
                                writeln!(t, "!").expect("term::writeln failed");
                            }
                        }
                    }
                }
            }
            _ => {
                if !options.upgradable_only {
                    if options.quiet > 0 {
                        t.fg(avg.severity.to_color()).expect("term::fg failed");
                        writeln!(t, "{}", pkg).expect("term::writeln failed");
                        t.reset().expect("term::stdout failed");
                    } else {
                        match options.format {
                            Some(ref f) => {
                                t.fg(avg.severity.to_color()).expect("term::fg failed");
                                writeln!(
                                    t,
                                    "{}",
                                    f.replace("%n", pkg.as_str())
                                        .replace("%c", avg.issues.iter().join(",").as_str(),)
                                )
                                .expect("term::writeln failed");
                                t.reset().expect("term::stdout failed");
                            }
                            None => {
                                print_avg_colored(&mut t, pkg, avg);
                                writeln!(t).expect("term::writeln failed");
                            }
                        }
                    }
                }
            }
        }
    }
}

/// Prints "Package {pkg} is affected by {issues}. {severity}!" colored
fn print_avg_colored(t: &mut Box<term::StdoutTerminal>, pkg: &String, avg: &avg::AVG) {
    // Bold package
    write!(t, "Package").expect("term::write failed");
    t.attr(term::Attr::Bold).expect("term::attr failed");
    write!(t, " {}", pkg).expect("term::writeln failed");
    // Normal "is affected by {issues}"
    t.reset().expect("term::stdout failed");
    write!(t, " is affected by {}.", avg.issues.join(", ")).expect("term::write failed");
    // Colored severit
    t.fg(avg.severity.to_color()).expect("term::fg failed");
    write!(t, " {}!", avg.severity).expect("term::write failed");
    t.reset().expect("term::stdout failed");
}
<|endoftext|>