<commit_before>extern crate camera_controllers;
extern crate collada;
extern crate dev_menu;
extern crate env_logger;
extern crate gfx;
extern crate gfx_debug_draw;
extern crate gfx_device_gl;
extern crate gfx_gl as gl;
extern crate piston;
extern crate piston_window;
extern crate sdl2;
extern crate sdl2_window;
extern crate shader_version;
extern crate skeletal_animation;
extern crate vecmath;

use std::rc::Rc;
use std::cell::RefCell;
use std::collections::HashMap;

use gfx_debug_draw::DebugRenderer;

use piston::window::{
    WindowSettings,
};

use piston::event::*;

use vecmath::{mat4_id};

use sdl2_window::Sdl2Window;

use camera_controllers::{
    OrbitZoomCamera,
    OrbitZoomCameraSettings,
    CameraPerspective,
    model_view_projection
};

mod demo;
use demo::Settings;

fn main() {

    env_logger::init().unwrap();

    let (win_width, win_height) = (640, 480);
    let window = Rc::new(RefCell::new(Sdl2Window::new(
        shader_version::OpenGL::_3_2,
        WindowSettings::new(
            "Skeletal Animation Demo".to_string(),
            piston::window::Size { width: 640, height: 480 }
        ).exit_on_esc(true)
    )));

    let piston_window = piston_window::PistonWindow::new(window, piston_window::empty_app());

    let mut debug_renderer = DebugRenderer::from_canvas(
        &mut piston_window.canvas.borrow_mut(),
        64,
        None,
        None,
    ).ok().unwrap();

    let model = mat4_id();
    let mut projection = CameraPerspective {
        fov: 90.0f32,
        near_clip: 0.1,
        far_clip: 1000.0,
        aspect_ratio: (win_width as f32) / (win_height as f32)
    }.projection();

    let mut orbit_zoom_camera: OrbitZoomCamera<f32> = OrbitZoomCamera::new(
        [0.0, 0.0, 0.0],
        OrbitZoomCameraSettings::default()
    );

    // Start event loop

    let mut settings = Settings {

        use_dlb: true,
        draw_skeleton: true,
        draw_labels: false,
        draw_mesh: true,
        playback_speed: 1.0,

        params: HashMap::new(),
    };

    let mut menu = dev_menu::Menu::<Settings>::new();

    menu.add_item(dev_menu::MenuItem::action_item(
        "Toggle DLB/LBS Skinning",
        Box::new( |ref mut settings| {
            settings.use_dlb = !settings.use_dlb;
        })
    ));

    menu.add_item(dev_menu::MenuItem::action_item(
        "Toggle Skeleton",
        Box::new( |ref mut settings| { settings.draw_skeleton = !settings.draw_skeleton; })
    ));

    menu.add_item(dev_menu::MenuItem::action_item(
        "Toggle Joint Labels",
        Box::new( |ref mut settings| { settings.draw_labels = !settings.draw_labels; })
    ));

    menu.add_item(dev_menu::MenuItem::action_item(
        "Toggle Mesh",
        Box::new( |ref mut settings| { settings.draw_mesh = !settings.draw_mesh; })
    ));

    menu.add_item(dev_menu::MenuItem::slider_item(
        "Playback Speed = ",
        [-5.0, 5.0],
        0.01,
        Box::new( |ref settings| { settings.playback_speed }),
        Box::new( |ref mut settings, value| { settings.playback_speed = value }),
    ));

    let mut lbs_demo = demo::dlb_demo(&mut piston_window.canvas.borrow_mut());
    let mut dlb_demo = demo::lbs_demo(&mut piston_window.canvas.borrow_mut());

    for (param, &value) in dlb_demo.controller.get_parameters().iter() {
        settings.params.insert(param.clone(), value);

        // Apparently need to make our own string copies to move into each closure..
        let param_copy_1 = param.clone();
        let param_copy_2 = param.clone();

        menu.add_item(dev_menu::MenuItem::slider_item(
            &format!("Param[{}] = ", param)[..],
            [0.0, 1.0],
            0.01,
            Box::new( move |ref settings| {
                settings.params[&param_copy_1[..]]
            }),
            Box::new( move |ref mut settings, value| {
                settings.params.insert(param_copy_2.clone(), value);
            }),
        ));
    }

    // set head look controller params to nice initial values..
    settings.params.insert("head-look-level".to_string(), 1.0);
    settings.params.insert("head-look-sideways-level".to_string(), 1.0);
    settings.params.insert("head-down-to-up".to_string(), 0.5);
    settings.params.insert("head-left-to-right".to_string(), 0.5);

    for e in piston_window {

        orbit_zoom_camera.event(&e);
        menu.event(&e, &mut settings);

        e.resize(|width, height| {
            // Update projection matrix
            projection = CameraPerspective {
                fov: 90.0f32,
                near_clip: 0.1,
                far_clip: 1000.0,
                aspect_ratio: (width as f32) / (height as f32)
            }.projection();
        });

        e.update(|args| {
            dlb_demo.update(&settings, args.dt);
            lbs_demo.update(&settings, args.dt);
        });

        e.draw_3d(|canvas| {

            let args = e.render_args().unwrap();

            canvas.renderer.clear(
                gfx::ClearData {
                    color: [0.3, 0.3, 0.3, 1.0],
                    depth: 1.0,
                    stencil: 0,
                },
                gfx::COLOR | gfx::DEPTH,
                &canvas.output
            );

            let camera_view = orbit_zoom_camera.camera(args.ext_dt).orthogonal();

            let camera_projection = model_view_projection(
                model,
                camera_view,
                projection
            );

            // Draw axes
            debug_renderer.draw_line([0.0, 0.0, 0.0], [5.0, 0.0, 0.0], [1.0, 0.0, 0.0, 1.0]);
            debug_renderer.draw_line([0.0, 0.0, 0.0], [0.0, 5.0, 0.0], [0.0, 1.0, 0.0, 1.0]);
            debug_renderer.draw_line([0.0, 0.0, 0.0], [0.0, 0.0, 5.0], [0.0, 0.0, 1.0, 1.0]);

            debug_renderer.draw_text_at_position(
                "X",
                [6.0, 0.0, 0.0],
                [1.0, 0.0, 0.0, 1.0],
            );

            debug_renderer.draw_text_at_position(
                "Y",
                [0.0, 6.0, 0.0],
                [0.0, 1.0, 0.0, 1.0],
            );

            debug_renderer.draw_text_at_position(
                "Z",
                [0.0, 0.0, 6.0],
                [0.0, 0.0, 1.0, 1.0],
            );

            dlb_demo.render(&settings, &mut debug_renderer, canvas, camera_view, camera_projection, args.ext_dt, settings.use_dlb);

            lbs_demo.render(&settings, &mut debug_renderer, canvas, camera_view, camera_projection, args.ext_dt, !settings.use_dlb);

            menu.draw(&settings, &mut debug_renderer);

            debug_renderer.render_canvas(canvas, camera_projection);
        });
    }
}
<commit_msg>fixed lbs/dlb mixup<commit_after>extern crate camera_controllers;
extern crate collada;
extern crate dev_menu;
extern crate env_logger;
extern crate gfx;
extern crate gfx_debug_draw;
extern crate gfx_device_gl;
extern crate gfx_gl as gl;
extern crate piston;
extern crate piston_window;
extern crate sdl2;
extern crate sdl2_window;
extern crate shader_version;
extern crate skeletal_animation;
extern crate vecmath;

use std::rc::Rc;
use std::cell::RefCell;
use std::collections::HashMap;

use gfx_debug_draw::DebugRenderer;

use piston::window::{
    WindowSettings,
};

use piston::event::*;

use vecmath::{mat4_id};

use sdl2_window::Sdl2Window;

use camera_controllers::{
    OrbitZoomCamera,
    OrbitZoomCameraSettings,
    CameraPerspective,
    model_view_projection
};

mod demo;
use demo::Settings;

