<commit_before>use crate::messages::{Notification, Request};
use crate::worker::common::{EventHandlers, SubscriptionTarget, WeakEventHandlers};
use crate::worker::utils::{PreparedPayloadChannelRead, PreparedPayloadChannelWrite};
use crate::worker::{utils, RequestError, SubscriptionHandler};
use futures_lite::future;
use log::{debug, error, trace, warn};
use mediasoup_sys::UvAsyncT;
use parking_lot::Mutex;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::{HashMap, VecDeque};
use std::fmt::Debug;
use std::sync::{Arc, Weak};
use std::time::Duration;
use thiserror::Error;

const PAYLOAD_MAX_LEN: usize = 4_194_304;

#[derive(Debug)]
pub(super) enum InternalMessage {
    /// Unknown data
    UnexpectedData(Vec<u8>),
}

#[derive(Debug, Deserialize)]
#[serde(untagged)]
enum PayloadChannelReceiveMessage {
    ResponseSuccess {
        id: u32,
        accepted: bool,
        data: Option<Value>,
    },
    ResponseError {
        id: u32,
        error: Value,
        reason: String,
    },
    #[serde(rename_all = "camelCase")]
    Notification { target_id: SubscriptionTarget },
    /// Unknown data
    #[serde(skip)]
    Internal(InternalMessage),
}

fn deserialize_message(bytes: &[u8]) -> PayloadChannelReceiveMessage {
    match serde_json::from_slice(bytes) {
        Ok(message) => message,
        Err(error) => {
            warn!("Failed to deserialize message: {}", error);

            PayloadChannelReceiveMessage::Internal(InternalMessage::UnexpectedData(Vec::from(
                bytes,
            )))
        }
    }
}

#[derive(Debug, Error, Eq, PartialEq)]
pub enum NotificationError {
    #[error("Channel already closed")]
    ChannelClosed,
    #[error("Payload is too long")]
    PayloadTooLong,
}

struct ResponseError {
    reason: String,
}

type Response<T> = Result<Option<T>, ResponseError>;

#[derive(Default)]
struct RequestsContainer {
    next_id: u32,
    handlers: HashMap<u32, async_oneshot::Sender<Response<Value>>>,
}

struct OutgoingMessageBuffer {
    handle: Option<UvAsyncT>,
    messages: VecDeque<(Vec<u8>, Vec<u8>)>,
}

struct Inner {
    outgoing_message_buffer: Arc<Mutex<OutgoingMessageBuffer>>,
    internal_message_receiver: async_channel::Receiver<InternalMessage>,
    requests_container_weak: Weak<Mutex<RequestsContainer>>,
    #[allow(clippy::type_complexity)]
    event_handlers_weak: WeakEventHandlers<Arc<dyn Fn(&[u8], &[u8]) + Send + Sync + 'static>>,
}

impl Drop for Inner {
    fn drop(&mut self) {
        self.internal_message_receiver.close();
    }
}

#[derive(Clone)]
pub(crate) struct PayloadChannel {
    inner: Arc<Inner>,
}

impl PayloadChannel {
    pub(super) fn new() -> (
        Self,
        PreparedPayloadChannelRead,
        PreparedPayloadChannelWrite,
    ) {
        let outgoing_message_buffer = Arc::new(Mutex::new(OutgoingMessageBuffer {
            handle: None,
            messages: VecDeque::<(Vec<u8>, Vec<u8>)>::with_capacity(10),
        }));
        let requests_container = Arc::<Mutex<RequestsContainer>>::default();
        let requests_container_weak = Arc::downgrade(&requests_container);
        let event_handlers = EventHandlers::new();
        let event_handlers_weak = event_handlers.downgrade();

        let prepared_payload_channel_read = utils::prepare_payload_channel_read_fn({
            let outgoing_message_buffer = Arc::clone(&outgoing_message_buffer);

            move |handle| {
                let mut outgoing_message_buffer = outgoing_message_buffer.lock();
                if outgoing_message_buffer.handle.is_none() {
                    outgoing_message_buffer.handle.insert(handle);
                }

                outgoing_message_buffer.messages.pop_front()
            }
        });

        let (internal_message_sender, internal_message_receiver) = async_channel::bounded(1);

        let prepared_payload_channel_write =
            utils::prepare_payload_channel_write_fn(move |message, payload| {
                trace!("received raw message: {}", String::from_utf8_lossy(message));

                match deserialize_message(message) {
                    PayloadChannelReceiveMessage::ResponseSuccess {
                        id,
                        accepted: _,
                        data,
                    } => {
                        let sender = requests_container.lock().handlers.remove(&id);
                        if let Some(mut sender) = sender {
                            let _ = sender.send(Ok(data));
                        } else {
                            warn!(
                                "received success response does not match any sent request [id:{}]",
                                id,
                            );
                        }
                    }
                    PayloadChannelReceiveMessage::ResponseError {
                        id,
                        error: _,
                        reason,
                    } => {
                        let sender = requests_container.lock().handlers.remove(&id);
                        if let Some(mut sender) = sender {
                            let _ = sender.send(Err(ResponseError { reason }));
                        } else {
                            warn!(
                                "received error response does not match any sent request [id:{}]",
                                id,
                            );
                        }
                    }
                    PayloadChannelReceiveMessage::Notification { target_id } => {
                        trace!("received notification payload of {} bytes", payload.len());

                        event_handlers.call_callbacks_with_two_values(&target_id, message, payload);
                    }
                    PayloadChannelReceiveMessage::Internal(internal_message) => {
                        let _ = internal_message_sender.try_send(internal_message);
                    }
                }
            });

        let inner = Arc::new(Inner {
            outgoing_message_buffer,
            internal_message_receiver,
            requests_container_weak,
            event_handlers_weak,
        });

        (
            Self { inner },
            prepared_payload_channel_read,
            prepared_payload_channel_write,
        )
    }

    pub(super) fn get_internal_message_receiver(&self) -> async_channel::Receiver<InternalMessage> {
        self.inner.internal_message_receiver.clone()
    }

    pub(crate) async fn request<R>(
        &self,
        request: R,
        payload: Vec<u8>,
    ) -> Result<R::Response, RequestError>
    where
        R: Request,
    {
        let method = request.as_method();

        let id;
        let queue_len;
        let (result_sender, result_receiver) = async_oneshot::oneshot();

        {
            let requests_container_lock = self
                .inner
                .requests_container_weak
                .upgrade()
                .ok_or(RequestError::ChannelClosed)?;
            let mut requests_container = requests_container_lock.lock();

            id = requests_container.next_id;
            queue_len = requests_container.handlers.len();

            requests_container.next_id = requests_container.next_id.wrapping_add(1);
            requests_container.handlers.insert(id, result_sender);
        }

        debug!("request() [method:{}, id:{}]: {:?}", method, id, request);

        if payload.len() > PAYLOAD_MAX_LEN {
            self.inner
                .requests_container_weak
                .upgrade()
                .ok_or(RequestError::ChannelClosed)?
                .lock()
                .handlers
                .remove(&id);
            return Err(RequestError::PayloadTooLong);
        }

        {
            #[derive(Debug, Serialize)]
            struct RequestMessagePrivate<'a, R> {
                id: u32,
                method: &'static str,
                #[serde(flatten)]
                request: &'a R,
            }

            let message = serde_json::to_vec(&RequestMessagePrivate {
                id,
                method: request.as_method(),
                request: &request,
            })
            .unwrap();

            let mut outgoing_message_buffer = self.inner.outgoing_message_buffer.lock();
            outgoing_message_buffer
                .messages
                .push_back((message, payload));
            if let Some(handle) = &outgoing_message_buffer.handle {
                unsafe {
                    // Notify worker that there is something to read
                    let ret = mediasoup_sys::uv_async_send(*handle);
                    if ret != 0 {
                        error!("uv_async_send call failed with code {}", ret);
                        return Err(RequestError::ChannelClosed);
                    }
                }
            }
        }

        let result = future::or(
            async move {
                result_receiver
                    .await
                    .map_err(|_| RequestError::ChannelClosed {})
            },
            async move {
                async_io::Timer::after(Duration::from_millis(
                    (1000.0 * (15.0 + (0.1 * queue_len as f64))).round() as u64,
                ))
                .await;

                if let Some(requests_container) = self.inner.requests_container_weak.upgrade() {
                    requests_container.lock().handlers.remove(&id);
                }

                Err(RequestError::TimedOut)
            },
        )
        .await?;

        let data = match result {
            Ok(data) => {
                debug!("request succeeded [method:{}, id:{}]", method, id);
                data
            }
            Err(ResponseError { reason }) => {
                debug!("request failed [method:{}, id:{}]: {}", method, id, reason);

                return Err(RequestError::Response { reason });
            }
        };

        // Default will work for `()` response
        serde_json::from_value(data.unwrap_or_default()).map_err(|error| {
            RequestError::FailedToParse {
                error: error.to_string(),
            }
        })
    }

    pub(crate) async fn notify<N>(
        &self,
        notification: N,
        payload: Vec<u8>,
    ) -> Result<(), NotificationError>
    where
        N: Notification,
    {
        debug!("notify() [event:{}]", notification.as_event());

        if payload.len() > PAYLOAD_MAX_LEN {
            return Err(NotificationError::PayloadTooLong);
        }

        {
            #[derive(Debug, Serialize)]
            struct NotificationMessagePrivate<'a, N: Serialize> {
                event: &'static str,
                #[serde(flatten)]
                notification: &'a N,
            }

            let message = serde_json::to_vec(&NotificationMessagePrivate {
                event: notification.as_event(),
                notification: &notification,
            })
            .unwrap();

            let mut outgoing_message_buffer = self.inner.outgoing_message_buffer.lock();
            outgoing_message_buffer
                .messages
                .push_back((message, payload));
            if let Some(handle) = &outgoing_message_buffer.handle {
                unsafe {
                    // Notify worker that there is something to read
                    let ret = mediasoup_sys::uv_async_send(*handle);
                    if ret != 0 {
                        error!("uv_async_send call failed with code {}", ret);
                        return Err(NotificationError::ChannelClosed);
                    }
                }
            }
        }

        Ok(())
    }

    pub(crate) fn subscribe_to_notifications<F>(
        &self,
        target_id: SubscriptionTarget,
        callback: F,
    ) -> Option<SubscriptionHandler>
    where
        F: Fn(&[u8], &[u8]) + Send + Sync + 'static,
    {
        self.inner
            .event_handlers_weak
            .upgrade()
            .map(|event_handlers| event_handlers.add(target_id, Arc::new(callback)))
    }
}
<commit_msg>Optimize worker communication performance by making notification in payload channel the first item in the enum, accelerating deserialization for performance-sensitive case<commit_after>use crate::messages::{Notification, Request};
use crate::worker::common::{EventHandlers, SubscriptionTarget, WeakEventHandlers};
use crate::worker::utils::{PreparedPayloadChannelRead, PreparedPayloadChannelWrite};
use crate::worker::{utils, RequestError, SubscriptionHandler};
use futures_lite::future;
use log::{debug, error, trace, warn};
use mediasoup_sys::UvAsyncT;
use parking_lot::Mutex;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::{HashMap, VecDeque};
use std::fmt::Debug;
use std::sync::{Arc, Weak};
use std::time::Duration;
use thiserror::Error;

const PAYLOAD_MAX_LEN: usize = 4_194_304;

#[derive(Debug)]
pub(super) enum InternalMessage {
    /// Unknown data
    UnexpectedData(Vec<u8>),
}

#[derive(Debug, Deserialize)]
#[serde(untagged)]
enum PayloadChannelReceiveMessage {
    #[serde(rename_all = "camelCase")]
    Notification { target_id: SubscriptionTarget },
    ResponseSuccess {
        id: u32,
        accepted: bool,
        data: Option<Value>,
    },
    ResponseError {
        id: u32,
        error: Value,
        reason: String,
    },
    /// Unknown data
    #[serde(skip)]
    Internal(InternalMessage),
}

fn deserialize_message(bytes: &[u8]) -> PayloadChannelReceiveMessage {
    match serde_json::from_slice(bytes) {
        Ok(message) => message,
        Err(error) => {
            warn!("Failed to deserialize message: {}", error);

            PayloadChannelReceiveMessage::Internal(InternalMessage::UnexpectedData(Vec::from(
                bytes,
            )))
        }
    }
}

