<commit_before>use std::{fmt, ptr, slice};
use std::cmp::Ordering;
use std::collections::HashMap;
use std::ffi::CString;
use std::io::{Error, ErrorKind, Result};
use std::sync::Mutex;
use std::time::Duration;
use super::{LcmSubscription, handler_callback};
use message::Message;
use ffi::*;

/// A threadsafe version of the LCM instance. Because this can be used on
/// multiple threads, all callbacks must also be threadsafe.
pub struct ThreadsafeLcm {
    lcm: *mut lcm_t,

    // This is a solid type. I'm proud of what I've done here.
    subscriptions:
      Mutex<HashMap<*mut lcm_subscription_t, Box<Box<FnMut(*const lcm_recv_buf_t) + Sync + Send>>>>,
}
unsafe impl Sync for ThreadsafeLcm { }
unsafe impl Send for ThreadsafeLcm { }

impl ThreadsafeLcm {
    /// Creates a new `ThreadsafeLcm` instance.
    ///
    /// ```
    /// use lcm::ThreadsafeLcm;
    /// let mut lcm = ThreadsafeLcm::new().unwrap();
    /// ```
    pub fn new() -> Result<Self> {
        trace!("Creating LCM instance");
        let lcm = unsafe { lcm_create(ptr::null()) };
        match lcm.is_null() {
            true => Err(Error::new(ErrorKind::Other, "Failed to initialize LCM.")),
            false => {
                Ok(ThreadsafeLcm {
                    lcm: lcm,
                    subscriptions: Mutex::new(HashMap::new()),
                })
            }
        }
    }

    pub fn get_fileno(&self) -> ::std::os::raw::c_int {
        unsafe { lcm_get_fileno(self.lcm) }
    }

    /// Subscribes a callback to a particular topic.
    ///
    /// If the callback function panics, then it will poison the inner mutex
    /// which will then cause all other subscription operations to panic.
    ///
    /// ```
    /// # use lcm::ThreadsafeLcm;
    /// let mut lcm = ThreadsafeLcm::new().unwrap();
    /// lcm.subscribe("GREETINGS", |name: String| println!("Hello, {}!", name) );
    /// ```
    pub fn subscribe<M, F>(&mut self, channel: &str, mut callback: F) -> LcmSubscription
        where M: Message,
              F: FnMut(M) + Sync + Send + 'static
    {
        trace!("Subscribing handler to channel {}", channel);

        let channel = CString::new(channel).unwrap();

        // This is a double box for a reason
        let handler = {
            let handler: Box<FnMut(*const lcm_recv_buf_t) + Sync + Send> =
                Box::new(move |rbuf: *const lcm_recv_buf_t| {
                    trace!("Running handler");
                    let mut buf = unsafe {
                        let ref rbuf = *rbuf;
                        let data = rbuf.data as *mut u8;
                        let len = rbuf.data_size as usize;
                        slice::from_raw_parts(data, len)
                    };
                    trace!("Decoding buffer: {:?}", buf);
                    match M::decode_with_hash(&mut buf) {
                        Ok(msg) => callback(msg),
                        Err(_) => error!("Failed to decode buffer: {:?}", buf),
                    }
                });
            Box::new(handler)
        };

        let subscription = unsafe {
            lcm_subscribe(self.lcm,
                          channel.as_ptr(),
                          Some(handler_callback),
                          &*handler as *const _  as *mut _)
        };

        let mut subs = self.subscriptions.lock().unwrap();
        assert!(!subs.contains_key(&subscription));
        subs.insert(subscription, handler);

        LcmSubscription { subscription }
    }

    /// Unsubscribes a message handler.
    ///
    /// ```
    /// # use lcm::ThreadsafeLcm;
    /// # let handler_function = |name: String| println!("Hello, {}!", name);
    /// # let mut lcm = ThreadsafeLcm::new().unwrap();
    /// let subscription = lcm.subscribe("GREETINGS", handler_function);
    /// // ...
    /// lcm.unsubscribe(subscription);
    /// ```
    pub fn unsubscribe(&mut self, subscription: LcmSubscription) -> Result<()> {
        trace!("Unsubscribing handler {:?}", subscription.subscription);
        let result = unsafe { lcm_unsubscribe(self.lcm, subscription.subscription) };

        match result {
            0 => {
                self.subscriptions.lock().unwrap().remove(&subscription.subscription);
                Ok(())
            },
            _ => Err(Error::new(ErrorKind::Other, "LCM: Failed to unsubscribe")),
        }
    }

    /// Publishes a message on the specified channel.
    ///
    /// ```
    /// # use lcm::ThreadsafeLcm;
    /// let mut lcm = ThreadsafeLcm::new().unwrap();
    /// lcm.publish("GREETINGS", &"Charles".to_string()).unwrap();
    /// ```
    pub fn publish<M>(&mut self, channel: &str, message: &M) -> Result<()>
        where M: Message
    {
        let channel = CString::new(channel).unwrap();
        let buffer = message.encode_with_hash()?;
        let result = unsafe {
            lcm_publish(self.lcm,
                        channel.as_ptr(),
                        buffer.as_ptr() as *mut _,
                        buffer.len() as _)
        };
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM Error")),
        }
    }

    /// Waits for and dispatches the next incoming message.
    ///
    /// ```
    /// # use lcm::ThreadsafeLcm;
    /// # let handler_function = |name: String| println!("Hello, {}!", name);
    /// let mut lcm = ThreadsafeLcm::new().unwrap();
    /// lcm.subscribe("POSITION", handler_function);
    /// loop {
    /// # break;
    ///     lcm.handle().unwrap();
    /// }
    /// ```
    pub fn handle(&mut self) -> Result<()> {
        let result = unsafe { lcm_handle(self.lcm) };
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM Error")),
        }
    }

    /// Waits for and dispatches the next incoming message, up to a time limit.
    ///
    /// ```
    /// # use std::time::Duration;
    /// # use lcm::ThreadsafeLcm;
    /// # let handler_function = |name: String| println!("Hello, {}!", name);
    /// let mut lcm = ThreadsafeLcm::new().unwrap();
    /// lcm.subscribe("POSITION", handler_function);
    /// let wait_dur = Duration::from_millis(100);
    /// loop {
    /// # break;
    ///     lcm.handle_timeout(Duration::from_millis(1000)).unwrap();
    /// }
    /// ```
    pub fn handle_timeout(&mut self, timeout: Duration) -> Result<()> {
        let result = unsafe { lcm_handle_timeout(self.lcm, (timeout.as_secs() * 1000) as i32 + (timeout.subsec_nanos() / 1000_000) as i32) };
        match result.cmp(&0) {
            Ordering::Less => Err(Error::new(ErrorKind::Other, "LCM Error")),
            Ordering::Equal => Err(Error::new(ErrorKind::Other, "LCM Timeout")),
            Ordering::Greater => Ok(()),
        }
    }

    /// Adjusts the maximum number of received messages that can be queued up for a subscription.
    /// The default is `30`.
    ///
    /// ```
    /// # use lcm::ThreadsafeLcm;
    /// # let handler_function = |name: String| println!("Hello, {}!", name);
    /// # let mut lcm = ThreadsafeLcm::new().unwrap();
    /// let subscription = lcm.subscribe("POSITION", handler_function);
    /// lcm.subscription_set_queue_capacity(subscription, 30);
    /// ```
    pub fn subscription_set_queue_capacity(&self, subscription: LcmSubscription, num_messages: usize) {
        let handler = subscription.subscription;
        let num_messages = num_messages as _;
        unsafe { lcm_subscription_set_queue_capacity(handler, num_messages) };
    }

}

impl Drop for ThreadsafeLcm {
    fn drop(&mut self) {
        trace!("Destroying Lcm instance");
        unsafe { lcm_destroy(self.lcm) };
    }
}

impl fmt::Debug for ThreadsafeLcm {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "ThreadsafeLcm {{ lcm: {:?} }}", self.lcm)
    }
}

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

    #[test]
    fn initialized() {
        let _lcm = ThreadsafeLcm::new().unwrap();
    }

    #[test]
    fn test_subscribe() {
        let mut lcm = ThreadsafeLcm::new().unwrap();
        lcm.subscribe("channel", |_: String| {});
        let subs_len = lcm.subscriptions.lock().unwrap().len();
        assert_eq!(subs_len, 1);
    }

    #[test]
    fn test_unsubscribe() {
        let mut lcm = ThreadsafeLcm::new().unwrap();
        let sub = lcm.subscribe("channel", |_: String| {});
        lcm.unsubscribe(sub).unwrap();
        let subs_len = lcm.subscriptions.lock().unwrap().len();
        assert_eq!(subs_len, 0);
    }
}
<commit_msg>Make threadsafe usable on multiple threads<commit_after>use std::{fmt, ptr, slice};
use std::cmp::Ordering;
use std::collections::HashMap;
use std::ffi::CString;
use std::io::{Error, ErrorKind, Result};
use std::sync::Mutex;
use std::time::Duration;
use super::{LcmSubscription, handler_callback};
use message::Message;
use ffi::*;

/// A threadsafe version of the LCM instance. Because this can be used on
/// multiple threads, all callbacks must also be threadsafe.
pub struct ThreadsafeLcm {
    lcm: *mut lcm_t,

    // This is a solid type. I'm proud of what I've done here.
    subscriptions:
      Mutex<HashMap<*mut lcm_subscription_t, Box<Box<FnMut(*const lcm_recv_buf_t) + Sync + Send>>>>,
}
unsafe impl Sync for ThreadsafeLcm { }
unsafe impl Send for ThreadsafeLcm { }

impl ThreadsafeLcm {
    /// Creates a new `ThreadsafeLcm` instance.
    ///
    /// ```
    /// use lcm::ThreadsafeLcm;
    /// let mut lcm = ThreadsafeLcm::new().unwrap();
    /// ```
    pub fn new() -> Result<Self> {
        trace!("Creating LCM instance");
        let lcm = unsafe { lcm_create(ptr::null()) };
        match lcm.is_null() {
            true => Err(Error::new(ErrorKind::Other, "Failed to initialize LCM.")),
            false => {
                Ok(ThreadsafeLcm {
                    lcm: lcm,
                    subscriptions: Mutex::new(HashMap::new()),
                })
            }
        }
    }

    pub fn get_fileno(&self) -> ::std::os::raw::c_int {
        unsafe { lcm_get_fileno(self.lcm) }
    }

