<commit_before>#![cfg(target_os = "windows")]

pub use api::win32;
pub use api::win32::{MonitorID, get_available_monitors, get_primary_monitor};
pub use api::win32::{WindowProxy, PollEventsIterator, WaitEventsIterator};

use libc;

use Api;
use BuilderAttribs;
use ContextError;
use CreationError;
use PixelFormat;
use GlContext;

use api::egl::ffi::egl::Egl;

use std::ffi::CString;
use std::ops::{Deref, DerefMut};
use kernel32;

/// Stupid wrapper because `*const libc::c_void` doesn't implement `Sync`.
struct EglWrapper(Egl);
unsafe impl Sync for EglWrapper {}

lazy_static! {
    // An EGL implementation available on the system.
    static ref EGL: Option<EglWrapper> = {
        // the ATI drivers provide an EGL implementation in their DLLs
        let dll_name = if cfg!(target_pointer_width = "64") {
            b"atio6axx.dll\0"
        } else {
            b"atioglxx.dll\0"
        };

        let dll = unsafe { kernel32::LoadLibraryA(dll_name.as_ptr() as *const _) };

        if !dll.is_null() {
            let egl = Egl::load_with(|name| {
                let name = CString::new(name).unwrap();
                unsafe { kernel32::GetProcAddress(dll, name.as_ptr()) as *const _ }
            });

            Some(EglWrapper(egl))

        } else {
            None
        }
    };
}

/// The Win32 implementation of the main `Window` object.
pub struct Window(win32::Window);

impl Window {
    /// See the docs in the crate root file.
    pub fn new(builder: BuilderAttribs) -> Result<Window, CreationError> {
        win32::Window::new(builder, EGL.as_ref().map(|w| &w.0)).map(|w| Window(w))
    }
}

impl Deref for Window {
    type Target = win32::Window;

    fn deref(&self) -> &win32::Window {
        &self.0
    }
}

impl DerefMut for Window {
    fn deref_mut(&mut self) -> &mut win32::Window {
        &mut self.0
    }
}

///
pub struct HeadlessContext(win32::Window);

impl HeadlessContext {
    pub fn new(mut builder: BuilderAttribs) -> Result<HeadlessContext, CreationError> {
        builder.visible = false;
        win32::Window::new(builder, EGL.as_ref().map(|w| &w.0)).map(|w| HeadlessContext(w))
    }
}

impl GlContext for HeadlessContext {
    unsafe fn make_current(&self) -> Result<(), ContextError> {
        self.0.make_current()
    }

    fn is_current(&self) -> bool {
        self.0.is_current()
    }

    fn get_proc_address(&self, addr: &str) -> *const libc::c_void {
        self.0.get_proc_address(addr)
    }

    fn swap_buffers(&self) -> Result<(), ContextError> {
        self.0.swap_buffers()
    }

    fn get_api(&self) -> Api {
        self.0.get_api()
    }

    fn get_pixel_format(&self) -> PixelFormat {
        self.0.get_pixel_format()
    }
}
<commit_msg>Headless contexts now create an EGL pbuffer on win32<commit_after>#![cfg(target_os = "windows")]

pub use api::win32;
pub use api::win32::{MonitorID, get_available_monitors, get_primary_monitor};
pub use api::win32::{WindowProxy, PollEventsIterator, WaitEventsIterator};

use libc;

use Api;
use BuilderAttribs;
use ContextError;
use CreationError;
use PixelFormat;
use GlContext;

use api::egl::ffi::egl::Egl;
use api::egl::Context as EglContext;

use std::ffi::CString;
use std::ops::{Deref, DerefMut};
use kernel32;

/// Stupid wrapper because `*const libc::c_void` doesn't implement `Sync`.
struct EglWrapper(Egl);
unsafe impl Sync for EglWrapper {}

lazy_static! {
    // An EGL implementation available on the system.
    static ref EGL: Option<EglWrapper> = {
        // the ATI drivers provide an EGL implementation in their DLLs
        let dll_name = if cfg!(target_pointer_width = "64") {
            b"atio6axx.dll\0"
        } else {
            b"atioglxx.dll\0"
        };

        let dll = unsafe { kernel32::LoadLibraryA(dll_name.as_ptr() as *const _) };

        if !dll.is_null() {
            let egl = Egl::load_with(|name| {
                let name = CString::new(name).unwrap();
                unsafe { kernel32::GetProcAddress(dll, name.as_ptr()) as *const _ }
            });

            Some(EglWrapper(egl))

        } else {
            None
        }
    };
}

/// The Win32 implementation of the main `Window` object.
pub struct Window(win32::Window);

impl Window {
    /// See the docs in the crate root file.
    pub fn new(builder: BuilderAttribs) -> Result<Window, CreationError> {
        win32::Window::new(builder, EGL.as_ref().map(|w| &w.0)).map(|w| Window(w))
    }
}

impl Deref for Window {
    type Target = win32::Window;

    fn deref(&self) -> &win32::Window {
        &self.0
    }
}

impl DerefMut for Window {
    fn deref_mut(&mut self) -> &mut win32::Window {
        &mut self.0
    }
}

///
pub enum HeadlessContext {
    /// A regular window, but invisible.
    HiddenWindow(win32::Window),
    /// An EGL pbuffer.
    EglPbuffer(EglContext),
}

impl HeadlessContext {
    pub fn new(mut builder: BuilderAttribs) -> Result<HeadlessContext, CreationError> {
        builder.visible = false;

        // if EGL is available, we try using EGL first
        // if EGL returns an error, we try the hidden window method
        if let &Some(ref egl) = &*EGL {
            let context = EglContext::new(egl.0.clone(), &builder, None)
                                .and_then(|prototype| prototype.finish_pbuffer())
                                .map(|ctxt| HeadlessContext::EglPbuffer(ctxt));

            if let Ok(context) = context {
                return Ok(context);
            }
        }

        let window = try!(win32::Window::new(builder, EGL.as_ref().map(|w| &w.0)));
        Ok(HeadlessContext::HiddenWindow(window))
    }
}

impl GlContext for HeadlessContext {
    unsafe fn make_current(&self) -> Result<(), ContextError> {
        match self {
            &HeadlessContext::HiddenWindow(ref ctxt) => ctxt.make_current(),
            &HeadlessContext::EglPbuffer(ref ctxt) => ctxt.make_current(),
        }
    }

    fn is_current(&self) -> bool {
        match self {
            &HeadlessContext::HiddenWindow(ref ctxt) => ctxt.is_current(),
            &HeadlessContext::EglPbuffer(ref ctxt) => ctxt.is_current(),
        }
    }

    fn get_proc_address(&self, addr: &str) -> *const libc::c_void {
        match self {
            &HeadlessContext::HiddenWindow(ref ctxt) => ctxt.get_proc_address(addr),
            &HeadlessContext::EglPbuffer(ref ctxt) => ctxt.get_proc_address(addr),
        }
    }

    fn swap_buffers(&self) -> Result<(), ContextError> {
        match self {
            &HeadlessContext::HiddenWindow(ref ctxt) => ctxt.swap_buffers(),
            &HeadlessContext::EglPbuffer(ref ctxt) => ctxt.swap_buffers(),
        }
    }

    fn get_api(&self) -> Api {
        match self {
            &HeadlessContext::HiddenWindow(ref ctxt) => ctxt.get_api(),
            &HeadlessContext::EglPbuffer(ref ctxt) => ctxt.get_api(),
        }
    }

    fn get_pixel_format(&self) -> PixelFormat {
        match self {
            &HeadlessContext::HiddenWindow(ref ctxt) => ctxt.get_pixel_format(),
            &HeadlessContext::EglPbuffer(ref ctxt) => ctxt.get_pixel_format(),
        }
    }
}
<|endoftext|>


<commit_before>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
        };

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

        // Make sure `length` is not too big (fits into `usize`).
        if length > usize::max_value() as u64 {
            return Err(Error::Capacity(format!("Message length too big: {}", length).into()));
        }

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

        // Size is checked above, so it won't be truncated here.
        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);
            }
        }

        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::parse(&mut raw); // should not crash
    }
}
<commit_msg>frame: eliminate an unsafe when reading message data into the buffer<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 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
        };

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

        // Make sure `length` is not too big (fits into `usize`).
        if length > usize::max_value() as u64 {
            return Err(Error::Capacity(format!("Message length too big: {}", length).into()));
        }

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

        // Size is checked above, so it won't be truncated here.
        let mut data = Vec::with_capacity(length as usize);
        if length > 0 {
            try!(cursor.take(length).read_to_end(&mut data));
        }
        debug_assert_eq!(data.len() as u64, length);

        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::parse(&mut raw); // should not crash
    }
}
<|endoftext|>