#[derive(Debug, Error, Eq, PartialEq)]
pub enum NotificationError {
    #[error("Channel already closed")]
    ChannelClosed,
    #[error("Payload is too long")]
    PayloadTooLong,
}

struct ResponseError {
    reason: String,
}

type Response<T> = Result<Option<T>, ResponseError>;

#[derive(Default)]
struct RequestsContainer {
    next_id: u32,
    handlers: HashMap<u32, async_oneshot::Sender<Response<Value>>>,
}

struct OutgoingMessageBuffer {
    handle: Option<UvAsyncT>,
    messages: VecDeque<(Vec<u8>, Vec<u8>)>,
}

struct Inner {
    outgoing_message_buffer: Arc<Mutex<OutgoingMessageBuffer>>,
    internal_message_receiver: async_channel::Receiver<InternalMessage>,
    requests_container_weak: Weak<Mutex<RequestsContainer>>,
    #[allow(clippy::type_complexity)]
    event_handlers_weak: WeakEventHandlers<Arc<dyn Fn(&[u8], &[u8]) + Send + Sync + 'static>>,
}

impl Drop for Inner {
    fn drop(&mut self) {
        self.internal_message_receiver.close();
    }
}

#[derive(Clone)]
pub(crate) struct PayloadChannel {
    inner: Arc<Inner>,
}

impl PayloadChannel {
    pub(super) fn new() -> (
        Self,
        PreparedPayloadChannelRead,
        PreparedPayloadChannelWrite,
    ) {
        let outgoing_message_buffer = Arc::new(Mutex::new(OutgoingMessageBuffer {
            handle: None,
            messages: VecDeque::<(Vec<u8>, Vec<u8>)>::with_capacity(10),
        }));
        let requests_container = Arc::<Mutex<RequestsContainer>>::default();
        let requests_container_weak = Arc::downgrade(&requests_container);
        let event_handlers = EventHandlers::new();
        let event_handlers_weak = event_handlers.downgrade();

        let prepared_payload_channel_read = utils::prepare_payload_channel_read_fn({
            let outgoing_message_buffer = Arc::clone(&outgoing_message_buffer);

            move |handle| {
                let mut outgoing_message_buffer = outgoing_message_buffer.lock();
                if outgoing_message_buffer.handle.is_none() {
                    outgoing_message_buffer.handle.insert(handle);
                }

                outgoing_message_buffer.messages.pop_front()
            }
        });

        let (internal_message_sender, internal_message_receiver) = async_channel::bounded(1);

        let prepared_payload_channel_write =
            utils::prepare_payload_channel_write_fn(move |message, payload| {
                trace!("received raw message: {}", String::from_utf8_lossy(message));

                match deserialize_message(message) {
                    PayloadChannelReceiveMessage::Notification { target_id } => {
                        trace!("received notification payload of {} bytes", payload.len());

                        event_handlers.call_callbacks_with_two_values(&target_id, message, payload);
                    }
                    PayloadChannelReceiveMessage::ResponseSuccess {
                        id,
                        accepted: _,
                        data,
                    } => {
                        let sender = requests_container.lock().handlers.remove(&id);
                        if let Some(mut sender) = sender {
                            let _ = sender.send(Ok(data));
                        } else {
                            warn!(
                                "received success response does not match any sent request [id:{}]",
                                id,
                            );
                        }
                    }
                    PayloadChannelReceiveMessage::ResponseError {
                        id,
                        error: _,
                        reason,
                    } => {
                        let sender = requests_container.lock().handlers.remove(&id);
                        if let Some(mut sender) = sender {
                            let _ = sender.send(Err(ResponseError { reason }));
                        } else {
                            warn!(
                                "received error response does not match any sent request [id:{}]",
                                id,
                            );
                        }
                    }
                    PayloadChannelReceiveMessage::Internal(internal_message) => {
                        let _ = internal_message_sender.try_send(internal_message);
                    }
                }
            });

        let inner = Arc::new(Inner {
            outgoing_message_buffer,
            internal_message_receiver,
            requests_container_weak,
            event_handlers_weak,
        });

        (
            Self { inner },
            prepared_payload_channel_read,
            prepared_payload_channel_write,
        )
    }

    pub(super) fn get_internal_message_receiver(&self) -> async_channel::Receiver<InternalMessage> {
        self.inner.internal_message_receiver.clone()
    }

    pub(crate) async fn request<R>(
        &self,
        request: R,
        payload: Vec<u8>,
    ) -> Result<R::Response, RequestError>
    where
        R: Request,
    {
        let method = request.as_method();

        let id;
        let queue_len;
        let (result_sender, result_receiver) = async_oneshot::oneshot();

        {
            let requests_container_lock = self
                .inner
                .requests_container_weak
                .upgrade()
                .ok_or(RequestError::ChannelClosed)?;
            let mut requests_container = requests_container_lock.lock();

            id = requests_container.next_id;
            queue_len = requests_container.handlers.len();

            requests_container.next_id = requests_container.next_id.wrapping_add(1);
            requests_container.handlers.insert(id, result_sender);
        }

        debug!("request() [method:{}, id:{}]: {:?}", method, id, request);

        if payload.len() > PAYLOAD_MAX_LEN {
            self.inner
                .requests_container_weak
                .upgrade()
                .ok_or(RequestError::ChannelClosed)?
                .lock()
                .handlers
                .remove(&id);
            return Err(RequestError::PayloadTooLong);
        }

        {
            #[derive(Debug, Serialize)]
            struct RequestMessagePrivate<'a, R> {
                id: u32,
                method: &'static str,
                #[serde(flatten)]
                request: &'a R,
            }

            let message = serde_json::to_vec(&RequestMessagePrivate {
                id,
                method: request.as_method(),
                request: &request,
            })
            .unwrap();

            let mut outgoing_message_buffer = self.inner.outgoing_message_buffer.lock();
            outgoing_message_buffer
                .messages
                .push_back((message, payload));
            if let Some(handle) = &outgoing_message_buffer.handle {
                unsafe {
                    // Notify worker that there is something to read
                    let ret = mediasoup_sys::uv_async_send(*handle);
                    if ret != 0 {
                        error!("uv_async_send call failed with code {}", ret);
                        return Err(RequestError::ChannelClosed);
                    }
                }
            }
        }

        let result = future::or(
            async move {
                result_receiver
                    .await
                    .map_err(|_| RequestError::ChannelClosed {})
            },
            async move {
                async_io::Timer::after(Duration::from_millis(
                    (1000.0 * (15.0 + (0.1 * queue_len as f64))).round() as u64,
                ))
                .await;

                if let Some(requests_container) = self.inner.requests_container_weak.upgrade() {
                    requests_container.lock().handlers.remove(&id);
                }

                Err(RequestError::TimedOut)
            },
        )
        .await?;

        let data = match result {
            Ok(data) => {
                debug!("request succeeded [method:{}, id:{}]", method, id);
                data
            }
            Err(ResponseError { reason }) => {
                debug!("request failed [method:{}, id:{}]: {}", method, id, reason);

                return Err(RequestError::Response { reason });
            }
        };

        // Default will work for `()` response
        serde_json::from_value(data.unwrap_or_default()).map_err(|error| {
            RequestError::FailedToParse {
                error: error.to_string(),
            }
        })
    }

    pub(crate) async fn notify<N>(
        &self,
        notification: N,
        payload: Vec<u8>,
    ) -> Result<(), NotificationError>
    where
        N: Notification,
    {
        debug!("notify() [event:{}]", notification.as_event());

        if payload.len() > PAYLOAD_MAX_LEN {
            return Err(NotificationError::PayloadTooLong);
        }

        {
            #[derive(Debug, Serialize)]
            struct NotificationMessagePrivate<'a, N: Serialize> {
                event: &'static str,
                #[serde(flatten)]
                notification: &'a N,
            }

            let message = serde_json::to_vec(&NotificationMessagePrivate {
                event: notification.as_event(),
                notification: &notification,
            })
            .unwrap();

            let mut outgoing_message_buffer = self.inner.outgoing_message_buffer.lock();
            outgoing_message_buffer
                .messages
                .push_back((message, payload));
            if let Some(handle) = &outgoing_message_buffer.handle {
                unsafe {
                    // Notify worker that there is something to read
                    let ret = mediasoup_sys::uv_async_send(*handle);
                    if ret != 0 {
                        error!("uv_async_send call failed with code {}", ret);
                        return Err(NotificationError::ChannelClosed);
                    }
                }
            }
        }

        Ok(())
    }

    pub(crate) fn subscribe_to_notifications<F>(
        &self,
        target_id: SubscriptionTarget,
        callback: F,
    ) -> Option<SubscriptionHandler>
    where
        F: Fn(&[u8], &[u8]) + Send + Sync + 'static,
    {
        self.inner
            .event_handlers_weak
            .upgrade()
            .map(|event_handlers| event_handlers.add(target_id, Arc::new(callback)))
    }
}
<|endoftext|>


<commit_before>//
// Copyright (c) 2016, Boris Popov <popov@whitekefir.ru>
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//

use logger;
use env;
use std::process::exit;
use blkid;
use mount;

//////////////////////////////////////////////////////////////////
pub struct Mounter {
    part_label: String,
    part_type:  String,
    target:     String
}

impl Mounter {
    pub fn new() -> Mounter {
        Mounter{
            part_label: "".to_string(),
            part_type:  "".to_string(),
            target:     "".to_string()
        }
    }

    pub fn mount(&mut self, l: &logger::Logger, e: &env::Env) {
        self.load_part_info(e.getDevice(), l);
        let result = self.verify_label(e.getLabel(), l);
        if ! result {
            let i = format!("Label does not match {}.", e.getLabel());
            l.info(&i);
            exit(0);
        }
        self.process(l, e);
    }

    fn process(&mut self, l: &logger::Logger, e: &env::Env) {
        let source = e.getDevice().to_string();
        self.target = e.getDir().to_string();
        let res = mount::mount(source.as_str(),
                               self.target.as_str(),
                               self.part_type.as_str());
        match res {
            Err(e) => {
                l.info(e.as_str());
                exit(0);
            },
            _ => {
                let i = format!("Mount {} to {} OK!",
                                source.as_str(),
                                self.target.as_str());
                l.info(i.as_str());
                return;
            }
        }
    }

    fn verify_label(&self, templ: &str, l: &logger::Logger) -> bool {
        if self.part_label.len() == 0 {
            return false;
        }
        return self.part_label.contains(templ);
    }

    fn load_part_info(&mut self, dev: &str, l: &logger::Logger) {
        let prober = blkid::Prober::new(dev);
        let res = prober.probe();
        if let Err(e) = res {
            l.error(&e);
            exit(1);
        }
        self.loadLabel(&prober, l);
        self.loadType(&prober, l);
    }

    fn loadLabel(&mut self, prober: &blkid::Prober,
                 l: &logger::Logger) {
        let lr = prober.getLabel();
        match lr {
            Ok(o) => {
                self.part_label = o;
            },
            Err(e) => {
                l.error(&e);
            }
        }
    }

    fn loadType(&mut self, prober: &blkid::Prober,
                l: &logger::Logger) {
        let lr = prober.getType();
        match lr {
            Ok(o) => {
                self.part_type = o;
            },
            Err(e) => {
                l.error(&e);
            }
        }
    }
}

impl Drop for Mounter {
    fn drop(&mut self) {
        if self.target.len() != 0 {
            //
            //TODO
            //
            //mount::umount(&self.target);
        }
    }
}
//////////////////////////////////////////////////////////////////
<commit_msg>some fixes for Mounter<commit_after>//
// Copyright (c) 2016, Boris Popov <popov@whitekefir.ru>
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//

use logger;
use env;
use std::process::exit;
use blkid;
use mount;

//////////////////////////////////////////////////////////////////
pub struct Mounter {
    part_label: String,
    part_type:  String,
    target:     String,
    umount:     bool
}

impl Mounter {
    pub fn new() -> Mounter {
        Mounter{
            part_label: "".to_string(),
            part_type:  "".to_string(),
            target:     "".to_string(),
            umount:     false
        }
    }

    pub fn mount(&mut self, l: &logger::Logger, e: &env::Env) {
        self.load_part_info(e.getDevice(), l);
        let result = self.verify_label(e.getLabel());
        if ! result {
            let i = format!("Label does not match {}.", e.getLabel());
            l.info(&i);
            exit(0);
        }
        self.process(l, e);
    }

