<commit_before>#![cfg(any(target_os = "linux", target_os = "dragonfly", target_os = "freebsd", target_os = "openbsd"))]

use std::collections::VecDeque;
use std::sync::Arc;
use std::env;

use {CreationError, CursorState, EventsLoopClosed, MouseCursor, ControlFlow};
use libc;

use self::x11::XConnection;
use self::x11::XError;
use self::x11::XNotSupported;
use self::x11::ffi::XVisualInfo;

mod dlopen;
pub mod wayland;
pub mod x11;

/// Environment variable that indicates a backend preference
/// `WINIT_UNIX_BACKEND=x11` : Will try to use an X11 backend
/// `WINIT_UNIX_BACKEND=wayland` : Will try to use a wayland backend
/// When used is selected backends will not fallback to one another
/// (Current default behaviour is wayland, fallback to X)
const BACKEND_PREFERENCE_ENV_VAR: &str = "WINIT_UNIX_BACKEND";

#[derive(Clone, Default)]
pub struct PlatformSpecificWindowBuilderAttributes {
    pub visual_infos: Option<XVisualInfo>,
    pub screen_id: Option<i32>,
}

pub enum UnixBackend {
    X(Arc<XConnection>),
    Wayland(Arc<wayland::WaylandContext>),
    Error(Option<XNotSupported>, Option<String>),
}

lazy_static!(
    pub static ref UNIX_BACKEND: UnixBackend = {
        #[inline]
        fn x_backend() -> Result<UnixBackend, XNotSupported> {
            match XConnection::new(Some(x_error_callback)) {
                Ok(x) => Ok(UnixBackend::X(Arc::new(x))),
                Err(e) => Err(e),
            }
        }
        #[inline]
        fn wayland_backend() -> Result<UnixBackend, ()> {
             wayland::WaylandContext::init()
                .map(|ctx| UnixBackend::Wayland(Arc::new(ctx)))
                .ok_or(())
        }
        match env::var(BACKEND_PREFERENCE_ENV_VAR) {
            Ok(s) => match s.as_str() {
                "x11" => x_backend().unwrap_or_else(|e| UnixBackend::Error(Some(e), None)),
                "wayland" => wayland_backend().unwrap_or_else(|_| {
                    UnixBackend::Error(None, Some("Wayland not available".into()))
                }),
                _ => panic!("Unknown environment variable value for {}, try one of `x11`,`wayland`",
                            BACKEND_PREFERENCE_ENV_VAR),
            },
            Err(_) => {
                // Try wayland, fallback to X11
                wayland_backend().unwrap_or_else(|_| {
                    x_backend().unwrap_or_else(|x_err| {
                        UnixBackend::Error(Some(x_err), Some("Wayland not available".into()))
                    })
                })
            },
        }
    };
);