<commit_before>#![feature(const_fn)]

const X : usize = 2;

const fn f(x: usize) -> usize {
    let mut sum = 0;
    for i in 0..x {
        //~^ ERROR E0015
        //~| ERROR E0015
        //~| ERROR E0658
        //~| ERROR E0080
        //~| ERROR E0744
        sum += i;
    }
    sum
}

#[allow(unused_variables)]
fn main() {
    let a : [i32; f(X)];
}
<commit_msg>Bless compile-fail<commit_after>#![feature(const_fn)]

const X : usize = 2;

const fn f(x: usize) -> usize {
    let mut sum = 0;
    for i in 0..x {
        //~^ ERROR mutable references
        //~| ERROR E0080
        //~| ERROR E0744
        sum += i;
    }
    sum
}

#[allow(unused_variables)]
fn main() {
    let a : [i32; f(X)];
}
<|endoftext|>


<commit_before>mod bindings {
    ::windows::include_bindings!();
}

#[cfg(windows)]
use bindings::{
    Windows::Win32::Security::{
        GetTokenInformation, OpenProcessToken, PSID, TOKEN_ACCESS_MASK, TOKEN_INFORMATION_CLASS,
        TOKEN_USER,
    },
    Windows::Win32::SystemServices::{
        GetCurrentProcess, OpenEventLogA, ReportEventA, ReportEvent_wType, HANDLE, PSTR,
    },
};

#[cfg(windows)]
fn main() -> windows::Result<()> {
    let ph = unsafe { GetCurrentProcess() };
    let mut th: HANDLE = HANDLE(0);
    unsafe { OpenProcessToken(ph, TOKEN_ACCESS_MASK::TOKEN_QUERY, &mut th) }.ok()?;

    // Determine the required buffer size, ignore ERROR_INSUFFICIENT_BUFFER
    let mut length = 0_u32;
    unsafe {
        GetTokenInformation(
            th,
            TOKEN_INFORMATION_CLASS::TokenUser,
            std::ptr::null_mut(),
            0,
            &mut length,
        )
    }
    .ok()
    .unwrap_err();

    // Retrieve the user token.
    let mut token_user_bytes = vec![0u8; length as usize];
    unsafe {
        GetTokenInformation(
            th,
            TOKEN_INFORMATION_CLASS::TokenUser,
            token_user_bytes.as_mut_ptr().cast(),
            length,
            &mut length,
        )
    }
    .ok()?;

    // Extract the pointer to the user SID.
    let user_sid: PSID = unsafe { (*token_user_bytes.as_ptr().cast::<TOKEN_USER>()).User.Sid };

    // use the Application event log
    let event_log_handle = unsafe { OpenEventLogA(PSTR::default(), "Application") };

    let mut event_msg = PSTR(b"Hello in the event log\0".as_ptr() as _);
    unsafe {
        ReportEventA(
            HANDLE(event_log_handle.0),               //h_event_log: T0__,
            ReportEvent_wType::EVENTLOG_WARNING_TYPE, // for type use EVENTLOG_WARNING_TYPE w_type: u16,
            5,                                        // for category use "Shell" w_category: u16,
            1,                                        // for ID use 1  dw_event_id: u32,
            user_sid,                                 // lp_user_sid: *mut c_void,
            1,                                        // w_num_strings: u16,
            0,                                        // dw_data_size: u32,
            &mut event_msg,                           // lp_strings: *mut PSTR,
            std::ptr::null_mut(),                     // lp_raw_data: *mut c_void,
        )
    }
    .ok()?;

    Ok(())
}

#[cfg(not(windows))]
fn main() {
    println!("Not implemented");
}
<commit_msg>Fix cfg<commit_after>#[cfg(windows)]
mod bindings {
    ::windows::include_bindings!();
}

#[cfg(windows)]
use bindings::{
    Windows::Win32::Security::{
        GetTokenInformation, OpenProcessToken, PSID, TOKEN_ACCESS_MASK, TOKEN_INFORMATION_CLASS,
        TOKEN_USER,
    },
    Windows::Win32::SystemServices::{
        GetCurrentProcess, OpenEventLogA, ReportEventA, ReportEvent_wType, HANDLE, PSTR,
    },
};

#[cfg(windows)]
fn main() -> windows::Result<()> {
    let ph = unsafe { GetCurrentProcess() };
    let mut th: HANDLE = HANDLE(0);
    unsafe { OpenProcessToken(ph, TOKEN_ACCESS_MASK::TOKEN_QUERY, &mut th) }.ok()?;

    // Determine the required buffer size, ignore ERROR_INSUFFICIENT_BUFFER
    let mut length = 0_u32;
    unsafe {
        GetTokenInformation(
            th,
            TOKEN_INFORMATION_CLASS::TokenUser,
            std::ptr::null_mut(),
            0,
            &mut length,
        )
    }
    .ok()
    .unwrap_err();

    // Retrieve the user token.
    let mut token_user_bytes = vec![0u8; length as usize];
    unsafe {
        GetTokenInformation(
            th,
            TOKEN_INFORMATION_CLASS::TokenUser,
            token_user_bytes.as_mut_ptr().cast(),
            length,
            &mut length,
        )
    }
    .ok()?;

    // Extract the pointer to the user SID.
    let user_sid: PSID = unsafe { (*token_user_bytes.as_ptr().cast::<TOKEN_USER>()).User.Sid };

    // use the Application event log
    let event_log_handle = unsafe { OpenEventLogA(PSTR::default(), "Application") };

    let mut event_msg = PSTR(b"Hello in the event log\0".as_ptr() as _);
    unsafe {
        ReportEventA(
            HANDLE(event_log_handle.0),               //h_event_log: T0__,
            ReportEvent_wType::EVENTLOG_WARNING_TYPE, // for type use EVENTLOG_WARNING_TYPE w_type: u16,
            5,                                        // for category use "Shell" w_category: u16,
            1,                                        // for ID use 1  dw_event_id: u32,
            user_sid,                                 // lp_user_sid: *mut c_void,
            1,                                        // w_num_strings: u16,
            0,                                        // dw_data_size: u32,
            &mut event_msg,                           // lp_strings: *mut PSTR,
            std::ptr::null_mut(),                     // lp_raw_data: *mut c_void,
        )
    }
    .ok()?;

    Ok(())
}

#[cfg(not(windows))]
fn main() {
    println!("Not implemented");
}
<|endoftext|>


<commit_before>use std::fmt;
use std::borrow::Cow;
use super::token::is_token_delim;

#[derive(Debug, PartialEq, Clone)]
pub enum Atom {
    // String
    Str(String),

    // Unsigned 64-bit integer
    UInt(u64),

    // Signed 64-bit integer
    SInt(i64),

    // 64-bit floating point number
    Float(f64),
}

// Str
impl From<String> for Atom {
    fn from(s: String) -> Atom {
        Atom::Str(s)
    }
}

impl<'a> From<&'a str> for Atom {
    fn from(s: &'a str) -> Atom {
        Atom::Str(s.to_string())
    }
}

// Uint
impl From<usize> for Atom {
    fn from(u: usize) -> Atom {
        Atom::UInt(u as u64)
    }
}

impl From<u64> for Atom {
    fn from(u: u64) -> Atom {
        Atom::UInt(u.into())
    }
}
impl From<u32> for Atom {
    fn from(u: u32) -> Atom {
        Atom::UInt(u.into())
    }
}

impl From<u16> for Atom {
    fn from(u: u16) -> Atom {
        Atom::UInt(u.into())
    }
}

impl From<u8> for Atom {
    fn from(u: u8) -> Atom {
        Atom::UInt(u.into())
    }
}

// SInt
impl From<isize> for Atom {
    fn from(i: isize) -> Atom {
        Atom::SInt(i as i64)
    }
}

impl From<i64> for Atom {
    fn from(i: i64) -> Atom {
        Atom::SInt(i.into())
    }
}

impl From<i32> for Atom {
    fn from(i: i32) -> Atom {
        Atom::SInt(i.into())
    }
}

impl From<i16> for Atom {
    fn from(i: i16) -> Atom {
        Atom::SInt(i.into())
    }
}

impl From<i8> for Atom {
    fn from(i: i8) -> Atom {
        Atom::SInt(i.into())
    }
}

// Float
impl From<f64> for Atom {
    fn from(f: f64) -> Atom {
        Atom::Float(f.into())
    }
}