    fn process(&mut self, l: &logger::Logger, e: &env::Env) {
        let source = e.getDevice().to_string();
        self.target = e.getDir().to_string();
        self.umount = e.getUmount();
        l.info("Trying mount part...");
        let res = mount::mount(source.as_str(),
                               self.target.as_str(),
                               self.part_type.as_str());
        match res {
            Err(e) => {
                l.info(e.as_str());
                exit(0);
            },
            _ => {
                let i = format!("... {} to {} OK!",
                                source.as_str(),
                                self.target.as_str());
                l.info(i.as_str());
                return;
            }
        }
    }

    fn verify_label(&self, templ: &str) -> bool {
        if self.part_label.len() == 0 {
            return false;
        }
        return self.part_label.contains(templ);
    }

    fn load_part_info(&mut self, dev: &str, l: &logger::Logger) {
        let prober = blkid::Prober::new(dev);
        let res = prober.probe();
        if let Err(e) = res {
            l.error(&e);
            exit(1);
        }
        self.loadLabel(&prober, l);
        self.loadType(&prober, l);
    }

    fn loadLabel(&mut self, prober: &blkid::Prober,
                 l: &logger::Logger) {
        let lr = prober.getLabel();
        match lr {
            Ok(o) => {
                self.part_label = o;
                let s = format!("Label of part: {}",
                                self.part_label);
                l.info(s.as_str());
            },
            Err(e) => {
                l.error(&e);
            }
        }
    }

    fn loadType(&mut self, prober: &blkid::Prober,
                l: &logger::Logger) {
        let lr = prober.getType();
        match lr {
            Ok(o) => {
                self.part_type = o;
                let s = format!("Type of part: {}",
                                self.part_type);
                l.info(s.as_str());
            },
            Err(e) => {
                l.error(&e);
            }
        }
    }
}

impl Drop for Mounter {
    fn drop(&mut self) {
        if self.target.len() != 0 {
            if self.umount {
                mount::umount(&self.target);
            }
        }
    }
}
//////////////////////////////////////////////////////////////////
<|endoftext|>


<commit_before>#[derive(Debug, serde_derive::Deserialize)]
pub struct Config {
    pub name: String,
    pub package_key: Option<String>,
    pub repo_key: Option<String>,
    pub srcdest: String,
    pub logdest: String,
    pub pkgbuild: String,
    pub builds: std::collections::HashMap<super::builder::Arch, BuildConfig>,
    pub s3: Option<S3Config>,
}

#[derive(Debug, serde_derive::Deserialize)]
pub struct BuildConfig {
    pub chroot: String,
}

#[derive(Debug, serde_derive::Deserialize)]
pub struct S3Config {
    pub bucket: String,
    pub region: Region,
}

#[derive(Debug)]
pub struct Region(rusoto_core::Region);

impl<'de> serde::Deserialize<'de> for Region {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        struct Visitor;
        impl<'de> serde::de::Visitor<'de> for Visitor {
            type Value = Region;

            fn expecting(
                &self,
                formatter: &mut std::fmt::Formatter,
            ) -> Result<(), std::fmt::Error> {
                write!(formatter, "a valid AWS region name")
            }

            fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
            where
                E: serde::de::Error,
            {
                use std::str::FromStr;

                match rusoto_core::Region::from_str(v) {
                    Ok(r) => Ok(Region(r)),
                    Err(e) => Err(serde::de::Error::invalid_value(
                        serde::de::Unexpected::Str(v),
                        &format!("{}", e).as_str(),
                    )),
                }
            }
        }
        deserializer.deserialize_str(Visitor {})
    }
}

impl Config {
    pub fn from_reader<R>(reader: R) -> serde_yaml::Result<Self>
    where
        R: std::io::Read,
    {
        serde_yaml::from_reader(reader)
    }

    pub fn repo_dir(&self, arch: &super::builder::Arch) -> std::path::PathBuf {
        std::path::PathBuf::from(&self.name)
            .join("os")
            .join(format!("{}", arch))
    }

    pub fn db_path(&self, arch: &super::builder::Arch) -> std::path::PathBuf {
        let mut path = self.repo_dir(arch).join(&self.name).into_os_string();
        path.push(".db");
        std::path::PathBuf::from(path)
    }

    pub fn files_path(&self, arch: &super::builder::Arch) -> std::path::PathBuf {
        let mut path = self.repo_dir(arch).join(&self.name).into_os_string();
        path.push(".files");
        std::path::PathBuf::from(path)
    }

    pub fn package_dir(&self, package_name: &str) -> std::path::PathBuf {
        std::path::PathBuf::from(&self.pkgbuild).join(package_name)
    }
}

pub struct S3 {
    client: rusoto_s3::S3Client,
    bucket: String,
}

impl S3 {
    pub fn new(config: &S3Config) -> Self {
        let Region(ref region) = config.region;
        let client = rusoto_s3::S3Client::new(region.clone());
        S3 {
            client,
            bucket: config.bucket.to_owned(),
        }
    }

    pub async fn download_repository(
        &self,
        config: &Config,
        arch: &super::builder::Arch,
    ) -> Result<(), anyhow::Error> {
        self.get(config.db_path(arch)).await?;
        self.get(config.files_path(arch)).await
    }

    pub async fn upload_repository<P>(
        &self,
        config: &Config,
        arch: &super::builder::Arch,
        package_paths: &[P],
    ) -> Result<(), anyhow::Error>
    where
        P: AsRef<std::path::Path>,
    {
        const XZ_MIME_TYPE: &str = "application/x-xz";
        const SIG_MIME_TYPE: &str = "application/pgp-signature";
        const GZIP_MIME_TYPE: &str = "application/gzip";

        for package_path in package_paths {
            self.put(package_path, XZ_MIME_TYPE).await?;
            if config.package_key.is_some() {
                let mut sig_path = package_path.as_ref().as_os_str().to_os_string();
                sig_path.push(".sig");
                self.put(&sig_path, SIG_MIME_TYPE).await?;
            }
        }
        self.put(config.files_path(arch), GZIP_MIME_TYPE).await?;
        let db_path = config.db_path(arch);
        self.put(&db_path, GZIP_MIME_TYPE).await?;
        if config.repo_key.is_some() {
            let mut sig_path = db_path.clone().into_os_string();
            sig_path.push(".sig");
            self.put(sig_path, SIG_MIME_TYPE).await?;
        }
        Ok(())
    }

    async fn get<P>(&self, path: P) -> Result<(), anyhow::Error>
    where
        P: AsRef<std::path::Path>,
    {
        use rusoto_s3::S3;

        let path = path.as_ref();
        let request = rusoto_s3::GetObjectRequest {
            bucket: self.bucket.to_owned(),
            key: path.to_string_lossy().into_owned(),
            ..rusoto_s3::GetObjectRequest::default()
        };
        println!("Download {}", path.display());
        match self.client.get_object(request).await {
            Ok(output) => {
                if let Some(mut body) = output.body {
                    use futures::StreamExt as _;
                    use tokio::io::AsyncWriteExt as _;

                    let file = tokio::fs::File::create(path).await?;
                    let mut writer = tokio::io::BufWriter::new(file);
                    while let Some(item) = body.next().await {
                        writer.write_all(&item?).await?;
                    }
                    writer.shutdown().await?;
                }
                Ok(())
            }
            Err(rusoto_core::RusotoError::Service(rusoto_s3::GetObjectError::NoSuchKey(_))) => {
                Ok(())
            }
            Err(e) => Err(anyhow::Error::from(e)),
        }
    }

    async fn put<P>(&self, path: P, content_type: &str) -> Result<(), anyhow::Error>
    where
        P: AsRef<std::path::Path>,
    {
        use futures::FutureExt as _;
        use futures::TryStreamExt as _;
        use rusoto_s3::S3;

        let path = path.as_ref();
        let metadata = tokio::fs::metadata(path).await?;
        let stream = rusoto_s3::StreamingBody::new(
            tokio::fs::read(path.to_owned())
                .into_stream()
                .map_ok(bytes::Bytes::from),
        );
        let request = rusoto_s3::PutObjectRequest {
            bucket: self.bucket.to_owned(),
            key: path.to_string_lossy().into_owned(),
            content_type: Some(content_type.to_owned()),
            content_length: Some(metadata.len() as i64),
            body: Some(stream),
            ..Default::default()
        };
        println!("Upload {}", path.display());
        self.client.put_object(request).await?;
        Ok(())
    }
}
<commit_msg>Fix content-type of package files<commit_after>#[derive(Debug, serde_derive::Deserialize)]
pub struct Config {
    pub name: String,
    pub package_key: Option<String>,
    pub repo_key: Option<String>,
    pub srcdest: String,
    pub logdest: String,
    pub pkgbuild: String,
    pub builds: std::collections::HashMap<super::builder::Arch, BuildConfig>,
    pub s3: Option<S3Config>,
}

#[derive(Debug, serde_derive::Deserialize)]
pub struct BuildConfig {
    pub chroot: String,
}

#[derive(Debug, serde_derive::Deserialize)]
pub struct S3Config {
    pub bucket: String,
    pub region: Region,
}

#[derive(Debug)]
pub struct Region(rusoto_core::Region);

impl<'de> serde::Deserialize<'de> for Region {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        struct Visitor;
        impl<'de> serde::de::Visitor<'de> for Visitor {
            type Value = Region;

            fn expecting(
                &self,
                formatter: &mut std::fmt::Formatter,
            ) -> Result<(), std::fmt::Error> {
                write!(formatter, "a valid AWS region name")
            }

            fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
            where
                E: serde::de::Error,
            {
                use std::str::FromStr;

                match rusoto_core::Region::from_str(v) {
                    Ok(r) => Ok(Region(r)),
                    Err(e) => Err(serde::de::Error::invalid_value(
                        serde::de::Unexpected::Str(v),
                        &format!("{}", e).as_str(),
                    )),
                }
            }
        }
        deserializer.deserialize_str(Visitor {})
    }
}

impl Config {
    pub fn from_reader<R>(reader: R) -> serde_yaml::Result<Self>
    where
        R: std::io::Read,
    {
        serde_yaml::from_reader(reader)
    }

    pub fn repo_dir(&self, arch: &super::builder::Arch) -> std::path::PathBuf {
        std::path::PathBuf::from(&self.name)
            .join("os")
            .join(format!("{}", arch))
    }

    pub fn db_path(&self, arch: &super::builder::Arch) -> std::path::PathBuf {
        let mut path = self.repo_dir(arch).join(&self.name).into_os_string();
        path.push(".db");
        std::path::PathBuf::from(path)
    }

    pub fn files_path(&self, arch: &super::builder::Arch) -> std::path::PathBuf {
        let mut path = self.repo_dir(arch).join(&self.name).into_os_string();
        path.push(".files");
        std::path::PathBuf::from(path)
    }

    pub fn package_dir(&self, package_name: &str) -> std::path::PathBuf {
        std::path::PathBuf::from(&self.pkgbuild).join(package_name)
    }
}

pub struct S3 {
    client: rusoto_s3::S3Client,
    bucket: String,
}

impl S3 {
    pub fn new(config: &S3Config) -> Self {
        let Region(ref region) = config.region;
        let client = rusoto_s3::S3Client::new(region.clone());
        S3 {
            client,
            bucket: config.bucket.to_owned(),
        }
    }

    pub async fn download_repository(
        &self,
        config: &Config,
        arch: &super::builder::Arch,
    ) -> Result<(), anyhow::Error> {
        self.get(config.db_path(arch)).await?;
        self.get(config.files_path(arch)).await
    }