pub enum Window2 {
    #[doc(hidden)]
    X(x11::Window2),
    #[doc(hidden)]
    Wayland(wayland::Window)
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum WindowId {
    #[doc(hidden)]
    X(x11::WindowId),
    #[doc(hidden)]
    Wayland(wayland::WindowId)
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum DeviceId {
    #[doc(hidden)]
    X(x11::DeviceId),
    #[doc(hidden)]
    Wayland(wayland::DeviceId)
}

#[derive(Clone)]
pub enum MonitorId {
    #[doc(hidden)]
    X(x11::MonitorId),
    #[doc(hidden)]
    Wayland(wayland::MonitorId),
    #[doc(hidden)]
    None,
}

#[inline]
pub fn get_available_monitors() -> VecDeque<MonitorId> {
    match *UNIX_BACKEND {
        UnixBackend::Wayland(ref ctxt) => wayland::get_available_monitors(ctxt)
                                .into_iter()
                                .map(MonitorId::Wayland)
                                .collect(),
        UnixBackend::X(ref connec) => x11::get_available_monitors(connec)
                                    .into_iter()
                                    .map(MonitorId::X)
                                    .collect(),
        UnixBackend::Error(..) => { let mut d = VecDeque::new(); d.push_back(MonitorId::None); d},
    }
}

#[inline]
pub fn get_primary_monitor() -> MonitorId {
    match *UNIX_BACKEND {
        UnixBackend::Wayland(ref ctxt) => MonitorId::Wayland(wayland::get_primary_monitor(ctxt)),
        UnixBackend::X(ref connec) => MonitorId::X(x11::get_primary_monitor(connec)),
        UnixBackend::Error(..) => MonitorId::None,
    }
}

impl MonitorId {
    #[inline]
    pub fn get_name(&self) -> Option<String> {
        match self {
            &MonitorId::X(ref m) => m.get_name(),
            &MonitorId::Wayland(ref m) => m.get_name(),
            &MonitorId::None => None,
        }
    }

    #[inline]
    pub fn get_native_identifier(&self) -> ::native_monitor::NativeMonitorId {
        match self {
            &MonitorId::X(ref m) => m.get_native_identifier(),
            &MonitorId::Wayland(ref m) => m.get_native_identifier(),
            &MonitorId::None => unimplemented!()        // FIXME:
        }
    }

    #[inline]
    pub fn get_dimensions(&self) -> (u32, u32) {
        match self {
            &MonitorId::X(ref m) => m.get_dimensions(),
            &MonitorId::Wayland(ref m) => m.get_dimensions(),
            &MonitorId::None => (800, 600),     // FIXME:
        }
    }
}

impl Window2 {
    #[inline]
    pub fn new(events_loop: &EventsLoop,
               window: &::WindowAttributes,
               pl_attribs: &PlatformSpecificWindowBuilderAttributes)
               -> Result<Self, CreationError>
    {
        match *UNIX_BACKEND {
            UnixBackend::Wayland(ref ctxt) => {
                if let EventsLoop::Wayland(ref evlp) = *events_loop {
                    wayland::Window::new(evlp, ctxt.clone(), window).map(Window2::Wayland)
                } else {
                    // It is not possible to instanciate an EventsLoop not matching its backend
                    unreachable!()
                }
            },

            UnixBackend::X(_) => {
                x11::Window2::new(events_loop, window, pl_attribs).map(Window2::X)
            },
            UnixBackend::Error(..) => {
                // If the Backend is Error(), it is not possible to instanciate an EventsLoop at all,
                // thus this function cannot be called!
                unreachable!()
            }
        }
    }

    #[inline]
    pub fn id(&self) -> WindowId {
        match self {
            &Window2::X(ref w) => WindowId::X(w.id()),
            &Window2::Wayland(ref w) => WindowId::Wayland(w.id())
        }
    }

    #[inline]
    pub fn set_title(&self, title: &str) {
        match self {
            &Window2::X(ref w) => w.set_title(title),
            &Window2::Wayland(ref w) => w.set_title(title)
        }
    }

    #[inline]
    pub fn show(&self) {
        match self {
            &Window2::X(ref w) => w.show(),
            &Window2::Wayland(ref w) => w.show()
        }
    }

    #[inline]
    pub fn hide(&self) {
        match self {
            &Window2::X(ref w) => w.hide(),
            &Window2::Wayland(ref w) => w.hide()
        }
    }

    #[inline]
    pub fn get_position(&self) -> Option<(i32, i32)> {
        match self {
            &Window2::X(ref w) => w.get_position(),
            &Window2::Wayland(ref w) => w.get_position()
        }
    }

    #[inline]
    pub fn set_position(&self, x: i32, y: i32) {
        match self {
            &Window2::X(ref w) => w.set_position(x, y),
            &Window2::Wayland(ref w) => w.set_position(x, y)
        }
    }

    #[inline]
    pub fn get_inner_size(&self) -> Option<(u32, u32)> {
        match self {
            &Window2::X(ref w) => w.get_inner_size(),
            &Window2::Wayland(ref w) => w.get_inner_size()
        }
    }

    #[inline]
    pub fn get_outer_size(&self) -> Option<(u32, u32)> {
        match self {
            &Window2::X(ref w) => w.get_outer_size(),
            &Window2::Wayland(ref w) => w.get_outer_size()
        }
    }

    #[inline]
    pub fn set_inner_size(&self, x: u32, y: u32) {
        match self {
            &Window2::X(ref w) => w.set_inner_size(x, y),
            &Window2::Wayland(ref w) => w.set_inner_size(x, y)
        }
    }

    #[inline]
    pub fn set_cursor(&self, cursor: MouseCursor) {
        match self {
            &Window2::X(ref w) => w.set_cursor(cursor),
            &Window2::Wayland(ref w) => w.set_cursor(cursor)
        }
    }

    #[inline]
    pub fn set_cursor_state(&self, state: CursorState) -> Result<(), String> {
        match self {
            &Window2::X(ref w) => w.set_cursor_state(state),
            &Window2::Wayland(ref w) => w.set_cursor_state(state)
        }
    }

    #[inline]
    pub fn hidpi_factor(&self) -> f32 {
       match self {
            &Window2::X(ref w) => w.hidpi_factor(),
            &Window2::Wayland(ref w) => w.hidpi_factor()
        }
    }

    #[inline]
    pub fn set_cursor_position(&self, x: i32, y: i32) -> Result<(), ()> {
        match self {
            &Window2::X(ref w) => w.set_cursor_position(x, y),
            &Window2::Wayland(ref w) => w.set_cursor_position(x, y)
        }
    }

    #[inline]
    pub fn platform_display(&self) -> *mut libc::c_void {
        use wayland_client::Proxy;
        match self {
            &Window2::X(ref w) => w.platform_display(),
            &Window2::Wayland(ref w) => w.get_display().ptr() as *mut _
        }
    }

    #[inline]
    pub fn platform_window(&self) -> *mut libc::c_void {
        use wayland_client::Proxy;
        match self {
            &Window2::X(ref w) => w.platform_window(),
            &Window2::Wayland(ref w) => w.get_surface().ptr() as *mut _
        }
    }
}

unsafe extern "C" fn x_error_callback(dpy: *mut x11::ffi::Display, event: *mut x11::ffi::XErrorEvent)
                                      -> libc::c_int
{
    use std::ffi::CStr;

    if let UnixBackend::X(ref x) = *UNIX_BACKEND {
        let mut buff: Vec<u8> = Vec::with_capacity(1024);
        (x.xlib.XGetErrorText)(dpy, (*event).error_code as i32, buff.as_mut_ptr() as *mut libc::c_char, buff.capacity() as i32);
        let description = CStr::from_ptr(buff.as_mut_ptr() as *const libc::c_char).to_string_lossy();

        let error = XError {
            description: description.into_owned(),
            error_code: (*event).error_code,
            request_code: (*event).request_code,
            minor_code: (*event).minor_code,
        };

        *x.latest_error.lock().unwrap() = Some(error);
    }

    0
}

pub enum EventsLoop {
    #[doc(hidden)]
    Wayland(wayland::EventsLoop),
    #[doc(hidden)]
    X(x11::EventsLoop)
}

pub enum EventsLoopProxy {
    X(x11::EventsLoopProxy),
    Wayland(wayland::EventsLoopProxy),
}

impl EventsLoop {
    pub fn new() -> EventsLoop {
        match *UNIX_BACKEND {
            UnixBackend::Wayland(ref ctxt) => {
                EventsLoop::Wayland(wayland::EventsLoop::new(ctxt.clone()))
            },

            UnixBackend::X(ref ctxt) => {
                EventsLoop::X(x11::EventsLoop::new(ctxt.clone()))
            },

            UnixBackend::Error(..) => {
                panic!("Attempted to create an EventsLoop while no backend was available.")
            }
        }
    }

    pub fn create_proxy(&self) -> EventsLoopProxy {
        match *self {
            EventsLoop::Wayland(ref evlp) => EventsLoopProxy::Wayland(evlp.create_proxy()),
            EventsLoop::X(ref evlp) => EventsLoopProxy::X(evlp.create_proxy()),
        }
    }

    pub fn poll_events<F>(&mut self, callback: F)
        where F: FnMut(::Event)
    {
        match *self {
            EventsLoop::Wayland(ref mut evlp) => evlp.poll_events(callback),
            EventsLoop::X(ref mut evlp) => evlp.poll_events(callback)
        }
    }

    pub fn run_forever<F>(&mut self, callback: F)
        where F: FnMut(::Event) -> ControlFlow
    {
        match *self {
            EventsLoop::Wayland(ref mut evlp) => evlp.run_forever(callback),
            EventsLoop::X(ref mut evlp) => evlp.run_forever(callback)
        }
    }
}

impl EventsLoopProxy {
    pub fn wakeup(&self) -> Result<(), EventsLoopClosed> {
        match *self {
            EventsLoopProxy::Wayland(ref proxy) => proxy.wakeup(),
            EventsLoopProxy::X(ref proxy) => proxy.wakeup(),
        }
    }
}
<commit_msg>Improve unix backend env var docs<commit_after>#![cfg(any(target_os = "linux", target_os = "dragonfly", target_os = "freebsd", target_os = "openbsd"))]

use std::collections::VecDeque;
use std::sync::Arc;
use std::env;

use {CreationError, CursorState, EventsLoopClosed, MouseCursor, ControlFlow};
use libc;

use self::x11::XConnection;
use self::x11::XError;
use self::x11::XNotSupported;
use self::x11::ffi::XVisualInfo;

mod dlopen;
pub mod wayland;
pub mod x11;

/// Environment variable that indicates a backend preference
/// `WINIT_UNIX_BACKEND=x11` : Will try to use an X11 backend
/// `WINIT_UNIX_BACKEND=wayland` : Will try to use a wayland backend
/// When the variable is present the indicated backend will not fallback to another
/// (Current default behaviour, without the env var, is try wayland then fallback to X)
const BACKEND_PREFERENCE_ENV_VAR: &str = "WINIT_UNIX_BACKEND";

#[derive(Clone, Default)]
pub struct PlatformSpecificWindowBuilderAttributes {
    pub visual_infos: Option<XVisualInfo>,
    pub screen_id: Option<i32>,
}

pub enum UnixBackend {
    X(Arc<XConnection>),
    Wayland(Arc<wayland::WaylandContext>),
    Error(Option<XNotSupported>, Option<String>),
}

lazy_static!(
    pub static ref UNIX_BACKEND: UnixBackend = {
        #[inline]
        fn x_backend() -> Result<UnixBackend, XNotSupported> {
            match XConnection::new(Some(x_error_callback)) {
                Ok(x) => Ok(UnixBackend::X(Arc::new(x))),
                Err(e) => Err(e),
            }
        }
        #[inline]
        fn wayland_backend() -> Result<UnixBackend, ()> {
             wayland::WaylandContext::init()
                .map(|ctx| UnixBackend::Wayland(Arc::new(ctx)))
                .ok_or(())
        }
        match env::var(BACKEND_PREFERENCE_ENV_VAR) {
            Ok(s) => match s.as_str() {
                "x11" => x_backend().unwrap_or_else(|e| UnixBackend::Error(Some(e), None)),
                "wayland" => wayland_backend().unwrap_or_else(|_| {
                    UnixBackend::Error(None, Some("Wayland not available".into()))
                }),
                _ => panic!("Unknown environment variable value for {}, try one of `x11`,`wayland`",
                            BACKEND_PREFERENCE_ENV_VAR),
            },
            Err(_) => {
                // Try wayland, fallback to X11
                wayland_backend().unwrap_or_else(|_| {
                    x_backend().unwrap_or_else(|x_err| {
                        UnixBackend::Error(Some(x_err), Some("Wayland not available".into()))
                    })
                })
            },
        }
    };
);

pub enum Window2 {
    #[doc(hidden)]
    X(x11::Window2),
    #[doc(hidden)]
    Wayland(wayland::Window)
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum WindowId {
    #[doc(hidden)]
    X(x11::WindowId),
    #[doc(hidden)]
    Wayland(wayland::WindowId)
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum DeviceId {
    #[doc(hidden)]
    X(x11::DeviceId),
    #[doc(hidden)]
    Wayland(wayland::DeviceId)
}

#[derive(Clone)]
pub enum MonitorId {
    #[doc(hidden)]
    X(x11::MonitorId),
    #[doc(hidden)]
    Wayland(wayland::MonitorId),
    #[doc(hidden)]
    None,
}

#[inline]
pub fn get_available_monitors() -> VecDeque<MonitorId> {
    match *UNIX_BACKEND {
        UnixBackend::Wayland(ref ctxt) => wayland::get_available_monitors(ctxt)
                                .into_iter()
                                .map(MonitorId::Wayland)
                                .collect(),
        UnixBackend::X(ref connec) => x11::get_available_monitors(connec)
                                    .into_iter()
                                    .map(MonitorId::X)
                                    .collect(),
        UnixBackend::Error(..) => { let mut d = VecDeque::new(); d.push_back(MonitorId::None); d},
    }
}

#[inline]
pub fn get_primary_monitor() -> MonitorId {
    match *UNIX_BACKEND {
        UnixBackend::Wayland(ref ctxt) => MonitorId::Wayland(wayland::get_primary_monitor(ctxt)),
        UnixBackend::X(ref connec) => MonitorId::X(x11::get_primary_monitor(connec)),
        UnixBackend::Error(..) => MonitorId::None,
    }
}

impl MonitorId {
    #[inline]
    pub fn get_name(&self) -> Option<String> {
        match self {
            &MonitorId::X(ref m) => m.get_name(),
            &MonitorId::Wayland(ref m) => m.get_name(),
            &MonitorId::None => None,
        }
    }

    #[inline]
    pub fn get_native_identifier(&self) -> ::native_monitor::NativeMonitorId {
        match self {
            &MonitorId::X(ref m) => m.get_native_identifier(),
            &MonitorId::Wayland(ref m) => m.get_native_identifier(),
            &MonitorId::None => unimplemented!()        // FIXME:
        }
    }

    #[inline]
    pub fn get_dimensions(&self) -> (u32, u32) {
        match self {
            &MonitorId::X(ref m) => m.get_dimensions(),
            &MonitorId::Wayland(ref m) => m.get_dimensions(),
            &MonitorId::None => (800, 600),     // FIXME:
        }
    }
}

impl Window2 {
    #[inline]
    pub fn new(events_loop: &EventsLoop,
               window: &::WindowAttributes,
               pl_attribs: &PlatformSpecificWindowBuilderAttributes)
               -> Result<Self, CreationError>
    {
        match *UNIX_BACKEND {
            UnixBackend::Wayland(ref ctxt) => {
                if let EventsLoop::Wayland(ref evlp) = *events_loop {
                    wayland::Window::new(evlp, ctxt.clone(), window).map(Window2::Wayland)
                } else {
                    // It is not possible to instanciate an EventsLoop not matching its backend
                    unreachable!()
                }
            },

            UnixBackend::X(_) => {
                x11::Window2::new(events_loop, window, pl_attribs).map(Window2::X)
            },
            UnixBackend::Error(..) => {
                // If the Backend is Error(), it is not possible to instanciate an EventsLoop at all,
                // thus this function cannot be called!
                unreachable!()
            }
        }
    }

    #[inline]
    pub fn id(&self) -> WindowId {
        match self {
            &Window2::X(ref w) => WindowId::X(w.id()),
            &Window2::Wayland(ref w) => WindowId::Wayland(w.id())
        }
    }

    #[inline]
    pub fn set_title(&self, title: &str) {
        match self {
            &Window2::X(ref w) => w.set_title(title),
            &Window2::Wayland(ref w) => w.set_title(title)
        }
    }

    #[inline]
    pub fn show(&self) {
        match self {
            &Window2::X(ref w) => w.show(),
            &Window2::Wayland(ref w) => w.show()
        }
    }

    #[inline]
    pub fn hide(&self) {
        match self {
            &Window2::X(ref w) => w.hide(),
            &Window2::Wayland(ref w) => w.hide()
        }
    }

    #[inline]
    pub fn get_position(&self) -> Option<(i32, i32)> {
        match self {
            &Window2::X(ref w) => w.get_position(),
            &Window2::Wayland(ref w) => w.get_position()
        }
    }

    #[inline]
    pub fn set_position(&self, x: i32, y: i32) {
        match self {
            &Window2::X(ref w) => w.set_position(x, y),
            &Window2::Wayland(ref w) => w.set_position(x, y)
        }
    }

    #[inline]
    pub fn get_inner_size(&self) -> Option<(u32, u32)> {
        match self {
            &Window2::X(ref w) => w.get_inner_size(),
            &Window2::Wayland(ref w) => w.get_inner_size()
        }
    }

    #[inline]
    pub fn get_outer_size(&self) -> Option<(u32, u32)> {
        match self {
            &Window2::X(ref w) => w.get_outer_size(),
            &Window2::Wayland(ref w) => w.get_outer_size()
        }
    }

    #[inline]
    pub fn set_inner_size(&self, x: u32, y: u32) {
        match self {
            &Window2::X(ref w) => w.set_inner_size(x, y),
            &Window2::Wayland(ref w) => w.set_inner_size(x, y)
        }
    }

    #[inline]
    pub fn set_cursor(&self, cursor: MouseCursor) {
        match self {
            &Window2::X(ref w) => w.set_cursor(cursor),
            &Window2::Wayland(ref w) => w.set_cursor(cursor)
        }
    }

    #[inline]
    pub fn set_cursor_state(&self, state: CursorState) -> Result<(), String> {
        match self {
            &Window2::X(ref w) => w.set_cursor_state(state),
            &Window2::Wayland(ref w) => w.set_cursor_state(state)
        }
    }

    #[inline]
    pub fn hidpi_factor(&self) -> f32 {
       match self {
            &Window2::X(ref w) => w.hidpi_factor(),
            &Window2::Wayland(ref w) => w.hidpi_factor()
        }
    }

    #[inline]
    pub fn set_cursor_position(&self, x: i32, y: i32) -> Result<(), ()> {
        match self {
            &Window2::X(ref w) => w.set_cursor_position(x, y),
            &Window2::Wayland(ref w) => w.set_cursor_position(x, y)
        }
    }

    #[inline]
    pub fn platform_display(&self) -> *mut libc::c_void {
        use wayland_client::Proxy;
        match self {
            &Window2::X(ref w) => w.platform_display(),
            &Window2::Wayland(ref w) => w.get_display().ptr() as *mut _
        }
    }

    #[inline]
    pub fn platform_window(&self) -> *mut libc::c_void {
        use wayland_client::Proxy;
        match self {
            &Window2::X(ref w) => w.platform_window(),
            &Window2::Wayland(ref w) => w.get_surface().ptr() as *mut _
        }
    }
}

unsafe extern "C" fn x_error_callback(dpy: *mut x11::ffi::Display, event: *mut x11::ffi::XErrorEvent)
                                      -> libc::c_int
{
    use std::ffi::CStr;

    if let UnixBackend::X(ref x) = *UNIX_BACKEND {
        let mut buff: Vec<u8> = Vec::with_capacity(1024);
        (x.xlib.XGetErrorText)(dpy, (*event).error_code as i32, buff.as_mut_ptr() as *mut libc::c_char, buff.capacity() as i32);
        let description = CStr::from_ptr(buff.as_mut_ptr() as *const libc::c_char).to_string_lossy();

        let error = XError {
            description: description.into_owned(),
            error_code: (*event).error_code,
            request_code: (*event).request_code,
            minor_code: (*event).minor_code,
        };

        *x.latest_error.lock().unwrap() = Some(error);
    }

    0
}

pub enum EventsLoop {
    #[doc(hidden)]
    Wayland(wayland::EventsLoop),
    #[doc(hidden)]
    X(x11::EventsLoop)
}

pub enum EventsLoopProxy {
    X(x11::EventsLoopProxy),
    Wayland(wayland::EventsLoopProxy),
}

impl EventsLoop {
    pub fn new() -> EventsLoop {
        match *UNIX_BACKEND {
            UnixBackend::Wayland(ref ctxt) => {
                EventsLoop::Wayland(wayland::EventsLoop::new(ctxt.clone()))
            },

            UnixBackend::X(ref ctxt) => {
                EventsLoop::X(x11::EventsLoop::new(ctxt.clone()))
            },

            UnixBackend::Error(..) => {
                panic!("Attempted to create an EventsLoop while no backend was available.")
            }
        }
    }

    pub fn create_proxy(&self) -> EventsLoopProxy {
        match *self {
            EventsLoop::Wayland(ref evlp) => EventsLoopProxy::Wayland(evlp.create_proxy()),
            EventsLoop::X(ref evlp) => EventsLoopProxy::X(evlp.create_proxy()),
        }
    }

    pub fn poll_events<F>(&mut self, callback: F)
        where F: FnMut(::Event)
    {
        match *self {
            EventsLoop::Wayland(ref mut evlp) => evlp.poll_events(callback),
            EventsLoop::X(ref mut evlp) => evlp.poll_events(callback)
        }
    }

    pub fn run_forever<F>(&mut self, callback: F)
        where F: FnMut(::Event) -> ControlFlow
    {
        match *self {
            EventsLoop::Wayland(ref mut evlp) => evlp.run_forever(callback),
            EventsLoop::X(ref mut evlp) => evlp.run_forever(callback)
        }
    }
}

impl EventsLoopProxy {
    pub fn wakeup(&self) -> Result<(), EventsLoopClosed> {
        match *self {
            EventsLoopProxy::Wayland(ref proxy) => proxy.wakeup(),
            EventsLoopProxy::X(ref proxy) => proxy.wakeup(),
        }
    }
}
<|endoftext|>


<commit_before>use antidote::RwLock;
use irc::client::prelude::*;
use std::thread::{sleep, spawn};
use std::{fmt, sync::Arc, time::Duration};

use chrono::{self, NaiveDateTime};
use time;

use plugin::*;
pub mod database;
mod parser;
use self::database::Database;
use self::parser::CommandParser;

use self::error::*;
use error::ErrorKind as FrippyErrorKind;
use error::FrippyError;
use failure::ResultExt;

fn get_time() -> NaiveDateTime {
    let tm = time::now().to_timespec();
    NaiveDateTime::from_timestamp(tm.sec, 0u32)
}

fn get_events<T: Database>(db: &RwLock<T>, in_next: chrono::Duration) -> Vec<database::Event> {
    loop {
        let before = get_time() + in_next;
        match db.read().get_events_before(&before) {
            Ok(events) => return events,
            Err(e) => {
                if e.kind() != ErrorKind::NotFound {
                    error!("Failed to get events: {}", e);
                }
            }
        }

        sleep(in_next.to_std().expect("Failed to convert look ahead time"));
    }
}

fn run<T: Database>(client: &IrcClient, db: Arc<RwLock<T>>) {
    let look_ahead = chrono::Duration::minutes(2);

    let mut events = get_events(&db, look_ahead);

    let mut sleep_time = look_ahead
        .to_std()
        .expect("Failed to convert look ahead time");

    loop {
        let now = get_time();
        for event in events {
            if event.time <= now {
                let msg = format!("Reminder from {}: {}", event.author, event.content);
                if let Err(e) = client.send_notice(&event.receiver, &msg) {
                    error!("Failed to send reminder: {}", e);
                } else {
                    debug!("Sent reminder {:?}", event);

                    if let Some(repeat) = event.repeat {
                        let next_time = event.time + chrono::Duration::seconds(repeat);

                        if let Err(e) = db.write().update_event_time(event.id, &next_time) {
                            error!("Failed to update reminder: {}", e);
                        } else {
                            debug!("Updated time");
                        }
                    } else if let Err(e) = db.write().delete_event(event.id) {
                        error!("Failed to delete reminder: {}", e);
                    }
                }
            } else {
                let until_event = (event.time - now)
                    .to_std()
                    .expect("Failed to convert until event time");

                if until_event < sleep_time {
                    sleep_time = until_event + Duration::from_secs(1);
                }
            }
        }

        sleep(sleep_time);
        sleep_time = Duration::from_secs(120);

        events = get_events(&db, look_ahead);
    }
}

#[derive(PluginName)]
pub struct Remind<T: 'static + Database> {
    events: Arc<RwLock<T>>,
    has_reminder: RwLock<bool>,
}

impl<T: 'static + Database> Remind<T> {
    pub fn new(db: T) -> Self {
        let events = Arc::new(RwLock::new(db));

        Remind {
            events: events,
            has_reminder: RwLock::new(false),
        }
    }

    fn user_cmd(&self, command: PluginCommand) -> Result<String, RemindError> {
        let parser = CommandParser::parse_target(command.tokens)?;

        self.set(parser, &command.source)
    }

    fn me_cmd(&self, command: PluginCommand) -> Result<String, RemindError> {
        let source = command.source.clone();
        let parser = CommandParser::with_target(command.tokens, command.source)?;

        self.set(parser, &source)
    }

    fn set(&self, parser: CommandParser, author: &str) -> Result<String, RemindError> {
        debug!("parser: {:?}", parser);

        let target = parser.get_target();
        let time = parser.get_time(Duration::from_secs(120))?;

        let event = database::NewEvent {
            receiver: target,
            content: &parser.get_message(),
            author: author,
            time: &time,
            repeat: parser
                .get_repeat(Duration::from_secs(600))?
                .map(|d| d.as_secs() as i64),
        };

        debug!("New event: {:?}", event);

        Ok(self.events
            .write()
            .insert_event(&event)
            .map(|id| format!("Created reminder with id {} at {} UTC", id, time))?)
    }

    fn list(&self, user: &str) -> Result<String, RemindError> {
        let mut events = self.events.read().get_user_events(user)?;

        if events.is_empty() {
            Err(ErrorKind::NotFound)?;
        }

        let mut list = events.remove(0).to_string();
        for ev in events {
            list.push_str("\r\n");
            list.push_str(&ev.to_string());
        }

        Ok(list)
    }

    fn delete(&self, mut command: PluginCommand) -> Result<&str, RemindError> {
        let id = command
            .tokens
            .remove(0)
            .parse::<i64>()
            .context(ErrorKind::Parsing)?;
        let event = self.events
            .read()
            .get_event(id)
            .context(ErrorKind::NotFound)?;

        if event.receiver.eq_ignore_ascii_case(&command.source)
            || event.author.eq_ignore_ascii_case(&command.source)
        {
            self.events
                .write()
                .delete_event(id)
                .map(|()| "Successfully deleted")
        } else {
            Ok("Only the author or receiver can delete a reminder")
        }
    }

    fn help(&self) -> &str {
        "usage: remind <subcommand>\r\n\
         subcommands: user, list, delete, help\r\n\
         examples\r\n\
         remind user foo to sleep in 1 hour\r\n\
         remind user bar to leave early on 1.1 at 16:00 every week"
    }
}

impl<T: Database> Plugin for Remind<T> {
    fn execute(&self, client: &IrcClient, msg: &Message) -> ExecutionStatus {
        if let Command::JOIN(_, _, _) = msg.command {
            let mut has_reminder = self.has_reminder.write();

            if !*has_reminder {
                let events = Arc::clone(&self.events);
                let client = client.clone();

                spawn(move || run(&client, events));

                *has_reminder = true;
            }
        }

        ExecutionStatus::Done
    }

    fn execute_threaded(&self, _: &IrcClient, _: &Message) -> Result<(), FrippyError> {
        panic!("Remind should not use frippy's threading")
    }

    fn command(&self, client: &IrcClient, mut command: PluginCommand) -> Result<(), FrippyError> {
        if command.tokens.is_empty() {
            return Ok(client
                .send_notice(&command.source, &ErrorKind::InvalidCommand.to_string())
                .context(FrippyErrorKind::Connection)?);
        }

        let source = command.source.clone();

        let sub_command = command.tokens.remove(0);
        let response = match sub_command.as_ref() {
            "user" => self.user_cmd(command),
            "me" => self.me_cmd(command),
            "delete" => self.delete(command).map(|s| s.to_owned()),
            "list" => self.list(&source),
            "help" => Ok(self.help().to_owned()),
            _ => Err(ErrorKind::InvalidCommand.into()),
        };

        let result = match response {
            Ok(msg) => client
                .send_notice(&source, &msg)
                .context(FrippyErrorKind::Connection)?,
            Err(e) => {
                let message = e.to_string();

                client
                    .send_notice(&source, &message)
                    .context(FrippyErrorKind::Connection)?;

                Err(e).context(FrippyErrorKind::Remind)?
            }
        };

        Ok(result)
    }

    fn evaluate(&self, _: &IrcClient, _: PluginCommand) -> Result<String, String> {
        Err(String::from(
            "Evaluation of commands is not implemented for remind at this time",
        ))
    }
}

impl<T: Database> fmt::Debug for Remind<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Remind {{ ... }}")
    }
}