impl From<f32> for Atom {
    fn from(f: f32) -> Atom {
        Atom::Float(f.into())
    }
}

fn is_num_string(s: &str) -> bool {
    match s.slice_shift_char() {
        None => false,
        Some((c, rest)) => {
            if char::is_digit(c, 10) {
                true
            } else if c == '-' || c == '+' {
                match rest.slice_shift_char() {
                    Some((c, _)) if char::is_digit(c, 10) => true,
                    _ => false,
                }
            } else {
                false
            }
        }
    }
}

#[test]
fn test_is_num_string() {
    assert_eq!(true, is_num_string("0"));
    assert_eq!(true, is_num_string("0123"));
    assert_eq!(true, is_num_string("0a"));
    assert_eq!(true, is_num_string("0a11"));
    assert_eq!(true, is_num_string("1"));
    assert_eq!(true, is_num_string("1.123"));
    assert_eq!(true, is_num_string("+1.123"));
    assert_eq!(true, is_num_string("-1.123"));
    assert_eq!(true, is_num_string("-1"));
    assert_eq!(true, is_num_string("+1"));
    assert_eq!(true, is_num_string("+1abc"));
    assert_eq!(true, is_num_string("-1abc"));

    assert_eq!(false, is_num_string("-a"));
    assert_eq!(false, is_num_string("+a"));
    assert_eq!(false, is_num_string("+"));
    assert_eq!(false, is_num_string("-"));
    assert_eq!(false, is_num_string("a"));
    assert_eq!(false, is_num_string("abc"));
    assert_eq!(false, is_num_string("a0123"));
    assert_eq!(false, is_num_string("a+0123"));
}

fn quote(s: &str) -> Cow<str> {
    if s.is_empty() {
        Cow::Borrowed("\"\"")
    } else if is_num_string(s) || s.contains(is_token_delim) || s.contains("\"") || s.contains("\\") {
        // XXX: Replace \n and others
        let mut r: String = "\"".to_string();
        r.push_str(&s.replace("\\", "\\\\").replace("\"", "\\\""));
        r.push_str("\"");
        Cow::Owned(r)
    } else {
        Cow::Borrowed(s)
    }
}

#[test]
fn test_quote() {
    assert_eq!("\"\"", quote(""));
    assert_eq!("abc", quote("abc"));
    assert_eq!("\"abc\\\\def\"", quote("abc\\def"));
}

impl fmt::Display for Atom {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            Atom::Str(ref s) => write!(f, "{}", quote(s)),
            Atom::UInt(u) => write!(f, "{}", u),
            Atom::SInt(s) => {
                if s >= 0 {
                    write!(f, "+{}", s)
                } else {
                    write!(f, "{}", s)
                }
            }
            Atom::Float(n) => write!(f, "{:.1}", n),
        }
    }
}

#[test]
fn test_display() {
    assert_eq!("12345", &format!("{}", Atom::UInt(12345)));
    assert_eq!("+12345", &format!("{}", Atom::SInt(12345)));
    assert_eq!("-12345", &format!("{}", Atom::SInt(-12345)));
    assert_eq!("-12345.5", &format!("{}", Atom::Float(-12345.5)));
    assert_eq!("abc", &format!("{}", Atom::Str("abc".to_string())));
    assert_eq!("\"(\"", &format!("{}", Atom::Str("(".to_string())));
    assert_eq!("\"\\\"\"", &format!("{}", Atom::Str("\"".to_string())));
    assert_eq!("\"123\"", &format!("{}", Atom::Str("123".to_string())));
    assert_eq!("\"123abc\"",
               &format!("{}", Atom::Str("123abc".to_string())));
    assert_eq!("\"+123\"", &format!("{}", Atom::Str("+123".to_string())));
    assert_eq!("+", &format!("{}", Atom::Str("+".to_string())));
    assert_eq!("-", &format!("{}", Atom::Str("-".to_string())));
    assert_eq!("-b", &format!("{}", Atom::Str("-b".to_string())));
    assert_eq!("abc123", &format!("{}", Atom::Str("abc123".to_string())));
}
<commit_msg>Do not limit precision to one digit. But write floats always with a "."<commit_after>use std::fmt;
use std::borrow::Cow;
use super::token::is_token_delim;

#[derive(Debug, PartialEq, Clone)]
pub enum Atom {
    // String
    Str(String),

    // Unsigned 64-bit integer
    UInt(u64),

    // Signed 64-bit integer
    SInt(i64),

    // 64-bit floating point number
    Float(f64),
}

// Str
impl From<String> for Atom {
    fn from(s: String) -> Atom {
        Atom::Str(s)
    }
}

impl<'a> From<&'a str> for Atom {
    fn from(s: &'a str) -> Atom {
        Atom::Str(s.to_string())
    }
}

// Uint
impl From<usize> for Atom {
    fn from(u: usize) -> Atom {
        Atom::UInt(u as u64)
    }
}

impl From<u64> for Atom {
    fn from(u: u64) -> Atom {
        Atom::UInt(u.into())
    }
}
impl From<u32> for Atom {
    fn from(u: u32) -> Atom {
        Atom::UInt(u.into())
    }
}

impl From<u16> for Atom {
    fn from(u: u16) -> Atom {
        Atom::UInt(u.into())
    }
}

impl From<u8> for Atom {
    fn from(u: u8) -> Atom {
        Atom::UInt(u.into())
    }
}

// SInt
impl From<isize> for Atom {
    fn from(i: isize) -> Atom {
        Atom::SInt(i as i64)
    }
}

impl From<i64> for Atom {
    fn from(i: i64) -> Atom {
        Atom::SInt(i.into())
    }
}

impl From<i32> for Atom {
    fn from(i: i32) -> Atom {
        Atom::SInt(i.into())
    }
}

impl From<i16> for Atom {
    fn from(i: i16) -> Atom {
        Atom::SInt(i.into())
    }
}

impl From<i8> for Atom {
    fn from(i: i8) -> Atom {
        Atom::SInt(i.into())
    }
}

// Float
impl From<f64> for Atom {
    fn from(f: f64) -> Atom {
        Atom::Float(f.into())
    }
}

impl From<f32> for Atom {
    fn from(f: f32) -> Atom {
        Atom::Float(f.into())
    }
}

fn is_num_string(s: &str) -> bool {
    match s.slice_shift_char() {
        None => false,
        Some((c, rest)) => {
            if char::is_digit(c, 10) {
                true
            } else if c == '-' || c == '+' {
                match rest.slice_shift_char() {
                    Some((c, _)) if char::is_digit(c, 10) => true,
                    _ => false,
                }
            } else {
                false
            }
        }
    }
}

#[test]
fn test_is_num_string() {
    assert_eq!(true, is_num_string("0"));
    assert_eq!(true, is_num_string("0123"));
    assert_eq!(true, is_num_string("0a"));
    assert_eq!(true, is_num_string("0a11"));
    assert_eq!(true, is_num_string("1"));
    assert_eq!(true, is_num_string("1.123"));
    assert_eq!(true, is_num_string("+1.123"));
    assert_eq!(true, is_num_string("-1.123"));
    assert_eq!(true, is_num_string("-1"));
    assert_eq!(true, is_num_string("+1"));
    assert_eq!(true, is_num_string("+1abc"));
    assert_eq!(true, is_num_string("-1abc"));

    assert_eq!(false, is_num_string("-a"));
    assert_eq!(false, is_num_string("+a"));
    assert_eq!(false, is_num_string("+"));
    assert_eq!(false, is_num_string("-"));
    assert_eq!(false, is_num_string("a"));
    assert_eq!(false, is_num_string("abc"));
    assert_eq!(false, is_num_string("a0123"));
    assert_eq!(false, is_num_string("a+0123"));
}

fn quote(s: &str) -> Cow<str> {
    if s.is_empty() {
        Cow::Borrowed("\"\"")
    } else if is_num_string(s) || s.contains(is_token_delim) || s.contains("\"") || s.contains("\\") {
        // XXX: Replace \n and others
        let mut r: String = "\"".to_string();
        r.push_str(&s.replace("\\", "\\\\").replace("\"", "\\\""));
        r.push_str("\"");
        Cow::Owned(r)
    } else {
        Cow::Borrowed(s)
    }
}

#[test]
fn test_quote() {
    assert_eq!("\"\"", quote(""));
    assert_eq!("abc", quote("abc"));
    assert_eq!("\"abc\\\\def\"", quote("abc\\def"));
}