    pub async fn upload_repository<P>(
        &self,
        config: &Config,
        arch: &super::builder::Arch,
        package_paths: &[P],
    ) -> Result<(), anyhow::Error>
    where
        P: AsRef<std::path::Path>,
    {
        const XZ_MIME_TYPE: &str = "application/x-xz";
        const ZST_MIME_TYPE: &str = "application/zstd";
        const OCTET_STREAM_MIME_TYPE: &str = "application/octet-stream";
        const SIG_MIME_TYPE: &str = "application/pgp-signature";
        const GZIP_MIME_TYPE: &str = "application/gzip";

        for package_path in package_paths {
            let mime_type = if let Some(ext) = package_path.as_ref().extension() {
                if ext == "zst" {
                    ZST_MIME_TYPE
                } else if ext == "xz" {
                    XZ_MIME_TYPE
                } else {
                    OCTET_STREAM_MIME_TYPE
                }
            } else {
                OCTET_STREAM_MIME_TYPE
            };
            self.put(package_path, mime_type).await?;
            if config.package_key.is_some() {
                let mut sig_path = package_path.as_ref().as_os_str().to_os_string();
                sig_path.push(".sig");
                self.put(&sig_path, SIG_MIME_TYPE).await?;
            }
        }
        self.put(config.files_path(arch), GZIP_MIME_TYPE).await?;
        let db_path = config.db_path(arch);
        self.put(&db_path, GZIP_MIME_TYPE).await?;
        if config.repo_key.is_some() {
            let mut sig_path = db_path.clone().into_os_string();
            sig_path.push(".sig");
            self.put(sig_path, SIG_MIME_TYPE).await?;
        }
        Ok(())
    }

    async fn get<P>(&self, path: P) -> Result<(), anyhow::Error>
    where
        P: AsRef<std::path::Path>,
    {
        use rusoto_s3::S3;

        let path = path.as_ref();
        let request = rusoto_s3::GetObjectRequest {
            bucket: self.bucket.to_owned(),
            key: path.to_string_lossy().into_owned(),
            ..rusoto_s3::GetObjectRequest::default()
        };
        println!("Download {}", path.display());
        match self.client.get_object(request).await {
            Ok(output) => {
                if let Some(mut body) = output.body {
                    use futures::StreamExt as _;
                    use tokio::io::AsyncWriteExt as _;

                    let file = tokio::fs::File::create(path).await?;
                    let mut writer = tokio::io::BufWriter::new(file);
                    while let Some(item) = body.next().await {
                        writer.write_all(&item?).await?;
                    }
                    writer.shutdown().await?;
                }
                Ok(())
            }
            Err(rusoto_core::RusotoError::Service(rusoto_s3::GetObjectError::NoSuchKey(_))) => {
                Ok(())
            }
            Err(e) => Err(anyhow::Error::from(e)),
        }
    }

    async fn put<P>(&self, path: P, content_type: &str) -> Result<(), anyhow::Error>
    where
        P: AsRef<std::path::Path>,
    {
        use futures::FutureExt as _;
        use futures::TryStreamExt as _;
        use rusoto_s3::S3;

        let path = path.as_ref();
        let metadata = tokio::fs::metadata(path).await?;
        let stream = rusoto_s3::StreamingBody::new(
            tokio::fs::read(path.to_owned())
                .into_stream()
                .map_ok(bytes::Bytes::from),
        );
        let request = rusoto_s3::PutObjectRequest {
            bucket: self.bucket.to_owned(),
            key: path.to_string_lossy().into_owned(),
            content_type: Some(content_type.to_owned()),
            content_length: Some(metadata.len() as i64),
            body: Some(stream),
            ..Default::default()
        };
        println!("Upload {}", path.display());
        self.client.put_object(request).await?;
        Ok(())
    }
}
<|endoftext|>


<commit_before>use helper::*;
use nom::{self,is_digit};

use super::{Time, DateTime, Date};

macro_rules! empty_or(
    ($i:expr, $submac:ident!( $($args:tt)* )) => (
        if $i.len() == 0 {
            nom::IResult::Done($i, None)
        } else {
            match $submac!($i, $($args)*) {
                nom::IResult::Done(i,o)     => nom::IResult::Done(i, Some(o)),
                nom::IResult::Error(_)      => nom::IResult::Done($i, None),
                nom::IResult::Incomplete(i) => nom::IResult::Incomplete(i)

            }
        }
    );
);

macro_rules! check(
  ($input:expr, $submac:ident!( $($args:tt)* )) => (

    {
      let mut failed = false;
      for idx in 0..$input.len() {
        if !$submac!($input[idx], $($args)*) {
            failed = true;
            break;
        }
      }
      if failed {
        nom::IResult::Error(nom::Err::Position(nom::ErrorCode::Filter as u32,$input))
      } else {
        nom::IResult::Done(&b""[..], $input)
      }
    }
  );
  ($input:expr, $f:expr) => (
    check!($input, call!($f));
  );
);

macro_rules! char_between(
    ($input:expr, $min:expr, $max:expr) => (
        {
        fn f(c: u8) -> bool { c >= ($min as u8) && c <= ($max as u8)}
        flat_map!($input, take!(1), check!(f))
        }
    );
);

named!(pub take_4_digits, flat_map!(take!(4), check!(is_digit)));

// year
named!(year_prefix, alt!(tag!("+") | tag!("-")));
named!(pub year <i32>, chain!(
        pref: opt!(year_prefix) ~
        year: call!(take_4_digits)
        ,
        || {
            match pref {
                Some(b"-") => -buf_to_i32(year),
                _ => buf_to_i32(year)
            }
        }));

// MM
named!(lower_month <u32>, chain!(tag!("0") ~ s:char_between!('1', '9') , || buf_to_u32(s)));
named!(upper_month <u32>, chain!(tag!("1") ~ s:char_between!('0', '2') , || 10+buf_to_u32(s)));
named!(pub month <u32>, alt!(lower_month | upper_month));

// DD
named!(day_zero  <u32>, chain!(tag!("0") ~ s:char_between!('1', '9') , || buf_to_u32(s)));
named!(day_one   <u32>, chain!(tag!("1") ~ s:char_between!('0', '9') , || 10+buf_to_u32(s)));
named!(day_two   <u32>, chain!(tag!("2") ~ s:char_between!('0', '9') , || 20+buf_to_u32(s)));
named!(day_three <u32>, chain!(tag!("3") ~ s:char_between!('0', '1') , || 30+buf_to_u32(s)));
named!(pub day <u32>, alt!(day_zero | day_one | day_two | day_three));

// WW
// reusing day_N parsers, sorry
named!(week_three <u32>, chain!(tag!("3") ~ s:char_between!('0', '9') , || 30+buf_to_u32(s)));
named!(week_four  <u32>, chain!(tag!("4") ~ s:char_between!('0', '9') , || 40+buf_to_u32(s)));
named!(week_five  <u32>, chain!(tag!("5") ~ s:char_between!('0', '3') , || 50+buf_to_u32(s)));

named!(week <u32>, alt!(day_zero | day_one | day_two | week_three| week_four | week_five ));
named!(week_day <u32>, chain!(s:char_between!('1', '7') , || buf_to_u32(s)));

// ordinal DDD
named!(ord_day <u32>, chain!(
        a:char_between!('0','3') ~
        b:char_between!('0','9') ~
        c:char_between!('0','9')
        ,
        || { buf_to_u32(a)*100 + buf_to_u32(b)*10 + buf_to_u32(c) }
        ));

// YYYY-MM-DD
named!(pub ymd_date <Date>, chain!(
        y: year ~
        opt!(tag!("-")) ~
        m: month ~
        opt!(tag!("-")) ~
        d: day
        ,
        || { Date::YMD{ year: y, month: m, day: d } }
        ));

// YYYY-MM-DD
named!(pub ordinal_date <Date>, chain!(
        y: year ~
        opt!(tag!("-")) ~
        d: ord_day
        ,
        || { Date::Ordinal{ year: y, ddd: d } }
        ));

// YYYY-"W"WW-D
named!(pub iso_week_date <Date>, chain!(
        y: year ~
        opt!(tag!("-")) ~
        tag!("W") ~
        w: week ~
        opt!(tag!("-")) ~
        d: week_day
        ,
        || { Date::Week{ year: y, ww: w, d: d } }
        ));

named!(pub parse_date <Date>, alt!( ymd_date | iso_week_date | ordinal_date ) );

//    TIME

// HH
named!(lower_hour <u32>, chain!(f:char_between!('0','1') ~ s:char_between!('0','9') ,
                                       || { buf_to_u32(f)*10 + buf_to_u32(s) } ));
named!(upper_hour <u32>, chain!(tag!("2") ~ s:char_between!('0','4') , || 20+buf_to_u32(s)));
named!(pub hour <u32>, alt!(lower_hour | upper_hour));

// MM
named!(below_sixty <u32>, chain!(f:char_between!('0','5') ~ s:char_between!('0','9'), || { buf_to_u32(f)*10 + buf_to_u32(s) } ));
named!(upto_sixty <u32>, alt!(below_sixty | map!(tag!("60"), |_| 60)));

named!(pub minute <u32>, call!(below_sixty));
named!(pub second <u32>, call!(upto_sixty));
named!(pub millisecond <u32>, chain!( ms: is_a!("0123456789"), || buf_to_u32(ms) ) );

// HH:MM:[SS][.(m*)][(Z|+...|-...)]
named!(pub parse_time <Time>, chain!(
        h: hour ~
        opt!(tag!(":")) ~
        m: minute ~
        s:  opt!( chain!( opt!(tag!(":")) ~ s:second, || s)) ~
        ms: opt!( chain!( tag!(".") ~ ms:millisecond, || ms)) ~
        z:  opt!( alt!( timezone_hour | timezone_utc) ) ,
        || {
            Time {
                hour: h,
                minute: m,
                second: s.unwrap_or(0),
                millisecond: ms.unwrap_or(0),
                tz_offset_hours: z.unwrap_or((0,0)).0,
                tz_offset_minutes: z.unwrap_or((0,0)).1
            }
        }
        ));

named!(sign <i32>, alt!(
        tag!("-") => { |_| -1 } |
        tag!("+") => { |_| 1 }
        )
    );

named!(timezone_hour <(i32,i32)>, chain!(
        s: sign ~
        h: hour ~
        m: empty_or!(
            chain!(
                tag!(":")? ~ m: minute , || { m }
            ))
        ,
        || { (s * (h as i32) , s * (m.unwrap_or(0) as i32)) }
        ));

named!(tz_z, tag!("Z")); // TODO inline below
named!(timezone_utc <(i32,i32)>, map!(tz_z, |_| (0,0)));

// Full ISO8601
named!(pub parse_datetime <DateTime>, chain!(
        d: parse_date ~
        tag!("T") ~
        t: parse_time
        ,
        || {
            DateTime{
                date: d,
                time: t,
            }
        }
        ));

<commit_msg>little inlining<commit_after>use helper::*;
use nom::{self,is_digit};
use super::{Time, DateTime, Date};

macro_rules! empty_or(
    ($i:expr, $submac:ident!( $($args:tt)* )) => (
        if $i.len() == 0 {
            nom::IResult::Done($i, None)
        } else {
            match $submac!($i, $($args)*) {
                nom::IResult::Done(i,o)     => nom::IResult::Done(i, Some(o)),
                nom::IResult::Error(_)      => nom::IResult::Done($i, None),
                nom::IResult::Incomplete(i) => nom::IResult::Incomplete(i)

            }
        }
    );
);

macro_rules! check(
  ($input:expr, $submac:ident!( $($args:tt)* )) => (

    {
      let mut failed = false;
      for idx in 0..$input.len() {
        if !$submac!($input[idx], $($args)*) {
            failed = true;
            break;
        }
      }
      if failed {
        nom::IResult::Error(nom::Err::Position(nom::ErrorCode::Filter as u32,$input))
      } else {
        nom::IResult::Done(&b""[..], $input)
      }
    }
  );
  ($input:expr, $f:expr) => (
    check!($input, call!($f));
  );
);

macro_rules! char_between(
    ($input:expr, $min:expr, $max:expr) => (
        {
        fn f(c: u8) -> bool { c >= ($min as u8) && c <= ($max as u8)}
        flat_map!($input, take!(1), check!(f))
        }
    );
);

named!(pub take_4_digits, flat_map!(take!(4), check!(is_digit)));

// year
named!(year_prefix, alt!(tag!("+") | tag!("-")));
named!(pub year <i32>, chain!(
        pref: opt!(year_prefix) ~
        year: call!(take_4_digits)
        ,
        || {
            match pref {
                Some(b"-") => -buf_to_i32(year),
                _ => buf_to_i32(year)
            }
        }));

// MM
named!(lower_month <u32>, chain!(tag!("0") ~ s:char_between!('1', '9') , || buf_to_u32(s)));
named!(upper_month <u32>, chain!(tag!("1") ~ s:char_between!('0', '2') , || 10+buf_to_u32(s)));
named!(pub month <u32>, alt!(lower_month | upper_month));