pub mod error {
    #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)]
    #[error = "RemindError"]
    pub enum ErrorKind {
        /// Invalid command error
        #[fail(display = "Incorrect Command. Send \"remind help\" for help.")]
        InvalidCommand,

        /// Missing message error
        #[fail(display = "Reminder needs to have a description")]
        MissingMessage,

        /// Missing receiver error
        #[fail(display = "Specify who to remind")]
        MissingReceiver,

        /// Missing time error
        #[fail(display = "Reminder needs to have a time")]
        MissingTime,

        /// Invalid time error
        #[fail(display = "Could not parse time")]
        InvalidTime,

        /// Invalid date error
        #[fail(display = "Could not parse date")]
        InvalidDate,

        /// Parse error
        #[fail(display = "Could not parse integers")]
        Parsing,

        /// Ambigous time error
        #[fail(display = "Time specified is ambiguous")]
        AmbiguousTime,

        /// Time too short error
        #[fail(display = "Reminder needs to be in over 2 minutes")]
        TimeShort,

        /// Repeat time too short error
        #[fail(display = "Repeat time needs to be over 5 minutes")]
        RepeatTimeShort,

        /// Duplicate error
        #[fail(display = "Entry already exists")]
        Duplicate,

        /// Not found error
        #[fail(display = "No events found")]
        NotFound,

        /// MySQL error
        #[cfg(feature = "mysql")]
        #[fail(display = "Failed to execute MySQL Query")]
        MysqlError,

        /// No connection error
        #[cfg(feature = "mysql")]
        #[fail(display = "No connection to the database")]
        NoConnection,
    }
}
<commit_msg>Remind: Fix help message<commit_after>use antidote::RwLock;
use irc::client::prelude::*;
use std::thread::{sleep, spawn};
use std::{fmt, sync::Arc, time::Duration};

use chrono::{self, NaiveDateTime};
use time;

use plugin::*;
pub mod database;
mod parser;
use self::database::Database;
use self::parser::CommandParser;

use self::error::*;
use error::ErrorKind as FrippyErrorKind;
use error::FrippyError;
use failure::ResultExt;

fn get_time() -> NaiveDateTime {
    let tm = time::now().to_timespec();
    NaiveDateTime::from_timestamp(tm.sec, 0u32)
}

fn get_events<T: Database>(db: &RwLock<T>, in_next: chrono::Duration) -> Vec<database::Event> {
    loop {
        let before = get_time() + in_next;
        match db.read().get_events_before(&before) {
            Ok(events) => return events,
            Err(e) => {
                if e.kind() != ErrorKind::NotFound {
                    error!("Failed to get events: {}", e);
                }
            }
        }

        sleep(in_next.to_std().expect("Failed to convert look ahead time"));
    }
}