    /// Subscribes a callback to a particular topic.
    ///
    /// If the callback function panics, then it will poison the inner mutex
    /// which will then cause all other subscription operations to panic.
    ///
    /// ```
    /// # use lcm::ThreadsafeLcm;
    /// let mut lcm = ThreadsafeLcm::new().unwrap();
    /// lcm.subscribe("GREETINGS", |name: String| println!("Hello, {}!", name) );
    /// ```
    pub fn subscribe<M, F>(&self, channel: &str, mut callback: F) -> LcmSubscription
        where M: Message,
              F: FnMut(M) + Sync + Send + 'static
    {
        trace!("Subscribing handler to channel {}", channel);

        let channel = CString::new(channel).unwrap();

        // This is a double box for a reason
        let handler = {
            let handler: Box<FnMut(*const lcm_recv_buf_t) + Sync + Send> =
                Box::new(move |rbuf: *const lcm_recv_buf_t| {
                    trace!("Running handler");
                    let mut buf = unsafe {
                        let ref rbuf = *rbuf;
                        let data = rbuf.data as *mut u8;
                        let len = rbuf.data_size as usize;
                        slice::from_raw_parts(data, len)
                    };
                    trace!("Decoding buffer: {:?}", buf);
                    match M::decode_with_hash(&mut buf) {
                        Ok(msg) => callback(msg),
                        Err(_) => error!("Failed to decode buffer: {:?}", buf),
                    }
                });
            Box::new(handler)
        };

        let subscription = unsafe {
            lcm_subscribe(self.lcm,
                          channel.as_ptr(),
                          Some(handler_callback),
                          &*handler as *const _  as *mut _)
        };

        let mut subs = self.subscriptions.lock().unwrap();
        assert!(!subs.contains_key(&subscription));
        subs.insert(subscription, handler);

        LcmSubscription { subscription }
    }

    /// Unsubscribes a message handler.
    ///
    /// ```
    /// # use lcm::ThreadsafeLcm;
    /// # let handler_function = |name: String| println!("Hello, {}!", name);
    /// # let mut lcm = ThreadsafeLcm::new().unwrap();
    /// let subscription = lcm.subscribe("GREETINGS", handler_function);
    /// // ...
    /// lcm.unsubscribe(subscription);
    /// ```
    pub fn unsubscribe(&self, subscription: LcmSubscription) -> Result<()> {
        trace!("Unsubscribing handler {:?}", subscription.subscription);
        let result = unsafe { lcm_unsubscribe(self.lcm, subscription.subscription) };

        match result {
            0 => {
                self.subscriptions.lock().unwrap().remove(&subscription.subscription);
                Ok(())
            },
            _ => Err(Error::new(ErrorKind::Other, "LCM: Failed to unsubscribe")),
        }
    }

    /// Publishes a message on the specified channel.
    ///
    /// ```
    /// # use lcm::ThreadsafeLcm;
    /// let mut lcm = ThreadsafeLcm::new().unwrap();
    /// lcm.publish("GREETINGS", &"Charles".to_string()).unwrap();
    /// ```
    pub fn publish<M>(&self, channel: &str, message: &M) -> Result<()>
        where M: Message
    {
        let channel = CString::new(channel).unwrap();
        let buffer = message.encode_with_hash()?;
        let result = unsafe {
            lcm_publish(self.lcm,
                        channel.as_ptr(),
                        buffer.as_ptr() as *mut _,
                        buffer.len() as _)
        };
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM Error")),
        }
    }

    /// Waits for and dispatches the next incoming message.
    ///
    /// ```
    /// # use lcm::ThreadsafeLcm;
    /// # let handler_function = |name: String| println!("Hello, {}!", name);
    /// let mut lcm = ThreadsafeLcm::new().unwrap();
    /// lcm.subscribe("POSITION", handler_function);
    /// loop {
    /// # break;
    ///     lcm.handle().unwrap();
    /// }
    /// ```
    pub fn handle(&self) -> Result<()> {
        let result = unsafe { lcm_handle(self.lcm) };
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM Error")),
        }
    }

    /// Waits for and dispatches the next incoming message, up to a time limit.
    ///
    /// ```
    /// # use std::time::Duration;
    /// # use lcm::ThreadsafeLcm;
    /// # let handler_function = |name: String| println!("Hello, {}!", name);
    /// let mut lcm = ThreadsafeLcm::new().unwrap();
    /// lcm.subscribe("POSITION", handler_function);
    /// let wait_dur = Duration::from_millis(100);
    /// loop {
    /// # break;
    ///     lcm.handle_timeout(Duration::from_millis(1000)).unwrap();
    /// }
    /// ```
    pub fn handle_timeout(&self, timeout: Duration) -> Result<()> {
        let result = unsafe { lcm_handle_timeout(self.lcm, (timeout.as_secs() * 1000) as i32 + (timeout.subsec_nanos() / 1000_000) as i32) };
        match result.cmp(&0) {
            Ordering::Less => Err(Error::new(ErrorKind::Other, "LCM Error")),
            Ordering::Equal => Err(Error::new(ErrorKind::Other, "LCM Timeout")),
            Ordering::Greater => Ok(()),
        }
    }

    /// Adjusts the maximum number of received messages that can be queued up for a subscription.
    /// The default is `30`.
    ///
    /// ```
    /// # use lcm::ThreadsafeLcm;
    /// # let handler_function = |name: String| println!("Hello, {}!", name);
    /// # let mut lcm = ThreadsafeLcm::new().unwrap();
    /// let subscription = lcm.subscribe("POSITION", handler_function);
    /// lcm.subscription_set_queue_capacity(subscription, 30);
    /// ```
    pub fn subscription_set_queue_capacity(&self, subscription: LcmSubscription, num_messages: usize) {
        let handler = subscription.subscription;
        let num_messages = num_messages as _;
        unsafe { lcm_subscription_set_queue_capacity(handler, num_messages) };
    }

}

impl Drop for ThreadsafeLcm {
    fn drop(&mut self) {
        trace!("Destroying Lcm instance");
        unsafe { lcm_destroy(self.lcm) };
    }
}

impl fmt::Debug for ThreadsafeLcm {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "ThreadsafeLcm {{ lcm: {:?} }}", self.lcm)
    }
}

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

    #[test]
    fn initialized() {
        let _lcm = ThreadsafeLcm::new().unwrap();
    }

    #[test]
    fn test_subscribe() {
        let mut lcm = ThreadsafeLcm::new().unwrap();
        lcm.subscribe("channel", |_: String| {});
        let subs_len = lcm.subscriptions.lock().unwrap().len();
        assert_eq!(subs_len, 1);
    }

    #[test]
    fn test_unsubscribe() {
        let mut lcm = ThreadsafeLcm::new().unwrap();
        let sub = lcm.subscribe("channel", |_: String| {});
        lcm.unsubscribe(sub).unwrap();
        let subs_len = lcm.subscriptions.lock().unwrap().len();
        assert_eq!(subs_len, 0);
    }
}
<|endoftext|>


<commit_before>// Copyright (C) 2018 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

use super::{Idle, Poll, PrepareDownload, SharedState, State, StateChangeImpl, StateMachine};
use crate::client::Api;
use slog_scope::{debug, error, info};

#[derive(Debug, PartialEq, Clone)]
pub(super) enum ServerAddress {
    Default,
    Custom(String),
}

#[derive(Debug, PartialEq)]
pub(super) struct Probe {
    pub(super) server_address: ServerAddress,
}

create_state_step!(Probe => Idle);
create_state_step!(Probe => Poll);

/// Implements the state change for State<Probe>.
impl StateChangeImpl for State<Probe> {
    fn name(&self) -> &'static str {
        "probe"
    }

    fn handle_trigger_probe(&self) -> actor::probe::Response {
        actor::probe::Response::RequestAccepted(self.name().to_owned())
    }

    fn handle(self, shared_state: &mut SharedState) -> Result<StateMachine, failure::Error> {
        use crate::client::ProbeResponse;
        use chrono::{Duration, Utc};
        use std::thread;

        let server_address = match self.0.server_address.clone() {
            ServerAddress::Default => shared_state.settings.network.server_address.clone(),
            ServerAddress::Custom(s) => s,
        };

        let r = loop {
            let probe = Api::new(&server_address)
                .probe(&shared_state.runtime_settings, &shared_state.firmware);
            if let Err(e) = probe {
                error!("{}", e);
                shared_state.runtime_settings.inc_retries();
                thread::sleep(Duration::seconds(1).to_std().unwrap());
            } else {
                shared_state.runtime_settings.clear_retries();
                break probe?;
            }
        };

        if let ProbeResponse::ExtraPoll(s) = r {
            info!("Delaying the probing as requested by the server.");
            shared_state
                .runtime_settings
                .set_polling_extra_interval(Duration::seconds(s))?;
        };

        match r {
            ProbeResponse::NoUpdate => {
                debug!("Moving to Idle state as no update is available.");

                // Store timestamp of last polling
                shared_state.runtime_settings.set_last_polling(Utc::now())?;
                Ok(StateMachine::Idle(self.into()))
            }

            ProbeResponse::ExtraPoll(_) => {
                debug!("Moving to Poll state due the extra polling interval.");
                Ok(StateMachine::Poll(self.into()))
            }

            ProbeResponse::Update(u) => {
                // Ensure the package is compatible
                u.compatible_with(&shared_state.firmware)?;
                // Store timestamp of last polling
                shared_state.runtime_settings.set_last_polling(Utc::now())?;

                if Some(u.package_uid()) == shared_state.runtime_settings.applied_package_uid() {
                    info!(
                        "Not applying the update package. Same package has already been installed."
                    );
                    debug!("Moving to Idle state as this update package is already installed.");
                    Ok(StateMachine::Idle(self.into()))
                } else {
                    debug!("Moving to PrepareDownload state to process the update package.");
                    Ok(StateMachine::PrepareDownload(State(PrepareDownload {
                        update_package: u,
                    })))
                }
            }
        }
    }
}

#[test]
fn update_not_available() {
    use super::*;
    use crate::{
        client::tests::{create_mock_server, FakeServer},
        firmware::tests::{create_fake_metadata, FakeDevice},
    };
    use std::fs;
    use tempfile::NamedTempFile;

    let tmpfile = NamedTempFile::new().unwrap();
    let tmpfile = tmpfile.path();
    fs::remove_file(&tmpfile).unwrap();

    let mock = create_mock_server(FakeServer::NoUpdate);

    let settings = Settings::default();
    let runtime_settings = RuntimeSettings::new()
        .load(tmpfile.to_str().unwrap())
        .unwrap();
    let firmware = Metadata::from_path(&create_fake_metadata(FakeDevice::NoUpdate)).unwrap();
    let mut shared_state = SharedState {
        settings,
        runtime_settings,
        firmware,
    };

    let machine = StateMachine::Probe(State(Probe {
        server_address: ServerAddress::Default,
    }))
    .move_to_next_state(&mut shared_state)
    .unwrap();

    mock.assert();

    assert_state!(machine, Idle);
}

#[test]
fn update_available() {
    use super::*;
    use crate::{
        client::tests::{create_mock_server, FakeServer},
        firmware::tests::{create_fake_metadata, FakeDevice},
    };
    use std::fs;
    use tempfile::NamedTempFile;

    let tmpfile = NamedTempFile::new().unwrap();
    let tmpfile = tmpfile.path();
    fs::remove_file(&tmpfile).unwrap();

    let mock = create_mock_server(FakeServer::HasUpdate);

    let settings = Settings::default();
    let runtime_settings = RuntimeSettings::new()
        .load(tmpfile.to_str().unwrap())
        .unwrap();
    let firmware = Metadata::from_path(&create_fake_metadata(FakeDevice::HasUpdate)).unwrap();
    let mut shared_state = SharedState {
        settings,
        runtime_settings,
        firmware,
    };

    let machine = StateMachine::Probe(State(Probe {
        server_address: ServerAddress::Default,
    }))
    .move_to_next_state(&mut shared_state)
    .unwrap();

    mock.assert();

    assert_state!(machine, PrepareDownload);
}