// DD
named!(day_zero  <u32>, chain!(tag!("0") ~ s:char_between!('1', '9') , || buf_to_u32(s)));
named!(day_one   <u32>, chain!(tag!("1") ~ s:char_between!('0', '9') , || 10+buf_to_u32(s)));
named!(day_two   <u32>, chain!(tag!("2") ~ s:char_between!('0', '9') , || 20+buf_to_u32(s)));
named!(day_three <u32>, chain!(tag!("3") ~ s:char_between!('0', '1') , || 30+buf_to_u32(s)));
named!(pub day <u32>, alt!(day_zero | day_one | day_two | day_three));

// WW
// reusing day_N parsers, sorry
named!(week_three <u32>, chain!(tag!("3") ~ s:char_between!('0', '9') , || 30+buf_to_u32(s)));
named!(week_four  <u32>, chain!(tag!("4") ~ s:char_between!('0', '9') , || 40+buf_to_u32(s)));
named!(week_five  <u32>, chain!(tag!("5") ~ s:char_between!('0', '3') , || 50+buf_to_u32(s)));

named!(week <u32>, alt!(day_zero | day_one | day_two | week_three| week_four | week_five ));
named!(week_day <u32>, chain!(s:char_between!('1', '7') , || buf_to_u32(s)));

// ordinal DDD
named!(ord_day <u32>, chain!(
        a:char_between!('0','3') ~
        b:char_between!('0','9') ~
        c:char_between!('0','9')
        ,
        || { buf_to_u32(a)*100 + buf_to_u32(b)*10 + buf_to_u32(c) }
        ));

// YYYY-MM-DD
named!(pub ymd_date <Date>, chain!(
        y: year ~
        opt!(tag!("-")) ~
        m: month ~
        opt!(tag!("-")) ~
        d: day
        ,
        || { Date::YMD{ year: y, month: m, day: d } }
        ));

// YYYY-MM-DD
named!(pub ordinal_date <Date>, chain!(
        y: year ~
        opt!(tag!("-")) ~
        d: ord_day
        ,
        || { Date::Ordinal{ year: y, ddd: d } }
        ));

// YYYY-"W"WW-D
named!(pub iso_week_date <Date>, chain!(
        y: year ~
        opt!(tag!("-")) ~
        tag!("W") ~
        w: week ~
        opt!(tag!("-")) ~
        d: week_day
        ,
        || { Date::Week{ year: y, ww: w, d: d } }
        ));

named!(pub parse_date <Date>, alt!( ymd_date | iso_week_date | ordinal_date ) );

//    TIME

// HH
named!(lower_hour <u32>, chain!(f:char_between!('0','1') ~ s:char_between!('0','9') ,
                                       || { buf_to_u32(f)*10 + buf_to_u32(s) } ));
named!(upper_hour <u32>, chain!(tag!("2") ~ s:char_between!('0','4') , || 20+buf_to_u32(s)));
named!(pub hour <u32>, alt!(lower_hour | upper_hour));

// MM
named!(below_sixty <u32>, chain!(f:char_between!('0','5') ~ s:char_between!('0','9'), || { buf_to_u32(f)*10 + buf_to_u32(s) } ));
named!(upto_sixty <u32>, alt!(below_sixty | map!(tag!("60"), |_| 60)));

named!(pub minute <u32>, call!(below_sixty));
named!(pub second <u32>, call!(upto_sixty));
named!(pub millisecond <u32>, chain!( ms: is_a!("0123456789"), || buf_to_u32(ms) ) );

// HH:MM:[SS][.(m*)][(Z|+...|-...)]
named!(pub parse_time <Time>, chain!(
        h: hour ~
        opt!(tag!(":")) ~
        m: minute ~
        s:  opt!( chain!( opt!(tag!(":")) ~ s:second, || s)) ~
        ms: opt!( chain!( tag!(".") ~ ms:millisecond, || ms)) ~
        z:  opt!( alt!( timezone_hour | timezone_utc) ) ,
        || {
            Time {
                hour: h,
                minute: m,
                second: s.unwrap_or(0),
                millisecond: ms.unwrap_or(0),
                tz_offset_hours: z.unwrap_or((0,0)).0,
                tz_offset_minutes: z.unwrap_or((0,0)).1
            }
        }
        ));

named!(sign <i32>, alt!(
        tag!("-") => { |_| -1 } |
        tag!("+") => { |_| 1 }
        )
    );

named!(timezone_hour <(i32,i32)>, chain!(
        s: sign ~
        h: hour ~
        m: empty_or!(
            chain!(
                tag!(":")? ~ m: minute , || { m }
            ))
        ,
        || { (s * (h as i32) , s * (m.unwrap_or(0) as i32)) }
        ));

named!(timezone_utc <(i32,i32)>, map!(tag!("Z"), |_| (0,0)));

// Full ISO8601
named!(pub parse_datetime <DateTime>, chain!(
        d: parse_date ~
        tag!("T") ~
        t: parse_time
        ,
        || {
            DateTime{
                date: d,
                time: t,
            }
        }
        ));

<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::hash::BuildHasherDefault;

use fnv::FnvHasher;

use {Entity, Index, Generation};
use bitset::BitSet;

/// Base trait for a component storage that is used as a trait object.
/// Doesn't depent on the actual component type.
pub trait StorageBase {
    /// Delete a particular entity from the storage.
    fn del(&mut self, Entity);
}

/// Typed component storage trait.
pub trait Storage<T>: StorageBase + Sized {
    /// Used during iterator
    type UnprotectedStorage: UnprotectedStorage<T>;

    /// Create a new storage. This is called when you register a new
    /// component type within the world.
    fn new() -> Self;
    /// Insert a new data for a given entity.
    fn insert(&mut self, Entity, T);
    /// Try reading the data associated with an entity.
    fn get(&self, Entity) -> Option<&T>;
    /// Try mutating the data associated with an entity.
    fn get_mut(&mut self, Entity) -> Option<&mut T>;
    /// Remove the data associated with an entity.
    fn remove(&mut self, Entity) -> Option<T>;
    /// splits the bitset from the storage for use
    /// by the join iterator.
    fn open(&self) -> (&BitSet, &Self::UnprotectedStorage);
    /// splits the bitset from the storage for use
    /// by the join iterator.
    fn open_mut(&mut self) -> (&BitSet, &mut Self::UnprotectedStorage);
}

/// Used by the framework to quickly join componets
pub trait UnprotectedStorage<T> {
    /// Try reading the data associated with an entity.
    /// this is unsafe because the external set used
    /// to protect this storage is absent
    unsafe fn get(&self, id: Index) -> &T;
    /// Try mutating the data associated with an entity.
    /// this is unsafe because the external set used
    /// to protect this storage is absent
    unsafe fn get_mut(&mut self, id: Index) -> &mut T;
}

pub struct InnerHashMap<T>(HashMap<Index, GenerationData<T>, BuildHasherDefault<FnvHasher>>);

/// HashMap-based storage. Best suited for rare components.
pub struct HashMapStorage<T>{
    set: BitSet,
    map: InnerHashMap<T>
}

impl<T> StorageBase for HashMapStorage<T> {
    fn del(&mut self, entity: Entity) {
        if self.set.remove(entity.get_id() as u32) {
            self.map.0.remove(&(entity.get_id() as u32));
        }
    }
}

impl<T> Storage<T> for HashMapStorage<T> {
    type UnprotectedStorage = InnerHashMap<T>;

    fn new() -> Self {
        let fnv = BuildHasherDefault::<FnvHasher>::default();
        HashMapStorage {
            set: BitSet::new(),
            map: InnerHashMap(HashMap::with_hasher(fnv))
        }
    }
    fn get(&self, entity: Entity) -> Option<&T> {
        self.map.0.get(&(entity.get_id() as u32))
            .and_then(|x| if x.generation == entity.get_gen() { Some(&x.data) } else { None })

    }
    fn get_mut(&mut self, entity: Entity) -> Option<&mut T> {
        self.map.0.get_mut(&(entity.get_id() as u32))
            .and_then(|x| if x.generation == entity.get_gen() { Some(&mut x.data) } else { None })

    }
    fn insert(&mut self, entity: Entity, value: T) {
        let value = GenerationData{
            data: value,
            generation: entity.get_gen()
        };
        if self.map.0.insert(entity.get_id() as u32, value).is_none() {
            self.set.add(entity.get_id() as u32);
        }
    }
    fn remove(&mut self, entity: Entity) -> Option<T> {
        if self.set.remove(entity.get_id() as u32) {
            let value = self.map.0.remove(&(entity.get_id() as u32)).unwrap();
            if value.generation == entity.get_gen() {
                Some(value.data)
            } else {
                // it should be unlikely that the generation mismatch was
                // wrong, so this is a slow path to re-add the value back
                self.insert(entity, value.data);
                None
            }
        } else {
            None
        }
    }
    fn open(&self) -> (&BitSet, &Self::UnprotectedStorage) {
        (&self.set, &self.map)
    }
    fn open_mut(&mut self) -> (&BitSet, &mut Self::UnprotectedStorage) {
        (&self.set, &mut self.map)
    }
}

impl<T> UnprotectedStorage<T> for InnerHashMap<T> {
    unsafe fn get(&self, e: Index) -> &T {
        &self.0.get(&e).unwrap().data
    }
    unsafe fn get_mut(&mut self, e: Index) -> &mut T {
        &mut self.0.get_mut(&e).unwrap().data
    }
}

pub struct InnerVec<T>(Vec<GenerationData<T>>);

pub struct GenerationData<T> {
    pub generation: Generation,
    pub data: T
}

/// Vec-based storage, actually wraps data into options and stores the generations
/// of the data in order to match with given entities. Supposed to have maximum
/// performance for the components mostly present in entities.
pub struct VecStorage<T> {
    set: BitSet,
    values: InnerVec<T>,
}

impl<T> VecStorage<T> {
    fn extend(&mut self, id: usize) {
        unsafe {
            let delta = (id + 1) - self.values.0.len();
            self.values.0.reserve(delta);
            self.values.0.set_len(id + 1);
        }
    }
}

impl<T> Drop for VecStorage<T> {
    fn drop(&mut self) {
        use std::mem;

        for (i, v) in self.values.0.drain(..).enumerate() {
            // if v was not in the set the data is invalid
            // and we must forget it instead of dropping it
            if !self.set.remove(i as u32) {
                mem::forget(v);
            }
        }
    }
}

impl<T> super::StorageBase for VecStorage<T> {
    fn del(&mut self, e: Entity) {
        self.remove(e);
    }
}

impl<T> super::Storage<T> for VecStorage<T> {
    type UnprotectedStorage = InnerVec<T>;

    fn new() -> Self {
        VecStorage {
            set: BitSet::new(),
            values: InnerVec(Vec::new()),
        }
    }
    fn get(&self, e: Entity) -> Option<&T> {
        let id = e.get_id();
        if self.set.contains(id as u32) {
            let v = unsafe { self.values.0.get_unchecked(id) };
            if v.generation == e.get_gen() {
                return Some(&v.data);
            }
        }
        None

    }
    fn get_mut(&mut self, e: Entity) -> Option<&mut T> {
        let id = e.get_id();
        if self.set.contains(id as u32) {
            let v = unsafe { self.values.0.get_unchecked_mut(id) };
            if v.generation == e.get_gen() {
                return Some(&mut v.data);
            }
        }
        None
    }
    fn insert(&mut self, e: Entity, mut v: T) {
        use std::{ptr, mem};

        let id = e.get_id();
        if self.set.contains(id as u32) {
            let mut data = &mut self.values.0[id];
            data.generation = e.get_gen();
            mem::swap(&mut data.data, &mut v);
            Some(v)
        } else {
            self.set.add(id as u32);
            if self.values.0.len() <= id {
                self.extend(id);
            }
            unsafe {
                ptr::write(
                    &mut self.values.0[id],
                    GenerationData{
                        generation: e.get_gen(),
                        data: v
                    }
                );
            }
            None
        };
    }
    fn remove(&mut self, e: Entity) -> Option<T> {
        use std::ptr;

        let (id, gen) = (e.get_id(), e.get_gen());
        let gen_matches = self.values.0.get(id)
            .map(|x| x.generation == gen).unwrap_or(false);

        if gen_matches && self.set.remove(id as u32) {
            let value = unsafe { ptr::read(&self.values.0[id]) };
            Some(value.data)
        } else {
            None
        }
    }
    fn open(&self) -> (&BitSet, &InnerVec<T>) {
        (&self.set, &self.values)
    }
    fn open_mut(&mut self) -> (&BitSet, &mut InnerVec<T>) {
        (&self.set, &mut self.values)
    }
}