fn run<T: Database>(client: &IrcClient, db: Arc<RwLock<T>>) {
    let look_ahead = chrono::Duration::minutes(2);

    let mut events = get_events(&db, look_ahead);

    let mut sleep_time = look_ahead
        .to_std()
        .expect("Failed to convert look ahead time");

    loop {
        let now = get_time();
        for event in events {
            if event.time <= now {
                let msg = format!("Reminder from {}: {}", event.author, event.content);
                if let Err(e) = client.send_notice(&event.receiver, &msg) {
                    error!("Failed to send reminder: {}", e);
                } else {
                    debug!("Sent reminder {:?}", event);

                    if let Some(repeat) = event.repeat {
                        let next_time = event.time + chrono::Duration::seconds(repeat);

                        if let Err(e) = db.write().update_event_time(event.id, &next_time) {
                            error!("Failed to update reminder: {}", e);
                        } else {
                            debug!("Updated time");
                        }
                    } else if let Err(e) = db.write().delete_event(event.id) {
                        error!("Failed to delete reminder: {}", e);
                    }
                }
            } else {
                let until_event = (event.time - now)
                    .to_std()
                    .expect("Failed to convert until event time");

                if until_event < sleep_time {
                    sleep_time = until_event + Duration::from_secs(1);
                }
            }
        }

        sleep(sleep_time);
        sleep_time = Duration::from_secs(120);

        events = get_events(&db, look_ahead);
    }
}

#[derive(PluginName)]
pub struct Remind<T: 'static + Database> {
    events: Arc<RwLock<T>>,
    has_reminder: RwLock<bool>,
}

impl<T: 'static + Database> Remind<T> {
    pub fn new(db: T) -> Self {
        let events = Arc::new(RwLock::new(db));

        Remind {
            events: events,
            has_reminder: RwLock::new(false),
        }
    }

    fn user_cmd(&self, command: PluginCommand) -> Result<String, RemindError> {
        let parser = CommandParser::parse_target(command.tokens)?;

        self.set(parser, &command.source)
    }

    fn me_cmd(&self, command: PluginCommand) -> Result<String, RemindError> {
        let source = command.source.clone();
        let parser = CommandParser::with_target(command.tokens, command.source)?;

        self.set(parser, &source)
    }

    fn set(&self, parser: CommandParser, author: &str) -> Result<String, RemindError> {
        debug!("parser: {:?}", parser);

        let target = parser.get_target();
        let time = parser.get_time(Duration::from_secs(120))?;

        let event = database::NewEvent {
            receiver: target,
            content: &parser.get_message(),
            author: author,
            time: &time,
            repeat: parser
                .get_repeat(Duration::from_secs(600))?
                .map(|d| d.as_secs() as i64),
        };

        debug!("New event: {:?}", event);

        Ok(self.events
            .write()
            .insert_event(&event)
            .map(|id| format!("Created reminder with id {} at {} UTC", id, time))?)
    }

    fn list(&self, user: &str) -> Result<String, RemindError> {
        let mut events = self.events.read().get_user_events(user)?;

        if events.is_empty() {
            Err(ErrorKind::NotFound)?;
        }

        let mut list = events.remove(0).to_string();
        for ev in events {
            list.push_str("\r\n");
            list.push_str(&ev.to_string());
        }

        Ok(list)
    }

    fn delete(&self, mut command: PluginCommand) -> Result<&str, RemindError> {
        let id = command
            .tokens
            .remove(0)
            .parse::<i64>()
            .context(ErrorKind::Parsing)?;
        let event = self.events
            .read()
            .get_event(id)
            .context(ErrorKind::NotFound)?;

        if event.receiver.eq_ignore_ascii_case(&command.source)
            || event.author.eq_ignore_ascii_case(&command.source)
        {
            self.events
                .write()
                .delete_event(id)
                .map(|()| "Successfully deleted")
        } else {
            Ok("Only the author or receiver can delete a reminder")
        }
    }

    fn help(&self) -> &str {
        "usage: remind <subcommand>\r\n\
         subcommands: user, me, list, delete, help\r\n\
         examples\r\n\
         remind user foo to sleep in 1 hour\r\n\
         remind me to leave early on 1.1 at 16:00 every week"
    }
}

impl<T: Database> Plugin for Remind<T> {
    fn execute(&self, client: &IrcClient, msg: &Message) -> ExecutionStatus {
        if let Command::JOIN(_, _, _) = msg.command {
            let mut has_reminder = self.has_reminder.write();

            if !*has_reminder {
                let events = Arc::clone(&self.events);
                let client = client.clone();

                spawn(move || run(&client, events));

                *has_reminder = true;
            }
        }

        ExecutionStatus::Done
    }

    fn execute_threaded(&self, _: &IrcClient, _: &Message) -> Result<(), FrippyError> {
        panic!("Remind should not use frippy's threading")
    }

    fn command(&self, client: &IrcClient, mut command: PluginCommand) -> Result<(), FrippyError> {
        if command.tokens.is_empty() {
            return Ok(client
                .send_notice(&command.source, &ErrorKind::InvalidCommand.to_string())
                .context(FrippyErrorKind::Connection)?);
        }

        let source = command.source.clone();

        let sub_command = command.tokens.remove(0);
        let response = match sub_command.as_ref() {
            "user" => self.user_cmd(command),
            "me" => self.me_cmd(command),
            "delete" => self.delete(command).map(|s| s.to_owned()),
            "list" => self.list(&source),
            "help" => Ok(self.help().to_owned()),
            _ => Err(ErrorKind::InvalidCommand.into()),
        };

        let result = match response {
            Ok(msg) => client
                .send_notice(&source, &msg)
                .context(FrippyErrorKind::Connection)?,
            Err(e) => {
                let message = e.to_string();

                client
                    .send_notice(&source, &message)
                    .context(FrippyErrorKind::Connection)?;

                Err(e).context(FrippyErrorKind::Remind)?
            }
        };

        Ok(result)
    }

    fn evaluate(&self, _: &IrcClient, _: PluginCommand) -> Result<String, String> {
        Err(String::from(
            "Evaluation of commands is not implemented for remind at this time",
        ))
    }
}

impl<T: Database> fmt::Debug for Remind<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Remind {{ ... }}")
    }
}

pub mod error {
    #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)]
    #[error = "RemindError"]
    pub enum ErrorKind {
        /// Invalid command error
        #[fail(display = "Incorrect Command. Send \"remind help\" for help.")]
        InvalidCommand,

        /// Missing message error
        #[fail(display = "Reminder needs to have a description")]
        MissingMessage,

        /// Missing receiver error
        #[fail(display = "Specify who to remind")]
        MissingReceiver,

        /// Missing time error
        #[fail(display = "Reminder needs to have a time")]
        MissingTime,

        /// Invalid time error
        #[fail(display = "Could not parse time")]
        InvalidTime,

        /// Invalid date error
        #[fail(display = "Could not parse date")]
        InvalidDate,

        /// Parse error
        #[fail(display = "Could not parse integers")]
        Parsing,

        /// Ambigous time error
        #[fail(display = "Time specified is ambiguous")]
        AmbiguousTime,

        /// Time too short error
        #[fail(display = "Reminder needs to be in over 2 minutes")]
        TimeShort,

        /// Repeat time too short error
        #[fail(display = "Repeat time needs to be over 5 minutes")]
        RepeatTimeShort,

        /// Duplicate error
        #[fail(display = "Entry already exists")]
        Duplicate,

        /// Not found error
        #[fail(display = "No events found")]
        NotFound,

        /// MySQL error
        #[cfg(feature = "mysql")]
        #[fail(display = "Failed to execute MySQL Query")]
        MysqlError,

        /// No connection error
        #[cfg(feature = "mysql")]
        #[fail(display = "No connection to the database")]
        NoConnection,
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::borrow::Cow;
use std::mem::transmute;
use std::io::{Cursor, Read, Write, ErrorKind};
use std::default::Default;
use std::string::{String, FromUtf8Error};
use std::result::Result as StdResult;
use byteorder::{ByteOrder, ReadBytesExt, NetworkEndian};
use bytes::BufMut;

use error::{Error, Result};
use super::coding::{OpCode, Control, Data, CloseCode};
use super::mask::{generate_mask, apply_mask};

/// A struct representing the close command.
#[derive(Debug, Clone)]
pub struct CloseFrame<'t> {
    /// The reason as a code.
    pub code: CloseCode,
    /// The reason as text string.
    pub reason: Cow<'t, str>,
}

impl<'t> CloseFrame<'t> {
    /// Convert into a owned string.
    pub fn into_owned(self) -> CloseFrame<'static> {
        CloseFrame {
            code: self.code,
            reason: self.reason.into_owned().into(),
        }
    }
}

impl<'t> fmt::Display for CloseFrame<'t> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} ({})", self.reason, self.code)
    }
}

/// A struct representing a WebSocket frame.
#[derive(Debug, Clone)]
pub struct Frame {
    finished: bool,
    rsv1: bool,
    rsv2: bool,
    rsv3: bool,
    opcode: OpCode,

    mask: Option<[u8; 4]>,

    payload: Vec<u8>,
}

impl Frame {

    /// Get the length of the frame.
    /// This is the length of the header + the length of the payload.
    #[inline]
    pub fn len(&self) -> usize {
        let mut header_length = 2;
        let payload_len = self.payload().len();
        if payload_len > 125 {
            if payload_len <= u16::max_value() as usize {
                header_length += 2;
            } else {
                header_length += 8;
            }
        }

        if self.is_masked() {
            header_length += 4;
        }

        header_length + payload_len
    }

    /// Test whether the frame is a final frame.
    #[inline]
    pub fn is_final(&self) -> bool {
        self.finished
    }

    /// Test whether the first reserved bit is set.
    #[inline]
    pub fn has_rsv1(&self) -> bool {
        self.rsv1
    }

    /// Test whether the second reserved bit is set.
    #[inline]
    pub fn has_rsv2(&self) -> bool {
        self.rsv2
    }

    /// Test whether the third reserved bit is set.
    #[inline]
    pub fn has_rsv3(&self) -> bool {
        self.rsv3
    }

    /// Get the OpCode of the frame.
    #[inline]
    pub fn opcode(&self) -> OpCode {
        self.opcode
    }

    /// Get a reference to the frame's payload.
    #[inline]
    pub fn payload(&self) -> &Vec<u8> {
        &self.payload
    }

    // Test whether the frame is masked.
    #[doc(hidden)]
    #[inline]
    pub fn is_masked(&self) -> bool {
        self.mask.is_some()
    }

    // Get an optional reference to the frame's mask.
    #[doc(hidden)]
    #[allow(dead_code)]
    #[inline]
    pub fn mask(&self) -> Option<&[u8; 4]> {
        self.mask.as_ref()
    }

    /// Make this frame a final frame.
    #[allow(dead_code)]
    #[inline]
    pub fn set_final(&mut self, is_final: bool) -> &mut Frame {
        self.finished = is_final;
        self
    }

    /// Set the first reserved bit.
    #[inline]
    pub fn set_rsv1(&mut self, has_rsv1: bool) -> &mut Frame {
        self.rsv1 = has_rsv1;
        self
    }

    /// Set the second reserved bit.
    #[inline]
    pub fn set_rsv2(&mut self, has_rsv2: bool) -> &mut Frame {
        self.rsv2 = has_rsv2;
        self
    }

    /// Set the third reserved bit.
    #[inline]
    pub fn set_rsv3(&mut self, has_rsv3: bool) -> &mut Frame {
        self.rsv3 = has_rsv3;
        self
    }