fn main() {

    env_logger::init().unwrap();

    let (win_width, win_height) = (640, 480);
    let window = Rc::new(RefCell::new(Sdl2Window::new(
        shader_version::OpenGL::_3_2,
        WindowSettings::new(
            "Skeletal Animation Demo".to_string(),
            piston::window::Size { width: 640, height: 480 }
        ).exit_on_esc(true)
    )));

    let piston_window = piston_window::PistonWindow::new(window, piston_window::empty_app());

    let mut debug_renderer = DebugRenderer::from_canvas(
        &mut piston_window.canvas.borrow_mut(),
        64,
        None,
        None,
    ).ok().unwrap();

    let model = mat4_id();
    let mut projection = CameraPerspective {
        fov: 90.0f32,
        near_clip: 0.1,
        far_clip: 1000.0,
        aspect_ratio: (win_width as f32) / (win_height as f32)
    }.projection();

    let mut orbit_zoom_camera: OrbitZoomCamera<f32> = OrbitZoomCamera::new(
        [0.0, 0.0, 0.0],
        OrbitZoomCameraSettings::default()
    );

    // Start event loop

    let mut settings = Settings {

        use_dlb: true,
        draw_skeleton: true,
        draw_labels: false,
        draw_mesh: true,
        playback_speed: 1.0,

        params: HashMap::new(),
    };

    let mut menu = dev_menu::Menu::<Settings>::new();

    menu.add_item(dev_menu::MenuItem::action_item(
        "Toggle DLB/LBS Skinning",
        Box::new( |ref mut settings| {
            settings.use_dlb = !settings.use_dlb;
        })
    ));

    menu.add_item(dev_menu::MenuItem::action_item(
        "Toggle Skeleton",
        Box::new( |ref mut settings| { settings.draw_skeleton = !settings.draw_skeleton; })
    ));

    menu.add_item(dev_menu::MenuItem::action_item(
        "Toggle Joint Labels",
        Box::new( |ref mut settings| { settings.draw_labels = !settings.draw_labels; })
    ));

    menu.add_item(dev_menu::MenuItem::action_item(
        "Toggle Mesh",
        Box::new( |ref mut settings| { settings.draw_mesh = !settings.draw_mesh; })
    ));

    menu.add_item(dev_menu::MenuItem::slider_item(
        "Playback Speed = ",
        [-5.0, 5.0],
        0.01,
        Box::new( |ref settings| { settings.playback_speed }),
        Box::new( |ref mut settings, value| { settings.playback_speed = value }),
    ));

    let mut lbs_demo = demo::lbs_demo(&mut piston_window.canvas.borrow_mut());
    let mut dlb_demo = demo::dlb_demo(&mut piston_window.canvas.borrow_mut());

    for (param, &value) in dlb_demo.controller.get_parameters().iter() {
        settings.params.insert(param.clone(), value);

        // Apparently need to make our own string copies to move into each closure..
        let param_copy_1 = param.clone();
        let param_copy_2 = param.clone();

        menu.add_item(dev_menu::MenuItem::slider_item(
            &format!("Param[{}] = ", param)[..],
            [0.0, 1.0],
            0.01,
            Box::new( move |ref settings| {
                settings.params[&param_copy_1[..]]
            }),
            Box::new( move |ref mut settings, value| {
                settings.params.insert(param_copy_2.clone(), value);
            }),
        ));
    }

    // set head look controller params to nice initial values..
    settings.params.insert("head-look-level".to_string(), 1.0);
    settings.params.insert("head-look-sideways-level".to_string(), 1.0);
    settings.params.insert("head-down-to-up".to_string(), 0.5);
    settings.params.insert("head-left-to-right".to_string(), 0.5);

    for e in piston_window {

        orbit_zoom_camera.event(&e);
        menu.event(&e, &mut settings);

        e.resize(|width, height| {
            // Update projection matrix
            projection = CameraPerspective {
                fov: 90.0f32,
                near_clip: 0.1,
                far_clip: 1000.0,
                aspect_ratio: (width as f32) / (height as f32)
            }.projection();
        });

        e.update(|args| {
            dlb_demo.update(&settings, args.dt);
            lbs_demo.update(&settings, args.dt);
        });

        e.draw_3d(|canvas| {

            let args = e.render_args().unwrap();

            canvas.renderer.clear(
                gfx::ClearData {
                    color: [0.3, 0.3, 0.3, 1.0],
                    depth: 1.0,
                    stencil: 0,
                },
                gfx::COLOR | gfx::DEPTH,
                &canvas.output
            );

            let camera_view = orbit_zoom_camera.camera(args.ext_dt).orthogonal();

            let camera_projection = model_view_projection(
                model,
                camera_view,
                projection
            );

            // Draw axes
            debug_renderer.draw_line([0.0, 0.0, 0.0], [5.0, 0.0, 0.0], [1.0, 0.0, 0.0, 1.0]);
            debug_renderer.draw_line([0.0, 0.0, 0.0], [0.0, 5.0, 0.0], [0.0, 1.0, 0.0, 1.0]);
            debug_renderer.draw_line([0.0, 0.0, 0.0], [0.0, 0.0, 5.0], [0.0, 0.0, 1.0, 1.0]);

            debug_renderer.draw_text_at_position(
                "X",
                [6.0, 0.0, 0.0],
                [1.0, 0.0, 0.0, 1.0],
            );

            debug_renderer.draw_text_at_position(
                "Y",
                [0.0, 6.0, 0.0],
                [0.0, 1.0, 0.0, 1.0],
            );

            debug_renderer.draw_text_at_position(
                "Z",
                [0.0, 0.0, 6.0],
                [0.0, 0.0, 1.0, 1.0],
            );

            dlb_demo.render(&settings, &mut debug_renderer, canvas, camera_view, camera_projection, args.ext_dt, settings.use_dlb);

            lbs_demo.render(&settings, &mut debug_renderer, canvas, camera_view, camera_projection, args.ext_dt, !settings.use_dlb);

            menu.draw(&settings, &mut debug_renderer);

            debug_renderer.render_canvas(canvas, camera_projection);
        });
    }
}
<|endoftext|>


<commit_before>#![recursion_limit = "1024"] // error_chain

extern crate websocket;
extern crate env_logger;
#[macro_use]
extern crate log;
#[macro_use]
extern crate error_chain;
extern crate url;
extern crate clap;

const BUFSIZ : usize = 8192;

use std::thread;
use std::io::{stdin,stdout};

use websocket::{Message, Sender, Receiver, DataFrame, Server as WsServer};
use websocket::message::Type;
use websocket::client::request::Url;
use websocket::Client;

use std::borrow::Borrow;
use std::io::{Error as IoError, ErrorKind as IoErrorKind, Write, Read};

error_chain! {
    foreign_links {
        ::std::io::Error, Io;
        log::SetLoggerError, Log;
        ::url::ParseError, Url;
        ::websocket::result::WebSocketError, Ws;
        ::std::env::VarError, Ev;
    }
    errors {
        InvalidSpecifier(t : String) {
            description("invalid specifier")
            display("Invalid client or server specifier `{}`", t)
        }
    }
}

// Initialize logger with default "info" log level:
fn init_logger() -> Result<()> {
    let mut builder = env_logger::LogBuilder::new();
    builder.filter(None, log::LogLevelFilter::Info);
    if ::std::env::var("RUST_LOG").is_ok() {
       builder.parse(&::std::env::var("RUST_LOG")?);
    }
    builder.init()?;
    Ok(())
}

struct SenderWrapper<T: Sender> (T);

impl<T: Sender> ::std::io::Write for SenderWrapper<T> {
    fn write(&mut self, buf: &[u8]) -> ::std::io::Result<usize> {
        let message = Message::binary(buf);
        let ret;
        let len = buf.len();
        if len > 0 {
            debug!("Sending message of {} bytes", len);
            ret = self.0.send_message(&message);
        } else {
            // Interpret zero length buffer is request
            // to close communication

            debug!("Sending the closing message");
            ret = self.0.send_message(&Message::close());
        }
        ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;
        Ok(len)
    }
    fn flush(&mut self) -> ::std::io::Result<()> {
        Ok(())
    }
}

struct ReceiverWrapper<T: Receiver<DataFrame>> (T);

impl<T:Receiver<DataFrame>> ::std::io::Read for ReceiverWrapper<T> {
    fn read(&mut self, buf: &mut [u8]) -> ::std::io::Result<usize> {
        let ret = self.0.recv_message();
        let msg : Message = ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;

        match msg.opcode {
            Type::Close => {
                Ok(0)
            }
            Type::Ping => {
                // Sender used to be in a separate thread with a channel
                // now there's no channel, so trickier to combine ping replies
                // and usual data exchange
                unimplemented!();
            }
            _ => {
                let msgpayload : &[u8] = msg.payload.borrow();
                let len = msgpayload.len();
                debug!("Received message of {} bytes", len);

                assert!(buf.len() >= len);

                buf[0..len].copy_from_slice(msgpayload);

                Ok(len)
            }
        }
    }
}

struct Peer<R, W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    reader: R,
    writer: W,
}

type IPeer = Peer<Box<Read+Send>, Box<Write+Send>>;

struct DataExchangeSession<R1, R2, W1, W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    peer1: Peer<R1, W1>,
    peer2: Peer<R2, W2>,
}

// Derived from https://doc.rust-lang.org/src/std/up/src/libstd/io/util.rs.html#46-61
pub fn copy_with_flushes<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W) -> ::std::io::Result<u64>
    where R: Read, W: Write
{
    let mut buf = [0; BUFSIZ];
    let mut written = 0;
    loop {
        let len = match reader.read(&mut buf) {
            Ok(0) => return Ok(written),
            Ok(len) => len,
            Err(ref e) if e.kind() == IoErrorKind::Interrupted => continue,
            Err(ref e) if e.kind() == IoErrorKind::WouldBlock => continue,
            Err(e) => return Err(e),
        };
        writer.write_all(&buf[..len])?;
        writer.flush()?;
        written += len as u64;
    }
}

impl<R1,R2,W1,W2> DataExchangeSession<R1,R2,W1,W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    fn data_exchange(self) -> Result<()> {

        let mut reader1 = self.peer1.reader;
        let mut writer1 = self.peer1.writer;
        let mut reader2 = self.peer2.reader;
        let mut writer2 = self.peer2.writer;