impl<T> super::UnprotectedStorage<T> for InnerVec<T> {
    unsafe fn get(&self, e: u32) -> &T {
        &self.0.get_unchecked(e as usize).data
    }
    unsafe fn get_mut(&mut self, e: u32) -> &mut T {
        &mut self.0.get_unchecked_mut(e as usize).data
    }
}

#[cfg(test)]
mod map_test {
    use {Storage, Entity, Generation};
    use super::VecStorage;

    #[test]
    fn insert() {
        let mut c = VecStorage::new();
        for i in 0..1_000 {
            c.insert(Entity::new(i, Generation(0)), i);
        }

        for i in 0..1_000 {
            assert_eq!(c.get(Entity::new(i, Generation(0))).unwrap(), &i);
        }
    }

    #[test]
    fn insert_100k() {
        let mut c = VecStorage::new();
        for i in 0..100_000 {
            c.insert(Entity::new(i, Generation(0)), i);
        }

        for i in 0..100_000 {
            assert_eq!(c.get(Entity::new(i, Generation(0))).unwrap(), &i);
        }
    }

    #[test]
    fn remove() {
        let mut c = VecStorage::new();
        for i in 0..1_000 {
            c.insert(Entity::new(i, Generation(0)), i);
        }

        for i in 0..1_000 {
            assert_eq!(c.get(Entity::new(i, Generation(0))).unwrap(), &i);
        }

        for i in 0..1_000 {
            c.remove(Entity::new(i, Generation(0)));
        }

        for i in 0..1_000 {
            assert!(c.get(Entity::new(i, Generation(0))).is_none());
        }
    }

    #[test]
    fn test_gen() {
        let mut c = VecStorage::new();
        for i in 0..1_000i32 {
            c.insert(Entity::new(i as u32, Generation(0)), i);
            c.insert(Entity::new(i as u32, Generation(0)), -i);
        }

        for i in 0..1_000i32 {
            assert_eq!(c.get(Entity::new(i as u32, Generation(0))).unwrap(), &-i);
        }
    }

    #[test]
    fn insert_same_key() {
        let mut c = VecStorage::new();
        for i in 0..10_000 {
            c.insert(Entity::new(i as u32, Generation(0)), i);
            assert_eq!(c.get(Entity::new(i as u32, Generation(0))).unwrap(), &i);
        }
    }

    #[should_panic]
    #[test]
    fn wrap() {
        let mut c = VecStorage::new();
        c.insert(Entity::new(1 << 21, Generation(0)), 7);
    }
}

#[cfg(test)]
mod test {
    use {Entity, Generation, Storage, VecStorage, HashMapStorage};

    fn test_add<S>() where S: Storage<u32> {
        let mut s = S::new();
        for i in 0..1_000 {
            s.insert(Entity::new(i, Generation(1)), i + 2718);
        }

        for i in 0..1_000 {
            assert_eq!(*s.get(Entity::new(i, Generation(1))).unwrap(), i + 2718);
        }
    }

    fn test_sub<S>() where S: Storage<u32> {
        let mut s = S::new();
        for i in 0..1_000 {
            s.insert(Entity::new(i, Generation(1)), i + 2718);
        }

        for i in 0..1_000 {
            assert_eq!(s.remove(Entity::new(i, Generation(1))).unwrap(), i + 2718);
            assert!(s.remove(Entity::new(i, Generation(1))).is_none());
        }
    }

    fn test_get_mut<S>() where S: Storage<u32> {
        let mut s = S::new();
        for i in 0..1_000 {
            s.insert(Entity::new(i, Generation(1)), i + 2718);
        }

        for i in 0..1_000 {
            *s.get_mut(Entity::new(i, Generation(1))).unwrap() -= 718;
        }

        for i in 0..1_000 {
            assert_eq!(*s.get(Entity::new(i, Generation(1))).unwrap(), i + 2000);
        }
    }

    fn test_add_gen<S>() where S: Storage<u32> {
        let mut s = S::new();
        for i in 0..1_000 {
            s.insert(Entity::new(i, Generation(1)), i + 2718);
            s.insert(Entity::new(i, Generation(2)), i + 31415);
        }

        for i in 0..1_000 {
            // this is removed since vec and hashmap disagree
            // on how this behavior should work...
            //assert!(s.get(Entity::new(i, 1)).is_none());
            assert_eq!(*s.get(Entity::new(i, Generation(2))).unwrap(), i + 31415);
        }
    }

    fn test_sub_gen<S>() where S: Storage<u32> {
        let mut s = S::new();
        for i in 0..1_000 {
            s.insert(Entity::new(i, Generation(2)), i + 2718);
        }

        for i in 0..1_000 {
            assert!(s.remove(Entity::new(i, Generation(1))).is_none());
        }
    }

    #[test] fn vec_test_add() { test_add::<VecStorage<u32>>(); }
    #[test] fn vec_test_sub() { test_sub::<VecStorage<u32>>(); }
    #[test] fn vec_test_get_mut() { test_get_mut::<VecStorage<u32>>(); }
    #[test] fn vec_test_add_gen() { test_add_gen::<VecStorage<u32>>(); }
    #[test] fn vec_test_sub_gen() { test_sub_gen::<VecStorage<u32>>(); }

    #[test] fn hash_test_add() { test_add::<HashMapStorage<u32>>(); }
    #[test] fn hash_test_sub() { test_sub::<HashMapStorage<u32>>(); }
    #[test] fn hash_test_get_mut() { test_get_mut::<HashMapStorage<u32>>(); }
    #[test] fn hash_test_add_gen() { test_add_gen::<HashMapStorage<u32>>(); }
    #[test] fn hash_test_sub_gen() { test_sub_gen::<HashMapStorage<u32>>(); }
}

<commit_msg>move safe logic out of unsafe block<commit_after>use std::collections::HashMap;
use std::hash::BuildHasherDefault;

use fnv::FnvHasher;

use {Entity, Index, Generation};
use bitset::BitSet;

/// Base trait for a component storage that is used as a trait object.
/// Doesn't depent on the actual component type.
pub trait StorageBase {
    /// Delete a particular entity from the storage.
    fn del(&mut self, Entity);
}

/// Typed component storage trait.
pub trait Storage<T>: StorageBase + Sized {
    /// Used during iterator
    type UnprotectedStorage: UnprotectedStorage<T>;

    /// Create a new storage. This is called when you register a new
    /// component type within the world.
    fn new() -> Self;
    /// Insert a new data for a given entity.
    fn insert(&mut self, Entity, T);
    /// Try reading the data associated with an entity.
    fn get(&self, Entity) -> Option<&T>;
    /// Try mutating the data associated with an entity.
    fn get_mut(&mut self, Entity) -> Option<&mut T>;
    /// Remove the data associated with an entity.
    fn remove(&mut self, Entity) -> Option<T>;
    /// splits the bitset from the storage for use
    /// by the join iterator.
    fn open(&self) -> (&BitSet, &Self::UnprotectedStorage);
    /// splits the bitset from the storage for use
    /// by the join iterator.
    fn open_mut(&mut self) -> (&BitSet, &mut Self::UnprotectedStorage);
}

/// Used by the framework to quickly join componets
pub trait UnprotectedStorage<T> {
    /// Try reading the data associated with an entity.
    /// this is unsafe because the external set used
    /// to protect this storage is absent
    unsafe fn get(&self, id: Index) -> &T;
    /// Try mutating the data associated with an entity.
    /// this is unsafe because the external set used
    /// to protect this storage is absent
    unsafe fn get_mut(&mut self, id: Index) -> &mut T;
}

pub struct InnerHashMap<T>(HashMap<Index, GenerationData<T>, BuildHasherDefault<FnvHasher>>);

/// HashMap-based storage. Best suited for rare components.
pub struct HashMapStorage<T>{
    set: BitSet,
    map: InnerHashMap<T>
}

impl<T> StorageBase for HashMapStorage<T> {
    fn del(&mut self, entity: Entity) {
        if self.set.remove(entity.get_id() as u32) {
            self.map.0.remove(&(entity.get_id() as u32));
        }
    }
}

impl<T> Storage<T> for HashMapStorage<T> {
    type UnprotectedStorage = InnerHashMap<T>;

    fn new() -> Self {
        let fnv = BuildHasherDefault::<FnvHasher>::default();
        HashMapStorage {
            set: BitSet::new(),
            map: InnerHashMap(HashMap::with_hasher(fnv))
        }
    }
    fn get(&self, entity: Entity) -> Option<&T> {
        self.map.0.get(&(entity.get_id() as u32))
            .and_then(|x| if x.generation == entity.get_gen() { Some(&x.data) } else { None })

    }
    fn get_mut(&mut self, entity: Entity) -> Option<&mut T> {
        self.map.0.get_mut(&(entity.get_id() as u32))
            .and_then(|x| if x.generation == entity.get_gen() { Some(&mut x.data) } else { None })

    }
    fn insert(&mut self, entity: Entity, value: T) {
        let value = GenerationData{
            data: value,
            generation: entity.get_gen()
        };
        if self.map.0.insert(entity.get_id() as u32, value).is_none() {
            self.set.add(entity.get_id() as u32);
        }
    }
    fn remove(&mut self, entity: Entity) -> Option<T> {
        if self.set.remove(entity.get_id() as u32) {
            let value = self.map.0.remove(&(entity.get_id() as u32)).unwrap();
            if value.generation == entity.get_gen() {
                Some(value.data)
            } else {
                // it should be unlikely that the generation mismatch was
                // wrong, so this is a slow path to re-add the value back
                self.insert(entity, value.data);
                None
            }
        } else {
            None
        }
    }
    fn open(&self) -> (&BitSet, &Self::UnprotectedStorage) {
        (&self.set, &self.map)
    }
    fn open_mut(&mut self) -> (&BitSet, &mut Self::UnprotectedStorage) {
        (&self.set, &mut self.map)
    }
}

impl<T> UnprotectedStorage<T> for InnerHashMap<T> {
    unsafe fn get(&self, e: Index) -> &T {
        &self.0.get(&e).unwrap().data
    }
    unsafe fn get_mut(&mut self, e: Index) -> &mut T {
        &mut self.0.get_mut(&e).unwrap().data
    }
}

pub struct InnerVec<T>(Vec<GenerationData<T>>);

pub struct GenerationData<T> {
    pub generation: Generation,
    pub data: T
}

/// Vec-based storage, actually wraps data into options and stores the generations
/// of the data in order to match with given entities. Supposed to have maximum
/// performance for the components mostly present in entities.
pub struct VecStorage<T> {
    set: BitSet,
    values: InnerVec<T>,
}

impl<T> VecStorage<T> {
    fn extend(&mut self, id: usize) {
        debug_assert!(id >= self.values.0.len());
        let delta = (id + 1) - self.values.0.len();
        self.values.0.reserve(delta);
        unsafe {
            self.values.0.set_len(id + 1);
        }
    }
}

impl<T> Drop for VecStorage<T> {
    fn drop(&mut self) {
        use std::mem;

        for (i, v) in self.values.0.drain(..).enumerate() {
            // if v was not in the set the data is invalid
            // and we must forget it instead of dropping it
            if !self.set.remove(i as u32) {
                mem::forget(v);
            }
        }
    }
}

impl<T> super::StorageBase for VecStorage<T> {
    fn del(&mut self, e: Entity) {
        self.remove(e);
    }
}

impl<T> super::Storage<T> for VecStorage<T> {
    type UnprotectedStorage = InnerVec<T>;