    /// Set the OpCode.
    #[allow(dead_code)]
    #[inline]
    pub fn set_opcode(&mut self, opcode: OpCode) -> &mut Frame {
        self.opcode = opcode;
        self
    }

    /// Edit the frame's payload.
    #[allow(dead_code)]
    #[inline]
    pub fn payload_mut(&mut self) -> &mut Vec<u8> {
        &mut self.payload
    }

    // Generate a new mask for this frame.
    //
    // This method simply generates and stores the mask. It does not change the payload data.
    // Instead, the payload data will be masked with the generated mask when the frame is sent
    // to the other endpoint.
    #[doc(hidden)]
    #[inline]
    pub fn set_mask(&mut self) -> &mut Frame {
        self.mask = Some(generate_mask());
        self
    }

    // This method unmasks the payload and should only be called on frames that are actually
    // masked. In other words, those frames that have just been received from a client endpoint.
    #[doc(hidden)]
    #[inline]
    pub fn remove_mask(&mut self) {
        self.mask.and_then(|mask| {
            Some(apply_mask(&mut self.payload, &mask))
        });
        self.mask = None;
    }

    /// Consume the frame into its payload as binary.
    #[inline]
    pub fn into_data(self) -> Vec<u8> {
        self.payload
    }

    /// Consume the frame into its payload as string.
    #[inline]
    pub fn into_string(self) -> StdResult<String, FromUtf8Error> {
        String::from_utf8(self.payload)
    }

     /// Consume the frame into a closing frame.
    #[inline]
    pub fn into_close(self) -> Result<Option<CloseFrame<'static>>> {
        match self.payload.len() {
            0 => Ok(None),
            1 => Err(Error::Protocol("Invalid close sequence".into())),
            _ => {
                let mut data = self.payload;
                let code = NetworkEndian::read_u16(&data[0..2]).into();
                data.drain(0..2);
                let text = String::from_utf8(data)?;
                Ok(Some(CloseFrame { code: code, reason: text.into() }))
            }
        }
    }

    /// Create a new data frame.
    #[inline]
    pub fn message(data: Vec<u8>, code: OpCode, finished: bool) -> Frame {
        debug_assert!(match code {
            OpCode::Data(_) => true,
            _ => false,
        }, "Invalid opcode for data frame.");

        Frame {
            finished: finished,
            opcode: code,
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Pong control frame.
    #[inline]
    pub fn pong(data: Vec<u8>) -> Frame {
        Frame {
            opcode: OpCode::Control(Control::Pong),
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Ping control frame.
    #[inline]
    pub fn ping(data: Vec<u8>) -> Frame {
        Frame {
            opcode: OpCode::Control(Control::Ping),
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Close control frame.
    #[inline]
    pub fn close(msg: Option<CloseFrame>) -> Frame {
        let payload = if let Some(CloseFrame { code, reason }) = msg {
            let raw: [u8; 2] = unsafe {
                let u: u16 = code.into();
                transmute(u.to_be())
            };
            [&raw[..], reason.as_bytes()].concat()
        } else {
            Vec::new()
        };

        Frame {
            payload: payload,
            .. Frame::default()
        }
    }

    /// Parse the input stream into a frame.
    pub fn parse(cursor: &mut Cursor<Vec<u8>>) -> Result<Option<Frame>> {
        let size = cursor.get_ref().len() as u64 - cursor.position();
        let initial = cursor.position();
        trace!("Position in buffer {}", initial);

        let mut head = [0u8; 2];
        if try!(cursor.read(&mut head)) != 2 {
            cursor.set_position(initial);
            return Ok(None)
        }

        trace!("Parsed headers {:?}", head);

        let first = head[0];
        let second = head[1];
        trace!("First: {:b}", first);
        trace!("Second: {:b}", second);

        let finished = first & 0x80 != 0;

        let rsv1 = first & 0x40 != 0;
        let rsv2 = first & 0x20 != 0;
        let rsv3 = first & 0x10 != 0;

        let opcode = OpCode::from(first & 0x0F);
        trace!("Opcode: {:?}", opcode);

        let masked = second & 0x80 != 0;
        trace!("Masked: {:?}", masked);

        let mut header_length = 2;

        let mut length = (second & 0x7F) as u64;

        if let Some(length_nbytes) = match length {
            126 => Some(2),
            127 => Some(8),
            _ => None,
        } {
            match cursor.read_uint::<NetworkEndian>(length_nbytes) {
                Err(ref err) if err.kind() == ErrorKind::UnexpectedEof => {
                    cursor.set_position(initial);
                    return Ok(None);
                }
                Err(err) => {
                    return Err(Error::from(err));
                }
                Ok(read) => {
                    length = read;
                }
            };
            header_length += length_nbytes as u64;
        }
        trace!("Payload length: {}", length);

        let mask = if masked {
            let mut mask_bytes = [0u8; 4];
            if try!(cursor.read(&mut mask_bytes)) != 4 {
                cursor.set_position(initial);
                return Ok(None)
            } else {
                header_length += 4;
                Some(mask_bytes)
            }
        } else {
            None
        };

        if size < header_length || size - header_length < length {
            cursor.set_position(initial);
            return Ok(None)
        }

        let mut data = Vec::with_capacity(length as usize);
        if length > 0 {
            unsafe {
                try!(cursor.read_exact(data.bytes_mut()));
                data.advance_mut(length as usize);
            }
        }

        // Disallow bad opcode
        match opcode {
            OpCode::Control(Control::Reserved(_)) | OpCode::Data(Data::Reserved(_)) => {
                return Err(Error::Protocol(format!("Encountered invalid opcode: {}", first & 0x0F).into()))
            }
            _ => ()
        }

        let frame = Frame {
            finished: finished,
            rsv1: rsv1,
            rsv2: rsv2,
            rsv3: rsv3,
            opcode: opcode,
            mask: mask,
            payload: data,
        };

        Ok(Some(frame))
    }

    /// Write a frame out to a buffer
    pub fn format<W>(mut self, w: &mut W) -> Result<()>
        where W: Write
    {
        let mut one = 0u8;
        let code: u8 = self.opcode.into();
        if self.is_final() {
            one |= 0x80;
        }
        if self.has_rsv1() {
            one |= 0x40;
        }
        if self.has_rsv2() {
            one |= 0x20;
        }
        if self.has_rsv3() {
            one |= 0x10;
        }
        one |= code;

        let mut two = 0u8;

        if self.is_masked() {
            two |= 0x80;
        }

        if self.payload.len() < 126 {
            two |= self.payload.len() as u8;
            let headers = [one, two];
            try!(w.write(&headers));
        } else if self.payload.len() <= 65535 {
            two |= 126;
            let mut length_bytes = [0u8; 2];
            NetworkEndian::write_u16(&mut length_bytes, self.payload.len() as u16);
            let headers = [one, two, length_bytes[0], length_bytes[1]];
            try!(w.write(&headers));
        } else {
            two |= 127;
            let mut length_bytes = [0u8; 8];
            NetworkEndian::write_u64(&mut length_bytes, self.payload.len() as u64);
            let headers = [
                one,
                two,
                length_bytes[0],
                length_bytes[1],
                length_bytes[2],
                length_bytes[3],
                length_bytes[4],
                length_bytes[5],
                length_bytes[6],
                length_bytes[7],
            ];
            try!(w.write(&headers));
        }

        if self.is_masked() {
            let mask = self.mask.take().unwrap();
            apply_mask(&mut self.payload, &mask);
            try!(w.write(&mask));
        }

        try!(w.write(&self.payload));
        Ok(())
    }
}

impl Default for Frame {
    fn default() -> Frame {
        Frame {
            finished: true,
            rsv1: false,
            rsv2: false,
            rsv3: false,
            opcode: OpCode::Control(Control::Close),
            mask: None,
            payload: Vec::new(),
        }
    }
}

impl fmt::Display for Frame {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
            "
<FRAME>
final: {}
reserved: {} {} {}
opcode: {}
length: {}
payload length: {}
payload: 0x{}
            ",
            self.finished,
            self.rsv1,
            self.rsv2,
            self.rsv3,
            self.opcode,
            // self.mask.map(|mask| format!("{:?}", mask)).unwrap_or("NONE".into()),
            self.len(),
            self.payload.len(),
            self.payload.iter().map(|byte| format!("{:x}", byte)).collect::<String>())
    }
}

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

    use super::super::coding::{OpCode, Data};
    use std::io::Cursor;

    #[test]
    fn parse() {
        let mut raw: Cursor<Vec<u8>> = Cursor::new(vec![
            0x82, 0x07, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07
        ]);
        let frame = Frame::parse(&mut raw).unwrap().unwrap();
        assert_eq!(frame.into_data(), vec![ 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07 ]);
    }

    #[test]
    fn format() {
        let frame = Frame::ping(vec![0x01, 0x02]);
        let mut buf = Vec::with_capacity(frame.len());
        frame.format(&mut buf).unwrap();
        assert_eq!(buf, vec![0x89, 0x02, 0x01, 0x02]);
    }

    #[test]
    fn display() {
        let f = Frame::message("hi there".into(), OpCode::Data(Data::Text), true);
        let view = format!("{}", f);
        assert!(view.contains("payload:"));
    }

    #[test]
    fn parse_overflow() {
        let mut raw: Cursor<Vec<u8>> = Cursor::new(vec![
            0x83, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
            0xff, 0xff, 0x00, 0x00, 0x00, 0x00,
        ]);
        let frame_none = Frame::parse(&mut raw).unwrap();
        assert!(frame_none.is_none());
    }
}
<commit_msg>refactor: remove transmute in favor of safe code<commit_after>use std::fmt;
use std::borrow::Cow;
use std::io::{Cursor, Read, Write, ErrorKind};
use std::default::Default;
use std::string::{String, FromUtf8Error};
use std::result::Result as StdResult;
use byteorder::{ByteOrder, ReadBytesExt, WriteBytesExt, NetworkEndian};
use bytes::BufMut;

use error::{Error, Result};
use super::coding::{OpCode, Control, Data, CloseCode};
use super::mask::{generate_mask, apply_mask};

/// A struct representing the close command.
#[derive(Debug, Clone)]
pub struct CloseFrame<'t> {
    /// The reason as a code.
    pub code: CloseCode,
    /// The reason as text string.
    pub reason: Cow<'t, str>,
}

impl<'t> CloseFrame<'t> {
    /// Convert into a owned string.
    pub fn into_owned(self) -> CloseFrame<'static> {
        CloseFrame {
            code: self.code,
            reason: self.reason.into_owned().into(),
        }
    }
}

impl<'t> fmt::Display for CloseFrame<'t> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} ({})", self.reason, self.code)
    }
}

/// A struct representing a WebSocket frame.
#[derive(Debug, Clone)]
pub struct Frame {
    finished: bool,
    rsv1: bool,
    rsv2: bool,
    rsv3: bool,
    opcode: OpCode,

    mask: Option<[u8; 4]>,

    payload: Vec<u8>,
}

impl Frame {

    /// Get the length of the frame.
    /// This is the length of the header + the length of the payload.
    #[inline]
    pub fn len(&self) -> usize {
        let mut header_length = 2;
        let payload_len = self.payload().len();
        if payload_len > 125 {
            if payload_len <= u16::max_value() as usize {
                header_length += 2;
            } else {
                header_length += 8;
            }
        }

        if self.is_masked() {
            header_length += 4;
        }

        header_length + payload_len
    }

    /// Test whether the frame is a final frame.
    #[inline]
    pub fn is_final(&self) -> bool {
        self.finished
    }

    /// Test whether the first reserved bit is set.
    #[inline]
    pub fn has_rsv1(&self) -> bool {
        self.rsv1
    }

    /// Test whether the second reserved bit is set.
    #[inline]
    pub fn has_rsv2(&self) -> bool {
        self.rsv2
    }