        let receive_loop = thread::Builder::new().spawn(move || -> Result<()> {
            // Actual data transfer happens here
            copy_with_flushes(&mut reader1, &mut writer2)?;
            writer2.write(b"")?; // signal close
            Ok(())
        })?;

        // Actual data transfer happens here
        copy_with_flushes(&mut reader2, &mut writer1)?;
        writer1.write(b"")?; // Signal close

        debug!("Waiting for receiver side to exit");

        receive_loop.join().map_err(|x|format!("{:?}",x))?
    }
}

fn get_websocket_peer(urlstr: &str) -> Result<
        Peer<
            ReceiverWrapper<websocket::client::Receiver<websocket::WebSocketStream>>,
            SenderWrapper<websocket::client::Sender<websocket::WebSocketStream>>>
        > {
    let url = Url::parse(urlstr)?;

    info!("Connecting to {}", url);

    let request = Client::connect(url)?;

    let response = request.send()?; // Send the request and retrieve a response

    info!("Validating response...");

    response.validate()?; // Validate the response

    info!("Successfully connected");

    let (sender, receiver) = response.begin().split();

    let peer = Peer {
        reader : ReceiverWrapper(receiver),
        writer : SenderWrapper(sender),
    };
    Ok(peer)
}

fn get_tcp_peer(addr: &str) -> Result<
        Peer<
            ::std::net::TcpStream,
            ::std::net::TcpStream,
        >> {
    let sock = ::std::net::TcpStream::connect(addr)?;

    let peer = Peer {
        reader : sock.try_clone()?,
        writer : sock.try_clone()?,
    };
    info!("Connected to TCP {}", addr);
    Ok(peer)
}

fn get_stdio_peer() -> Result<Peer<std::io::Stdin, std::io::Stdout>> {
    Ok(
        Peer {
            reader : stdin(),
            writer : stdout(),
        }
    )
}

struct TcpServer(::std::net::TcpListener);

impl TcpServer {
    fn new(addr: &str) -> Result<Self> {
        Ok(TcpServer(::std::net::TcpListener::bind(addr)?))
    }
}

impl Server for TcpServer {
    fn accept_client(&mut self) -> Result<IPeer> {
        let (sock, addr) = self.0.accept()?;
        info!("TCP client connection from {}", addr);
        let peer = Peer {
            reader : sock.try_clone()?,
            writer : sock.try_clone()?,
        };
        Ok(peer.upcast())
    }
}

struct WebsockServer<'a>(WsServer<'a>);

impl<'a> WebsockServer<'a> {
    fn new(addr: &str) -> Result<Self> {
        Ok(WebsockServer(WsServer::bind(addr)?))
    }
}

impl<'a> Server for WebsockServer<'a> {
    fn accept_client(&mut self) -> Result<IPeer> {
        let connection = self.0.accept()?;
        info!("WebSocket client connection ...");
        let request = connection.read_request()?;
        request.validate()?;
        let response = request.accept(); // Form a response
        let mut client = response.send()?; // Send the response

        let ip = client.get_mut_sender()
            .get_mut()
            .peer_addr()
            .unwrap();

        info!("... from IP {}", ip);

        let (sender, receiver) = client.split();

        let peer = Peer {
            reader : ReceiverWrapper(receiver),
            writer : SenderWrapper(sender),
        };
        Ok(peer.upcast())
    }
}

impl<R,W> Peer<R,W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    fn upcast(self) -> IPeer  {
        Peer {
            reader: Box::new(self.reader) as Box<Read +Send>,
            writer: Box::new(self.writer) as Box<Write+Send>,
        }
    }
}

trait Server
{
    fn accept_client(&mut self) -> Result<IPeer>;

    fn start_serving(&mut self, spec2: &str, once: bool) -> Result<()> {
        let spec2s = spec2.to_string();
        let closure = move |peer, spec2 : String|{
            let spec2_ = get_peer_by_spec(spec2.as_str())?;
            let peer2 = match spec2_ {
                Spec::Server(mut x) => {
                    x.accept_client()?
                }
                Spec::Client(p1) => {
                    p1
                }
            };
            let des = DataExchangeSession {
                peer1 : peer,
                peer2 : peer2,
            };

            des.data_exchange()
        };
        if once {
            let peer = self.accept_client()?;
            closure(peer, spec2s)
        } else {
            let cl2 = ::std::sync::Arc::new(closure);
            loop {
                let ret = self.accept_client();
                let peer = match ret {
                    Ok(x) => x,
                    Err(er) => {
                        warn!("Can't accept client: {}", er);
                        continue;
                    }
                };
                let cl3 = cl2.clone();
                let spec2s2 = spec2s.clone();
                if let Err(x) = thread::Builder::new().spawn(move|| {
                    if let Err(x) = cl3(peer, spec2s2) {
                        warn!("Error while serving: {}", x);
                    }
                }) {
                    warn!("Error creating thread: {}", x);
                    thread::sleep(::std::time::Duration::from_millis(200));
                }
            }
        }
    }

    fn upcast(self) -> Box<Server+Send>
        where Self : Sized + Send + 'static
        { Box::new(self) as Box<Server+Send> }
}

fn main2(spec1: &str, spec2: &str, once: bool) -> Result<()> {
    let spec1_ = get_peer_by_spec(spec1)?;

    match spec1_ {
        Spec::Server(mut x) => {
            x.start_serving(spec2, once)
        }
        Spec::Client(p1) => {
            let spec2_ = get_peer_by_spec(spec2)?;

            let otherpeer = match spec2_ {
                Spec::Server(mut x) => {
                    let t = x.accept_client()?;
                    t
                }
                Spec::Client(p2) => {
                    p2
                }
            };

            let des = DataExchangeSession {
                peer1 : p1,
                peer2 : otherpeer,
            };

            des.data_exchange()
        }
    }
}

enum Spec {
    Server(Box<Server + Send>),
    Client(IPeer)
}

fn get_peer_by_spec(specifier: &str) -> Result<Spec> {
    use Spec::{Server,Client};
    match specifier {
        x if x == "-"               => Ok(Client(get_stdio_peer()?.upcast())),
        x if x.starts_with("ws:")   => Ok(Client(get_websocket_peer(x)?.upcast())),
        x if x.starts_with("wss:")  => Ok(Client(get_websocket_peer(x)?.upcast())),
        x if x.starts_with("tcp:")  => Ok(Client(get_tcp_peer(&x[4..])?.upcast())),
        x if x.starts_with("l-tcp:")  => Ok(Server(TcpServer::new(&x[6..])?.upcast())),
        x if x.starts_with("l-ws:")  => Ok(Server(WebsockServer::new(&x[5..])?.upcast())),
        x => Err(ErrorKind::InvalidSpecifier(x.to_string()).into()),
    }
}