#[test]
fn invalid_hardware() {
    use super::*;
    use crate::{
        client::tests::{create_mock_server, FakeServer},
        firmware::tests::{create_fake_metadata, FakeDevice},
    };
    use std::fs;
    use tempfile::NamedTempFile;

    let tmpfile = NamedTempFile::new().unwrap();
    let tmpfile = tmpfile.path();
    fs::remove_file(&tmpfile).unwrap();

    let mock = create_mock_server(FakeServer::InvalidHardware);

    let settings = Settings::default();
    let runtime_settings = RuntimeSettings::new()
        .load(tmpfile.to_str().unwrap())
        .unwrap();
    let firmware = Metadata::from_path(&create_fake_metadata(FakeDevice::InvalidHardware)).unwrap();
    let mut shared_state = SharedState {
        settings,
        runtime_settings,
        firmware,
    };

    let machine = StateMachine::Probe(State(Probe {
        server_address: ServerAddress::Default,
    }))
    .move_to_next_state(&mut shared_state);

    mock.assert();

    assert!(machine.is_err(), "Did not catch an incompatible hardware");
}

#[test]
fn extra_poll_interval() {
    use super::*;
    use crate::{
        client::tests::{create_mock_server, FakeServer},
        firmware::tests::{create_fake_metadata, FakeDevice},
    };
    use std::fs;
    use tempfile::NamedTempFile;

    let tmpfile = NamedTempFile::new().unwrap();
    let tmpfile = tmpfile.path();
    fs::remove_file(&tmpfile).unwrap();

    let mock = create_mock_server(FakeServer::ExtraPoll);

    let settings = Settings::default();
    let runtime_settings = RuntimeSettings::new()
        .load(tmpfile.to_str().unwrap())
        .unwrap();
    let firmware = Metadata::from_path(&create_fake_metadata(FakeDevice::ExtraPoll)).unwrap();
    let mut shared_state = SharedState {
        settings,
        runtime_settings,
        firmware,
    };

    let machine = StateMachine::Probe(State(Probe {
        server_address: ServerAddress::Default,
    }))
    .move_to_next_state(&mut shared_state)
    .unwrap();

    mock.assert();

    assert_state!(machine, Poll);
}

#[test]
fn skip_same_package_uid() {
    use super::*;
    use crate::{
        client::{
            tests::{create_mock_server, FakeServer},
            ProbeResponse,
        },
        firmware::tests::{create_fake_metadata, FakeDevice},
    };
    use std::fs;
    use tempfile::NamedTempFile;

    let tmpfile = NamedTempFile::new().unwrap();
    let tmpfile = tmpfile.path();
    fs::remove_file(&tmpfile).unwrap();

    let mock = create_mock_server(FakeServer::HasUpdate).expect(2);

    let mut runtime_settings = RuntimeSettings::new()
        .load(tmpfile.to_str().unwrap())
        .unwrap();

    // We first get the package_uid that will be returned so we can
    // use it for the upcoming test.
    //
    // This has been done so we don't need to manually update it every
    // time we change the package payload.
    let probe = Api::new(&Settings::default().network.server_address)
        .probe(
            &RuntimeSettings::default(),
            &Metadata::from_path(&create_fake_metadata(FakeDevice::HasUpdate)).unwrap(),
        )
        .unwrap();

    if let ProbeResponse::Update(u) = probe {
        runtime_settings
            .set_applied_package_uid(&u.package_uid())
            .unwrap();
    }

    let settings = Settings::default();
    let firmware = Metadata::from_path(&create_fake_metadata(FakeDevice::HasUpdate)).unwrap();
    let mut shared_state = SharedState {
        settings,
        runtime_settings,
        firmware,
    };

    let machine = StateMachine::Probe(State(Probe {
        server_address: ServerAddress::Default,
    }))
    .move_to_next_state(&mut shared_state)
    .unwrap();

    mock.assert();

    assert_state!(machine, Idle);
}

#[test]
fn error() {
    use super::*;
    use crate::{
        client::tests::{create_mock_server, FakeServer},
        firmware::tests::{create_fake_metadata, FakeDevice},
    };
    use std::fs;
    use tempfile::NamedTempFile;

    let tmpfile = NamedTempFile::new().unwrap();
    let tmpfile = tmpfile.path();
    fs::remove_file(&tmpfile).unwrap();

    // The server here waits for the second request which includes the
    // retries to succeed.
    let mock = create_mock_server(FakeServer::ErrorOnce);

    let settings = Settings::default();
    let runtime_settings = RuntimeSettings::new()
        .load(tmpfile.to_str().unwrap())
        .unwrap();
    let firmware = Metadata::from_path(&create_fake_metadata(FakeDevice::NoUpdate)).unwrap();
    let mut shared_state = SharedState {
        settings,
        runtime_settings,
        firmware,
    };

    let machine = StateMachine::Probe(State(Probe {
        server_address: ServerAddress::Default,
    }))
    .move_to_next_state(&mut shared_state)
    .unwrap();

    mock.assert();

    assert_state!(machine, Idle);
}
<commit_msg>updatehub: Probe test moved to a test module<commit_after>// Copyright (C) 2018 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

use super::{
    actor, Idle, Poll, PrepareDownload, SharedState, State, StateChangeImpl, StateMachine,
};
use crate::client::Api;
use slog_scope::{debug, error, info};

#[derive(Debug, PartialEq, Clone)]
pub(super) enum ServerAddress {
    Default,
    Custom(String),
}

#[derive(Debug, PartialEq)]
pub(super) struct Probe {
    pub(super) server_address: ServerAddress,
}

create_state_step!(Probe => Idle);
create_state_step!(Probe => Poll);

/// Implements the state change for State<Probe>.
impl StateChangeImpl for State<Probe> {
    fn name(&self) -> &'static str {
        "probe"
    }

    fn handle_trigger_probe(&self) -> actor::probe::Response {
        actor::probe::Response::RequestAccepted(self.name().to_owned())
    }

    fn handle(self, shared_state: &mut SharedState) -> Result<StateMachine, failure::Error> {
        use crate::client::ProbeResponse;
        use chrono::{Duration, Utc};
        use std::thread;

        let server_address = match self.0.server_address.clone() {
            ServerAddress::Default => shared_state.settings.network.server_address.clone(),
            ServerAddress::Custom(s) => s,
        };

        let r = loop {
            let probe = Api::new(&server_address)
                .probe(&shared_state.runtime_settings, &shared_state.firmware);
            if let Err(e) = probe {
                error!("{}", e);
                shared_state.runtime_settings.inc_retries();
                thread::sleep(Duration::seconds(1).to_std().unwrap());
            } else {
                shared_state.runtime_settings.clear_retries();
                break probe?;
            }
        };

        if let ProbeResponse::ExtraPoll(s) = r {
            info!("Delaying the probing as requested by the server.");
            shared_state
                .runtime_settings
                .set_polling_extra_interval(Duration::seconds(s))?;
        };

        match r {
            ProbeResponse::NoUpdate => {
                debug!("Moving to Idle state as no update is available.");

                // Store timestamp of last polling
                shared_state.runtime_settings.set_last_polling(Utc::now())?;
                Ok(StateMachine::Idle(self.into()))
            }

            ProbeResponse::ExtraPoll(_) => {
                debug!("Moving to Poll state due the extra polling interval.");
                Ok(StateMachine::Poll(self.into()))
            }

            ProbeResponse::Update(u) => {
                // Ensure the package is compatible
                u.compatible_with(&shared_state.firmware)?;
                // Store timestamp of last polling
                shared_state.runtime_settings.set_last_polling(Utc::now())?;

                if Some(u.package_uid()) == shared_state.runtime_settings.applied_package_uid() {
                    info!(
                        "Not applying the update package. Same package has already been installed."
                    );
                    debug!("Moving to Idle state as this update package is already installed.");
                    Ok(StateMachine::Idle(self.into()))
                } else {
                    debug!("Moving to PrepareDownload state to process the update package.");
                    Ok(StateMachine::PrepareDownload(State(PrepareDownload {
                        update_package: u,
                    })))
                }
            }
        }
    }
}

#[cfg(tets)]
mod tests {
    use super::*;
    use crate::{
        client::{
            tests::{create_mock_server, FakeServer},
            ProbeResponse,
        },
        firmware::tests::{create_fake_metadata, FakeDevice},
    };
    use std::fs;
    use tempfile::NamedTempFile;

    #[test]
    fn update_not_available() {
        let tmpfile = NamedTempFile::new().unwrap();
        let tmpfile = tmpfile.path();
        fs::remove_file(&tmpfile).unwrap();

        let mock = create_mock_server(FakeServer::NoUpdate);

        let settings = Settings::default();
        let runtime_settings = RuntimeSettings::new()
            .load(tmpfile.to_str().unwrap())
            .unwrap();
        let firmware = Metadata::from_path(&create_fake_metadata(FakeDevice::NoUpdate)).unwrap();
        let mut shared_state = SharedState {
            settings,
            runtime_settings,
            firmware,
        };

        let machine = StateMachine::Probe(State(Probe {
            server_address: ServerAddress::Default,
        }))
        .move_to_next_state(&mut shared_state)
        .unwrap();

        mock.assert();

        assert_state!(machine, Idle);
    }

    #[test]
    fn update_available() {
        let tmpfile = NamedTempFile::new().unwrap();
        let tmpfile = tmpfile.path();
        fs::remove_file(&tmpfile).unwrap();

        let mock = create_mock_server(FakeServer::HasUpdate);

        let settings = Settings::default();
        let runtime_settings = RuntimeSettings::new()
            .load(tmpfile.to_str().unwrap())
            .unwrap();
        let firmware = Metadata::from_path(&create_fake_metadata(FakeDevice::HasUpdate)).unwrap();
        let mut shared_state = SharedState {
            settings,
            runtime_settings,
            firmware,
        };

        let machine = StateMachine::Probe(State(Probe {
            server_address: ServerAddress::Default,
        }))
        .move_to_next_state(&mut shared_state)
        .unwrap();

        mock.assert();

        assert_state!(machine, PrepareDownload);
    }