    /// Test whether the third reserved bit is set.
    #[inline]
    pub fn has_rsv3(&self) -> bool {
        self.rsv3
    }

    /// Get the OpCode of the frame.
    #[inline]
    pub fn opcode(&self) -> OpCode {
        self.opcode
    }

    /// Get a reference to the frame's payload.
    #[inline]
    pub fn payload(&self) -> &Vec<u8> {
        &self.payload
    }

    // Test whether the frame is masked.
    #[doc(hidden)]
    #[inline]
    pub fn is_masked(&self) -> bool {
        self.mask.is_some()
    }

    // Get an optional reference to the frame's mask.
    #[doc(hidden)]
    #[allow(dead_code)]
    #[inline]
    pub fn mask(&self) -> Option<&[u8; 4]> {
        self.mask.as_ref()
    }

    /// Make this frame a final frame.
    #[allow(dead_code)]
    #[inline]
    pub fn set_final(&mut self, is_final: bool) -> &mut Frame {
        self.finished = is_final;
        self
    }

    /// Set the first reserved bit.
    #[inline]
    pub fn set_rsv1(&mut self, has_rsv1: bool) -> &mut Frame {
        self.rsv1 = has_rsv1;
        self
    }

    /// Set the second reserved bit.
    #[inline]
    pub fn set_rsv2(&mut self, has_rsv2: bool) -> &mut Frame {
        self.rsv2 = has_rsv2;
        self
    }

    /// Set the third reserved bit.
    #[inline]
    pub fn set_rsv3(&mut self, has_rsv3: bool) -> &mut Frame {
        self.rsv3 = has_rsv3;
        self
    }

    /// Set the OpCode.
    #[allow(dead_code)]
    #[inline]
    pub fn set_opcode(&mut self, opcode: OpCode) -> &mut Frame {
        self.opcode = opcode;
        self
    }

    /// Edit the frame's payload.
    #[allow(dead_code)]
    #[inline]
    pub fn payload_mut(&mut self) -> &mut Vec<u8> {
        &mut self.payload
    }

    // Generate a new mask for this frame.
    //
    // This method simply generates and stores the mask. It does not change the payload data.
    // Instead, the payload data will be masked with the generated mask when the frame is sent
    // to the other endpoint.
    #[doc(hidden)]
    #[inline]
    pub fn set_mask(&mut self) -> &mut Frame {
        self.mask = Some(generate_mask());
        self
    }

    // This method unmasks the payload and should only be called on frames that are actually
    // masked. In other words, those frames that have just been received from a client endpoint.
    #[doc(hidden)]
    #[inline]
    pub fn remove_mask(&mut self) {
        self.mask.and_then(|mask| {
            Some(apply_mask(&mut self.payload, &mask))
        });
        self.mask = None;
    }

    /// Consume the frame into its payload as binary.
    #[inline]
    pub fn into_data(self) -> Vec<u8> {
        self.payload
    }

    /// Consume the frame into its payload as string.
    #[inline]
    pub fn into_string(self) -> StdResult<String, FromUtf8Error> {
        String::from_utf8(self.payload)
    }

     /// Consume the frame into a closing frame.
    #[inline]
    pub fn into_close(self) -> Result<Option<CloseFrame<'static>>> {
        match self.payload.len() {
            0 => Ok(None),
            1 => Err(Error::Protocol("Invalid close sequence".into())),
            _ => {
                let mut data = self.payload;
                let code = NetworkEndian::read_u16(&data[0..2]).into();
                data.drain(0..2);
                let text = String::from_utf8(data)?;
                Ok(Some(CloseFrame { code: code, reason: text.into() }))
            }
        }
    }

    /// Create a new data frame.
    #[inline]
    pub fn message(data: Vec<u8>, code: OpCode, finished: bool) -> Frame {
        debug_assert!(match code {
            OpCode::Data(_) => true,
            _ => false,
        }, "Invalid opcode for data frame.");

        Frame {
            finished: finished,
            opcode: code,
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Pong control frame.
    #[inline]
    pub fn pong(data: Vec<u8>) -> Frame {
        Frame {
            opcode: OpCode::Control(Control::Pong),
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Ping control frame.
    #[inline]
    pub fn ping(data: Vec<u8>) -> Frame {
        Frame {
            opcode: OpCode::Control(Control::Ping),
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Close control frame.
    #[inline]
    pub fn close(msg: Option<CloseFrame>) -> Frame {
        let payload = if let Some(CloseFrame { code, reason }) = msg {
            let mut p = Vec::with_capacity(reason.as_bytes().len() + 2);
            p.write_u16::<NetworkEndian>(code.into()).unwrap(); // can't fail
            p.extend_from_slice(reason.as_bytes());
            p
        } else {
            Vec::new()
        };

        Frame {
            payload: payload,
            .. Frame::default()
        }
    }

    /// Parse the input stream into a frame.
    pub fn parse(cursor: &mut Cursor<Vec<u8>>) -> Result<Option<Frame>> {
        let size = cursor.get_ref().len() as u64 - cursor.position();
        let initial = cursor.position();
        trace!("Position in buffer {}", initial);

        let mut head = [0u8; 2];
        if try!(cursor.read(&mut head)) != 2 {
            cursor.set_position(initial);
            return Ok(None)
        }

        trace!("Parsed headers {:?}", head);

        let first = head[0];
        let second = head[1];
        trace!("First: {:b}", first);
        trace!("Second: {:b}", second);

        let finished = first & 0x80 != 0;

        let rsv1 = first & 0x40 != 0;
        let rsv2 = first & 0x20 != 0;
        let rsv3 = first & 0x10 != 0;

        let opcode = OpCode::from(first & 0x0F);
        trace!("Opcode: {:?}", opcode);

        let masked = second & 0x80 != 0;
        trace!("Masked: {:?}", masked);

        let mut header_length = 2;

        let mut length = (second & 0x7F) as u64;

        if let Some(length_nbytes) = match length {
            126 => Some(2),
            127 => Some(8),
            _ => None,
        } {
            match cursor.read_uint::<NetworkEndian>(length_nbytes) {
                Err(ref err) if err.kind() == ErrorKind::UnexpectedEof => {
                    cursor.set_position(initial);
                    return Ok(None);
                }
                Err(err) => {
                    return Err(Error::from(err));
                }
                Ok(read) => {
                    length = read;
                }
            };
            header_length += length_nbytes as u64;
        }
        trace!("Payload length: {}", length);

        let mask = if masked {
            let mut mask_bytes = [0u8; 4];
            if try!(cursor.read(&mut mask_bytes)) != 4 {
                cursor.set_position(initial);
                return Ok(None)
            } else {
                header_length += 4;
                Some(mask_bytes)
            }
        } else {
            None
        };

        if size < header_length || size - header_length < length {
            cursor.set_position(initial);
            return Ok(None)
        }

        let mut data = Vec::with_capacity(length as usize);
        if length > 0 {
            unsafe {
                try!(cursor.read_exact(data.bytes_mut()));
                data.advance_mut(length as usize);
            }
        }

        // Disallow bad opcode
        match opcode {
            OpCode::Control(Control::Reserved(_)) | OpCode::Data(Data::Reserved(_)) => {
                return Err(Error::Protocol(format!("Encountered invalid opcode: {}", first & 0x0F).into()))
            }
            _ => ()
        }

        let frame = Frame {
            finished: finished,
            rsv1: rsv1,
            rsv2: rsv2,
            rsv3: rsv3,
            opcode: opcode,
            mask: mask,
            payload: data,
        };

        Ok(Some(frame))
    }

    /// Write a frame out to a buffer
    pub fn format<W>(mut self, w: &mut W) -> Result<()>
        where W: Write
    {
        let mut one = 0u8;
        let code: u8 = self.opcode.into();
        if self.is_final() {
            one |= 0x80;
        }
        if self.has_rsv1() {
            one |= 0x40;
        }
        if self.has_rsv2() {
            one |= 0x20;
        }
        if self.has_rsv3() {
            one |= 0x10;
        }
        one |= code;

        let mut two = 0u8;

        if self.is_masked() {
            two |= 0x80;
        }

        if self.payload.len() < 126 {
            two |= self.payload.len() as u8;
            let headers = [one, two];
            try!(w.write(&headers));
        } else if self.payload.len() <= 65535 {
            two |= 126;
            let mut length_bytes = [0u8; 2];
            NetworkEndian::write_u16(&mut length_bytes, self.payload.len() as u16);
            let headers = [one, two, length_bytes[0], length_bytes[1]];
            try!(w.write(&headers));
        } else {
            two |= 127;
            let mut length_bytes = [0u8; 8];
            NetworkEndian::write_u64(&mut length_bytes, self.payload.len() as u64);
            let headers = [
                one,
                two,
                length_bytes[0],
                length_bytes[1],
                length_bytes[2],
                length_bytes[3],
                length_bytes[4],
                length_bytes[5],
                length_bytes[6],
                length_bytes[7],
            ];
            try!(w.write(&headers));
        }

        if self.is_masked() {
            let mask = self.mask.take().unwrap();
            apply_mask(&mut self.payload, &mask);
            try!(w.write(&mask));
        }

        try!(w.write(&self.payload));
        Ok(())
    }
}

impl Default for Frame {
    fn default() -> Frame {
        Frame {
            finished: true,
            rsv1: false,
            rsv2: false,
            rsv3: false,
            opcode: OpCode::Control(Control::Close),
            mask: None,
            payload: Vec::new(),
        }
    }
}

impl fmt::Display for Frame {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
            "
<FRAME>
final: {}
reserved: {} {} {}
opcode: {}
length: {}
payload length: {}
payload: 0x{}
            ",
            self.finished,
            self.rsv1,
            self.rsv2,
            self.rsv3,
            self.opcode,
            // self.mask.map(|mask| format!("{:?}", mask)).unwrap_or("NONE".into()),
            self.len(),
            self.payload.len(),
            self.payload.iter().map(|byte| format!("{:x}", byte)).collect::<String>())
    }
}

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

    use super::super::coding::{OpCode, Data};
    use std::io::Cursor;

    #[test]
    fn parse() {
        let mut raw: Cursor<Vec<u8>> = Cursor::new(vec![
            0x82, 0x07, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07
        ]);
        let frame = Frame::parse(&mut raw).unwrap().unwrap();
        assert_eq!(frame.into_data(), vec![ 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07 ]);
    }

    #[test]
    fn format() {
        let frame = Frame::ping(vec![0x01, 0x02]);
        let mut buf = Vec::with_capacity(frame.len());
        frame.format(&mut buf).unwrap();
        assert_eq!(buf, vec![0x89, 0x02, 0x01, 0x02]);
    }

    #[test]
    fn display() {
        let f = Frame::message("hi there".into(), OpCode::Data(Data::Text), true);
        let view = format!("{}", f);
        assert!(view.contains("payload:"));
    }

    #[test]
    fn parse_overflow() {
        let mut raw: Cursor<Vec<u8>> = Cursor::new(vec![
            0x83, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
            0xff, 0xff, 0x00, 0x00, 0x00, 0x00,
        ]);
        let frame_none = Frame::parse(&mut raw).unwrap();
        assert!(frame_none.is_none());
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::iter;
use std::f64;

use adivon::bag::Bag;
use adivon::stack::Stack;
use adivon::priority_queue::IndexMinPQ;

/// Weighted directed edge
#[derive(Clone, Copy)]
pub struct DirectedEdge {
    v: usize,
    w: usize,
    weight: f64
}

impl DirectedEdge {
    pub fn new(v: usize, w: usize, weight: f64) -> DirectedEdge {
        assert!(!weight.is_nan(), "weight is NaN");
        DirectedEdge {
            v: v,
            w: w,
            weight: weight
        }
    }

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

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

    #[inline]
    pub fn weight(&self) -> f64 {
        self.weight
    }
}

impl fmt::Debug for DirectedEdge {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} -> {} {:5.2}", self.v, self.w, self.weight)
    }
}

#[test]
fn test_directed_edge() {
    let e = DirectedEdge::new(12, 24, 3.14);
    assert_eq!(format!("{:?}", e), "12 -> 24  3.14");
}

/// Edge-weighted digraph, implemented using adjacency lists
#[derive(Clone)]
pub struct EdgeWeightedDigraph {
    v: usize,
    e: usize,
    adj: Vec<Bag<DirectedEdge>>
}

impl EdgeWeightedDigraph {
    pub fn new(v: usize) -> EdgeWeightedDigraph {
        EdgeWeightedDigraph {
            v: v,
            e: 0,
            adj: iter::repeat(Bag::new()).take(v).collect()
        }
    }

    pub fn v(&self) -> usize {
        self.v
    }

    pub fn e(&self) -> usize {
        self.e
    }

    #[inline]
    fn validate_vertex(&self, v: usize) {
        assert!(v < self.v, "vertex must be between 0 and V");
    }

    pub fn add_edge(&mut self, e: DirectedEdge) {
        let v = e.from();
        let w = e.to();
        self.validate_vertex(v);
        self.validate_vertex(w);
        self.adj[v].add(e);
        self.e += 1;
    }

    pub fn adj(&self, v: usize) -> ::std::vec::IntoIter<DirectedEdge> {
        self.validate_vertex(v);
        self.adj[v].iter().map(|e| e.clone()).collect::<Vec<DirectedEdge>>().into_iter()
    }

    pub fn outdegree(&self, v: usize) -> usize {
        self.validate_vertex(v);
        self.adj[v].len()
    }

    pub fn edges(&self) -> ::std::vec::IntoIter<DirectedEdge> {
        self.adj.iter()
            .flat_map(|adj| {
                adj.iter().map(|e| e.clone()).collect::<Vec<DirectedEdge>>().into_iter()
            })
            .collect::<Vec<DirectedEdge>>()
            .into_iter()
    }

    pub fn to_dot(&self) -> String {
        let mut dot = String::new();

        dot.push_str("digraph G {\n");
        for i in 0 .. self.v {
            dot.push_str(&format!("  {};\n", i));
        }

        for e in self.edges() {
            let v = e.from();
            let w = e.to();
            dot.push_str(&format!("  {} -> {} [ label=\"{}\" ];\n",
                                  v, w, e.weight))
        }
        dot.push_str("}\n");
        dot
    }
}

// Single-source shortest paths API
pub struct DijkstraSP<'a> {
    graph: &'a EdgeWeightedDigraph,
    dist_to: Vec<f64>,
    edge_to: Vec<Option<DirectedEdge>>,
    pq: IndexMinPQ<f64>,
    s: usize
}

impl<'a> DijkstraSP<'a> {
    fn new<'b>(graph: &'b EdgeWeightedDigraph, s: usize) -> DijkstraSP<'b> {
        let n = graph.v();
        for e in graph.edges() {
            if e.weight() < 0.0 {
                panic!("edge has negative weight in DijkstraSP");
            }
        }
        let dist_to = iter::repeat(f64::INFINITY).take(n).collect();
        let edge_to = iter::repeat(None).take(n).collect();
        let pq = IndexMinPQ::with_capacity(n);
        let mut sp = DijkstraSP {
            graph: graph,
            s: s,
            dist_to: dist_to,
            edge_to: edge_to,
            pq: pq
        };
        // alogrithm
        sp.dist_to[s] = 0.0;

        sp.pq.insert(s, 0.0);
        while !sp.pq.is_empty() {
            let v = sp.pq.del_min().unwrap();
            for e in graph.adj(v) {
                sp.relax(e);
            }
        }

        sp
    }

    fn relax(&mut self, e: DirectedEdge) {
        let v = e.from();
        let w = e.to();
        if self.dist_to[w] > self.dist_to[v] + e.weight() {
            self.dist_to[w] = self.dist_to[v] + e.weight();
            self.edge_to[w] = Some(e);
            if self.pq.contains(w) {
                self.pq.decrease_key(w, self.dist_to[w]);
            } else {
                self.pq.insert(w, self.dist_to[w]);
            }
        }
    }

    // length of shortest path from s to v
    pub fn dist_to(&self, v: usize) -> f64 {
        self.dist_to[v]
    }

    pub fn has_path_to(&self, v: usize) -> bool {
        self.dist_to[v] < f64::INFINITY
    }

    // shortest path from s to v
    pub fn path_to(&self, v: usize) -> ::std::vec::IntoIter<DirectedEdge> {
        if !self.has_path_to(v) {
            vec!().into_iter()
        } else {
            let mut path = Stack::new();
            let mut e = self.edge_to[v];
            while e.is_some() {
                path.push(e.unwrap());
                e = self.edge_to[e.unwrap().from()]
            }
            path.into_iter().collect::<Vec<DirectedEdge>>().into_iter()
        }
    }

    // FIXME: should this be public?
    pub fn check(&self) -> bool {
        let s = self.s;
        for e in self.graph.edges() {
            if e.weight() < 0.0 {
                return false;
            }
        }

        if self.dist_to[s] != 0.0 || self.edge_to[s].is_some() {
            return false;
        }

        for v in 0 .. self.graph.v() {
            if v == s { continue }
            if self.edge_to[v].is_none() && self.dist_to[v] != f64::INFINITY {
                // dist_to[] edge_to[] inconsistent
                return false;
            }
        }

        for v in 0 .. self.graph.v() {
            for e in self.graph.adj(v) {
                let w  = e.to();
                if self.dist_to[v] + e.weight() < self.dist_to[w] {
                    // edge not relaxed
                    return false;
                }
            }
        }

        for w in 0 .. self.graph.v() {
            if self.edge_to[w].is_none() { continue }
            let e = self.edge_to[w].unwrap();
            let v = e.from();
            if w != e.to() {
                return false;
            }
            if self.dist_to[v] + e.weight() != self.dist_to[w] {
                // edge on shortest path not tight
                return false;
            }
        }

        true
    }
}