fn try_main() -> Result<()> {
    //env_logger::init()?;
    init_logger()?;

    // setup command line arguments
    let matches = ::clap::App::new("websocat")
        .version("0.1")
        .author("Vitaly \"_Vi\" Shukela <vi0oss@gmail.com>")
        .about("Exchange binary data between binary websocket and something.\nSocat analogue with websockets.")
        .arg(::clap::Arg::with_name("spec1")
             .help("First specifier.")
             .required(true)
             .index(1))
        .arg(::clap::Arg::with_name("spec2")
             .help("Second specifier.")
             .required(true)
             .index(2))
        .after_help(r#"
Specifiers can be:
  ws[s]://<rest of websocket URL>   Connect to websocket
  l-ws:host:port                    Listen unencrypted websocket
  tcp:host:port                     Connect to TCP
  l-tcp:host:port                   Listen TCP connections
  -                                 stdin/stdout
  (more to be implemented)

Examples:
  websocat l-tcp:0.0.0.0:9559 ws://echo.websocket.org/
    Listen port 9959 on address :: and forward
    all connections to a public loopback websocket
  websocat l-ws:127.0.0.1:7878 tcp:127.0.0.1:1194
    Listen websocket and forward connections to local tcp
    Use nginx proxy for SSL if you want
  websocat - wss://myserver/mysocket
    Connect stdin/stdout to a secure web socket.
    Like netcat, but for websocket.
    `ssh user@host -o ProxyHommand "websocat - ws://..."`
  websocat ws://localhost:1234/ tcp:localhost:1235
    Connect both to websocket and to TCP and exchange data.

Specify listening part first, unless you want websocat to serve once.

IPv6 supported, just use specs like `l-ws:::1:4567`

Web socket usage is not obligatory, you can use any specs on both sides.
"#)
        .get_matches();

    let spec1  = matches.value_of("spec1") .ok_or("no listener_spec" )?;
    let spec2 = matches.value_of("spec2").ok_or("no connector_spec")?;

    main2(spec1, spec2, false)?;

    debug!("Exited");
    Ok(())
}

fn main() {
    if let Err(x) = try_main() {
        let _ = writeln!(::std::io::stderr(), "{:?}", x);
    }
}

<commit_msg>Bump version to 0.2 in other place.<commit_after>#![recursion_limit = "1024"] // error_chain

extern crate websocket;
extern crate env_logger;
#[macro_use]
extern crate log;
#[macro_use]
extern crate error_chain;
extern crate url;
extern crate clap;

const BUFSIZ : usize = 8192;

use std::thread;
use std::io::{stdin,stdout};

use websocket::{Message, Sender, Receiver, DataFrame, Server as WsServer};
use websocket::message::Type;
use websocket::client::request::Url;
use websocket::Client;

use std::borrow::Borrow;
use std::io::{Error as IoError, ErrorKind as IoErrorKind, Write, Read};

error_chain! {
    foreign_links {
        ::std::io::Error, Io;
        log::SetLoggerError, Log;
        ::url::ParseError, Url;
        ::websocket::result::WebSocketError, Ws;
        ::std::env::VarError, Ev;
    }
    errors {
        InvalidSpecifier(t : String) {
            description("invalid specifier")
            display("Invalid client or server specifier `{}`", t)
        }
    }
}

// Initialize logger with default "info" log level:
fn init_logger() -> Result<()> {
    let mut builder = env_logger::LogBuilder::new();
    builder.filter(None, log::LogLevelFilter::Info);
    if ::std::env::var("RUST_LOG").is_ok() {
       builder.parse(&::std::env::var("RUST_LOG")?);
    }
    builder.init()?;
    Ok(())
}

struct SenderWrapper<T: Sender> (T);

impl<T: Sender> ::std::io::Write for SenderWrapper<T> {
    fn write(&mut self, buf: &[u8]) -> ::std::io::Result<usize> {
        let message = Message::binary(buf);
        let ret;
        let len = buf.len();
        if len > 0 {
            debug!("Sending message of {} bytes", len);
            ret = self.0.send_message(&message);
        } else {
            // Interpret zero length buffer is request
            // to close communication

            debug!("Sending the closing message");
            ret = self.0.send_message(&Message::close());
        }
        ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;
        Ok(len)
    }
    fn flush(&mut self) -> ::std::io::Result<()> {
        Ok(())
    }
}

struct ReceiverWrapper<T: Receiver<DataFrame>> (T);

impl<T:Receiver<DataFrame>> ::std::io::Read for ReceiverWrapper<T> {
    fn read(&mut self, buf: &mut [u8]) -> ::std::io::Result<usize> {
        let ret = self.0.recv_message();
        let msg : Message = ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;

        match msg.opcode {
            Type::Close => {
                Ok(0)
            }
            Type::Ping => {
                // Sender used to be in a separate thread with a channel
                // now there's no channel, so trickier to combine ping replies
                // and usual data exchange
                unimplemented!();
            }
            _ => {
                let msgpayload : &[u8] = msg.payload.borrow();
                let len = msgpayload.len();
                debug!("Received message of {} bytes", len);

                assert!(buf.len() >= len);

                buf[0..len].copy_from_slice(msgpayload);

                Ok(len)
            }
        }
    }
}

struct Peer<R, W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    reader: R,
    writer: W,
}

type IPeer = Peer<Box<Read+Send>, Box<Write+Send>>;

struct DataExchangeSession<R1, R2, W1, W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    peer1: Peer<R1, W1>,
    peer2: Peer<R2, W2>,
}

// Derived from https://doc.rust-lang.org/src/std/up/src/libstd/io/util.rs.html#46-61
pub fn copy_with_flushes<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W) -> ::std::io::Result<u64>
    where R: Read, W: Write
{
    let mut buf = [0; BUFSIZ];
    let mut written = 0;
    loop {
        let len = match reader.read(&mut buf) {
            Ok(0) => return Ok(written),
            Ok(len) => len,
            Err(ref e) if e.kind() == IoErrorKind::Interrupted => continue,
            Err(ref e) if e.kind() == IoErrorKind::WouldBlock => continue,
            Err(e) => return Err(e),
        };
        writer.write_all(&buf[..len])?;
        writer.flush()?;
        written += len as u64;
    }
}

impl<R1,R2,W1,W2> DataExchangeSession<R1,R2,W1,W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    fn data_exchange(self) -> Result<()> {

        let mut reader1 = self.peer1.reader;
        let mut writer1 = self.peer1.writer;
        let mut reader2 = self.peer2.reader;
        let mut writer2 = self.peer2.writer;

        let receive_loop = thread::Builder::new().spawn(move || -> Result<()> {
            // Actual data transfer happens here
            copy_with_flushes(&mut reader1, &mut writer2)?;
            writer2.write(b"")?; // signal close
            Ok(())
        })?;

        // Actual data transfer happens here
        copy_with_flushes(&mut reader2, &mut writer1)?;
        writer1.write(b"")?; // Signal close

        debug!("Waiting for receiver side to exit");

        receive_loop.join().map_err(|x|format!("{:?}",x))?
    }
}

fn get_websocket_peer(urlstr: &str) -> Result<
        Peer<
            ReceiverWrapper<websocket::client::Receiver<websocket::WebSocketStream>>,
            SenderWrapper<websocket::client::Sender<websocket::WebSocketStream>>>
        > {
    let url = Url::parse(urlstr)?;

    info!("Connecting to {}", url);

    let request = Client::connect(url)?;

    let response = request.send()?; // Send the request and retrieve a response

    info!("Validating response...");

    response.validate()?; // Validate the response

    info!("Successfully connected");

    let (sender, receiver) = response.begin().split();

    let peer = Peer {
        reader : ReceiverWrapper(receiver),
        writer : SenderWrapper(sender),
    };
    Ok(peer)
}

fn get_tcp_peer(addr: &str) -> Result<
        Peer<
            ::std::net::TcpStream,
            ::std::net::TcpStream,
        >> {
    let sock = ::std::net::TcpStream::connect(addr)?;

    let peer = Peer {
        reader : sock.try_clone()?,
        writer : sock.try_clone()?,
    };
    info!("Connected to TCP {}", addr);
    Ok(peer)
}

fn get_stdio_peer() -> Result<Peer<std::io::Stdin, std::io::Stdout>> {
    Ok(
        Peer {
            reader : stdin(),
            writer : stdout(),
        }
    )
}

struct TcpServer(::std::net::TcpListener);

impl TcpServer {
    fn new(addr: &str) -> Result<Self> {
        Ok(TcpServer(::std::net::TcpListener::bind(addr)?))
    }
}

impl Server for TcpServer {
    fn accept_client(&mut self) -> Result<IPeer> {
        let (sock, addr) = self.0.accept()?;
        info!("TCP client connection from {}", addr);
        let peer = Peer {
            reader : sock.try_clone()?,
            writer : sock.try_clone()?,
        };
        Ok(peer.upcast())
    }
}

struct WebsockServer<'a>(WsServer<'a>);

impl<'a> WebsockServer<'a> {
    fn new(addr: &str) -> Result<Self> {
        Ok(WebsockServer(WsServer::bind(addr)?))
    }
}

impl<'a> Server for WebsockServer<'a> {
    fn accept_client(&mut self) -> Result<IPeer> {
        let connection = self.0.accept()?;
        info!("WebSocket client connection ...");
        let request = connection.read_request()?;
        request.validate()?;
        let response = request.accept(); // Form a response
        let mut client = response.send()?; // Send the response

        let ip = client.get_mut_sender()
            .get_mut()
            .peer_addr()
            .unwrap();

        info!("... from IP {}", ip);

        let (sender, receiver) = client.split();

        let peer = Peer {
            reader : ReceiverWrapper(receiver),
            writer : SenderWrapper(sender),
        };
        Ok(peer.upcast())
    }
}

impl<R,W> Peer<R,W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    fn upcast(self) -> IPeer  {
        Peer {
            reader: Box::new(self.reader) as Box<Read +Send>,
            writer: Box::new(self.writer) as Box<Write+Send>,
        }
    }
}

trait Server
{
    fn accept_client(&mut self) -> Result<IPeer>;

    fn start_serving(&mut self, spec2: &str, once: bool) -> Result<()> {
        let spec2s = spec2.to_string();
        let closure = move |peer, spec2 : String|{
            let spec2_ = get_peer_by_spec(spec2.as_str())?;
            let peer2 = match spec2_ {
                Spec::Server(mut x) => {
                    x.accept_client()?
                }
                Spec::Client(p1) => {
                    p1
                }
            };
            let des = DataExchangeSession {
                peer1 : peer,
                peer2 : peer2,
            };

            des.data_exchange()
        };
        if once {
            let peer = self.accept_client()?;
            closure(peer, spec2s)
        } else {
            let cl2 = ::std::sync::Arc::new(closure);
            loop {
                let ret = self.accept_client();
                let peer = match ret {
                    Ok(x) => x,
                    Err(er) => {
                        warn!("Can't accept client: {}", er);
                        continue;
                    }
                };
                let cl3 = cl2.clone();
                let spec2s2 = spec2s.clone();
                if let Err(x) = thread::Builder::new().spawn(move|| {
                    if let Err(x) = cl3(peer, spec2s2) {
                        warn!("Error while serving: {}", x);
                    }
                }) {
                    warn!("Error creating thread: {}", x);
                    thread::sleep(::std::time::Duration::from_millis(200));
                }
            }
        }
    }

    fn upcast(self) -> Box<Server+Send>
        where Self : Sized + Send + 'static
        { Box::new(self) as Box<Server+Send> }
}