impl fmt::Display for Atom {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            Atom::Str(ref s) => write!(f, "{}", quote(s)),
            Atom::UInt(u) => write!(f, "{}", u),
            Atom::SInt(s) => {
                if s >= 0 {
                    write!(f, "+{}", s)
                } else {
                    write!(f, "{}", s)
                }
            }
            Atom::Float(n) => {
                let s = format!("{}", n);
                if s.contains('.') {
                    write!(f, "{}", s)
                } else {
                   write!(f, "{}.0", s)
                }
            }
        }
    }
}

#[test]
fn test_display() {
    assert_eq!("12345", &format!("{}", Atom::UInt(12345)));
    assert_eq!("+12345", &format!("{}", Atom::SInt(12345)));
    assert_eq!("-12345", &format!("{}", Atom::SInt(-12345)));
    assert_eq!("-12345.5", &format!("{}", Atom::Float(-12345.5)));
    assert_eq!("-12345.5", &format!("{}", Atom::Float(-12345.500)));
    assert_eq!("-12345.55", &format!("{}", Atom::Float(-12345.55)));
    assert_eq!("abc", &format!("{}", Atom::Str("abc".to_string())));
    assert_eq!("\"(\"", &format!("{}", Atom::Str("(".to_string())));
    assert_eq!("\"\\\"\"", &format!("{}", Atom::Str("\"".to_string())));
    assert_eq!("\"123\"", &format!("{}", Atom::Str("123".to_string())));
    assert_eq!("\"123abc\"",
               &format!("{}", Atom::Str("123abc".to_string())));
    assert_eq!("\"+123\"", &format!("{}", Atom::Str("+123".to_string())));
    assert_eq!("+", &format!("{}", Atom::Str("+".to_string())));
    assert_eq!("-", &format!("{}", Atom::Str("-".to_string())));
    assert_eq!("-b", &format!("{}", Atom::Str("-b".to_string())));
    assert_eq!("abc123", &format!("{}", Atom::Str("abc123".to_string())));
}
<|endoftext|>


<commit_before>use super::*;

use std::iter;

/// Doc-comments are promoted to attributes that have `is_sugared_doc` = true
#[derive(Debug, Clone, Eq, PartialEq, Hash)]
pub struct Attribute {
    pub style: AttrStyle,
    pub value: MetaItem,
    pub is_sugared_doc: bool,
}

impl Attribute {
    pub fn name(&self) -> &str {
        self.value.name()
    }
}

/// Distinguishes between Attributes that decorate items and Attributes that
/// are contained as statements within items. These two cases need to be
/// distinguished for pretty-printing.
#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
pub enum AttrStyle {
    /// Attribute of the form `#![...]`.
    Outer,

    /// Attribute of the form `#[...]`.
    Inner,
}

/// A compile-time attribute item.
///
/// E.g. `#[test]`, `#[derive(..)]` or `#[feature = "foo"]`
#[derive(Debug, Clone, Eq, PartialEq, Hash)]
pub enum MetaItem {
    /// Word meta item.
    ///
    /// E.g. `test` as in `#[test]`
    Word(Ident),

    /// List meta item.
    ///
    /// E.g. `derive(..)` as in `#[derive(..)]`
    List(Ident, Vec<NestedMetaItem>),

    /// Name-value meta item.
    ///
    /// E.g. `feature = "foo"` as in `#[feature = "foo"]`
    NameValue(Ident, Lit),
}

impl MetaItem {
    /// Name of the item.
    ///
    /// E.g. `test` as in `#[test]`, `derive` as in `#[derive(..)]`, and
    /// `feature` as in `#[feature = "foo"]`.
    pub fn name(&self) -> &str {
        match *self {
            MetaItem::Word(ref name) |
            MetaItem::List(ref name, _) |
            MetaItem::NameValue(ref name, _) => name.as_ref(),
        }
    }
}

/// Possible values inside of compile-time attribute lists.
///
/// E.g. the '..' in `#[name(..)]`.
#[derive(Debug, Clone, Eq, PartialEq, Hash)]
pub enum NestedMetaItem {
    /// A full `MetaItem`.
    ///
    /// E.g. `Copy` in `#[derive(Copy)]` would be a `MetaItem::Word(Ident::from("Copy"))`.
    MetaItem(MetaItem),

    /// A Rust literal.
    ///
    /// E.g. `"name"` in `#[rename("name")]`.
    Literal(Lit),
}

pub trait FilterAttrs<'a> {
    type Ret: Iterator<Item = &'a Attribute>;

    fn outer(self) -> Self::Ret;
    fn inner(self) -> Self::Ret;
}

impl<'a, T> FilterAttrs<'a> for T
    where T: IntoIterator<Item = &'a Attribute>
{
    type Ret = iter::Filter<T::IntoIter, fn(&&Attribute) -> bool>;

    fn outer(self) -> Self::Ret {
        fn is_outer(attr: &&Attribute) -> bool {
            attr.style == AttrStyle::Outer
        }
        self.into_iter().filter(is_outer)
    }

    fn inner(self) -> Self::Ret {
        fn is_inner(attr: &&Attribute) -> bool {
            attr.style == AttrStyle::Inner
        }
        self.into_iter().filter(is_inner)
    }
}

#[cfg(feature = "parsing")]
pub mod parsing {
    use super::*;
    use ident::parsing::ident;
    use lit::parsing::lit;
    use synom::space::{block_comment, whitespace};

    #[cfg(feature = "full")]
    named!(pub inner_attr -> Attribute, alt!(
        do_parse!(
            punct!("#") >>
            punct!("!") >>
            punct!("[") >>
            meta_item: meta_item >>
            punct!("]") >>
            (Attribute {
                style: AttrStyle::Inner,
                value: meta_item,
                is_sugared_doc: false,
            })
        )
        |
        do_parse!(
            punct!("//!") >>
            content: take_until!("\n") >>
            (Attribute {
                style: AttrStyle::Inner,
                value: MetaItem::NameValue(
                    "doc".into(),
                    format!("//!{}", content).into(),
                ),
                is_sugared_doc: true,
            })
        )
        |
        do_parse!(
            option!(whitespace) >>
            peek!(tag!("/*!")) >>
            com: block_comment >>
            (Attribute {
                style: AttrStyle::Inner,
                value: MetaItem::NameValue(
                    "doc".into(),
                    com.into(),
                ),
                is_sugared_doc: true,
            })
        )
    ));

    named!(pub outer_attr -> Attribute, alt!(
        do_parse!(
            punct!("#") >>
            punct!("[") >>
            meta_item: meta_item >>
            punct!("]") >>
            (Attribute {
                style: AttrStyle::Outer,
                value: meta_item,
                is_sugared_doc: false,
            })
        )
        |
        do_parse!(
            punct!("///") >>
            not!(tag!("/")) >>
            content: take_until!("\n") >>
            (Attribute {
                style: AttrStyle::Outer,
                value: MetaItem::NameValue(
                    "doc".into(),
                    format!("///{}", content).into(),
                ),
                is_sugared_doc: true,
            })
        )
        |
        do_parse!(
            option!(whitespace) >>
            peek!(tuple!(tag!("/**"), not!(tag!("*")))) >>
            com: block_comment >>
            (Attribute {
                style: AttrStyle::Outer,
                value: MetaItem::NameValue(
                    "doc".into(),
                    com.into(),
                ),
                is_sugared_doc: true,
            })
        )
    ));

    named!(meta_item -> MetaItem, alt!(
        do_parse!(
            id: ident >>
            punct!("(") >>
            inner: terminated_list!(punct!(","), nested_meta_item) >>
            punct!(")") >>
            (MetaItem::List(id, inner))
        )
        |
        do_parse!(
            name: ident >>
            punct!("=") >>
            value: lit >>
            (MetaItem::NameValue(name, value))
        )
        |
        map!(ident, MetaItem::Word)
    ));

    named!(nested_meta_item -> NestedMetaItem, alt!(
        meta_item => { NestedMetaItem::MetaItem }
        |
        lit => { NestedMetaItem::Literal }
    ));
}

#[cfg(feature = "printing")]
mod printing {
    use super::*;
    use lit::{Lit, StrStyle};
    use quote::{Tokens, ToTokens};