impl EdgeWeightedDigraph {
    pub fn dijkstra_sp<'a>(&'a self, s: usize) -> DijkstraSP<'a> {
        DijkstraSP::new(self, s)
    }
}

#[test]
fn test_dijkstra_shortest_path() {
    let mut g = EdgeWeightedDigraph::new(6);
    g.add_edge(DirectedEdge::new(0, 1, 7.0));
    g.add_edge(DirectedEdge::new(1, 2, 10.0));
    g.add_edge(DirectedEdge::new(0, 2, 9.0));
    g.add_edge(DirectedEdge::new(0, 5, 14.0));
    g.add_edge(DirectedEdge::new(1, 3, 15.0));
    g.add_edge(DirectedEdge::new(2, 5, 2.0));
    g.add_edge(DirectedEdge::new(2, 3, 11.0));
    g.add_edge(DirectedEdge::new(4, 5, 9.0));
    g.add_edge(DirectedEdge::new(3, 4, 6.0));
    g.add_edge(DirectedEdge::new(2, 2, 1.0));

    assert_eq!(20.0, g.dijkstra_sp(0).dist_to(3));
    assert_eq!(26.0, g.dijkstra_sp(0).path_to(4).map(|e| e.weight()).sum());

    assert!(g.dijkstra_sp(0).check());
}
<commit_msg>shortest_paths: add EdgeWeightedDirectedCycle, DepthFirstOrder support for edge weighted graph<commit_after>use std::fmt;
use std::iter;
use std::f64;

use adivon::bag::Bag;
use adivon::stack::Stack;
use adivon::queue::Queue;
use adivon::priority_queue::IndexMinPQ;

/// Weighted directed edge
#[derive(Clone, Copy)]
pub struct DirectedEdge {
    v: usize,
    w: usize,
    weight: f64
}

impl DirectedEdge {
    pub fn new(v: usize, w: usize, weight: f64) -> DirectedEdge {
        assert!(!weight.is_nan(), "weight is NaN");
        DirectedEdge {
            v: v,
            w: w,
            weight: weight
        }
    }

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

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

    #[inline]
    pub fn weight(&self) -> f64 {
        self.weight
    }
}

impl fmt::Debug for DirectedEdge {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} -> {} {:5.2}", self.v, self.w, self.weight)
    }
}

#[test]
fn test_directed_edge() {
    let e = DirectedEdge::new(12, 24, 3.14);
    assert_eq!(format!("{:?}", e), "12 -> 24  3.14");
}

/// Edge-weighted digraph, implemented using adjacency lists
#[derive(Clone)]
pub struct EdgeWeightedDigraph {
    v: usize,
    e: usize,
    adj: Vec<Bag<DirectedEdge>>
}

impl EdgeWeightedDigraph {
    pub fn new(v: usize) -> EdgeWeightedDigraph {
        EdgeWeightedDigraph {
            v: v,
            e: 0,
            adj: iter::repeat(Bag::new()).take(v).collect()
        }
    }

    pub fn v(&self) -> usize {
        self.v
    }

    pub fn e(&self) -> usize {
        self.e
    }

    #[inline]
    fn validate_vertex(&self, v: usize) {
        assert!(v < self.v, "vertex must be between 0 and V");
    }

    pub fn add_edge(&mut self, e: DirectedEdge) {
        let v = e.from();
        let w = e.to();
        self.validate_vertex(v);
        self.validate_vertex(w);
        self.adj[v].add(e);
        self.e += 1;
    }

    pub fn adj(&self, v: usize) -> ::std::vec::IntoIter<DirectedEdge> {
        self.validate_vertex(v);
        self.adj[v].iter().map(|e| e.clone()).collect::<Vec<DirectedEdge>>().into_iter()
    }

    pub fn outdegree(&self, v: usize) -> usize {
        self.validate_vertex(v);
        self.adj[v].len()
    }

    pub fn edges(&self) -> ::std::vec::IntoIter<DirectedEdge> {
        self.adj.iter()
            .flat_map(|adj| {
                adj.iter().map(|e| e.clone()).collect::<Vec<DirectedEdge>>().into_iter()
            })
            .collect::<Vec<DirectedEdge>>()
            .into_iter()
    }

    pub fn to_dot(&self) -> String {
        let mut dot = String::new();

        dot.push_str("digraph G {\n");
        for i in 0 .. self.v {
            dot.push_str(&format!("  {};\n", i));
        }

        for e in self.edges() {
            let v = e.from();
            let w = e.to();
            dot.push_str(&format!("  {} -> {} [ label=\"{}\" ];\n",
                                  v, w, e.weight))
        }
        dot.push_str("}\n");
        dot
    }
}

// Single-source shortest paths API
pub struct DijkstraSP<'a> {
    graph: &'a EdgeWeightedDigraph,
    dist_to: Vec<f64>,
    edge_to: Vec<Option<DirectedEdge>>,
    pq: IndexMinPQ<f64>,
    s: usize
}