fn main2(spec1: &str, spec2: &str, once: bool) -> Result<()> {
    let spec1_ = get_peer_by_spec(spec1)?;

    match spec1_ {
        Spec::Server(mut x) => {
            x.start_serving(spec2, once)
        }
        Spec::Client(p1) => {
            let spec2_ = get_peer_by_spec(spec2)?;

            let otherpeer = match spec2_ {
                Spec::Server(mut x) => {
                    let t = x.accept_client()?;
                    t
                }
                Spec::Client(p2) => {
                    p2
                }
            };

            let des = DataExchangeSession {
                peer1 : p1,
                peer2 : otherpeer,
            };

            des.data_exchange()
        }
    }
}

enum Spec {
    Server(Box<Server + Send>),
    Client(IPeer)
}

fn get_peer_by_spec(specifier: &str) -> Result<Spec> {
    use Spec::{Server,Client};
    match specifier {
        x if x == "-"               => Ok(Client(get_stdio_peer()?.upcast())),
        x if x.starts_with("ws:")   => Ok(Client(get_websocket_peer(x)?.upcast())),
        x if x.starts_with("wss:")  => Ok(Client(get_websocket_peer(x)?.upcast())),
        x if x.starts_with("tcp:")  => Ok(Client(get_tcp_peer(&x[4..])?.upcast())),
        x if x.starts_with("l-tcp:")  => Ok(Server(TcpServer::new(&x[6..])?.upcast())),
        x if x.starts_with("l-ws:")  => Ok(Server(WebsockServer::new(&x[5..])?.upcast())),
        x => Err(ErrorKind::InvalidSpecifier(x.to_string()).into()),
    }
}

fn try_main() -> Result<()> {
    //env_logger::init()?;
    init_logger()?;

    // setup command line arguments
    let matches = ::clap::App::new("websocat")
        .version("0.2")
        .author("Vitaly \"_Vi\" Shukela <vi0oss@gmail.com>")
        .about("Exchange binary data between binary websocket and something.\nSocat analogue with websockets.")
        .arg(::clap::Arg::with_name("spec1")
             .help("First specifier.")
             .required(true)
             .index(1))
        .arg(::clap::Arg::with_name("spec2")
             .help("Second specifier.")
             .required(true)
             .index(2))
        .after_help(r#"
Specifiers can be:
  ws[s]://<rest of websocket URL>   Connect to websocket
  l-ws:host:port                    Listen unencrypted websocket
  tcp:host:port                     Connect to TCP
  l-tcp:host:port                   Listen TCP connections
  -                                 stdin/stdout
  (more to be implemented)

Examples:
  websocat l-tcp:0.0.0.0:9559 ws://echo.websocket.org/
    Listen port 9959 on address :: and forward
    all connections to a public loopback websocket
  websocat l-ws:127.0.0.1:7878 tcp:127.0.0.1:1194
    Listen websocket and forward connections to local tcp
    Use nginx proxy for SSL if you want
  websocat - wss://myserver/mysocket
    Connect stdin/stdout to a secure web socket.
    Like netcat, but for websocket.
    `ssh user@host -o ProxyHommand "websocat - ws://..."`
  websocat ws://localhost:1234/ tcp:localhost:1235
    Connect both to websocket and to TCP and exchange data.

Specify listening part first, unless you want websocat to serve once.

IPv6 supported, just use specs like `l-ws:::1:4567`

Web socket usage is not obligatory, you can use any specs on both sides.
"#)
        .get_matches();

    let spec1  = matches.value_of("spec1") .ok_or("no listener_spec" )?;
    let spec2 = matches.value_of("spec2").ok_or("no connector_spec")?;

    main2(spec1, spec2, false)?;

    debug!("Exited");
    Ok(())
}

fn main() {
    if let Err(x) = try_main() {
        let _ = writeln!(::std::io::stderr(), "{:?}", x);
    }
}

<|endoftext|>


<commit_before>#[macro_use]
extern crate log;
extern crate getopts;
extern crate json;
extern crate tic;
extern crate tiny_http;

mod logging;
mod metrics;
mod options;

use logging::set_log_level;
use metrics::Metric;
use options::{PROGRAM, VERSION};
use std::{io, thread};
use tic::Sample;
use tiny_http::{Method, Request, Response};

fn main() {
    let options = options::init();

    // initialize logging
    set_log_level(options.opt_count("verbose"));
    info!("{} {}", PROGRAM, VERSION);

    // initialize metrics
    let mut metrics = metrics::init(options.opt_str("metrics"));
    let mut stats = metrics.get_sender();
    let clock = metrics.get_clocksource();
    thread::spawn(move || { metrics.run(); });

    // initialize http listener
    let http = options.opt_str("http").unwrap_or("0.0.0.0:4567".to_owned());
    let server = tiny_http::Server::http(&http).unwrap();
    info!("listening HTTP {}", &http);

    // run the server
    loop {
        if let Ok(Some(request)) = server.try_recv() {
            debug!("handle request");
            let t0 = clock.counter();
            handle_http(request);
            let t1 = clock.counter();
            let _ = stats.send(Sample::new(t0, t1, Metric::Request));
        }
    }
}

// actually handle the http request
fn handle_http(mut request: Request) {
    let response = match *request.method() {
        Method::Post => {
            match request.url() {
                "/payload" => {
                    info!("payload received");
				    let mut content = String::new();
				    request.as_reader().read_to_string(&mut content).unwrap();
				    handle_payload(&content)
                }
                _ => Response::empty(404),
            }
        }
        _ => Response::empty(405),
    };

    let _ = request.respond(response);
}

// parse payload
fn handle_payload(payload: &str) -> Response<io::Empty> {
	info!("handle payload");
	let parsed = json::parse(payload).unwrap();
	Response::empty(200)
}<commit_msg>rustfmt<commit_after>#[macro_use]
extern crate log;
extern crate getopts;
extern crate json;
extern crate tic;
extern crate tiny_http;

mod logging;
mod metrics;
mod options;

use logging::set_log_level;
use metrics::Metric;
use options::{PROGRAM, VERSION};
use std::{io, thread};
use tic::Sample;
use tiny_http::{Method, Request, Response};

fn main() {
    let options = options::init();

    // initialize logging
    set_log_level(options.opt_count("verbose"));
    info!("{} {}", PROGRAM, VERSION);

    // initialize metrics
    let mut metrics = metrics::init(options.opt_str("metrics"));
    let mut stats = metrics.get_sender();
    let clock = metrics.get_clocksource();
    thread::spawn(move || { metrics.run(); });

    // initialize http listener
    let http = options.opt_str("http").unwrap_or("0.0.0.0:4567".to_owned());
    let server = tiny_http::Server::http(&http).unwrap();
    info!("listening HTTP {}", &http);

    // run the server
    loop {
        if let Ok(Some(request)) = server.try_recv() {
            debug!("handle request");
            let t0 = clock.counter();
            handle_http(request);
            let t1 = clock.counter();
            let _ = stats.send(Sample::new(t0, t1, Metric::Request));
        }
    }
}

// actually handle the http request
fn handle_http(mut request: Request) {
    let response = match *request.method() {
        Method::Post => {
            match request.url() {
                "/payload" => {
                    info!("payload received");
                    let mut content = String::new();
                    request.as_reader().read_to_string(&mut content).unwrap();
                    handle_payload(&content)
                }
                _ => Response::empty(404),
            }
        }
        _ => Response::empty(405),
    };

    let _ = request.respond(response);
}

// parse payload
fn handle_payload(payload: &str) -> Response<io::Empty> {
    info!("handle payload");
    let parsed = json::parse(payload).unwrap();
    Response::empty(200)
}
<|endoftext|>


<commit_before>use std::fmt;
use std::str::FromStr;

#[derive(Debug, PartialEq)]
enum Atom {
    Number(f64),
    Symbol(String)
}

#[derive(Debug, PartialEq)]
enum ParseAtomError {
    IncorrectSymbolName
}

impl fmt::Display for Atom {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Atom::Number(n) => {
                write!(f, "{}", n)
            },
            Atom::Symbol(ref name) => {
                write!(f, "'{}", name)
            }
        }
    }
}

impl FromStr for Atom {
    type Err = ParseAtomError;

    fn from_str(s: &str) -> Result<Atom, ParseAtomError> {
        match s.parse::<f64>() {
            Ok(f) => {
                Ok(Atom::Number(f))
            },
            Err(..) => {
                match s.chars().next() {
                    Some(c) if !c.is_numeric() => {
                        Ok(Atom::Symbol(s.to_string()))
                    },
                    _ => {
                        Err(ParseAtomError::IncorrectSymbolName)
                    }
                }
            }
        }
    }
}

fn tokenize(s: &str) -> Vec<String> {
    s.replace("("," ( ")
        .replace(")", " ) ")
        .replace("\n", " ")
        .split(' ')
        .filter(|s| { !s.is_empty() })
        .map(|s| { s.to_string() })
        .collect()
}

#[cfg_attr(test, allow(dead_code))]
fn main() {
}

#[cfg(test)]
mod tests {
    use super::{Atom, ParseAtomError, tokenize};