    #[test]
    fn invalid_hardware() {
        let tmpfile = NamedTempFile::new().unwrap();
        let tmpfile = tmpfile.path();
        fs::remove_file(&tmpfile).unwrap();

        let mock = create_mock_server(FakeServer::InvalidHardware);

        let settings = Settings::default();
        let runtime_settings = RuntimeSettings::new()
            .load(tmpfile.to_str().unwrap())
            .unwrap();
        let firmware =
            Metadata::from_path(&create_fake_metadata(FakeDevice::InvalidHardware)).unwrap();
        let mut shared_state = SharedState {
            settings,
            runtime_settings,
            firmware,
        };

        let machine = StateMachine::Probe(State(Probe {
            server_address: ServerAddress::Default,
        }))
        .move_to_next_state(&mut shared_state);

        mock.assert();

        assert!(machine.is_err(), "Did not catch an incompatible hardware");
    }

    #[test]
    fn extra_poll_interval() {
        let tmpfile = NamedTempFile::new().unwrap();
        let tmpfile = tmpfile.path();
        fs::remove_file(&tmpfile).unwrap();

        let mock = create_mock_server(FakeServer::ExtraPoll);

        let settings = Settings::default();
        let runtime_settings = RuntimeSettings::new()
            .load(tmpfile.to_str().unwrap())
            .unwrap();
        let firmware = Metadata::from_path(&create_fake_metadata(FakeDevice::ExtraPoll)).unwrap();
        let mut shared_state = SharedState {
            settings,
            runtime_settings,
            firmware,
        };

        let machine = StateMachine::Probe(State(Probe {
            server_address: ServerAddress::Default,
        }))
        .move_to_next_state(&mut shared_state)
        .unwrap();

        mock.assert();

        assert_state!(machine, Poll);
    }

    #[test]
    fn skip_same_package_uid() {
        let tmpfile = NamedTempFile::new().unwrap();
        let tmpfile = tmpfile.path();
        fs::remove_file(&tmpfile).unwrap();

        let mock = create_mock_server(FakeServer::HasUpdate).expect(2);

        let mut runtime_settings = RuntimeSettings::new()
            .load(tmpfile.to_str().unwrap())
            .unwrap();

        // We first get the package_uid that will be returned so we can
        // use it for the upcoming test.
        //
        // This has been done so we don't need to manually update it every
        // time we change the package payload.
        let probe = Api::new(&Settings::default().network.server_address)
            .probe(
                &RuntimeSettings::default(),
                &Metadata::from_path(&create_fake_metadata(FakeDevice::HasUpdate)).unwrap(),
            )
            .unwrap();

        if let ProbeResponse::Update(u) = probe {
            runtime_settings
                .set_applied_package_uid(&u.package_uid())
                .unwrap();
        }

        let settings = Settings::default();
        let firmware = Metadata::from_path(&create_fake_metadata(FakeDevice::HasUpdate)).unwrap();
        let mut shared_state = SharedState {
            settings,
            runtime_settings,
            firmware,
        };

        let machine = StateMachine::Probe(State(Probe {
            server_address: ServerAddress::Default,
        }))
        .move_to_next_state(&mut shared_state)
        .unwrap();

        mock.assert();

        assert_state!(machine, Idle);
    }

    #[test]
    fn error() {
        let tmpfile = NamedTempFile::new().unwrap();
        let tmpfile = tmpfile.path();
        fs::remove_file(&tmpfile).unwrap();

        // The server here waits for the second request which includes the
        // retries to succeed.
        let mock = create_mock_server(FakeServer::ErrorOnce);

        let settings = Settings::default();
        let runtime_settings = RuntimeSettings::new()
            .load(tmpfile.to_str().unwrap())
            .unwrap();
        let firmware = Metadata::from_path(&create_fake_metadata(FakeDevice::NoUpdate)).unwrap();
        let mut shared_state = SharedState {
            settings,
            runtime_settings,
            firmware,
        };

        let machine = StateMachine::Probe(State(Probe {
            server_address: ServerAddress::Default,
        }))
        .move_to_next_state(&mut shared_state)
        .unwrap();

        mock.assert();

        assert_state!(machine, Idle);
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2017-2018, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use criterion::*;
use crate::partition::*;
use crate::partition::BlockSize::*;
use crate::plane::*;
use rand::{ChaChaRng, Rng, SeedableRng};
use rav1e::me;
use rav1e::Pixel;

fn fill_plane<T: Pixel>(ra: &mut ChaChaRng, plane: &mut Plane<T>) {
  let stride = plane.cfg.stride;
  for row in plane.data.chunks_mut(stride) {
    for pixel in row {
      let v: u8 = ra.gen();
      *pixel = T::cast_from(v);
    }
  }
}

fn new_plane<T: Pixel>(ra: &mut ChaChaRng, width: usize, height: usize) -> Plane<T> {
  let mut p = Plane::new(width, height, 0, 0, 128 + 8, 128 + 8);

  fill_plane(ra, &mut p);

  p
}

fn bench_get_sad(b: &mut Bencher, bs: &BlockSize) {
  let mut ra = ChaChaRng::from_seed([0; 32]);
  let bsw = bs.width();
  let bsh = bs.height();
  let w = 640;
  let h = 480;
  let bit_depth = 10;
  let input_plane = new_plane::<u16>(&mut ra, w, h);
  let rec_plane = new_plane::<u16>(&mut ra, w, h);
  let po = PlaneOffset { x: 0, y: 0 };

  let plane_org = input_plane.slice(po);
  let plane_ref = rec_plane.slice(po);

  b.iter(|| {
    let _ =
      black_box(me::get_sad(&plane_org, &plane_ref, bsh, bsw, bit_depth));
  })
}

pub fn get_sad(c: &mut Criterion) {
  let blocks = vec![
    BLOCK_4X4,
    BLOCK_4X8,
    BLOCK_8X4,
    BLOCK_8X8,
    BLOCK_8X16,
    BLOCK_16X8,
    BLOCK_16X16,
    BLOCK_16X32,
    BLOCK_32X16,
    BLOCK_32X32,
    BLOCK_32X64,
    BLOCK_64X32,
    BLOCK_64X64,
    BLOCK_64X128,
    BLOCK_128X64,
    BLOCK_128X128,
    BLOCK_4X16,
    BLOCK_16X4,
    BLOCK_8X32,
    BLOCK_32X8,
    BLOCK_16X64,
    BLOCK_64X16,
  ];

  c.bench_function_over_inputs("get_sad", bench_get_sad, blocks);
}
<commit_msg>Bench the sad function for 8 bit and 10 bit<commit_after>// Copyright (c) 2017-2018, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use criterion::*;
use crate::partition::*;
use crate::partition::BlockSize::*;
use crate::plane::*;
use rand::{ChaChaRng, Rng, SeedableRng};
use rav1e::me;
use rav1e::Pixel;

fn fill_plane<T: Pixel>(ra: &mut ChaChaRng, plane: &mut Plane<T>) {
  let stride = plane.cfg.stride;
  for row in plane.data.chunks_mut(stride) {
    for pixel in row {
      let v: u8 = ra.gen();
      *pixel = T::cast_from(v);
    }
  }
}

fn new_plane<T: Pixel>(ra: &mut ChaChaRng, width: usize, height: usize) -> Plane<T> {
  let mut p = Plane::new(width, height, 0, 0, 128 + 8, 128 + 8);

  fill_plane(ra, &mut p);

  p
}

fn bench_get_sad(b: &mut Bencher, &(bs, bit_depth): &(BlockSize, usize)) {
  let mut ra = ChaChaRng::from_seed([0; 32]);
  let bsw = bs.width();
  let bsh = bs.height();
  let w = 640;
  let h = 480;
  let input_plane = new_plane::<u16>(&mut ra, w, h);
  let rec_plane = new_plane::<u16>(&mut ra, w, h);
  let po = PlaneOffset { x: 0, y: 0 };

  let plane_org = input_plane.slice(po);
  let plane_ref = rec_plane.slice(po);

  b.iter(|| {
    let _ =
      black_box(me::get_sad(&plane_org, &plane_ref, bsh, bsw, bit_depth));
  })
}

pub fn get_sad(c: &mut Criterion) {
  let blocks = vec![
    (BLOCK_4X4, 8),
    (BLOCK_4X8, 8),
    (BLOCK_8X4, 8),
    (BLOCK_8X8, 8),
    (BLOCK_8X16, 8),
    (BLOCK_16X8, 8),
    (BLOCK_16X16, 8),
    (BLOCK_16X32, 8),
    (BLOCK_32X16, 8),
    (BLOCK_32X32, 8),
    (BLOCK_32X64, 8),
    (BLOCK_64X32, 8),
    (BLOCK_64X64, 8),
    (BLOCK_64X128, 8),
    (BLOCK_128X64, 8),
    (BLOCK_128X128, 8),
    (BLOCK_4X16, 8),
    (BLOCK_16X4, 8),
    (BLOCK_8X32, 8),
    (BLOCK_32X8, 8),
    (BLOCK_16X64, 8),
    (BLOCK_64X16, 8),

    (BLOCK_4X4, 10),
    (BLOCK_4X8, 10),
    (BLOCK_8X4, 10),
    (BLOCK_8X8, 10),
    (BLOCK_8X16, 10),
    (BLOCK_16X8, 10),
    (BLOCK_16X16, 10),
    (BLOCK_16X32, 10),
    (BLOCK_32X16, 10),
    (BLOCK_32X32, 10),
    (BLOCK_32X64, 10),
    (BLOCK_64X32, 10),
    (BLOCK_64X64, 10),
    (BLOCK_64X128, 10),
    (BLOCK_128X64, 10),
    (BLOCK_128X128, 10),
    (BLOCK_4X16, 10),
    (BLOCK_16X4, 10),
    (BLOCK_8X32, 10),
    (BLOCK_32X8, 10),
    (BLOCK_16X64, 10),
    (BLOCK_64X16, 10)
  ];

  c.bench_function_over_inputs("get_sad", bench_get_sad, blocks);
}
<|endoftext|>


<commit_before>use color::PuyoColor;
use field::{self, BitField, FieldHeight};

pub struct CoreField {
    field: BitField,
    height: [i16; 8],
}

impl CoreField {
    pub fn new() -> CoreField {
        CoreField {
            field: BitField::new(),
            height: [0; 8],
        }
    }

    pub fn from_str(s: &str) -> CoreField {
        let mut cf = CoreField {
            field: BitField::from_str(s),
            height: [0; 8],
        };

        for x in 1 .. field::WIDTH + 1 {
            for y in 1 .. 15 {
                if cf.is_empty(x, y) {
                    cf.height[x] = (y - 1) as i16;
                    break
                }
            }
        }

        cf
    }

    pub fn color(&self, x: usize, y: usize) -> PuyoColor {
        self.field.color(x, y)
    }

    pub fn height(&self, x: usize) -> usize {
        self.height[x] as usize
    }

    pub fn is_color(&self, x: usize, y: usize, c: PuyoColor) -> bool {
        self.field.is_color(x, y, c)
    }

    pub fn is_normal_color(&self, x: usize, y: usize) -> bool {
        self.field.is_normal_color(x, y)
    }

    pub fn is_empty(&self, x: usize, y: usize) -> bool {
        self.field.is_empty(x, y)
    }

    pub fn field(&self) -> &BitField {
        &self.field
    }
}

impl FieldHeight for CoreField {
    fn height(&self, x: usize) -> usize {
        CoreField::height(self, x)
    }
}

#[cfg(test)]
mod tests {
    use super::CoreField;
    use color::PuyoColor;
    use field;

    #[test]
    fn test_constructor() {
        let cf = CoreField::new();

        for x in 1 .. field::WIDTH + 1 {
            for y in 1 .. field::HEIGHT + 1 {
                assert!(cf.is_empty(x, y));
                assert!(cf.is_color(x, y, PuyoColor::EMPTY));
                assert_eq!(PuyoColor::EMPTY, cf.color(x, y));
            }
        }

        for x in 1 .. field::WIDTH + 1 {
            assert_eq!(0, cf.height(x));
        }
    }

    #[test]
    fn test_from_str() {
        let cf = CoreField::from_str(concat!(
            "R.....",
            "RRRB.."));

        assert_eq!(PuyoColor::WALL, cf.color(0, 1));
        assert_eq!(PuyoColor::RED, cf.color(1, 1));
        assert_eq!(PuyoColor::RED, cf.color(2, 1));
        assert_eq!(PuyoColor::RED, cf.color(3, 1));
        assert_eq!(PuyoColor::BLUE, cf.color(4, 1));
        assert_eq!(PuyoColor::EMPTY, cf.color(5, 1));
        assert_eq!(PuyoColor::EMPTY, cf.color(6, 1));
        assert_eq!(PuyoColor::WALL, cf.color(7, 1));

        assert_eq!(2, cf.height(1));
        assert_eq!(1, cf.height(2));
        assert_eq!(1, cf.height(3));
        assert_eq!(1, cf.height(4));
        assert_eq!(0, cf.height(5));
        assert_eq!(0, cf.height(6));
    }
}
<commit_msg>Add several BitField adaptor to CoreField<commit_after>use color::PuyoColor;
use field::{self, BitField, FieldHeight};

pub struct CoreField {
    field: BitField,
    height: [i16; 8],
}

impl CoreField {
    pub fn new() -> CoreField {
        CoreField {
            field: BitField::new(),
            height: [0; 8],
        }
    }

    pub fn from_str(s: &str) -> CoreField {
        let mut cf = CoreField {
            field: BitField::from_str(s),
            height: [0; 8],
        };

        for x in 1 .. field::WIDTH + 1 {
            for y in 1 .. 15 {
                if cf.is_empty(x, y) {
                    cf.height[x] = (y - 1) as i16;
                    break
                }
            }
        }

        cf
    }

    pub fn color(&self, x: usize, y: usize) -> PuyoColor {
        self.field.color(x, y)
    }

    pub fn height(&self, x: usize) -> usize {
        self.height[x] as usize
    }

    pub fn is_color(&self, x: usize, y: usize, c: PuyoColor) -> bool {
        self.field.is_color(x, y, c)
    }

    pub fn is_normal_color(&self, x: usize, y: usize) -> bool {
        self.field.is_normal_color(x, y)
    }

    pub fn is_empty(&self, x: usize, y: usize) -> bool {
        self.field.is_empty(x, y)
    }

    pub fn field(&self) -> &BitField {
        &self.field
    }

    pub fn is_connected(&self, x: usize, y: usize) -> bool {
        self.field.is_connected(x, y)
    }

    pub fn count_connected_max4(&self, x: usize, y: usize) -> usize {
        self.field.count_connected_max4(x, y)
    }

    pub fn count_connected_max4_with_color(&self, x: usize, y: usize, c: PuyoColor) -> usize {
        self.field.count_connected_max4_with_color(x, y, c)
    }
}

impl FieldHeight for CoreField {
    fn height(&self, x: usize) -> usize {
        CoreField::height(self, x)
    }
}

#[cfg(test)]
mod tests {
    use super::CoreField;
    use color::PuyoColor;
    use field;

    #[test]
    fn test_constructor() {
        let cf = CoreField::new();

        for x in 1 .. field::WIDTH + 1 {
            for y in 1 .. field::HEIGHT + 1 {
                assert!(cf.is_empty(x, y));
                assert!(cf.is_color(x, y, PuyoColor::EMPTY));
                assert_eq!(PuyoColor::EMPTY, cf.color(x, y));
            }
        }

        for x in 1 .. field::WIDTH + 1 {
            assert_eq!(0, cf.height(x));
        }
    }

    #[test]
    fn test_from_str() {
        let cf = CoreField::from_str(concat!(
            "R.....",
            "RRRB.."));

        assert_eq!(PuyoColor::WALL, cf.color(0, 1));
        assert_eq!(PuyoColor::RED, cf.color(1, 1));
        assert_eq!(PuyoColor::RED, cf.color(2, 1));
        assert_eq!(PuyoColor::RED, cf.color(3, 1));
        assert_eq!(PuyoColor::BLUE, cf.color(4, 1));
        assert_eq!(PuyoColor::EMPTY, cf.color(5, 1));
        assert_eq!(PuyoColor::EMPTY, cf.color(6, 1));
        assert_eq!(PuyoColor::WALL, cf.color(7, 1));

        assert_eq!(2, cf.height(1));
        assert_eq!(1, cf.height(2));
        assert_eq!(1, cf.height(3));
        assert_eq!(1, cf.height(4));
        assert_eq!(0, cf.height(5));
        assert_eq!(0, cf.height(6));
    }
}
<|endoftext|>


<commit_before>//! Convert images to black and white.

use cast;
use image::Rgba;
use std::collections::HashMap;

use errors::*;
use pixmap::{Pixel, Pixmap};
#[cfg(test)]
use test_util::{rgba_hex, test_images};

/// Different kinds of colors we might find in an image.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum ColorType {
    /// This color is transparent.
    Transparent,
    /// This color appears to be a shadow color which we should treat as
    /// transparent to facilitate letter separation.
    Shadow,
    /// This color is opaque, and should be used for letter recognition.
    Opaque,
}

impl From<ColorType> for usize {
    /// Map `ColorType` to something we can use to index an array.
    fn from(ct: ColorType) -> usize {
        match ct {
            ColorType::Transparent => 0,
            ColorType::Shadow => 1,
            ColorType::Opaque => 2,
        }
    }
}

/// Information about the `ColorType` of pixels adjacent to a given color.
#[derive(Debug)]
struct AdjacentPixelInfo {
    data: [usize; 3],
}

impl AdjacentPixelInfo {
    fn new() -> AdjacentPixelInfo {
        AdjacentPixelInfo { data: [0; 3] }
    }