    impl ToTokens for Attribute {
        fn to_tokens(&self, tokens: &mut Tokens) {
            if let Attribute { style,
                               value: MetaItem::NameValue(ref name,
                                                   Lit::Str(ref value, StrStyle::Cooked)),
                               is_sugared_doc: true } = *self {
                if name == "doc" {
                    match style {
                        AttrStyle::Inner if value.starts_with("//!") => {
                            tokens.append(&format!("{}\n", value));
                            return;
                        }
                        AttrStyle::Inner if value.starts_with("/*!") => {
                            tokens.append(value);
                            return;
                        }
                        AttrStyle::Outer if value.starts_with("///") => {
                            tokens.append(&format!("{}\n", value));
                            return;
                        }
                        AttrStyle::Outer if value.starts_with("/**") => {
                            tokens.append(value);
                            return;
                        }
                        _ => {}
                    }
                }
            }

            tokens.append("#");
            if let AttrStyle::Inner = self.style {
                tokens.append("!");
            }
            tokens.append("[");
            self.value.to_tokens(tokens);
            tokens.append("]");
        }
    }

    impl ToTokens for MetaItem {
        fn to_tokens(&self, tokens: &mut Tokens) {
            match *self {
                MetaItem::Word(ref ident) => {
                    ident.to_tokens(tokens);
                }
                MetaItem::List(ref ident, ref inner) => {
                    ident.to_tokens(tokens);
                    tokens.append("(");
                    tokens.append_separated(inner, ",");
                    tokens.append(")");
                }
                MetaItem::NameValue(ref name, ref value) => {
                    name.to_tokens(tokens);
                    tokens.append("=");
                    value.to_tokens(tokens);
                }
            }
        }
    }

    impl ToTokens for NestedMetaItem {
        fn to_tokens(&self, tokens: &mut Tokens) {
            match *self {
                NestedMetaItem::MetaItem(ref nested) => {
                    nested.to_tokens(tokens);
                }
                NestedMetaItem::Literal(ref lit) => {
                    lit.to_tokens(tokens);
                }
            }
        }
    }
}
<commit_msg>Support parsing attributes with arbitrary tokens.<commit_after>use super::*;

use std::iter;

/// Doc-comments are promoted to attributes that have `is_sugared_doc` = true
#[derive(Debug, Clone, Eq, PartialEq, Hash)]
pub struct Attribute {
    pub style: AttrStyle,
    pub value: MetaItem,
    pub is_sugared_doc: bool,
}

impl Attribute {
    pub fn name(&self) -> &str {
        self.value.name()
    }
}

/// Distinguishes between Attributes that decorate items and Attributes that
/// are contained as statements within items. These two cases need to be
/// distinguished for pretty-printing.
#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
pub enum AttrStyle {
    /// Attribute of the form `#![...]`.
    Outer,

    /// Attribute of the form `#[...]`.
    Inner,
}

/// A compile-time attribute item.
///
/// E.g. `#[test]`, `#[derive(..)]` or `#[feature = "foo"]`
#[derive(Debug, Clone, Eq, PartialEq, Hash)]
pub enum MetaItem {
    /// Word meta item.
    ///
    /// E.g. `test` as in `#[test]`
    Word(Ident),

    /// List meta item.
    ///
    /// E.g. `derive(..)` as in `#[derive(..)]`
    List(Ident, Vec<NestedMetaItem>),

    /// Name-value meta item.
    ///
    /// E.g. `feature = "foo"` as in `#[feature = "foo"]`
    NameValue(Ident, Lit),

    /// Tokens meta item.
    ///
    /// E.g. `test foo bar` as in `#[test foo bar]`
    Tokens(Ident, Vec<TokenTree>),
}

impl MetaItem {
    /// Name of the item.
    ///
    /// E.g. `test` as in `#[test]`, `derive` as in `#[derive(..)]`, and
    /// `feature` as in `#[feature = "foo"]`.
    pub fn name(&self) -> &str {
        match *self {
            MetaItem::Word(ref name) |
            MetaItem::List(ref name, _) |
            MetaItem::NameValue(ref name, _) |
            MetaItem::Tokens(ref name, _) => name.as_ref(),
        }
    }
}

/// Possible values inside of compile-time attribute lists.
///
/// E.g. the '..' in `#[name(..)]`.
#[derive(Debug, Clone, Eq, PartialEq, Hash)]
pub enum NestedMetaItem {
    /// A full `MetaItem`.
    ///
    /// E.g. `Copy` in `#[derive(Copy)]` would be a `MetaItem::Word(Ident::from("Copy"))`.
    MetaItem(MetaItem),

    /// A Rust literal.
    ///
    /// E.g. `"name"` in `#[rename("name")]`.
    Literal(Lit),
}

pub trait FilterAttrs<'a> {
    type Ret: Iterator<Item = &'a Attribute>;

    fn outer(self) -> Self::Ret;
    fn inner(self) -> Self::Ret;
}

impl<'a, T> FilterAttrs<'a> for T
    where T: IntoIterator<Item = &'a Attribute>
{
    type Ret = iter::Filter<T::IntoIter, fn(&&Attribute) -> bool>;

    fn outer(self) -> Self::Ret {
        fn is_outer(attr: &&Attribute) -> bool {
            attr.style == AttrStyle::Outer
        }
        self.into_iter().filter(is_outer)
    }

    fn inner(self) -> Self::Ret {
        fn is_inner(attr: &&Attribute) -> bool {
            attr.style == AttrStyle::Inner
        }
        self.into_iter().filter(is_inner)
    }
}

#[cfg(feature = "parsing")]
pub mod parsing {
    use super::*;
    use ident::parsing::ident;
    use lit::parsing::lit;
    use mac::parsing::token_trees;
    use synom::space::{block_comment, whitespace};

    #[cfg(feature = "full")]
    named!(pub inner_attr -> Attribute, alt!(
        do_parse!(
            punct!("#") >>
            punct!("!") >>
            punct!("[") >>
            meta_item: meta_item >>
            punct!("]") >>
            (Attribute {
                style: AttrStyle::Inner,
                value: meta_item,
                is_sugared_doc: false,
            })
        )
        |
        do_parse!(
            punct!("//!") >>
            content: take_until!("\n") >>
            (Attribute {
                style: AttrStyle::Inner,
                value: MetaItem::NameValue(
                    "doc".into(),
                    format!("//!{}", content).into(),
                ),
                is_sugared_doc: true,
            })
        )
        |
        do_parse!(
            option!(whitespace) >>
            peek!(tag!("/*!")) >>
            com: block_comment >>
            (Attribute {
                style: AttrStyle::Inner,
                value: MetaItem::NameValue(
                    "doc".into(),
                    com.into(),
                ),
                is_sugared_doc: true,
            })
        )
    ));

    named!(pub outer_attr -> Attribute, alt!(
        do_parse!(
            punct!("#") >>
            punct!("[") >>
            meta_item: meta_item >>
            punct!("]") >>
            (Attribute {
                style: AttrStyle::Outer,
                value: meta_item,
                is_sugared_doc: false,
            })
        )
        |
        do_parse!(
            name_and_token_trees: preceded!(
                punct!("#"),
                delimited!(
                    punct!("["),
                    tuple!(ident, token_trees),
                    punct!("]")
                )
            ) >>
            (Attribute {
                style: AttrStyle::Outer,
                value: {
                    let (name, token_trees) = name_and_token_trees;
                    MetaItem::Tokens(name, token_trees)
                },
                is_sugared_doc: false,
            })
        )
        |
        do_parse!(
            punct!("///") >>
            not!(tag!("/")) >>
            content: take_until!("\n") >>
            (Attribute {
                style: AttrStyle::Outer,
                value: MetaItem::NameValue(
                    "doc".into(),
                    format!("///{}", content).into(),
                ),
                is_sugared_doc: true,
            })
        )
        |
        do_parse!(
            option!(whitespace) >>
            peek!(tuple!(tag!("/**"), not!(tag!("*")))) >>
            com: block_comment >>
            (Attribute {
                style: AttrStyle::Outer,
                value: MetaItem::NameValue(
                    "doc".into(),
                    com.into(),
                ),
                is_sugared_doc: true,
            })
        )
    ));

    named!(meta_item -> MetaItem, alt!(
        do_parse!(
            id: ident >>
            punct!("(") >>
            inner: terminated_list!(punct!(","), nested_meta_item) >>
            punct!(")") >>
            (MetaItem::List(id, inner))
        )
        |
        do_parse!(
            name: ident >>
            punct!("=") >>
            value: lit >>
            (MetaItem::NameValue(name, value))
        )
        |
        map!(ident, MetaItem::Word)
    ));

    named!(nested_meta_item -> NestedMetaItem, alt!(
        meta_item => { NestedMetaItem::MetaItem }
        |
        lit => { NestedMetaItem::Literal }
    ));
}