    #[test]
    fn test_parse_integer() {
        assert_eq!(Atom::Number(64f64), "64".parse::<Atom>().ok().unwrap())
    }

    #[test]
    fn test_parse_float() {
        assert_eq!(Atom::Number(64.5), "64.5".parse::<Atom>().ok().unwrap())
    }

    #[test]
    fn test_parse_symbol() {
        assert_eq!(Atom::Symbol("name".to_string()), "name".parse::<Atom>().ok().unwrap())
    }

    #[test]
    fn test_parse_incorrect_symbol_starting_with_digit() {
        assert_eq!(ParseAtomError::IncorrectSymbolName, "6name".parse::<Atom>().err().unwrap())
    }

    #[test]
    fn test_tokenize_dense_expression() {
        let expected_result = ["(", "def", "a", "1", ")"]
                                .iter()
                                .map(|s| { s.to_string() })
                                .collect();
        assert_eq!(expected_result, tokenize("(def a 1)"))
    }

    #[test]
    fn test_tokenize_sparse_expression() {
        let expected_result = ["(", "def", "a", "1", ")"]
                                .iter()
                                .map(|s| { s.to_string() })
                                .collect();
        assert_eq!(expected_result, tokenize(" ( \n def a\n1)   \n"))
    }
}
<commit_msg>Specify use section<commit_after>use std::fmt;
use std::str::FromStr;

#[derive(Debug, PartialEq)]
enum Atom {
    Number(f64),
    Symbol(String)
}

#[derive(Debug, PartialEq)]
enum ParseAtomError {
    IncorrectSymbolName
}

impl fmt::Display for Atom {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Atom::Number(n) => {
                write!(f, "{}", n)
            },
            Atom::Symbol(ref name) => {
                write!(f, "'{}", name)
            }
        }
    }
}

impl FromStr for Atom {
    type Err = ParseAtomError;

    fn from_str(s: &str) -> Result<Atom, ParseAtomError> {
        match s.parse::<f64>() {
            Ok(f) => {
                Ok(Atom::Number(f))
            },
            Err(..) => {
                match s.chars().next() {
                    Some(c) if !c.is_numeric() => {
                        Ok(Atom::Symbol(s.to_string()))
                    },
                    _ => {
                        Err(ParseAtomError::IncorrectSymbolName)
                    }
                }
            }
        }
    }
}

fn tokenize(s: &str) -> Vec<String> {
    s.replace("("," ( ")
        .replace(")", " ) ")
        .replace("\n", " ")
        .split(' ')
        .filter(|s| { !s.is_empty() })
        .map(|s| { s.to_string() })
        .collect()
}

#[cfg_attr(test, allow(dead_code))]
fn main() {
}

#[cfg(test)]
mod tests {
    use super::Atom::{self, Number, Symbol};
    use super::ParseAtomError::IncorrectSymbolName;
    use super::tokenize;

    #[test]
    fn test_parse_integer() {
        assert_eq!(Number(64f64), "64".parse::<Atom>().ok().unwrap())
    }

    #[test]
    fn test_parse_float() {
        assert_eq!(Number(64.5), "64.5".parse::<Atom>().ok().unwrap())
    }

    #[test]
    fn test_parse_symbol() {
        assert_eq!(Symbol("name".to_string()), "name".parse::<Atom>().ok().unwrap())
    }

    #[test]
    fn test_parse_incorrect_symbol_starting_with_digit() {
        assert_eq!(IncorrectSymbolName, "6name".parse::<Atom>().err().unwrap())
    }

    #[test]
    fn test_tokenize_dense_expression() {
        let expected_result = ["(", "def", "a", "1", ")"]
                                .iter()
                                .map(|s| { s.to_string() })
                                .collect();
        assert_eq!(expected_result, tokenize("(def a 1)"))
    }

    #[test]
    fn test_tokenize_sparse_expression() {
        let expected_result = ["(", "def", "a", "1", ")"]
                                .iter()
                                .map(|s| { s.to_string() })
                                .collect();
        assert_eq!(expected_result, tokenize(" ( \n def a\n1)   \n"))
    }
}
<|endoftext|>


<commit_before>#![feature(unique, core_intrinsics)]

extern crate llvm_sys;

use std::ffi::CString;

mod id;
mod owned;
mod llvm;

fn main() {
    id::with(|context_id| {
        let context = llvm::Context::new(context_id);
        id::with(|module_id| {
            let module = llvm::Module::new(module_id, &CString::new("mymodule").unwrap(), &context);
            id::with(|function_id| {
                let i32_ty = llvm::Type::i32(&context);
                let func_ty = llvm::Type::function(&[i32_ty], i32_ty);
                let mut function = module.add_function(function_id, &CString::new("square").unwrap(), func_ty);
                {
                    let mut function_builder = function.builder();
                    let mut builder = llvm::Builder::new(&context);

                    let entry = function_builder.append_basic_block(&CString::new("entry").unwrap(), &context);
                    let builder = builder.position_at_end(entry);

                    let ret = builder.mul(function_builder.param(0), function_builder.param(0), &CString::new("square").unwrap());
                    builder.ret(ret);
                }

                function.dump();
            });
        });
    });
}
<commit_msg>Add another function to main to demonstrate the use of the FunctionBuilder<commit_after>#![feature(unique, core_intrinsics)]

extern crate llvm_sys;

use std::ffi::CString;

mod id;
mod owned;
mod llvm;

fn main() {
    id::with(|context_id| {
        let context = llvm::Context::new(context_id);
        id::with(|module_id| {
            let module = llvm::Module::new(module_id, &CString::new("mymodule").unwrap(), &context);

            let i32_ty = llvm::Type::i32(&context);
            let func_ty = llvm::Type::function(&[i32_ty], i32_ty);
            let mut builder = llvm::Builder::new(&context);

            id::with(|function_id| {
                let mut function = module.add_function(function_id, &CString::new("square").unwrap(), func_ty);
                {
                    let mut function_builder = function.builder();

                    let entry = function_builder.append_basic_block(&CString::new("entry").unwrap(), &context);
                    let builder = builder.position_at_end(entry);

                    let ret = builder.mul(function_builder.param(0), function_builder.param(0), &CString::new("square").unwrap());
                    builder.ret(ret);
                }

                function.dump();
            });

            id::with(|function_id| {
                let mut function = module.add_function(function_id, &CString::new("jumpy").unwrap(), func_ty);
                {
                    let mut function_builder = function.builder();
                    let entry = function_builder.append_basic_block(&CString::new("entry").unwrap(), &context);
                    let exit = function_builder.append_basic_block(&CString::new("exit").unwrap(), &context);

                    builder.position_at_end(entry).br(exit);
                    builder.position_at_end(exit).ret(function_builder.param(0));
                }

                function.dump()
            });
        });
    });
}
<|endoftext|>