    fn count(&self, ct: ColorType) -> usize {
        self.data[usize::from(ct)]
    }

    fn incr_count(&mut self, ct: ColorType) {
        self.data[usize::from(ct)] += 1;
    }

    fn total(&self) -> usize {
        self.data[0] + self.data[1] + self.data[2]
    }

    fn fraction_adj_to(&self, ct: ColorType) -> f64 {
        cast::f64(self.count(ct)) / cast::f64(self.total())
    }

    fn looks_like_opaque_inside_shadow(&self) -> bool {
        self.fraction_adj_to(ColorType::Opaque) > 0.95
    }

    fn looks_like_shadow(&self) -> bool {
        self.fraction_adj_to(ColorType::Opaque) > 0.33 &&
            self.fraction_adj_to(ColorType::Transparent) > 0.33
    }
}

/// Classify the colors in an image as transparent or non-transparent.
fn classify_colors(image: &Pixmap) -> Result<HashMap<Rgba<u8>, ColorType>> {
    // First divide colors into transparent and opaque based on alpha.  We
    // ensure that we always have an entry for the default color to
    // simplify the logic later on.
    let mut classification = HashMap::new();
    classification.insert(Rgba::default_color(), ColorType::Transparent);
    for px in image.pixels() {
        if px.is_transparent() {
            classification.entry(px).or_insert(ColorType::Transparent);
        } else {
            classification.entry(px).or_insert(ColorType::Opaque);
        }
    }
    debug!("color classification (initial): {:?}", &classification);

    // Calculate which colors are adjacent to which color classifications.
    // We'll use this to detect "shadow" colors.
    let mut adjacent = HashMap::new();
    for c in classification.keys() {
        if !c.is_transparent() {
            adjacent.insert(*c, AdjacentPixelInfo::new());
        }
    }
    for (x, y, px) in image.enumerate_pixels() {
        // Don't compute adjacency info for transparent pxiels.
        if px.is_transparent() {
            continue;
        }
        // Look at the 3x3 grid around this pixel.
        for &dy in &[-1, 0, 1] {
            for &dx in &[-1, 0, 1] {
                // Don't compare to ourselves.
                if dx == 0 && dy == 0 {
                    continue;
                }

                // Get our neighboring pixel, if it's in bounds.
                let px_adj =
                    image.get_default(cast::isize(x)? + dx, cast::isize(y)? + dy);

                // Ignore adjacent pixels of the same color.
                if px == px_adj {
                    continue;
                }

                // Classify our neighboring color an increment our counter.
                let ct_adj = *classification.get(&px_adj)
                    .expect("unknown classification");
                adjacent.get_mut(&px)
                    .expect("unknown adjacent color")
                    .incr_count(ct_adj);
            }
        }
    }
    debug!("color adjacency: {:?}", &adjacent);

    // Check to see whether we have any opaque colors which seem to be
    // _surrounded_ by a shadow.
    let total_adj: usize = adjacent.values().map(|adj| adj.total()).sum();
    let mut have_opaque_inside_shadow = false;
    for adj in adjacent.values() {
        // Only check colors which make up a reasonable fraction of our
        // total.
        if adj.total() >= total_adj / 4 && adj.looks_like_opaque_inside_shadow() {
            have_opaque_inside_shadow = true;
        }
    }

    // If we have colors _surrounded_ by shadows, look for the shadow
    // colors.
    if have_opaque_inside_shadow {
        for (c, adj) in adjacent.iter() {
            if adj.looks_like_shadow() {
                classification.insert(c.to_owned(), ColorType::Shadow);
            }
        }
    }

    debug!("color classification (final): {:?}", &classification);
    Ok(classification)
}

#[test]
fn classify_colors_as_transparent_and_opaque() {
    //use env_logger;
    //env_logger::init().unwrap();

    let images = test_images().unwrap();
    let colors = classify_colors(&Pixmap::from(images[0].clone())).unwrap();
    assert_eq!(colors.len(), 4);
    assert_eq!(*colors.get(&rgba_hex(0x00000000)).unwrap(), ColorType::Transparent);
    assert_eq!(*colors.get(&rgba_hex(0x000000ff)).unwrap(), ColorType::Shadow);
    assert_eq!(*colors.get(&rgba_hex(0x999999ff)).unwrap(), ColorType::Opaque);
    assert_eq!(*colors.get(&rgba_hex(0xf0f0f0ff)).unwrap(), ColorType::Opaque);
}

/// Reduce an image to two "colors": Areas inside letters, and transparent
/// background.
pub fn binarize(pixmap: &Pixmap) -> Result<Pixmap<bool>> {
    let classification = classify_colors(pixmap)?;
    Ok(pixmap.map(|px| {
        let ct = classification.get(&px).expect("Color wasn't classified");
        ct == &ColorType::Opaque
    }))
}

#[test]
fn binarize_reduces_to_transparent_and_black() {
    use env_logger;
    env_logger::init().unwrap();

    let images = test_images().unwrap();
    let bitmap = binarize(&Pixmap::from(images[0].clone())).unwrap();

    debug_pixmap!(&bitmap, "binarize.png");
}
<commit_msg>Use all_neighbors iterator to flood fill<commit_after>//! Convert images to black and white.

use cast;
use image::Rgba;
use std::collections::HashMap;

use errors::*;
use pixmap::{Pixel, Pixmap};
#[cfg(test)]
use test_util::{rgba_hex, test_images};

/// Different kinds of colors we might find in an image.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum ColorType {
    /// This color is transparent.
    Transparent,
    /// This color appears to be a shadow color which we should treat as
    /// transparent to facilitate letter separation.
    Shadow,
    /// This color is opaque, and should be used for letter recognition.
    Opaque,
}