#[cfg(feature = "printing")]
mod printing {
    use super::*;
    use lit::{Lit, StrStyle};
    use quote::{Tokens, ToTokens};

    impl ToTokens for Attribute {
        fn to_tokens(&self, tokens: &mut Tokens) {
            if let Attribute { style,
                               value: MetaItem::NameValue(ref name,
                                                   Lit::Str(ref value, StrStyle::Cooked)),
                               is_sugared_doc: true } = *self {
                if name == "doc" {
                    match style {
                        AttrStyle::Inner if value.starts_with("//!") => {
                            tokens.append(&format!("{}\n", value));
                            return;
                        }
                        AttrStyle::Inner if value.starts_with("/*!") => {
                            tokens.append(value);
                            return;
                        }
                        AttrStyle::Outer if value.starts_with("///") => {
                            tokens.append(&format!("{}\n", value));
                            return;
                        }
                        AttrStyle::Outer if value.starts_with("/**") => {
                            tokens.append(value);
                            return;
                        }
                        _ => {}
                    }
                }
            }

            tokens.append("#");
            if let AttrStyle::Inner = self.style {
                tokens.append("!");
            }
            tokens.append("[");
            self.value.to_tokens(tokens);
            tokens.append("]");
        }
    }

    impl ToTokens for MetaItem {
        fn to_tokens(&self, tokens: &mut Tokens) {
            match *self {
                MetaItem::Word(ref ident) => {
                    ident.to_tokens(tokens);
                }
                MetaItem::List(ref ident, ref inner) => {
                    ident.to_tokens(tokens);
                    tokens.append("(");
                    tokens.append_separated(inner, ",");
                    tokens.append(")");
                }
                MetaItem::NameValue(ref name, ref value) => {
                    name.to_tokens(tokens);
                    tokens.append("=");
                    value.to_tokens(tokens);
                }
                MetaItem::Tokens(ref name, ref token_trees) => {
                    name.to_tokens(tokens);
                    tokens.append_all(token_trees);
                }
            }
        }
    }

    impl ToTokens for NestedMetaItem {
        fn to_tokens(&self, tokens: &mut Tokens) {
            match *self {
                NestedMetaItem::MetaItem(ref nested) => {
                    nested.to_tokens(tokens);
                }
                NestedMetaItem::Literal(ref lit) => {
                    lit.to_tokens(tokens);
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::collections::hash_map::Entry as HashMapEntry;
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, RwLock};
use std::thread;
use std::time::Duration;
use std::vec::Vec;

use data::{Entry, Hint, SequenceNumber};
use log::{Log, LogWriter};
use stats::Stats;

const COMPACTION_CHECK_FREQUENCY: u64 = 60;
const FRAGMENTATION_THRESHOLD: f64 = 0.6;

#[derive(Debug)]
pub struct IndexEntry {
    pub file_id: u32,
    entry_pos: u64,
    pub entry_size: u64,
    sequence: SequenceNumber,
}

struct Index {
    map: HashMap<Vec<u8>, IndexEntry>,
    stats: Stats,
}

impl Index {
    fn new() -> Index {
        Index {
            map: HashMap::new(),
            stats: Stats::new(),
        }
    }

    fn get(&self, key: &[u8]) -> Option<&IndexEntry> {
        self.map.get(key)
    }

    fn insert(&mut self, key: Vec<u8>, index_entry: IndexEntry) -> Option<IndexEntry> {
        self.stats.add_entry(&index_entry);
        self.map
            .insert(key, index_entry)
            .map(|entry| {
                self.stats.remove_entry(&entry);
                entry
            })
    }

    fn remove(&mut self, key: &[u8]) -> Option<IndexEntry> {
        self.map
            .remove(key)
            .map(|entry| {
                self.stats.remove_entry(&entry);
                entry
            })
    }

    fn update(&mut self, hint: Hint, file_id: u32) {
        let index_entry = IndexEntry {
            file_id: file_id,
            entry_pos: hint.entry_pos,
            entry_size: hint.entry_size(),
            sequence: hint.sequence,
        };

        match self.map.entry(hint.key.to_vec()) {
            HashMapEntry::Occupied(mut o) => {
                if o.get().sequence <= hint.sequence {
                    self.stats.remove_entry(o.get());
                    if hint.deleted {
                        o.remove();
                    } else {
                        self.stats.add_entry(&index_entry);
                        o.insert(index_entry);
                    }
                } else {
                    self.stats.add_entry(&index_entry);
                    self.stats.remove_entry(&index_entry);
                }
            }
            HashMapEntry::Vacant(e) => {
                if !hint.deleted {
                    self.stats.add_entry(&index_entry);
                    e.insert(index_entry);
                }
            }
        }
    }
}

struct CaskInner {
    current_sequence: SequenceNumber,
    index: Index,
    log: Log,
}

impl CaskInner {
    fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
        self.index.get(key).and_then(|index_entry| {
            let entry = self.log.read_entry(index_entry.file_id, index_entry.entry_pos).unwrap();
            if entry.deleted {
                warn!("Index pointed to dead entry: Entry {{ key: {:?}, sequence: {} }} at file: \
                       {}",
                      entry.key,
                      entry.sequence,
                      index_entry.file_id);
                None
            } else {
                Some(entry.value.into_owned())
            }
        })
    }

    fn put(&mut self, key: Vec<u8>, value: &[u8]) {
        let index_entry = {
            let entry = Entry::new(self.current_sequence, &*key, value).unwrap();

            let (file_id, file_pos) = self.log.append_entry(&entry).unwrap();

            self.current_sequence += 1;

            IndexEntry {
                file_id: file_id,
                entry_pos: file_pos,
                entry_size: entry.size(),
                sequence: entry.sequence,
            }
        };

        self.index.insert(key, index_entry);
    }

    fn delete(&mut self, key: &[u8]) {
        if self.index.remove(key).is_some() {
            let entry = Entry::deleted(self.current_sequence, key);
            let _ = self.log.append_entry(&entry);
            self.current_sequence += 1;
        }
    }
}

#[derive(Clone)]
pub struct Cask {
    path: PathBuf,
    dropped: Arc<AtomicBool>,
    inner: Arc<RwLock<CaskInner>>,
}

impl Cask {
    pub fn open(path: &str, sync: bool) -> Cask {
        info!("Opening database: {:?}", &path);
        let mut log = Log::open(path, sync).unwrap();
        let mut index = Index::new();

        let mut sequence = 0;

        for file_id in log.files() {
            let mut f = |hint: Hint| {
                if hint.sequence > sequence {
                    sequence = hint.sequence;
                }

                index.update(hint, file_id);
            };

            match log.hints(file_id).unwrap() {
                Some(hints) => {
                    for hint in hints {
                        f(hint.unwrap());
                    }
                }
                None => {
                    for hint in log.recreate_hints(file_id).unwrap() {
                        f(hint.unwrap());
                    }
                }
            };
        }

        info!("Opened database: {:?}", &path);
        info!("Current sequence number: {:?}", sequence);

        let cask = Cask {
            path: log.path.clone(),
            dropped: Arc::new(AtomicBool::new(false)),
            inner: Arc::new(RwLock::new(CaskInner {
                                            current_sequence: sequence + 1,
                                            log: log,
                                            index: index,
                                        })),
        };

        let caskt = cask.clone();
        thread::spawn(move || loop {
                          thread::sleep(Duration::new(COMPACTION_CHECK_FREQUENCY, 0));

                          info!("Compaction thread wake up");

                          if caskt.dropped.load(Ordering::SeqCst) {
                              info!("Cask has been dropped, background compaction \
                                     thread is exiting");
                              break;
                          }

                          caskt.compact();
                      });

        cask
    }

    fn compact_file_aux(&self, file_id: u32) -> Option<u32> {
        if file_id == self.inner.read().unwrap().log.active_file_id {
            return None;
        }

        let hints = {
            self.inner.read().unwrap().log.hints(file_id).unwrap()
        };

        match hints {
            Some(hints) => {
                let new_file_id = {
                    // FIXME: turn into error
                    self.inner.read().unwrap().log.current_file_id.increment()
                };

                info!("Compacting data file: {} into: {}", file_id, new_file_id);

                let mut log_writer = LogWriter::new(&self.path, new_file_id, false).unwrap();
                let mut deletes = HashMap::new();

                {
                    let mut inserts = Vec::new();

                    for hint in hints {
                        let hint = hint.unwrap();
                        let inner = self.inner.read().unwrap();
                        let index_entry = inner.index.get(&*hint.key);

                        if hint.deleted {
                            if index_entry.is_none() {
                                match deletes.entry(hint.key.to_vec()) {
                                    HashMapEntry::Occupied(mut o) => {
                                        if *o.get() < hint.sequence {
                                            o.insert(hint.sequence);
                                        }
                                    }
                                    HashMapEntry::Vacant(e) => {
                                        e.insert(hint.sequence);
                                    }
                                }
                            }
                        } else if index_entry.is_some() &&
                                  index_entry.unwrap().sequence == hint.sequence {
                            inserts.push(hint)
                        }
                    }

                    for hint in inserts {
                        // FIXME: turn into error
                        let log = &self.inner.read().unwrap().log;
                        log_writer
                            .write(&log.read_entry(file_id, hint.entry_pos).unwrap())
                            .unwrap();
                    }
                }

                for (key, sequence) in deletes {
                    log_writer.write(&Entry::deleted(sequence, key)).unwrap();
                }

                Some(new_file_id)
            }
            _ => None,
        }
    }

    pub fn compact_file(&self, file_id: u32) {
        let new_file_id = self.compact_file_aux(file_id);

        if let Some(new_file_id) = new_file_id {
            let hints = {
                self.inner.read().unwrap().log.hints(new_file_id).unwrap()
            };

            if let Some(hints) = hints {
                for hint in hints {
                    let hint = hint.unwrap();
                    self.inner.write().unwrap().index.update(hint, new_file_id);
                }

                self.inner
                    .write()
                    .unwrap()
                    .log
                    .swap_file(file_id, new_file_id)
                    .unwrap();

                info!("Finished compacting data file: {} into: {}",
                      file_id,
                      new_file_id);
            };
        }
    }

    pub fn compact(&self) {
        let iter = {
            self.inner.read().unwrap().index.stats.fragmentation()
        };

        for &(file_id, fragmentation) in iter.iter().filter(|e| e.1 >= FRAGMENTATION_THRESHOLD) {

            info!("File {} has fragmentation factor of {}%, adding for compaction",
                  file_id,
                  fragmentation * 100.0);

            self.compact_file(file_id);
        }
    }

    pub fn get<K: AsRef<[u8]>>(&self, key: K) -> Option<Vec<u8>> {
        self.inner.read().unwrap().get(key.as_ref())
    }

    pub fn put<K: Into<Vec<u8>>, V: AsRef<[u8]>>(&self, key: K, value: V) {
        self.inner.write().unwrap().put(key.into(), value.as_ref())
    }

    pub fn delete<K: AsRef<[u8]>>(&self, key: K) {
        self.inner.write().unwrap().delete(key.as_ref())
    }
}

impl Drop for Cask {
    fn drop(&mut self) {
        self.dropped.store(true, Ordering::SeqCst);
    }
}
<commit_msg>handle errors in cask<commit_after>use std::collections::HashMap;
use std::collections::hash_map::Entry as HashMapEntry;
use std::path::PathBuf;
use std::result::Result::Ok;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, RwLock};
use std::thread;
use std::time::Duration;
use std::vec::Vec;

use data::{Entry, Hint, SequenceNumber};
use errors::Result;
use log::{Log, LogWriter};
use stats::Stats;

const COMPACTION_CHECK_FREQUENCY: u64 = 60;
const FRAGMENTATION_THRESHOLD: f64 = 0.6;

#[derive(Debug)]
pub struct IndexEntry {
    pub file_id: u32,
    entry_pos: u64,
    pub entry_size: u64,
    sequence: SequenceNumber,
}

struct Index {
    map: HashMap<Vec<u8>, IndexEntry>,
    stats: Stats,
}

impl Index {
    fn new() -> Index {
        Index {
            map: HashMap::new(),
            stats: Stats::new(),
        }
    }