<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.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");
}
<commit_msg>Pass in a thunk to append actions.<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 FnMut(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 FnMut(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 FnMut(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(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 FnMut(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(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 |action| parent_actions.push_back(action));
            }
            // 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>use std::error;
use std::io::{self, Write};
use std::process;
use std::sync::{Arc, Mutex};
use std::time::Instant;

use ignore::WalkState;

use args::Args;
use subject::Subject;

#[macro_use]
mod messages;

mod app;
mod args;
mod config;
mod logger;
mod path_printer;
mod search;
mod subject;

// Since Rust no longer uses jemalloc by default, ripgrep will, by default,
// use the system allocator. On Linux, this would normally be glibc's
// allocator, which is pretty good. In particular, ripgrep does not have a
// particularly allocation heavy workload, so there really isn't much
// difference (for ripgrep's purposes) between glibc's allocator and jemalloc.
//
// However, when ripgrep is built with musl, this means ripgrep will use musl's
// allocator, which appears to be substantially worse. (musl's goal is not to
// have the fastest version of everything. Its goal is to be small and amenable
// to static compilation.) Even though ripgrep isn't particularly allocation
// heavy, musl's allocator appears to slow down ripgrep quite a bit. Therefore,
// when building with musl, we use jemalloc.
//
// We don't unconditionally use jemalloc because it can be nice to use the
// system's default allocator by default. Moreover, jemalloc seems to increase
// compilation times by a bit.
//
// Moreover, we only do this on 64-bit systems since jemalloc doesn't support
// i686.
#[cfg(all(target_env = "musl", target_pointer_width = "64"))]
#[global_allocator]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

type Result<T> = ::std::result::Result<T, Box<dyn error::Error>>;

fn main() {
    if let Err(err) = Args::parse().and_then(try_main) {
        eprintln!("{}", err);
        process::exit(2);
    }
}

fn try_main(args: Args) -> Result<()> {
    use args::Command::*;

    let matched =
        match args.command()? {
            Search => search(&args),
            SearchParallel => search_parallel(&args),
            SearchNever => Ok(false),
            Files => files(&args),
            FilesParallel => files_parallel(&args),
            Types => types(&args),
            PCRE2Version => pcre2_version(&args),
        }?;
    if matched && (args.quiet() || !messages::errored()) {
        process::exit(0)
    } else if messages::errored() {
        process::exit(2)
    } else {
        process::exit(1)
    }
}

/// The top-level entry point for single-threaded search. This recursively
/// steps through the file list (current directory by default) and searches
/// each file sequentially.
fn search(args: &Args) -> Result<bool> {
    let started_at = Instant::now();
    let quit_after_match = args.quit_after_match()?;
    let subject_builder = args.subject_builder();
    let mut stats = args.stats()?;
    let mut searcher = args.search_worker(args.stdout())?;
    let mut matched = false;

    for result in args.walker()? {
        let subject = match subject_builder.build_from_result(result) {
            Some(subject) => subject,
            None => continue,
        };
        let search_result = match searcher.search(&subject) {
            Ok(search_result) => search_result,
            Err(err) => {
                // A broken pipe means graceful termination.
                if err.kind() == io::ErrorKind::BrokenPipe {
                    break;
                }
                err_message!("{}: {}", subject.path().display(), err);
                continue;
            }
        };
        matched = matched || search_result.has_match();
        if let Some(ref mut stats) = stats {
            *stats += search_result.stats().unwrap();
        }
        if matched && quit_after_match {
            break;
        }
    }
    if let Some(ref stats) = stats {
        let elapsed = Instant::now().duration_since(started_at);
        // We don't care if we couldn't print this successfully.
        let _ = searcher.print_stats(elapsed, stats);
    }
    Ok(matched)
}

/// The top-level entry point for multi-threaded search. The parallelism is
/// itself achieved by the recursive directory traversal. All we need to do is
/// feed it a worker for performing a search on each file.
fn search_parallel(args: &Args) -> Result<bool> {
    use std::sync::atomic::AtomicBool;
    use std::sync::atomic::Ordering::SeqCst;

    let quit_after_match = args.quit_after_match()?;
    let started_at = Instant::now();
    let subject_builder = Arc::new(args.subject_builder());
    let bufwtr = Arc::new(args.buffer_writer()?);
    let stats = Arc::new(args.stats()?.map(Mutex::new));
    let matched = Arc::new(AtomicBool::new(false));
    let mut searcher_err = None;
    args.walker_parallel()?.run(|| {
        let args = args.clone();
        let bufwtr = Arc::clone(&bufwtr);
        let stats = Arc::clone(&stats);
        let matched = Arc::clone(&matched);
        let subject_builder = Arc::clone(&subject_builder);
        let mut searcher = match args.search_worker(bufwtr.buffer()) {
            Ok(searcher) => searcher,
            Err(err) => {
                searcher_err = Some(err);
                return Box::new(move |_| {
                    WalkState::Quit
                });
            }
        };

        Box::new(move |result| {
            let subject = match subject_builder.build_from_result(result) {
                Some(subject) => subject,
                None => return WalkState::Continue,
            };
            searcher.printer().get_mut().clear();
            let search_result = match searcher.search(&subject) {
                Ok(search_result) => search_result,
                Err(err) => {
                    err_message!("{}: {}", subject.path().display(), err);
                    return WalkState::Continue;
                }
            };
            if search_result.has_match() {
                matched.store(true, SeqCst);
            }
            if let Some(ref locked_stats) = *stats {
                let mut stats = locked_stats.lock().unwrap();
                *stats += search_result.stats().unwrap();
            }
            if let Err(err) = bufwtr.print(searcher.printer().get_mut()) {
                // A broken pipe means graceful termination.
                if err.kind() == io::ErrorKind::BrokenPipe {
                    return WalkState::Quit;
                }
                // Otherwise, we continue on our merry way.
                err_message!("{}: {}", subject.path().display(), err);
            }
            if matched.load(SeqCst) && quit_after_match {
                WalkState::Quit
            } else {
                WalkState::Continue
            }
        })
    });
    if let Some(err) = searcher_err.take() {
        return Err(err);
    }
    if let Some(ref locked_stats) = *stats {
        let elapsed = Instant::now().duration_since(started_at);
        let stats = locked_stats.lock().unwrap();
        let mut searcher = args.search_worker(args.stdout())?;
        // We don't care if we couldn't print this successfully.
        let _ = searcher.print_stats(elapsed, &stats);
    }
    Ok(matched.load(SeqCst))
}

/// The top-level entry point for listing files without searching them. This
/// recursively steps through the file list (current directory by default) and
/// prints each path sequentially using a single thread.
fn files(args: &Args) -> Result<bool> {
    let quit_after_match = args.quit_after_match()?;
    let subject_builder = args.subject_builder();
    let mut matched = false;
    let mut path_printer = args.path_printer(args.stdout())?;
    for result in args.walker()? {
        let subject = match subject_builder.build_from_result(result) {
            Some(subject) => subject,
            None => continue,
        };
        matched = true;
        if quit_after_match {
            break;
        }
        if let Err(err) = path_printer.write_path(subject.path()) {
            // A broken pipe means graceful termination.
            if err.kind() == io::ErrorKind::BrokenPipe {
                break;
            }
            // Otherwise, we have some other error that's preventing us from
            // writing to stdout, so we should bubble it up.
            return Err(err.into());
        }
    }
    Ok(matched)
}

/// The top-level entry point for listing files without searching them. This
/// recursively steps through the file list (current directory by default) and
/// prints each path sequentially using multiple threads.
fn files_parallel(args: &Args) -> Result<bool> {
    use std::sync::atomic::AtomicBool;
    use std::sync::atomic::Ordering::SeqCst;
    use std::sync::mpsc;
    use std::thread;

    let quit_after_match = args.quit_after_match()?;
    let subject_builder = Arc::new(args.subject_builder());
    let mut path_printer = args.path_printer(args.stdout())?;
    let matched = Arc::new(AtomicBool::new(false));
    let (tx, rx) = mpsc::channel::<Subject>();

    let print_thread = thread::spawn(move || -> io::Result<()> {
        for subject in rx.iter() {
            path_printer.write_path(subject.path())?;
        }
        Ok(())
    });
    args.walker_parallel()?.run(|| {
        let subject_builder = Arc::clone(&subject_builder);
        let matched = Arc::clone(&matched);
        let tx = tx.clone();

        Box::new(move |result| {
            let subject = match subject_builder.build_from_result(result) {
                Some(subject) => subject,
                None => return WalkState::Continue,
            };
            matched.store(true, SeqCst);
            if quit_after_match {
                WalkState::Quit
            } else {
                match tx.send(subject) {
                    Ok(_) => WalkState::Continue,
                    Err(_) => WalkState::Quit,
                }
            }
        })
    });
    drop(tx);
    if let Err(err) = print_thread.join().unwrap() {
        // A broken pipe means graceful termination, so fall through.
        // Otherwise, something bad happened while writing to stdout, so bubble
        // it up.
        if err.kind() != io::ErrorKind::BrokenPipe {
            return Err(err.into());
        }
    }
    Ok(matched.load(SeqCst))
}

/// The top-level entry point for --type-list.
fn types(args: &Args) -> Result<bool> {
    let mut count = 0;
    let mut stdout = args.stdout();
    for def in args.type_defs()? {
        count += 1;
        stdout.write_all(def.name().as_bytes())?;
        stdout.write_all(b": ")?;

        let mut first = true;
        for glob in def.globs() {
            if !first {
                stdout.write_all(b", ")?;
            }
            stdout.write_all(glob.as_bytes())?;
            first = false;
        }
        stdout.write_all(b"\n")?;
    }
    Ok(count > 0)
}

/// The top-level entry point for --pcre2-version.
fn pcre2_version(args: &Args) -> Result<bool> {
    #[cfg(feature = "pcre2")]
    fn imp(args: &Args) -> Result<bool> {
        use grep::pcre2;

        let mut stdout = args.stdout();

        let (major, minor) = pcre2::version();
        writeln!(stdout, "PCRE2 {}.{} is available", major, minor)?;

        if cfg!(target_pointer_width = "64") && pcre2::is_jit_available() {
            writeln!(stdout, "JIT is available")?;
        }
        Ok(true)
    }

    #[cfg(not(feature = "pcre2"))]
    fn imp(args: &Args) -> Result<bool> {
        let mut stdout = args.stdout();
        writeln!(stdout, "PCRE2 is not available in this build of ripgrep.")?;
        Ok(false)
    }

    imp(args)
}
<commit_msg>core: simplify parallel walking using borrows<commit_after>use std::error;
use std::io::{self, Write};
use std::process;
use std::sync::Mutex;
use std::time::Instant;

use ignore::WalkState;

use args::Args;
use subject::Subject;

#[macro_use]
mod messages;

mod app;
mod args;
mod config;
mod logger;
mod path_printer;
mod search;
mod subject;

// Since Rust no longer uses jemalloc by default, ripgrep will, by default,
// use the system allocator. On Linux, this would normally be glibc's
// allocator, which is pretty good. In particular, ripgrep does not have a
// particularly allocation heavy workload, so there really isn't much
// difference (for ripgrep's purposes) between glibc's allocator and jemalloc.
//
// However, when ripgrep is built with musl, this means ripgrep will use musl's
// allocator, which appears to be substantially worse. (musl's goal is not to
// have the fastest version of everything. Its goal is to be small and amenable
// to static compilation.) Even though ripgrep isn't particularly allocation
// heavy, musl's allocator appears to slow down ripgrep quite a bit. Therefore,
// when building with musl, we use jemalloc.
//
// We don't unconditionally use jemalloc because it can be nice to use the
// system's default allocator by default. Moreover, jemalloc seems to increase
// compilation times by a bit.
//
// Moreover, we only do this on 64-bit systems since jemalloc doesn't support
// i686.
#[cfg(all(target_env = "musl", target_pointer_width = "64"))]
#[global_allocator]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

type Result<T> = ::std::result::Result<T, Box<dyn error::Error>>;

fn main() {
    if let Err(err) = Args::parse().and_then(try_main) {
        eprintln!("{}", err);
        process::exit(2);
    }
}

fn try_main(args: Args) -> Result<()> {
    use args::Command::*;

    let matched = match args.command()? {
        Search => search(&args),
        SearchParallel => search_parallel(&args),
        SearchNever => Ok(false),
        Files => files(&args),
        FilesParallel => files_parallel(&args),
        Types => types(&args),
        PCRE2Version => pcre2_version(&args),
    }?;
    if matched && (args.quiet() || !messages::errored()) {
        process::exit(0)
    } else if messages::errored() {
        process::exit(2)
    } else {
        process::exit(1)
    }
}

/// The top-level entry point for single-threaded search. This recursively
/// steps through the file list (current directory by default) and searches
/// each file sequentially.
fn search(args: &Args) -> Result<bool> {
    let started_at = Instant::now();
    let quit_after_match = args.quit_after_match()?;
    let subject_builder = args.subject_builder();
    let mut stats = args.stats()?;
    let mut searcher = args.search_worker(args.stdout())?;
    let mut matched = false;

    for result in args.walker()? {
        let subject = match subject_builder.build_from_result(result) {
            Some(subject) => subject,
            None => continue,
        };
        let search_result = match searcher.search(&subject) {
            Ok(search_result) => search_result,
            Err(err) => {
                // A broken pipe means graceful termination.
                if err.kind() == io::ErrorKind::BrokenPipe {
                    break;
                }
                err_message!("{}: {}", subject.path().display(), err);
                continue;
            }
        };
        matched = matched || search_result.has_match();
        if let Some(ref mut stats) = stats {
            *stats += search_result.stats().unwrap();
        }
        if matched && quit_after_match {
            break;
        }
    }
    if let Some(ref stats) = stats {
        let elapsed = Instant::now().duration_since(started_at);
        // We don't care if we couldn't print this successfully.
        let _ = searcher.print_stats(elapsed, stats);
    }
    Ok(matched)
}

/// The top-level entry point for multi-threaded search. The parallelism is
/// itself achieved by the recursive directory traversal. All we need to do is
/// feed it a worker for performing a search on each file.
fn search_parallel(args: &Args) -> Result<bool> {
    use std::sync::atomic::AtomicBool;
    use std::sync::atomic::Ordering::SeqCst;

    let quit_after_match = args.quit_after_match()?;
    let started_at = Instant::now();
    let subject_builder = args.subject_builder();
    let bufwtr = args.buffer_writer()?;
    let stats = args.stats()?.map(Mutex::new);
    let matched = AtomicBool::new(false);
    let mut searcher_err = None;
    args.walker_parallel()?.run(|| {
        let bufwtr = &bufwtr;
        let stats = &stats;
        let matched = &matched;
        let subject_builder = &subject_builder;
        let mut searcher = match args.search_worker(bufwtr.buffer()) {
            Ok(searcher) => searcher,
            Err(err) => {
                searcher_err = Some(err);
                return Box::new(move |_| WalkState::Quit);
            }
        };

        Box::new(move |result| {
            let subject = match subject_builder.build_from_result(result) {
                Some(subject) => subject,
                None => return WalkState::Continue,
            };
            searcher.printer().get_mut().clear();
            let search_result = match searcher.search(&subject) {
                Ok(search_result) => search_result,
                Err(err) => {
                    err_message!("{}: {}", subject.path().display(), err);
                    return WalkState::Continue;
                }
            };
            if search_result.has_match() {
                matched.store(true, SeqCst);
            }
            if let Some(ref locked_stats) = *stats {
                let mut stats = locked_stats.lock().unwrap();
                *stats += search_result.stats().unwrap();
            }
            if let Err(err) = bufwtr.print(searcher.printer().get_mut()) {
                // A broken pipe means graceful termination.
                if err.kind() == io::ErrorKind::BrokenPipe {
                    return WalkState::Quit;
                }
                // Otherwise, we continue on our merry way.
                err_message!("{}: {}", subject.path().display(), err);
            }
            if matched.load(SeqCst) && quit_after_match {
                WalkState::Quit
            } else {
                WalkState::Continue
            }
        })
    });
    if let Some(err) = searcher_err.take() {
        return Err(err);
    }
    if let Some(ref locked_stats) = stats {
        let elapsed = Instant::now().duration_since(started_at);
        let stats = locked_stats.lock().unwrap();
        let mut searcher = args.search_worker(args.stdout())?;
        // We don't care if we couldn't print this successfully.
        let _ = searcher.print_stats(elapsed, &stats);
    }
    Ok(matched.load(SeqCst))
}

/// The top-level entry point for listing files without searching them. This
/// recursively steps through the file list (current directory by default) and
/// prints each path sequentially using a single thread.
fn files(args: &Args) -> Result<bool> {
    let quit_after_match = args.quit_after_match()?;
    let subject_builder = args.subject_builder();
    let mut matched = false;
    let mut path_printer = args.path_printer(args.stdout())?;
    for result in args.walker()? {
        let subject = match subject_builder.build_from_result(result) {
            Some(subject) => subject,
            None => continue,
        };
        matched = true;
        if quit_after_match {
            break;
        }
        if let Err(err) = path_printer.write_path(subject.path()) {
            // A broken pipe means graceful termination.
            if err.kind() == io::ErrorKind::BrokenPipe {
                break;
            }
            // Otherwise, we have some other error that's preventing us from
            // writing to stdout, so we should bubble it up.
            return Err(err.into());
        }
    }
    Ok(matched)
}

/// The top-level entry point for listing files without searching them. This
/// recursively steps through the file list (current directory by default) and
/// prints each path sequentially using multiple threads.
fn files_parallel(args: &Args) -> Result<bool> {
    use std::sync::atomic::AtomicBool;
    use std::sync::atomic::Ordering::SeqCst;
    use std::sync::mpsc;
    use std::thread;

    let quit_after_match = args.quit_after_match()?;
    let subject_builder = args.subject_builder();
    let mut path_printer = args.path_printer(args.stdout())?;
    let matched = AtomicBool::new(false);
    let (tx, rx) = mpsc::channel::<Subject>();

    let print_thread = thread::spawn(move || -> io::Result<()> {
        for subject in rx.iter() {
            path_printer.write_path(subject.path())?;
        }
        Ok(())
    });
    args.walker_parallel()?.run(|| {
        let subject_builder = &subject_builder;
        let matched = &matched;
        let tx = tx.clone();

        Box::new(move |result| {
            let subject = match subject_builder.build_from_result(result) {
                Some(subject) => subject,
                None => return WalkState::Continue,
            };
            matched.store(true, SeqCst);
            if quit_after_match {
                WalkState::Quit
            } else {
                match tx.send(subject) {
                    Ok(_) => WalkState::Continue,
                    Err(_) => WalkState::Quit,
                }
            }
        })
    });
    drop(tx);
    if let Err(err) = print_thread.join().unwrap() {
        // A broken pipe means graceful termination, so fall through.
        // Otherwise, something bad happened while writing to stdout, so bubble
        // it up.
        if err.kind() != io::ErrorKind::BrokenPipe {
            return Err(err.into());
        }
    }
    Ok(matched.load(SeqCst))
}

/// The top-level entry point for --type-list.
fn types(args: &Args) -> Result<bool> {
    let mut count = 0;
    let mut stdout = args.stdout();
    for def in args.type_defs()? {
        count += 1;
        stdout.write_all(def.name().as_bytes())?;
        stdout.write_all(b": ")?;

        let mut first = true;
        for glob in def.globs() {
            if !first {
                stdout.write_all(b", ")?;
            }
            stdout.write_all(glob.as_bytes())?;
            first = false;
        }
        stdout.write_all(b"\n")?;
    }
    Ok(count > 0)
}

/// The top-level entry point for --pcre2-version.
fn pcre2_version(args: &Args) -> Result<bool> {
    #[cfg(feature = "pcre2")]
    fn imp(args: &Args) -> Result<bool> {
        use grep::pcre2;

        let mut stdout = args.stdout();

        let (major, minor) = pcre2::version();
        writeln!(stdout, "PCRE2 {}.{} is available", major, minor)?;

        if cfg!(target_pointer_width = "64") && pcre2::is_jit_available() {
            writeln!(stdout, "JIT is available")?;
        }
        Ok(true)
    }

    #[cfg(not(feature = "pcre2"))]
    fn imp(args: &Args) -> Result<bool> {
        let mut stdout = args.stdout();
        writeln!(stdout, "PCRE2 is not available in this build of ripgrep.")?;
        Ok(false)
    }

    imp(args)
}
<|endoftext|>