impl From<ColorType> for usize {
    /// Map `ColorType` to something we can use to index an array.
    fn from(ct: ColorType) -> usize {
        match ct {
            ColorType::Transparent => 0,
            ColorType::Shadow => 1,
            ColorType::Opaque => 2,
        }
    }
}

/// Information about the `ColorType` of pixels adjacent to a given color.
#[derive(Debug)]
struct AdjacentPixelInfo {
    data: [usize; 3],
}

impl AdjacentPixelInfo {
    fn new() -> AdjacentPixelInfo {
        AdjacentPixelInfo { data: [0; 3] }
    }

    fn count(&self, ct: ColorType) -> usize {
        self.data[usize::from(ct)]
    }

    fn incr_count(&mut self, ct: ColorType) {
        self.data[usize::from(ct)] += 1;
    }

    fn total(&self) -> usize {
        self.data[0] + self.data[1] + self.data[2]
    }

    fn fraction_adj_to(&self, ct: ColorType) -> f64 {
        cast::f64(self.count(ct)) / cast::f64(self.total())
    }

    fn looks_like_opaque_inside_shadow(&self) -> bool {
        self.fraction_adj_to(ColorType::Opaque) > 0.95
    }

    fn looks_like_shadow(&self) -> bool {
        self.fraction_adj_to(ColorType::Opaque) > 0.33 &&
            self.fraction_adj_to(ColorType::Transparent) > 0.33
    }
}

/// Classify the colors in an image as transparent or non-transparent.
fn classify_colors(image: &Pixmap) -> Result<HashMap<Rgba<u8>, ColorType>> {
    // First divide colors into transparent and opaque based on alpha.  We
    // ensure that we always have an entry for the default color to
    // simplify the logic later on.
    let mut classification = HashMap::new();
    classification.insert(Rgba::default_color(), ColorType::Transparent);
    for px in image.pixels() {
        if px.is_transparent() {
            classification.entry(px).or_insert(ColorType::Transparent);
        } else {
            classification.entry(px).or_insert(ColorType::Opaque);
        }
    }
    debug!("color classification (initial): {:?}", &classification);

    // Calculate which colors are adjacent to which color classifications.
    // We'll use this to detect "shadow" colors.
    let mut adjacent = HashMap::new();
    for c in classification.keys() {
        if !c.is_transparent() {
            adjacent.insert(*c, AdjacentPixelInfo::new());
        }
    }
    for (x, y, px) in image.enumerate_pixels() {
        // Don't compute adjacency info for transparent pxiels.
        if px.is_transparent() {
            continue;
        }
        // Look at the 3x3 grid around this pixel.
        for (x2, y2) in image.all_neighbors(x, y) {
            // Get our neighboring pixel, if it's in bounds.
            let px_adj = image.get_default(x2, y2);

            // Ignore adjacent pixels of the same color.
            if px == px_adj {
                continue;
            }

            // Classify our neighboring color an increment our counter.
            let ct_adj = *classification.get(&px_adj)
                .expect("unknown classification");
            adjacent.get_mut(&px)
                .expect("unknown adjacent color")
                .incr_count(ct_adj);
        }
    }
    debug!("color adjacency: {:?}", &adjacent);

    // Check to see whether we have any opaque colors which seem to be
    // _surrounded_ by a shadow.
    let total_adj: usize = adjacent.values().map(|adj| adj.total()).sum();
    let mut have_opaque_inside_shadow = false;
    for adj in adjacent.values() {
        // Only check colors which make up a reasonable fraction of our
        // total.
        if adj.total() >= total_adj / 4 && adj.looks_like_opaque_inside_shadow() {
            have_opaque_inside_shadow = true;
        }
    }

    // If we have colors _surrounded_ by shadows, look for the shadow
    // colors.
    if have_opaque_inside_shadow {
        for (c, adj) in adjacent.iter() {
            if adj.looks_like_shadow() {
                classification.insert(c.to_owned(), ColorType::Shadow);
            }
        }
    }

    debug!("color classification (final): {:?}", &classification);
    Ok(classification)
}

#[test]
fn classify_colors_as_transparent_and_opaque() {
    //use env_logger;
    //env_logger::init().unwrap();

    let images = test_images().unwrap();
    let colors = classify_colors(&Pixmap::from(images[0].clone())).unwrap();
    assert_eq!(colors.len(), 4);
    assert_eq!(*colors.get(&rgba_hex(0x00000000)).unwrap(), ColorType::Transparent);
    assert_eq!(*colors.get(&rgba_hex(0x000000ff)).unwrap(), ColorType::Shadow);
    assert_eq!(*colors.get(&rgba_hex(0x999999ff)).unwrap(), ColorType::Opaque);
    assert_eq!(*colors.get(&rgba_hex(0xf0f0f0ff)).unwrap(), ColorType::Opaque);
}

/// Reduce an image to two "colors": Areas inside letters, and transparent
/// background.
pub fn binarize(pixmap: &Pixmap) -> Result<Pixmap<bool>> {
    let classification = classify_colors(pixmap)?;
    Ok(pixmap.map(|px| {
        let ct = classification.get(&px).expect("Color wasn't classified");
        ct == &ColorType::Opaque
    }))
}

#[test]
fn binarize_reduces_to_transparent_and_black() {
    use env_logger;
    env_logger::init().unwrap();

    let images = test_images().unwrap();
    let bitmap = binarize(&Pixmap::from(images[0].clone())).unwrap();

    debug_pixmap!(&bitmap, "binarize.png");
}
<|endoftext|>


<commit_before>//! ReQL is the RethinkDB query language. It offers a very powerful and
//! convenient way to manipulate JSON documents.
//!
//! # Start the server #
//!
//! ## Linux and OS X ##
//!
//! Start the server from a terminal window.
//!
//! ```bash
//! $ rethinkdb
//! ```
//!
//! ## Windows ##
//!
//! Start the server from the Windows command prompt.
//!
//! ```bash
//! C:\Path\To\RethinkDB\>rethinkdb.exe
//! ```
//!
//! # Import the driver #
//!
//! First, make sure you have `protoc` installed and in your `PATH`. See
//! [`prost-build` documentation](https://docs.rs/prost-build/0.7.0/prost_build/#sourcing-protoc)
//! for more details if it fails to compile.
//!
//! Add this crate (`reql`) and the `futures` crate to your dependencies in `Cargo.toml`.
//! You may also need to depend on a crate that provides an async implementation of a
//! `TcpStream`, like `async-std`, for example. You can use this crate
//! with any `TcpStream` that implements `AsyncRead` and `AsyncWrite` from
//! the futures crate. The implementation also needs to implement the same
//! traits for `&TcpStream` because we use unmutable references of the
//! connection to take advantage of RethinkDB's connection pipelining abilities.
//!
//! Now import the RethinkDB driver:
//!
//! ```
//! use reql::r;
//! ```
//!
//! You can now access RethinkDB commands through the [`r` struct](r).
//!
//! # Open a connection #
//!
//! When you first start RethinkDB, the server opens a port for the client
//! drivers (`28015` by default). Let's open a connection:
//!
//! ```
//! use async_std::net::TcpStream;
//! use reql::{r, DEFAULT_ADDR};
//!
//! # async fn connect() -> reql::Result<()> {
//! let stream = TcpStream::connect(DEFAULT_ADDR).await?;
//! let connection = r.connection(stream).await?;
//! # Ok(())
//! # }
//! ```
//!
//! The variable `connection` is now initialized and we can run queries.
//!
//! # Send a query to the database #
//!
//! ```
//! use async_std::net::TcpStream;
//! use futures::TryStreamExt;
//! use reql::{r, DEFAULT_ADDR};
//!
//! # async fn connect() -> reql::Result<()> {
//! let stream = TcpStream::connect(DEFAULT_ADDR).await?;
//! let conn = r.connection(stream).await?;
//! let mut query = r.expr("Hello world!").run(&conn);
//! # let _: Option<String> = query.try_next().await?;
//! # Ok(())
//! # }
//! ```
//!
//! [See the `r` struct for more available commands](r)

pub mod cmd;
mod err;
mod proto;

use cmd::run::Response;
use dashmap::DashMap;
use futures::channel::mpsc::Sender;
use futures::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use futures::lock::Mutex;
use ql2::term::TermType;
use std::borrow::Cow;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};

pub use err::*;
pub use proto::Query;

/// Default address of the RethinkDB server
pub const DEFAULT_ADDR: (&str, u16) = ("localhost", 28015);

/// Custom result returned by various ReQL commands
pub type Result<T> = std::result::Result<T, Error>;

/// The connection object returned by `r.connection()`
#[derive(Debug)]
pub struct Connection<'a, T> {
    db: Cow<'a, str>,
    stream: T,
    token: AtomicU64,
    broken: AtomicBool,
    change_feed: AtomicBool,
    buffer: usize,
    senders: DashMap<u64, Sender<Result<Response>>>,
    locker: Mutex<()>,
}

impl<T> Connection<'_, T> {
    /// Convert the connection into an instance you can move around
    pub fn into_owned(self) -> Connection<'static, T> {
        Connection {
            db: Cow::from(self.db.into_owned()),
            stream: self.stream,
            token: self.token,
            broken: self.broken,
            change_feed: self.change_feed,
            buffer: self.buffer,
            senders: self.senders,
            locker: self.locker,
        }
    }

    fn mark_broken(&self) {
        self.broken.store(true, Ordering::SeqCst);
    }

    fn broken(&self) -> Result<()> {
        if self.broken.load(Ordering::SeqCst) {
            return Err(err::Client::ConnectionBroken.into());
        }
        Ok(())
    }

    fn mark_change_feed(&self) {
        self.change_feed.store(true, Ordering::SeqCst);
    }

    fn unmark_change_feed(&self) {
        self.change_feed.store(false, Ordering::SeqCst);
    }

    fn change_feed(&self) -> Result<()> {
        if self.change_feed.load(Ordering::SeqCst) {
            return Err(err::Client::ConnectionLocked.into());
        }
        Ok(())
    }
}