    fn get(&self, key: &[u8]) -> Option<&IndexEntry> {
        self.map.get(key)
    }

    fn insert(&mut self, key: Vec<u8>, index_entry: IndexEntry) -> Option<IndexEntry> {
        self.stats.add_entry(&index_entry);
        self.map
            .insert(key, index_entry)
            .map(|entry| {
                self.stats.remove_entry(&entry);
                entry
            })
    }

    fn remove(&mut self, key: &[u8]) -> Option<IndexEntry> {
        self.map
            .remove(key)
            .map(|entry| {
                self.stats.remove_entry(&entry);
                entry
            })
    }

    fn update(&mut self, hint: Hint, file_id: u32) {
        let index_entry = IndexEntry {
            file_id: file_id,
            entry_pos: hint.entry_pos,
            entry_size: hint.entry_size(),
            sequence: hint.sequence,
        };

        match self.map.entry(hint.key.to_vec()) {
            HashMapEntry::Occupied(mut o) => {
                if o.get().sequence <= hint.sequence {
                    self.stats.remove_entry(o.get());
                    if hint.deleted {
                        o.remove();
                    } else {
                        self.stats.add_entry(&index_entry);
                        o.insert(index_entry);
                    }
                } else {
                    self.stats.add_entry(&index_entry);
                    self.stats.remove_entry(&index_entry);
                }
            }
            HashMapEntry::Vacant(e) => {
                if !hint.deleted {
                    self.stats.add_entry(&index_entry);
                    e.insert(index_entry);
                }
            }
        }
    }
}

struct CaskInner {
    current_sequence: SequenceNumber,
    index: Index,
    log: Log,
}

impl CaskInner {
    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
        let value = match self.index.get(key) {
            Some(index_entry) => {
                let entry = self.log
                    .read_entry(index_entry.file_id, index_entry.entry_pos)?;
                if entry.deleted {
                    warn!("Index pointed to dead entry: Entry {{ key: {:?}, sequence: {} }} at \
                           file: {}",
                          entry.key,
                          entry.sequence,
                          index_entry.file_id);
                    None
                } else {
                    Some(entry.value.into_owned())
                }
            }
            _ => None,
        };

        Ok(value)
    }

    fn put(&mut self, key: Vec<u8>, value: &[u8]) -> Result<()> {
        let index_entry = {
            let entry = Entry::new(self.current_sequence, &*key, value)?;

            let (file_id, file_pos) = self.log.append_entry(&entry)?;

            self.current_sequence += 1;

            IndexEntry {
                file_id: file_id,
                entry_pos: file_pos,
                entry_size: entry.size(),
                sequence: entry.sequence,
            }
        };

        self.index.insert(key, index_entry);

        Ok(())
    }

    fn delete(&mut self, key: &[u8]) -> Result<()> {
        if self.index.remove(key).is_some() {
            let entry = Entry::deleted(self.current_sequence, key);
            self.log.append_entry(&entry)?;
            self.current_sequence += 1;
        }

        Ok(())
    }
}

#[derive(Clone)]
pub struct Cask {
    path: PathBuf,
    dropped: Arc<AtomicBool>,
    inner: Arc<RwLock<CaskInner>>,
}

impl Cask {
    pub fn open(path: &str, sync: bool) -> Result<Cask> {
        info!("Opening database: {:?}", &path);
        let mut log = Log::open(path, sync)?;
        let mut index = Index::new();

        let mut sequence = 0;

        for file_id in log.files() {
            let mut f = |hint: Hint| {
                if hint.sequence > sequence {
                    sequence = hint.sequence;
                }

                index.update(hint, file_id);
            };

            match log.hints(file_id)? {
                Some(hints) => {
                    for hint in hints {
                        f(hint?);
                    }
                }
                None => {
                    for hint in log.recreate_hints(file_id)? {
                        f(hint?);
                    }
                }
            };
        }

        info!("Opened database: {:?}", &path);
        info!("Current sequence number: {:?}", sequence);

        let cask = Cask {
            path: log.path.clone(),
            dropped: Arc::new(AtomicBool::new(false)),
            inner: Arc::new(RwLock::new(CaskInner {
                                            current_sequence: sequence + 1,
                                            log: log,
                                            index: index,
                                        })),
        };

        let caskt = cask.clone();
        thread::spawn(move || loop {
                          thread::sleep(Duration::new(COMPACTION_CHECK_FREQUENCY, 0));

                          info!("Compaction thread wake up");

                          if caskt.dropped.load(Ordering::SeqCst) {
                              info!("Cask has been dropped, background compaction \
                                     thread is exiting");
                              break;
                          }

                          caskt.compact().unwrap();
                      });

        Ok(cask)
    }