impl<'a> DijkstraSP<'a> {
    fn new<'b>(graph: &'b EdgeWeightedDigraph, s: usize) -> DijkstraSP<'b> {
        let n = graph.v();
        for e in graph.edges() {
            if e.weight() < 0.0 {
                panic!("edge has negative weight in DijkstraSP");
            }
        }
        let dist_to = iter::repeat(f64::INFINITY).take(n).collect();
        let edge_to = iter::repeat(None).take(n).collect();
        let pq = IndexMinPQ::with_capacity(n);
        let mut sp = DijkstraSP {
            graph: graph,
            s: s,
            dist_to: dist_to,
            edge_to: edge_to,
            pq: pq
        };
        // alogrithm
        sp.dist_to[s] = 0.0;

        sp.pq.insert(s, 0.0);
        while !sp.pq.is_empty() {
            let v = sp.pq.del_min().unwrap();
            for e in graph.adj(v) {
                sp.relax(e);
            }
        }

        sp
    }

    fn relax(&mut self, e: DirectedEdge) {
        let v = e.from();
        let w = e.to();
        if self.dist_to[w] > self.dist_to[v] + e.weight() {
            self.dist_to[w] = self.dist_to[v] + e.weight();
            self.edge_to[w] = Some(e);
            if self.pq.contains(w) {
                self.pq.decrease_key(w, self.dist_to[w]);
            } else {
                self.pq.insert(w, self.dist_to[w]);
            }
        }
    }

    // length of shortest path from s to v
    pub fn dist_to(&self, v: usize) -> f64 {
        self.dist_to[v]
    }

    pub fn has_path_to(&self, v: usize) -> bool {
        self.dist_to[v] < f64::INFINITY
    }

    // shortest path from s to v
    pub fn path_to(&self, v: usize) -> ::std::vec::IntoIter<DirectedEdge> {
        if !self.has_path_to(v) {
            vec!().into_iter()
        } else {
            let mut path = Stack::new();
            let mut e = self.edge_to[v];
            while e.is_some() {
                path.push(e.unwrap());
                e = self.edge_to[e.unwrap().from()]
            }
            path.into_iter().collect::<Vec<DirectedEdge>>().into_iter()
        }
    }

    // FIXME: should this be public?
    pub fn check(&self) -> bool {
        let s = self.s;
        for e in self.graph.edges() {
            if e.weight() < 0.0 {
                return false;
            }
        }

        if self.dist_to[s] != 0.0 || self.edge_to[s].is_some() {
            return false;
        }

        for v in 0 .. self.graph.v() {
            if v == s { continue }
            if self.edge_to[v].is_none() && self.dist_to[v] != f64::INFINITY {
                // dist_to[] edge_to[] inconsistent
                return false;
            }
        }

        for v in 0 .. self.graph.v() {
            for e in self.graph.adj(v) {
                let w  = e.to();
                if self.dist_to[v] + e.weight() < self.dist_to[w] {
                    // edge not relaxed
                    return false;
                }
            }
        }

        for w in 0 .. self.graph.v() {
            if self.edge_to[w].is_none() { continue }
            let e = self.edge_to[w].unwrap();
            let v = e.from();
            if w != e.to() {
                return false;
            }
            if self.dist_to[v] + e.weight() != self.dist_to[w] {
                // edge on shortest path not tight
                return false;
            }
        }

        true
    }
}

impl EdgeWeightedDigraph {
    pub fn dijkstra_sp<'a>(&'a self, s: usize) -> DijkstraSP<'a> {
        DijkstraSP::new(self, s)
    }
}

/// Compute preorder and postorder for a digraph or edge-weighted digraph
pub struct DepthFirstOrder<'a> {
    graph: &'a EdgeWeightedDigraph,
    pre: Vec<usize>,
    post: Vec<usize>,
    preorder: Queue<usize>,
    postorder: Queue<usize>,
    marked: Vec<bool>,
    pre_counter: usize,
    post_counter: usize
}

impl<'a> DepthFirstOrder<'a> {
    fn new<'b>(graph: &'b EdgeWeightedDigraph) -> DepthFirstOrder<'b> {
        let n = graph.v();
        let mut ret = DepthFirstOrder {
            graph: graph,
            pre: iter::repeat(0).take(n).collect(),
            post: iter::repeat(0).take(n).collect(),
            preorder: Queue::new(),
            postorder: Queue::new(),
            marked: iter::repeat(false).take(n).collect(),
            pre_counter: 0,
            post_counter: 0
        };
        ret.init();
        ret
    }

    fn init(&mut self) {
        for v in 0 .. self.graph.v() {
            if !self.marked[v] {
                self.dfs(v)
            }
        }
    }

    fn dfs(&mut self, v: usize) {
        self.marked[v] = true;
        self.pre[v] = self.pre_counter;
        self.pre_counter += 1;
        self.preorder.enqueue(v);
        for e in self.graph.adj(v) {
            let w = e.to();
            if !self.marked[w] {
                self.dfs(w);
            }
        }
        self.postorder.enqueue(v);
        self.post[v] = self.post_counter;
        self.post_counter += 1;
    }

    // preorder number of vertex v
    pub fn preorder(&self, v: usize) -> usize {
        self.pre[v]
    }

    // postorder number of vertex v
    pub fn postorder(&self, v: usize) -> usize {
        self.post[v]
    }

    pub fn pre(&self) -> ::std::vec::IntoIter<usize> {
        self.preorder.clone().into_iter().collect::<Vec<usize>>().into_iter()
    }

    pub fn post(&self) -> ::std::vec::IntoIter<usize> {
        self.postorder.clone().into_iter().collect::<Vec<usize>>().into_iter()
    }

    pub fn reverse_post(&self) -> ::std::vec::IntoIter<usize> {
        let mut reverse = Stack::new();
        for v in self.postorder.iter() {
            reverse.push(*v);
        }
        reverse.into_iter().collect::<Vec<usize>>().into_iter()
    }

    fn check(&self) -> bool {
        let mut r = 0;
        for v in self.post() {
            if self.postorder(v) != r {
                // post(v) and post() inconsistent
                return false;
            }
            r += 1;
        }

        r = 0;
        for v in self.pre() {
            if self.preorder(v) != r {
                // preorder(v) and pre() inconsistent
                return false;
            }
            r += 1;
        }
        return true;
    }
}

impl EdgeWeightedDigraph {
    pub fn depth_first_order<'a>(&'a self) -> DepthFirstOrder<'a> {
        DepthFirstOrder::new(self)
    }
}

// Finds a directed cycle in an edge-weighted digraph
pub struct EdgeWeightedDirectedCycle<'a> {
    graph: &'a EdgeWeightedDigraph,
    marked: Vec<bool>,
    edge_to: Vec<Option<DirectedEdge>>,
    on_stack: Vec<bool>,
    // directed cycle (or empty)
    cycle: Option<Stack<DirectedEdge>>
}

impl<'a> EdgeWeightedDirectedCycle<'a> {
    fn new<'b>(graph: &'b EdgeWeightedDigraph) -> EdgeWeightedDirectedCycle<'b> {
        let n = graph.v();
        let mut ret = EdgeWeightedDirectedCycle {
            graph: graph,
            marked: iter::repeat(false).take(n).collect(),
            edge_to: iter::repeat(None).take(n).collect(),
            on_stack: iter::repeat(false).take(n).collect(),
            cycle: None
        };
        ret.init();
        ret
    }

    fn init(&mut self) {
        for v in 0 .. self.graph.v() {
            if !self.marked[v] {
                self.dfs(v)
            }
        }
    }

    fn dfs(&mut self, v: usize) {
        self.on_stack[v] = true;
        self.marked[v] = true;
        for e in self.graph.adj(v) {
            let w = e.to();

            if self.cycle.is_some() {
                return;
            } else if !self.marked[w] {
                self.edge_to[w] = Some(e);
                self.dfs(w);
            } else if self.on_stack[w] {
                self.cycle = Some(Stack::new());
                // scope local
                let mut e = e.clone();
                while e.from() != w {
                    self.cycle.as_mut().map(|s| s.push(e));
                    e = self.edge_to[e.from()].unwrap();
                }
                self.cycle.as_mut().map(|s| s.push(e));
            }
        }
        self.on_stack[v] = false;
    }

    pub fn has_cycle(&self) -> bool {
        self.cycle.is_some()
    }

    pub fn cycle(&self) -> ::std::vec::IntoIter<DirectedEdge> {
        self.cycle.iter().flat_map(|e| e.clone()).collect::<Vec<DirectedEdge>>().into_iter()
    }

    pub fn check(&self) -> bool {
        if self.has_cycle() {
            let first = self.cycle().next().unwrap();
            let last = self.cycle().last().unwrap();

            if first.from() == last.to() {
                return true;
            } else {
                return false;
            }
        }
        return true;
    }
}

impl EdgeWeightedDigraph {
    pub fn cycle(&self) -> ::std::vec::IntoIter<DirectedEdge> {
        EdgeWeightedDirectedCycle::new(self).cycle()
    }

    pub fn has_cycle(&self) -> bool {
        EdgeWeightedDirectedCycle::new(self).has_cycle()
    }
}

/// Compute topological ordering of a DAG or edge-weighted DAG
pub struct Topological {
    order: Vec<usize>
}

impl Topological {
    pub fn new(graph: EdgeWeightedDigraph) -> Topological {
        unimplemented!()
    }
}

// Computes shortest paths in an edge-weighted acyclic digraph
pub struct AcyclicSP<'a> {
    graph: &'a EdgeWeightedDigraph,
    dist_to: Vec<f64>,
    edge_to: Vec<Option<DirectedEdge>>
}

impl<'a> AcyclicSP<'a> {
    fn new<'b>(graph: &'b EdgeWeightedDigraph, s: usize) -> AcyclicSP<'b> {
        let n = graph.v();
        let mut dist_to: Vec<f64> = iter::repeat(f64::INFINITY).take(n).collect();
        let edge_to = iter::repeat(None).take(n).collect();

        dist_to[s] = 0.0;
        AcyclicSP {
            graph: graph,
            dist_to: dist_to,
            edge_to: edge_to
        }
    }
}

#[test]
fn test_dijkstra_shortest_path() {
    let mut g = EdgeWeightedDigraph::new(6);
    g.add_edge(DirectedEdge::new(0, 1, 7.0));
    g.add_edge(DirectedEdge::new(1, 2, 10.0));
    g.add_edge(DirectedEdge::new(0, 2, 9.0));
    g.add_edge(DirectedEdge::new(0, 5, 14.0));
    g.add_edge(DirectedEdge::new(1, 3, 15.0));
    g.add_edge(DirectedEdge::new(2, 5, 2.0));
    g.add_edge(DirectedEdge::new(2, 3, 11.0));
    g.add_edge(DirectedEdge::new(4, 5, 9.0));
    g.add_edge(DirectedEdge::new(3, 4, 6.0));
    g.add_edge(DirectedEdge::new(2, 2, 1.0));

    assert_eq!(20.0, g.dijkstra_sp(0).dist_to(3));
    assert_eq!(26.0, g.dijkstra_sp(0).path_to(4).map(|e| e.weight()).sum());

    assert!(g.dijkstra_sp(0).check());
}

#[test]
fn test_cyclic_edge_weighted_directed_graph() {
    let mut g = EdgeWeightedDigraph::new(4);
    g.add_edge(DirectedEdge::new(0, 1, 0.5));
    g.add_edge(DirectedEdge::new(0, 2, 0.5));
    g.add_edge(DirectedEdge::new(1, 2, 0.5));
    g.add_edge(DirectedEdge::new(2, 3, 0.5));
    g.add_edge(DirectedEdge::new(3, 1, 0.5));

    assert!(g.has_cycle());
    assert_eq!(3, g.cycle().count());
}

#[test]
fn test_acyclic_shortest_path() {
    let mut g = EdgeWeightedDigraph::new(6);
    g.add_edge(DirectedEdge::new(0, 1, 7.0));
    g.add_edge(DirectedEdge::new(1, 2, 10.0));
    g.add_edge(DirectedEdge::new(0, 2, 9.0));
    g.add_edge(DirectedEdge::new(0, 5, 14.0));
    g.add_edge(DirectedEdge::new(1, 3, 15.0));
    g.add_edge(DirectedEdge::new(2, 5, 2.0));
    g.add_edge(DirectedEdge::new(2, 3, 11.0));
    g.add_edge(DirectedEdge::new(4, 5, 9.0));
    g.add_edge(DirectedEdge::new(3, 4, 6.0));
    g.add_edge(DirectedEdge::new(2, 2, 1.0));

    assert!(g.depth_first_order().check());
}
<|endoftext|>