/// A generic, runtime independent, TcpStream
///
/// This crate can be used with any runtime that implements
/// the `AsyncRead` and the `AsyncWrite` traits from the
/// futures crate
pub trait TcpStream<'a>: AsyncReadExt + AsyncWriteExt
where
    Self: Unpin + 'a,
    &'a Self: AsyncRead + AsyncWrite,
{
}

impl<'a, T: AsyncRead + AsyncWrite + ?Sized> TcpStream<'a> for T
where
    Self: Unpin + 'a,
    &'a Self: AsyncRead + AsyncWrite,
{
}

/// The top-level ReQL namespace
///
/// # Example
///
/// Set up your top-level namespace.
///
/// ```
/// use reql::r;
/// ```
#[allow(non_camel_case_types)]
pub struct r;

impl<'a> r {
    /// Create a new connection to the database server
    ///
    /// # Example
    ///
    /// Open a connection using the default host and port, specifying the default database.
    ///
    /// ```
    /// use async_std::net::TcpStream;
    /// use reql::cmd::connection::Options;
    /// use reql::{r, DEFAULT_ADDR};
    ///
    /// # async fn connect() -> reql::Result<()> {
    /// let stream = TcpStream::connect(DEFAULT_ADDR).await?;
    /// let conn = r.connection((stream, Options::new().db("marvel"))).await?;
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// Read more about this command [cmd::connection]
    pub async fn connection<A, T>(self, options: A) -> Result<Connection<'a, T>>
    where
        A: cmd::connection::Arg<'a, T>,
        T: TcpStream<'a>,
        &'a T: AsyncRead + AsyncWrite,
    {
        cmd::connection::new(options.into()).await
    }

    pub fn db_create<T>(self, arg: T) -> Query
    where
        T: cmd::db_create::Arg,
    {
        arg.into_query()
    }

    pub fn db_drop<T>(self, arg: T) -> Query
    where
        T: cmd::db_drop::Arg,
    {
        arg.into_query()
    }

    pub fn db_list(self) -> Query {
        Query::new(TermType::DbList)
    }

    /// Reference a database
    ///
    /// The `db` command is optional. If it is not present in a query, the
    /// query will run against the default database for the connection,
    /// specified in the `db` argument to [r::connection].
    ///
    /// # Examples
    ///
    /// Explicitly specify a database for a query.
    ///
    /// ```
    /// # use async_std::net::TcpStream;
    /// # use futures::TryStreamExt;
    /// # use reql::{r, DEFAULT_ADDR};
    /// # async fn example() -> reql::Result<()> {
    /// # let stream = TcpStream::connect(DEFAULT_ADDR).await?;
    /// # let conn = r.connection(stream).await?;
    /// let mut query = r.db("heroes").table("marvel").run(&conn);
    /// # let _: Option<String> = query.try_next().await?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn db<T>(self, arg: T) -> Query
    where
        T: cmd::db::Arg,
    {
        arg.into_query()
    }

    pub fn table_create<T>(self, arg: T) -> Query
    where
        T: cmd::table_create::Arg,
    {
        arg.into_query()
    }

    pub fn table<T>(self, arg: T) -> Query
    where
        T: cmd::table::Arg,
    {
        arg.into_query()
    }

    pub fn map<T>(self, arg: T) -> Query
    where
        T: cmd::map::Arg,
    {
        arg.into_query()
    }

    pub fn union<T>(self, arg: T) -> Query
    where
        T: cmd::union::Arg,
    {
        arg.into_query()
    }

    pub fn group<T>(self, arg: T) -> Query
    where
        T: cmd::group::Arg,
    {
        arg.into_query()
    }

    pub fn reduce<T>(self, arg: T) -> Query
    where
        T: cmd::reduce::Arg,
    {
        arg.into_query()
    }

    pub fn count<T>(self, arg: T) -> Query
    where
        T: cmd::count::Arg,
    {
        arg.into_query()
    }

    pub fn sum<T>(self, arg: T) -> Query
    where
        T: cmd::sum::Arg,
    {
        arg.into_query()
    }

    pub fn avg<T>(self, arg: T) -> Query
    where
        T: cmd::avg::Arg,
    {
        arg.into_query()
    }

    pub fn min<T>(self, arg: T) -> Query
    where
        T: cmd::min::Arg,
    {
        arg.into_query()
    }

    pub fn max<T>(self, arg: T) -> Query
    where
        T: cmd::max::Arg,
    {
        arg.into_query()
    }

    pub fn distinct<T>(self, arg: T) -> Query
    where
        T: cmd::distinct::Arg,
    {
        arg.into_query()
    }

    pub fn contains<T>(self, arg: T) -> Query
    where
        T: cmd::contains::Arg,
    {
        arg.into_query()
    }

    pub fn literal<T>(self, arg: T) -> Query
    where
        T: cmd::literal::Arg,
    {
        arg.into_query()
    }

    pub fn object<T>(self, arg: T) -> Query
    where
        T: cmd::object::Arg,
    {
        arg.into_query()
    }

    pub fn random<T>(self, arg: T) -> Query
    where
        T: cmd::random::Arg,
    {
        arg.into_query()
    }

    pub fn round<T>(self, arg: T) -> Query
    where
        T: cmd::round::Arg,
    {
        arg.into_query()
    }

    pub fn ceil<T>(self, arg: T) -> Query
    where
        T: cmd::ceil::Arg,
    {
        arg.into_query()
    }

    pub fn floor<T>(self, arg: T) -> Query
    where
        T: cmd::floor::Arg,
    {
        arg.into_query()
    }

    pub fn now(self) -> Query {
        Query::new(TermType::Now)
    }

    pub fn time<T>(self, arg: T) -> Query
    where
        T: cmd::time::Arg,
    {
        arg.into_query()
    }

    pub fn epoch_time<T>(self, arg: T) -> Query
    where
        T: cmd::epoch_time::Arg,
    {
        arg.into_query()
    }

    pub fn iso8601<T>(self, arg: T) -> Query
    where
        T: cmd::iso8601::Arg,
    {
        arg.into_query()
    }

    pub fn r#do<T>(self, arg: T) -> Query
    where
        T: cmd::r#do::Arg,
    {
        arg.into_query()
    }

    pub fn branch<T>(self, arg: T) -> Query
    where
        T: cmd::branch::Arg,
    {
        arg.into_query()
    }

    pub fn range<T>(self, arg: T) -> Query
    where
        T: cmd::range::Arg,
    {
        arg.into_query()
    }

    pub fn error<T>(self, arg: T) -> Query
    where
        T: cmd::error::Arg,
    {
        arg.into_query()
    }

    pub fn expr<T>(self, arg: T) -> Query
    where
        T: cmd::expr::Arg,
    {
        arg.into_query()
    }

    pub fn js<T>(self, arg: T) -> Query
    where
        T: cmd::js::Arg,
    {
        arg.into_query()
    }

    pub fn info<T>(self, arg: T) -> Query
    where
        T: cmd::info::Arg,
    {
        arg.into_query()
    }

    pub fn json<T>(self, arg: T) -> Query
    where
        T: cmd::json::Arg,
    {
        arg.into_query()
    }

    pub fn http<T>(self, arg: T) -> Query
    where
        T: cmd::http::Arg,
    {
        arg.into_query()
    }

    pub fn uuid<T>(self, arg: T) -> Query
    where
        T: cmd::uuid::Arg,
    {
        arg.into_query()
    }

    pub fn circle<T>(self, arg: T) -> Query
    where
        T: cmd::circle::Arg,
    {
        arg.into_query()
    }

    pub fn distance<T>(self, arg: T) -> Query
    where
        T: cmd::distance::Arg,
    {
        arg.into_query()
    }

    pub fn geojson<T>(self, arg: T) -> Query
    where
        T: cmd::geojson::Arg,
    {
        arg.into_query()
    }

    pub fn intersects<T>(self, arg: T) -> Query
    where
        T: cmd::intersects::Arg,
    {
        arg.into_query()
    }

    pub fn line<T>(self, arg: T) -> Query
    where
        T: cmd::line::Arg,
    {
        arg.into_query()
    }

    pub fn point<T>(self, arg: T) -> Query
    where
        T: cmd::point::Arg,
    {
        arg.into_query()
    }

    pub fn polygon<T>(self, arg: T) -> Query
    where
        T: cmd::polygon::Arg,
    {
        arg.into_query()
    }

    pub fn grant<T>(self, arg: T) -> Query
    where
        T: cmd::grant::Arg,
    {
        arg.into_query()
    }

    pub fn wait<T>(self, arg: T) -> Query
    where
        T: cmd::wait::Arg,
    {
        arg.into_query()
    }
}
<commit_msg>return TcpStream to caller<commit_after>//! ReQL is the RethinkDB query language. It offers a very powerful and
//! convenient way to manipulate JSON documents.
//!
//! # Start the server #
//!
//! ## Linux and OS X ##
//!
//! Start the server from a terminal window.
//!
//! ```bash
//! $ rethinkdb
//! ```
//!
//! ## Windows ##
//!
//! Start the server from the Windows command prompt.
//!
//! ```bash
//! C:\Path\To\RethinkDB\>rethinkdb.exe
//! ```
//!
//! # Import the driver #
//!
//! First, make sure you have `protoc` installed and in your `PATH`. See
//! [`prost-build` documentation](https://docs.rs/prost-build/0.7.0/prost_build/#sourcing-protoc)
//! for more details if it fails to compile.
//!
//! Add this crate (`reql`) and the `futures` crate to your dependencies in `Cargo.toml`.
//! You may also need to depend on a crate that provides an async implementation of a
//! `TcpStream`, like `async-std`, for example. You can use this crate
//! with any `TcpStream` that implements `AsyncRead` and `AsyncWrite` from
//! the futures crate. The implementation also needs to implement the same
//! traits for `&TcpStream` because we use unmutable references of the
//! connection to take advantage of RethinkDB's connection pipelining abilities.
//!
//! Now import the RethinkDB driver:
//!
//! ```
//! use reql::r;
//! ```
//!
//! You can now access RethinkDB commands through the [`r` struct](r).
//!
//! # Open a connection #
//!
//! When you first start RethinkDB, the server opens a port for the client
//! drivers (`28015` by default). Let's open a connection:
//!
//! ```
//! use async_std::net::TcpStream;
//! use reql::{r, DEFAULT_ADDR};
//!
//! # async fn connect() -> reql::Result<()> {
//! let stream = TcpStream::connect(DEFAULT_ADDR).await?;
//! let connection = r.connection(stream).await?;
//! # Ok(())
//! # }
//! ```
//!
//! The variable `connection` is now initialized and we can run queries.
//!
//! # Send a query to the database #
//!
//! ```
//! use async_std::net::TcpStream;
//! use futures::TryStreamExt;
//! use reql::{r, DEFAULT_ADDR};
//!
//! # async fn connect() -> reql::Result<()> {
//! let stream = TcpStream::connect(DEFAULT_ADDR).await?;
//! let conn = r.connection(stream).await?;
//! let mut query = r.expr("Hello world!").run(&conn);
//! # let _: Option<String> = query.try_next().await?;
//! # Ok(())
//! # }
//! ```
//!
//! [See the `r` struct for more available commands](r)