    fn new() -> Self {
        VecStorage {
            set: BitSet::new(),
            values: InnerVec(Vec::new()),
        }
    }
    fn get(&self, e: Entity) -> Option<&T> {
        let id = e.get_id();
        if self.set.contains(id as u32) {
            let v = unsafe { self.values.0.get_unchecked(id) };
            if v.generation == e.get_gen() {
                return Some(&v.data);
            }
        }
        None

    }
    fn get_mut(&mut self, e: Entity) -> Option<&mut T> {
        let id = e.get_id();
        if self.set.contains(id as u32) {
            let v = unsafe { self.values.0.get_unchecked_mut(id) };
            if v.generation == e.get_gen() {
                return Some(&mut v.data);
            }
        }
        None
    }
    fn insert(&mut self, e: Entity, mut v: T) {
        use std::{ptr, mem};

        let id = e.get_id();
        if self.set.contains(id as u32) {
            let mut data = &mut self.values.0[id];
            data.generation = e.get_gen();
            mem::swap(&mut data.data, &mut v);
            Some(v)
        } else {
            self.set.add(id as u32);
            if self.values.0.len() <= id {
                self.extend(id);
            }
            unsafe {
                ptr::write(
                    &mut self.values.0[id],
                    GenerationData{
                        generation: e.get_gen(),
                        data: v
                    }
                );
            }
            None
        };
    }
    fn remove(&mut self, e: Entity) -> Option<T> {
        use std::ptr;

        let (id, gen) = (e.get_id(), e.get_gen());
        let gen_matches = self.values.0.get(id)
            .map(|x| x.generation == gen).unwrap_or(false);

        if gen_matches && self.set.remove(id as u32) {
            let value = unsafe { ptr::read(&self.values.0[id]) };
            Some(value.data)
        } else {
            None
        }
    }
    fn open(&self) -> (&BitSet, &InnerVec<T>) {
        (&self.set, &self.values)
    }
    fn open_mut(&mut self) -> (&BitSet, &mut InnerVec<T>) {
        (&self.set, &mut self.values)
    }
}

impl<T> super::UnprotectedStorage<T> for InnerVec<T> {
    unsafe fn get(&self, e: u32) -> &T {
        &self.0.get_unchecked(e as usize).data
    }
    unsafe fn get_mut(&mut self, e: u32) -> &mut T {
        &mut self.0.get_unchecked_mut(e as usize).data
    }
}

#[cfg(test)]
mod map_test {
    use {Storage, Entity, Generation};
    use super::VecStorage;

    #[test]
    fn insert() {
        let mut c = VecStorage::new();
        for i in 0..1_000 {
            c.insert(Entity::new(i, Generation(0)), i);
        }

        for i in 0..1_000 {
            assert_eq!(c.get(Entity::new(i, Generation(0))).unwrap(), &i);
        }
    }

    #[test]
    fn insert_100k() {
        let mut c = VecStorage::new();
        for i in 0..100_000 {
            c.insert(Entity::new(i, Generation(0)), i);
        }

        for i in 0..100_000 {
            assert_eq!(c.get(Entity::new(i, Generation(0))).unwrap(), &i);
        }
    }

    #[test]
    fn remove() {
        let mut c = VecStorage::new();
        for i in 0..1_000 {
            c.insert(Entity::new(i, Generation(0)), i);
        }

        for i in 0..1_000 {
            assert_eq!(c.get(Entity::new(i, Generation(0))).unwrap(), &i);
        }

        for i in 0..1_000 {
            c.remove(Entity::new(i, Generation(0)));
        }

        for i in 0..1_000 {
            assert!(c.get(Entity::new(i, Generation(0))).is_none());
        }
    }

    #[test]
    fn test_gen() {
        let mut c = VecStorage::new();
        for i in 0..1_000i32 {
            c.insert(Entity::new(i as u32, Generation(0)), i);
            c.insert(Entity::new(i as u32, Generation(0)), -i);
        }

        for i in 0..1_000i32 {
            assert_eq!(c.get(Entity::new(i as u32, Generation(0))).unwrap(), &-i);
        }
    }

    #[test]
    fn insert_same_key() {
        let mut c = VecStorage::new();
        for i in 0..10_000 {
            c.insert(Entity::new(i as u32, Generation(0)), i);
            assert_eq!(c.get(Entity::new(i as u32, Generation(0))).unwrap(), &i);
        }
    }

    #[should_panic]
    #[test]
    fn wrap() {
        let mut c = VecStorage::new();
        c.insert(Entity::new(1 << 21, Generation(0)), 7);
    }
}

#[cfg(test)]
mod test {
    use {Entity, Generation, Storage, VecStorage, HashMapStorage};

    fn test_add<S>() where S: Storage<u32> {
        let mut s = S::new();
        for i in 0..1_000 {
            s.insert(Entity::new(i, Generation(1)), i + 2718);
        }

        for i in 0..1_000 {
            assert_eq!(*s.get(Entity::new(i, Generation(1))).unwrap(), i + 2718);
        }
    }

    fn test_sub<S>() where S: Storage<u32> {
        let mut s = S::new();
        for i in 0..1_000 {
            s.insert(Entity::new(i, Generation(1)), i + 2718);
        }

        for i in 0..1_000 {
            assert_eq!(s.remove(Entity::new(i, Generation(1))).unwrap(), i + 2718);
            assert!(s.remove(Entity::new(i, Generation(1))).is_none());
        }
    }

    fn test_get_mut<S>() where S: Storage<u32> {
        let mut s = S::new();
        for i in 0..1_000 {
            s.insert(Entity::new(i, Generation(1)), i + 2718);
        }

        for i in 0..1_000 {
            *s.get_mut(Entity::new(i, Generation(1))).unwrap() -= 718;
        }

        for i in 0..1_000 {
            assert_eq!(*s.get(Entity::new(i, Generation(1))).unwrap(), i + 2000);
        }
    }

    fn test_add_gen<S>() where S: Storage<u32> {
        let mut s = S::new();
        for i in 0..1_000 {
            s.insert(Entity::new(i, Generation(1)), i + 2718);
            s.insert(Entity::new(i, Generation(2)), i + 31415);
        }

        for i in 0..1_000 {
            // this is removed since vec and hashmap disagree
            // on how this behavior should work...
            //assert!(s.get(Entity::new(i, 1)).is_none());
            assert_eq!(*s.get(Entity::new(i, Generation(2))).unwrap(), i + 31415);
        }
    }

    fn test_sub_gen<S>() where S: Storage<u32> {
        let mut s = S::new();
        for i in 0..1_000 {
            s.insert(Entity::new(i, Generation(2)), i + 2718);
        }

        for i in 0..1_000 {
            assert!(s.remove(Entity::new(i, Generation(1))).is_none());
        }
    }

    #[test] fn vec_test_add() { test_add::<VecStorage<u32>>(); }
    #[test] fn vec_test_sub() { test_sub::<VecStorage<u32>>(); }
    #[test] fn vec_test_get_mut() { test_get_mut::<VecStorage<u32>>(); }
    #[test] fn vec_test_add_gen() { test_add_gen::<VecStorage<u32>>(); }
    #[test] fn vec_test_sub_gen() { test_sub_gen::<VecStorage<u32>>(); }

    #[test] fn hash_test_add() { test_add::<HashMapStorage<u32>>(); }
    #[test] fn hash_test_sub() { test_sub::<HashMapStorage<u32>>(); }
    #[test] fn hash_test_get_mut() { test_get_mut::<HashMapStorage<u32>>(); }
    #[test] fn hash_test_add_gen() { test_add_gen::<HashMapStorage<u32>>(); }
    #[test] fn hash_test_sub_gen() { test_sub_gen::<HashMapStorage<u32>>(); }
}

<|endoftext|>


<commit_before>use super::methods::*;
use libp2p::core::{upgrade, InboundUpgrade, OutboundUpgrade, UpgradeInfo};
use ssz::{ssz_encode, Decodable, Encodable, SszStream};
use std::io;
use std::iter;
use tokio::io::{AsyncRead, AsyncWrite};

/// The maximum bytes that can be sent across the RPC.
const MAX_READ_SIZE: usize = 4_194_304; // 4M

/// Implementation of the `ConnectionUpgrade` for the rpc protocol.

#[derive(Debug, Clone)]
pub struct RPCProtocol;

impl UpgradeInfo for RPCProtocol {
    type Info = &'static [u8];
    type InfoIter = iter::Once<Self::Info>;

    #[inline]
    fn protocol_info(&self) -> Self::InfoIter {
        iter::once(b"/eth/serenity/rpc/1.0.0")
    }
}

impl Default for RPCProtocol {
    fn default() -> Self {
        RPCProtocol
    }
}

/// The RPC types which are sent/received in this protocol.
#[derive(Debug, Clone)]
pub enum RPCEvent {
    Request {
        id: u64,
        method_id: u16,
        body: RPCRequest,
    },
    Response {
        id: u64,
        method_id: u16, //TODO: Remove and process decoding upstream
        result: RPCResponse,
    },
}

impl UpgradeInfo for RPCEvent {
    type Info = &'static [u8];
    type InfoIter = iter::Once<Self::Info>;

    #[inline]
    fn protocol_info(&self) -> Self::InfoIter {
        iter::once(b"/eth/serenity/rpc/1.0.0")
    }
}

impl<TSocket> InboundUpgrade<TSocket> for RPCProtocol
where
    TSocket: AsyncRead + AsyncWrite,
{
    type Output = RPCEvent;
    type Error = DecodeError;
    type Future =
        upgrade::ReadOneThen<TSocket, (), fn(Vec<u8>, ()) -> Result<RPCEvent, DecodeError>>;

    fn upgrade_inbound(self, socket: TSocket, _: Self::Info) -> Self::Future {
        upgrade::read_one_then(socket, MAX_READ_SIZE, (), |packet, ()| Ok(decode(packet)?))
    }
}

fn decode(packet: Vec<u8>) -> Result<RPCEvent, DecodeError> {
    // decode the header of the rpc
    // request/response
    let (request, index) = bool::ssz_decode(&packet, 0)?;
    let (id, index) = u64::ssz_decode(&packet, index)?;
    let (method_id, index) = u16::ssz_decode(&packet, index)?;

    if request {
        let body = match RPCMethod::from(method_id) {
            RPCMethod::Hello => {
                let (hello_body, _index) = HelloMessage::ssz_decode(&packet, index)?;
                RPCRequest::Hello(hello_body)
            }
            RPCMethod::Goodbye => {
                let (goodbye_code, _index) = u64::ssz_decode(&packet, index)?;
                RPCRequest::Goodbye(goodbye_code)
            }
            RPCMethod::BeaconBlockRoots => {
                let (block_roots_request, _index) =
                    BeaconBlockRootsRequest::ssz_decode(&packet, index)?;
                RPCRequest::BeaconBlockRoots(block_roots_request)
            }
            RPCMethod::BeaconBlockHeaders => {
                let (block_headers_request, _index) =
                    BeaconBlockHeadersRequest::ssz_decode(&packet, index)?;
                RPCRequest::BeaconBlockHeaders(block_headers_request)
            }
            RPCMethod::BeaconBlockBodies => {
                let (block_bodies_request, _index) =
                    BeaconBlockBodiesRequest::ssz_decode(&packet, index)?;
                RPCRequest::BeaconBlockBodies(block_bodies_request)
            }
            RPCMethod::Unknown | _ => return Err(DecodeError::UnknownRPCMethod),
        };

        Ok(RPCEvent::Request {
            id,
            method_id,
            body,
        })
    }
    // we have received a response
    else {
        let result = match RPCMethod::from(method_id) {
            RPCMethod::Hello => {
                let (body, _index) = HelloMessage::ssz_decode(&packet, index)?;
                RPCResponse::Hello(body)
            }
            RPCMethod::Goodbye => unreachable!("Should never receive a goodbye response"),
            RPCMethod::BeaconBlockRoots => {
                let (body, _index) = BeaconBlockRootsResponse::ssz_decode(&packet, index)?;
                RPCResponse::BeaconBlockRoots(body)
            }
            RPCMethod::BeaconBlockHeaders => {
                let (body, _index) = BeaconBlockHeadersResponse::ssz_decode(&packet, index)?;
                RPCResponse::BeaconBlockHeaders(body)
            }
            RPCMethod::BeaconBlockBodies => {
                let (body, _index) = BeaconBlockBodiesResponse::ssz_decode(&packet, index)?;
                RPCResponse::BeaconBlockBodies(body)
            }
            RPCMethod::Unknown | _ => return Err(DecodeError::UnknownRPCMethod),
        };
        Ok(RPCEvent::Response {
            id,
            method_id,
            result,
        })
    }
}

impl<TSocket> OutboundUpgrade<TSocket> for RPCEvent
where
    TSocket: AsyncWrite,
{
    type Output = ();
    type Error = io::Error;
    type Future = upgrade::WriteOne<TSocket>;

    #[inline]
    fn upgrade_outbound(self, socket: TSocket, _: Self::Info) -> Self::Future {
        let bytes = ssz_encode(&self);
        upgrade::write_one(socket, bytes)
    }
}