    fn compact_file_aux(&self, file_id: u32) -> Result<Option<u32>> {
        if file_id == self.inner.read().unwrap().log.active_file_id {
            return Ok(None);
        }

        let hints = {
            self.inner.read().unwrap().log.hints(file_id)?
        };

        Ok(match hints {
               Some(hints) => {
                   let new_file_id = {
                       // FIXME: turn into error
                       self.inner.read().unwrap().log.current_file_id.increment()
                   };

                   info!("Compacting data file: {} into: {}", file_id, new_file_id);

                   let mut log_writer = LogWriter::new(&self.path, new_file_id, false)?;
                   let mut deletes = HashMap::new();

                   {
                       let mut inserts = Vec::new();

                       for hint in hints {
                           let hint = hint?;
                           let inner = self.inner.read().unwrap();
                           let index_entry = inner.index.get(&*hint.key);

                           if hint.deleted {
                               if index_entry.is_none() {
                                   match deletes.entry(hint.key.to_vec()) {
                                       HashMapEntry::Occupied(mut o) => {
                                           if *o.get() < hint.sequence {
                                               o.insert(hint.sequence);
                                           }
                                       }
                                       HashMapEntry::Vacant(e) => {
                                           e.insert(hint.sequence);
                                       }
                                   }
                               }
                           } else if index_entry.is_some() &&
                                     index_entry.unwrap().sequence == hint.sequence {
                               inserts.push(hint)
                           }
                       }

                       for hint in inserts {
                           // FIXME: turn into error
                           let log = &self.inner.read().unwrap().log;
                           log_writer.write(&log.read_entry(file_id, hint.entry_pos)?)?;
                       }
                   }

                   for (key, sequence) in deletes {
                       log_writer.write(&Entry::deleted(sequence, key))?;
                   }

                   Some(new_file_id)
               }
               _ => None,
           })
    }

    pub fn compact_file(&self, file_id: u32) -> Result<()> {
        let new_file_id = self.compact_file_aux(file_id)?;

        if let Some(new_file_id) = new_file_id {
            let hints = {
                self.inner.read().unwrap().log.hints(new_file_id)?
            };

            if let Some(hints) = hints {
                for hint in hints {
                    let hint = hint?;
                    self.inner.write().unwrap().index.update(hint, new_file_id);
                }

                self.inner
                    .write()
                    .unwrap()
                    .log
                    .swap_file(file_id, new_file_id)?;

                info!("Finished compacting data file: {} into: {}",
                      file_id,
                      new_file_id);
            };
        }

        Ok(())
    }

    pub fn compact(&self) -> Result<()> {
        let iter = {
            self.inner.read().unwrap().index.stats.fragmentation()
        };

        for &(file_id, fragmentation) in iter.iter().filter(|e| e.1 >= FRAGMENTATION_THRESHOLD) {

            info!("File {} has fragmentation factor of {}%, adding for compaction",
                  file_id,
                  fragmentation * 100.0);

            self.compact_file(file_id)?;
        }

        Ok(())
    }

    pub fn get<K: AsRef<[u8]>>(&self, key: K) -> Result<Option<Vec<u8>>> {
        self.inner.read().unwrap().get(key.as_ref())
    }

    pub fn put<K: Into<Vec<u8>>, V: AsRef<[u8]>>(&self, key: K, value: V) -> Result<()> {
        self.inner.write().unwrap().put(key.into(), value.as_ref())
    }

    pub fn delete<K: AsRef<[u8]>>(&self, key: K) -> Result<()> {
        self.inner.write().unwrap().delete(key.as_ref())
    }
}

impl Drop for Cask {
    fn drop(&mut self) {
        self.dropped.store(true, Ordering::SeqCst);
    }
}
<|endoftext|>


<commit_before>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::cmp;

/// An iterator over a vector in (non-overlapping) mutable chunks (`size`
/// elements at a time).
///
/// When the vector len is not evenly divided by the chunk size,
/// the last slice of the iteration will be the remainder.
pub struct MutChunkIter<'self, T> {
    priv v: &'self mut [T],
    priv len: uint,
    priv size: uint,
    priv pos: uint
}

impl<'self, T> Iterator<&'self mut [T]> for MutChunkIter<'self, T> {
    #[inline]
    fn next(&mut self) -> Option<&'self mut [T]> {
        if self.pos >= self.len {
            None
        } else {
            let chunksz = cmp::min(self.len - self.pos, self.size);
            let out = self.v.mut_slice(self.pos, self.pos + chunksz);
            self.pos += chunksz;
            Some(out)
        }
    }

    #[inline]
    fn size_hint(&self) -> (uint, Option<uint>) {
        if self.len == 0 {
            (0, Some(0))
        } else {
            let (n, rem) = self.len.div_rem(&self.size);
            let n = if rem > 0 { n + 1 } else { n };
            (n, Some(n))
        }
    }
}

pub trait MutChunkIterable<'self, T> {
    /**
     * Returns an iterator over `size` elements of the vector at a
     * time. The chunks do not overlap. If `size` does not divide the
     * length of the vector, then the last chunk will not have length
     * `size`. The chunk are mutable.
     */
    fn mut_chunk_iter(self, size: uint) -> MutChunkIter<'self, T>;
}

impl<'self, T> MutChunkIterable<'self, T> for &'self mut [T] {
    #[inline]
    fn mut_chunk_iter(self, size: uint) -> MutChunkIter<'self, T> {
        assert!(size != 0);
        let len = self.len();
        MutChunkIter { v: self, len: len, size: size, pos: 0 }
    }
}

#[test]
fn test_mut_chunk_iterator() {
    let mut v = [0u8, 1, 2, 3, 4, 5];

    for (i, chunk) in v.mut_chunk_iter(3).enumerate() {
        chunk[0] = i as u8;
        chunk[1] = i as u8;
        chunk[2] = i as u8;
    }

    let result = [0u8, 0, 0, 1, 1, 1];
    assert_eq!(v, result);
}

#[test]
#[should_fail]
fn test_mut_chunk_iterator_0() {
    let mut v = [1, 2, 3, 4];
    let _it = v.mut_chunk_iter(0);
}
<commit_msg>Use some unsafe code to get MutChunkIter to compile again after Rust issue #8624.<commit_after>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::cmp;

/// An iterator over a vector in (non-overlapping) mutable chunks (`size`
/// elements at a time).
///
/// When the vector len is not evenly divided by the chunk size,
/// the last slice of the iteration will be the remainder.
pub struct MutChunkIter<'self, T> {
    priv v: &'self mut [T],
    priv len: uint,
    priv size: uint,
    priv pos: uint
}

impl<'self, T> Iterator<&'self mut [T]> for MutChunkIter<'self, T> {
    #[inline]
    fn next(&mut self) -> Option<&'self mut [T]> {
        use std::cast::transmute;
        if self.pos >= self.len {
            None
        } else {
            let chunksz = cmp::min(self.len - self.pos, self.size);
            let out: &'self mut [T];
            unsafe {
                out = transmute(self.v.mut_slice(self.pos, self.pos + chunksz));
            }
            self.pos += chunksz;
            Some(out)
        }
    }

    #[inline]
    fn size_hint(&self) -> (uint, Option<uint>) {
        if self.len == 0 {
            (0, Some(0))
        } else {
            let (n, rem) = self.len.div_rem(&self.size);
            let n = if rem > 0 { n + 1 } else { n };
            (n, Some(n))
        }
    }
}

pub trait MutChunkIterable<'self, T> {
    /**
     * Returns an iterator over `size` elements of the vector at a
     * time. The chunks do not overlap. If `size` does not divide the
     * length of the vector, then the last chunk will not have length
     * `size`. The chunk are mutable.
     */
    fn mut_chunk_iter(self, size: uint) -> MutChunkIter<'self, T>;
}

impl<'self, T> MutChunkIterable<'self, T> for &'self mut [T] {
    #[inline]
    fn mut_chunk_iter(self, size: uint) -> MutChunkIter<'self, T> {
        assert!(size != 0);
        let len = self.len();
        MutChunkIter { v: self, len: len, size: size, pos: 0 }
    }
}

#[test]
fn test_mut_chunk_iterator() {
    let mut v = [0u8, 1, 2, 3, 4, 5];

    for (i, chunk) in v.mut_chunk_iter(3).enumerate() {
        chunk[0] = i as u8;
        chunk[1] = i as u8;
        chunk[2] = i as u8;
    }

    let result = [0u8, 0, 0, 1, 1, 1];
    assert_eq!(v, result);
}

#[test]
#[should_fail]
fn test_mut_chunk_iterator_0() {
    let mut v = [1, 2, 3, 4];
    let _it = v.mut_chunk_iter(0);
}
<|endoftext|>