pub mod cmd;
mod err;
mod proto;

use cmd::run::Response;
use dashmap::DashMap;
use futures::channel::mpsc::Sender;
use futures::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use futures::lock::Mutex;
use ql2::term::TermType;
use std::borrow::Cow;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};

pub use err::*;
pub use proto::Query;

/// Default address of the RethinkDB server
pub const DEFAULT_ADDR: (&str, u16) = ("localhost", 28015);

/// Custom result returned by various ReQL commands
pub type Result<T> = std::result::Result<T, Error>;

/// The connection object returned by `r.connection()`
#[derive(Debug)]
pub struct Connection<'a, T> {
    db: Cow<'a, str>,
    stream: T,
    token: AtomicU64,
    broken: AtomicBool,
    change_feed: AtomicBool,
    buffer: usize,
    senders: DashMap<u64, Sender<Result<Response>>>,
    locker: Mutex<()>,
}

impl<T> Connection<'_, T> {
    /// Convert the connection into an instance you can move around
    pub fn into_owned(self) -> Connection<'static, T> {
        Connection {
            db: Cow::from(self.db.into_owned()),
            stream: self.stream,
            token: self.token,
            broken: self.broken,
            change_feed: self.change_feed,
            buffer: self.buffer,
            senders: self.senders,
            locker: self.locker,
        }
    }

    pub fn take_tcp_stream(self) -> T {
        self.stream
    }

    fn mark_broken(&self) {
        self.broken.store(true, Ordering::SeqCst);
    }

    fn broken(&self) -> Result<()> {
        if self.broken.load(Ordering::SeqCst) {
            return Err(err::Client::ConnectionBroken.into());
        }
        Ok(())
    }

    fn mark_change_feed(&self) {
        self.change_feed.store(true, Ordering::SeqCst);
    }

    fn unmark_change_feed(&self) {
        self.change_feed.store(false, Ordering::SeqCst);
    }

    fn change_feed(&self) -> Result<()> {
        if self.change_feed.load(Ordering::SeqCst) {
            return Err(err::Client::ConnectionLocked.into());
        }
        Ok(())
    }
}

/// A generic, runtime independent, TcpStream
///
/// This crate can be used with any runtime that implements
/// the `AsyncRead` and the `AsyncWrite` traits from the
/// futures crate
pub trait TcpStream<'a>: AsyncReadExt + AsyncWriteExt
where
    Self: Unpin + 'a,
    &'a Self: AsyncRead + AsyncWrite,
{
}

impl<'a, T: AsyncRead + AsyncWrite + ?Sized> TcpStream<'a> for T
where
    Self: Unpin + 'a,
    &'a Self: AsyncRead + AsyncWrite,
{
}

/// The top-level ReQL namespace
///
/// # Example
///
/// Set up your top-level namespace.
///
/// ```
/// use reql::r;
/// ```
#[allow(non_camel_case_types)]
pub struct r;

impl<'a> r {
    /// Create a new connection to the database server
    ///
    /// # Example
    ///
    /// Open a connection using the default host and port, specifying the default database.
    ///
    /// ```
    /// use async_std::net::TcpStream;
    /// use reql::cmd::connection::Options;
    /// use reql::{r, DEFAULT_ADDR};
    ///
    /// # async fn connect() -> reql::Result<()> {
    /// let stream = TcpStream::connect(DEFAULT_ADDR).await?;
    /// let conn = r.connection((stream, Options::new().db("marvel"))).await?;
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// Read more about this command [cmd::connection]
    pub async fn connection<A, T>(self, options: A) -> Result<Connection<'a, T>>
    where
        A: cmd::connection::Arg<'a, T>,
        T: TcpStream<'a>,
        &'a T: AsyncRead + AsyncWrite,
    {
        cmd::connection::new(options.into()).await
    }

    pub fn db_create<T>(self, arg: T) -> Query
    where
        T: cmd::db_create::Arg,
    {
        arg.into_query()
    }

    pub fn db_drop<T>(self, arg: T) -> Query
    where
        T: cmd::db_drop::Arg,
    {
        arg.into_query()
    }

    pub fn db_list(self) -> Query {
        Query::new(TermType::DbList)
    }

    /// Reference a database
    ///
    /// The `db` command is optional. If it is not present in a query, the
    /// query will run against the default database for the connection,
    /// specified in the `db` argument to [r::connection].
    ///
    /// # Examples
    ///
    /// Explicitly specify a database for a query.
    ///
    /// ```
    /// # use async_std::net::TcpStream;
    /// # use futures::TryStreamExt;
    /// # use reql::{r, DEFAULT_ADDR};
    /// # async fn example() -> reql::Result<()> {
    /// # let stream = TcpStream::connect(DEFAULT_ADDR).await?;
    /// # let conn = r.connection(stream).await?;
    /// let mut query = r.db("heroes").table("marvel").run(&conn);
    /// # let _: Option<String> = query.try_next().await?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn db<T>(self, arg: T) -> Query
    where
        T: cmd::db::Arg,
    {
        arg.into_query()
    }

    pub fn table_create<T>(self, arg: T) -> Query
    where
        T: cmd::table_create::Arg,
    {
        arg.into_query()
    }

    pub fn table<T>(self, arg: T) -> Query
    where
        T: cmd::table::Arg,
    {
        arg.into_query()
    }

    pub fn map<T>(self, arg: T) -> Query
    where
        T: cmd::map::Arg,
    {
        arg.into_query()
    }

    pub fn union<T>(self, arg: T) -> Query
    where
        T: cmd::union::Arg,
    {
        arg.into_query()
    }

    pub fn group<T>(self, arg: T) -> Query
    where
        T: cmd::group::Arg,
    {
        arg.into_query()
    }

    pub fn reduce<T>(self, arg: T) -> Query
    where
        T: cmd::reduce::Arg,
    {
        arg.into_query()
    }

    pub fn count<T>(self, arg: T) -> Query
    where
        T: cmd::count::Arg,
    {
        arg.into_query()
    }

    pub fn sum<T>(self, arg: T) -> Query
    where
        T: cmd::sum::Arg,
    {
        arg.into_query()
    }

    pub fn avg<T>(self, arg: T) -> Query
    where
        T: cmd::avg::Arg,
    {
        arg.into_query()
    }

    pub fn min<T>(self, arg: T) -> Query
    where
        T: cmd::min::Arg,
    {
        arg.into_query()
    }

    pub fn max<T>(self, arg: T) -> Query
    where
        T: cmd::max::Arg,
    {
        arg.into_query()
    }

    pub fn distinct<T>(self, arg: T) -> Query
    where
        T: cmd::distinct::Arg,
    {
        arg.into_query()
    }

    pub fn contains<T>(self, arg: T) -> Query
    where
        T: cmd::contains::Arg,
    {
        arg.into_query()
    }

    pub fn literal<T>(self, arg: T) -> Query
    where
        T: cmd::literal::Arg,
    {
        arg.into_query()
    }

    pub fn object<T>(self, arg: T) -> Query
    where
        T: cmd::object::Arg,
    {
        arg.into_query()
    }

    pub fn random<T>(self, arg: T) -> Query
    where
        T: cmd::random::Arg,
    {
        arg.into_query()
    }

    pub fn round<T>(self, arg: T) -> Query
    where
        T: cmd::round::Arg,
    {
        arg.into_query()
    }

    pub fn ceil<T>(self, arg: T) -> Query
    where
        T: cmd::ceil::Arg,
    {
        arg.into_query()
    }

    pub fn floor<T>(self, arg: T) -> Query
    where
        T: cmd::floor::Arg,
    {
        arg.into_query()
    }

    pub fn now(self) -> Query {
        Query::new(TermType::Now)
    }

    pub fn time<T>(self, arg: T) -> Query
    where
        T: cmd::time::Arg,
    {
        arg.into_query()
    }

    pub fn epoch_time<T>(self, arg: T) -> Query
    where
        T: cmd::epoch_time::Arg,
    {
        arg.into_query()
    }

    pub fn iso8601<T>(self, arg: T) -> Query
    where
        T: cmd::iso8601::Arg,
    {
        arg.into_query()
    }

    pub fn r#do<T>(self, arg: T) -> Query
    where
        T: cmd::r#do::Arg,
    {
        arg.into_query()
    }

    pub fn branch<T>(self, arg: T) -> Query
    where
        T: cmd::branch::Arg,
    {
        arg.into_query()
    }

    pub fn range<T>(self, arg: T) -> Query
    where
        T: cmd::range::Arg,
    {
        arg.into_query()
    }

    pub fn error<T>(self, arg: T) -> Query
    where
        T: cmd::error::Arg,
    {
        arg.into_query()
    }

    pub fn expr<T>(self, arg: T) -> Query
    where
        T: cmd::expr::Arg,
    {
        arg.into_query()
    }

    pub fn js<T>(self, arg: T) -> Query
    where
        T: cmd::js::Arg,
    {
        arg.into_query()
    }

    pub fn info<T>(self, arg: T) -> Query
    where
        T: cmd::info::Arg,
    {
        arg.into_query()
    }

    pub fn json<T>(self, arg: T) -> Query
    where
        T: cmd::json::Arg,
    {
        arg.into_query()
    }

    pub fn http<T>(self, arg: T) -> Query
    where
        T: cmd::http::Arg,
    {
        arg.into_query()
    }

    pub fn uuid<T>(self, arg: T) -> Query
    where
        T: cmd::uuid::Arg,
    {
        arg.into_query()
    }

    pub fn circle<T>(self, arg: T) -> Query
    where
        T: cmd::circle::Arg,
    {
        arg.into_query()
    }

    pub fn distance<T>(self, arg: T) -> Query
    where
        T: cmd::distance::Arg,
    {
        arg.into_query()
    }

    pub fn geojson<T>(self, arg: T) -> Query
    where
        T: cmd::geojson::Arg,
    {
        arg.into_query()
    }

    pub fn intersects<T>(self, arg: T) -> Query
    where
        T: cmd::intersects::Arg,
    {
        arg.into_query()
    }

    pub fn line<T>(self, arg: T) -> Query
    where
        T: cmd::line::Arg,
    {
        arg.into_query()
    }

    pub fn point<T>(self, arg: T) -> Query
    where
        T: cmd::point::Arg,
    {
        arg.into_query()
    }

    pub fn polygon<T>(self, arg: T) -> Query
    where
        T: cmd::polygon::Arg,
    {
        arg.into_query()
    }

    pub fn grant<T>(self, arg: T) -> Query
    where
        T: cmd::grant::Arg,
    {
        arg.into_query()
    }

    pub fn wait<T>(self, arg: T) -> Query
    where
        T: cmd::wait::Arg,
    {
        arg.into_query()
    }
}
<|endoftext|>