impl Encodable for RPCEvent {
    fn ssz_append(&self, s: &mut SszStream) {
        match self {
            RPCEvent::Request {
                id,
                method_id,
                body,
            } => {
                s.append(&true);
                s.append(id);
                s.append(method_id);
                match body {
                    RPCRequest::Hello(body) => {
                        s.append(body);
                    }
                    _ => {}
                }
            }
            RPCEvent::Response {
                id,
                method_id,
                result,
            } => {
                s.append(&false);
                s.append(id);
                s.append(method_id);
                match result {
                    RPCResponse::Hello(response) => {
                        s.append(response);
                    }
                    _ => {}
                }
            }
        }
    }
}

#[derive(Debug)]
pub enum DecodeError {
    ReadError(upgrade::ReadOneError),
    SSZDecodeError(ssz::DecodeError),
    UnknownRPCMethod,
}

impl From<upgrade::ReadOneError> for DecodeError {
    #[inline]
    fn from(err: upgrade::ReadOneError) -> Self {
        DecodeError::ReadError(err)
    }
}

impl From<ssz::DecodeError> for DecodeError {
    #[inline]
    fn from(err: ssz::DecodeError) -> Self {
        DecodeError::SSZDecodeError(err)
    }
}
<commit_msg>Implement BeaconChainState RPC method<commit_after>use super::methods::*;
use libp2p::core::{upgrade, InboundUpgrade, OutboundUpgrade, UpgradeInfo};
use ssz::{ssz_encode, Decodable, Encodable, SszStream};
use std::io;
use std::iter;
use tokio::io::{AsyncRead, AsyncWrite};

/// The maximum bytes that can be sent across the RPC.
const MAX_READ_SIZE: usize = 4_194_304; // 4M

/// Implementation of the `ConnectionUpgrade` for the rpc protocol.

#[derive(Debug, Clone)]
pub struct RPCProtocol;

impl UpgradeInfo for RPCProtocol {
    type Info = &'static [u8];
    type InfoIter = iter::Once<Self::Info>;

    #[inline]
    fn protocol_info(&self) -> Self::InfoIter {
        iter::once(b"/eth/serenity/rpc/1.0.0")
    }
}

impl Default for RPCProtocol {
    fn default() -> Self {
        RPCProtocol
    }
}

/// The RPC types which are sent/received in this protocol.
#[derive(Debug, Clone)]
pub enum RPCEvent {
    Request {
        id: u64,
        method_id: u16,
        body: RPCRequest,
    },
    Response {
        id: u64,
        method_id: u16, //TODO: Remove and process decoding upstream
        result: RPCResponse,
    },
}

impl UpgradeInfo for RPCEvent {
    type Info = &'static [u8];
    type InfoIter = iter::Once<Self::Info>;

    #[inline]
    fn protocol_info(&self) -> Self::InfoIter {
        iter::once(b"/eth/serenity/rpc/1.0.0")
    }
}

impl<TSocket> InboundUpgrade<TSocket> for RPCProtocol
where
    TSocket: AsyncRead + AsyncWrite,
{
    type Output = RPCEvent;
    type Error = DecodeError;
    type Future =
        upgrade::ReadOneThen<TSocket, (), fn(Vec<u8>, ()) -> Result<RPCEvent, DecodeError>>;

    fn upgrade_inbound(self, socket: TSocket, _: Self::Info) -> Self::Future {
        upgrade::read_one_then(socket, MAX_READ_SIZE, (), |packet, ()| Ok(decode(packet)?))
    }
}

fn decode(packet: Vec<u8>) -> Result<RPCEvent, DecodeError> {
    // decode the header of the rpc
    // request/response
    let (request, index) = bool::ssz_decode(&packet, 0)?;
    let (id, index) = u64::ssz_decode(&packet, index)?;
    let (method_id, index) = u16::ssz_decode(&packet, index)?;

    if request {
        let body = match RPCMethod::from(method_id) {
            RPCMethod::Hello => {
                let (hello_body, _index) = HelloMessage::ssz_decode(&packet, index)?;
                RPCRequest::Hello(hello_body)
            }
            RPCMethod::Goodbye => {
                let (goodbye_code, _index) = u64::ssz_decode(&packet, index)?;
                RPCRequest::Goodbye(goodbye_code)
            }
            RPCMethod::BeaconBlockRoots => {
                let (block_roots_request, _index) =
                    BeaconBlockRootsRequest::ssz_decode(&packet, index)?;
                RPCRequest::BeaconBlockRoots(block_roots_request)
            }
            RPCMethod::BeaconBlockHeaders => {
                let (block_headers_request, _index) =
                    BeaconBlockHeadersRequest::ssz_decode(&packet, index)?;
                RPCRequest::BeaconBlockHeaders(block_headers_request)
            }
            RPCMethod::BeaconBlockBodies => {
                let (block_bodies_request, _index) =
                    BeaconBlockBodiesRequest::ssz_decode(&packet, index)?;
                RPCRequest::BeaconBlockBodies(block_bodies_request)
            }
            RPCMethod::BeaconChainState => {
                let (chain_state_request, _index) =
                    BeaconChainStateRequest::ssz_decode(&packet, index)?;
                RPCRequest::BeaconChainState(chain_state_request)
            }
            RPCMethod::Unknown => return Err(DecodeError::UnknownRPCMethod),
        };

        Ok(RPCEvent::Request {
            id,
            method_id,
            body,
        })
    }
    // we have received a response
    else {
        let result = match RPCMethod::from(method_id) {
            RPCMethod::Hello => {
                let (body, _index) = HelloMessage::ssz_decode(&packet, index)?;
                RPCResponse::Hello(body)
            }
            RPCMethod::Goodbye => unreachable!("Should never receive a goodbye response"),
            RPCMethod::BeaconBlockRoots => {
                let (body, _index) = BeaconBlockRootsResponse::ssz_decode(&packet, index)?;
                RPCResponse::BeaconBlockRoots(body)
            }
            RPCMethod::BeaconBlockHeaders => {
                let (body, _index) = BeaconBlockHeadersResponse::ssz_decode(&packet, index)?;
                RPCResponse::BeaconBlockHeaders(body)
            }
            RPCMethod::BeaconBlockBodies => {
                let (body, _index) = BeaconBlockBodiesResponse::ssz_decode(&packet, index)?;
                RPCResponse::BeaconBlockBodies(body)
            }
            RPCMethod::BeaconChainState => {
                let (body, _index) = BeaconChainStateResponse::ssz_decode(&packet, index)?;
                RPCResponse::BeaconChainState(body)
            }
            RPCMethod::Unknown => return Err(DecodeError::UnknownRPCMethod),
        };
        Ok(RPCEvent::Response {
            id,
            method_id,
            result,
        })
    }
}

impl<TSocket> OutboundUpgrade<TSocket> for RPCEvent
where
    TSocket: AsyncWrite,
{
    type Output = ();
    type Error = io::Error;
    type Future = upgrade::WriteOne<TSocket>;

    #[inline]
    fn upgrade_outbound(self, socket: TSocket, _: Self::Info) -> Self::Future {
        let bytes = ssz_encode(&self);
        upgrade::write_one(socket, bytes)
    }
}

impl Encodable for RPCEvent {
    fn ssz_append(&self, s: &mut SszStream) {
        match self {
            RPCEvent::Request {
                id,
                method_id,
                body,
            } => {
                s.append(&true);
                s.append(id);
                s.append(method_id);
                match body {
                    RPCRequest::Hello(body) => {
                        s.append(body);
                    }
                    _ => {}
                }
            }
            RPCEvent::Response {
                id,
                method_id,
                result,
            } => {
                s.append(&false);
                s.append(id);
                s.append(method_id);
                match result {
                    RPCResponse::Hello(response) => {
                        s.append(response);
                    }
                    _ => {}
                }
            }
        }
    }
}

#[derive(Debug)]
pub enum DecodeError {
    ReadError(upgrade::ReadOneError),
    SSZDecodeError(ssz::DecodeError),
    UnknownRPCMethod,
}

impl From<upgrade::ReadOneError> for DecodeError {
    #[inline]
    fn from(err: upgrade::ReadOneError) -> Self {
        DecodeError::ReadError(err)
    }
}

impl From<ssz::DecodeError> for DecodeError {
    #[inline]
    fn from(err: ssz::DecodeError) -> Self {
        DecodeError::SSZDecodeError(err)
    }
}
<|endoftext|>


<commit_before>use super::SecretKey;
use bls_aggregates::PublicKey as RawPublicKey;
use ssz::{decode_ssz_list, ssz_encode, Decodable, DecodeError, Encodable, SszStream};
use std::hash::{Hash, Hasher};

/// A single BLS signature.
///
/// This struct is a wrapper upon a base type and provides helper functions (e.g., SSZ
/// serialization).
#[derive(Debug, PartialEq, Clone, Eq)]
pub struct PublicKey(RawPublicKey);

impl PublicKey {
    pub fn from_secret_key(secret_key: &SecretKey) -> Self {
        PublicKey(RawPublicKey::from_secret_key(secret_key.as_raw()))
    }

    /// Returns the underlying signature.
    pub fn as_raw(&self) -> &RawPublicKey {
        &self.0
    }
}

impl Encodable for PublicKey {
    fn ssz_append(&self, s: &mut SszStream) {
        s.append_vec(&self.0.as_bytes());
    }
}

impl Decodable for PublicKey {
    fn ssz_decode(bytes: &[u8], i: usize) -> Result<(Self, usize), DecodeError> {
        let (sig_bytes, i) = decode_ssz_list(bytes, i)?;
        let raw_sig = RawPublicKey::from_bytes(&sig_bytes).map_err(|_| DecodeError::TooShort)?;
        Ok((PublicKey(raw_sig), i))
    }
}

impl Hash for PublicKey {
    fn hash<H: Hasher>(&self, state: &mut H) {
        ssz_encode(self).hash(state)
    }
}

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

    #[test]
    pub fn test_ssz_round_trip() {
        let sk = SecretKey::random();
        let original = PublicKey::from_secret_key(&sk);

        let bytes = ssz_encode(&original);
        let (decoded, _) = PublicKey::ssz_decode(&bytes, 0).unwrap();

        assert_eq!(original, decoded);
    }
}
<commit_msg>Manually impl PartialEq for bls::PublicKey<commit_after>use super::SecretKey;
use bls_aggregates::PublicKey as RawPublicKey;
use ssz::{decode_ssz_list, ssz_encode, Decodable, DecodeError, Encodable, SszStream};
use std::hash::{Hash, Hasher};

/// A single BLS signature.
///
/// This struct is a wrapper upon a base type and provides helper functions (e.g., SSZ
/// serialization).
#[derive(Debug, Clone, Eq)]
pub struct PublicKey(RawPublicKey);

impl PublicKey {
    pub fn from_secret_key(secret_key: &SecretKey) -> Self {
        PublicKey(RawPublicKey::from_secret_key(secret_key.as_raw()))
    }

    /// Returns the underlying signature.
    pub fn as_raw(&self) -> &RawPublicKey {
        &self.0
    }
}

impl Encodable for PublicKey {
    fn ssz_append(&self, s: &mut SszStream) {
        s.append_vec(&self.0.as_bytes());
    }
}

impl Decodable for PublicKey {
    fn ssz_decode(bytes: &[u8], i: usize) -> Result<(Self, usize), DecodeError> {
        let (sig_bytes, i) = decode_ssz_list(bytes, i)?;
        let raw_sig = RawPublicKey::from_bytes(&sig_bytes).map_err(|_| DecodeError::TooShort)?;
        Ok((PublicKey(raw_sig), i))
    }
}

impl PartialEq for PublicKey {
    fn eq(&self, other: &PublicKey) -> bool {
        ssz_encode(self) == ssz_encode(other)
    }
}

impl Hash for PublicKey {
    fn hash<H: Hasher>(&self, state: &mut H) {
        ssz_encode(self).hash(state)
    }
}

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

    #[test]
    pub fn test_ssz_round_trip() {
        let sk = SecretKey::random();
        let original = PublicKey::from_secret_key(&sk);

        let bytes = ssz_encode(&original);
        let (decoded, _) = PublicKey::ssz_decode(&bytes, 0).unwrap();

        assert_eq!(original, decoded);
    }
}
<|endoftext|>