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

use std::collections::{btree_map::Entry, BTreeMap};
use std::net::SocketAddr;

// use futures::{
//     future,
//     stream::Fuse,
//     sync::{mpsc, oneshot},
//     Async, AsyncSink, Future, Poll, Sink, Stream,
// };

// use super::connect::{connect, RespConnection};

// use crate::{
//     error::{self, ConnectionReason},
//     reconnect::{reconnect, Reconnect},
//     resp::{self, FromResp},
// };

// #[derive(Debug)]
// enum PubsubEvent {
//     /// The: topic, sink to send messages through, and a oneshot to signal subscription has
//     /// occurred.
//     Subscribe(String, PubsubSink, oneshot::Sender<()>),
//     /// The name of the topic to unsubscribe from. Unsubscription will be signaled by the stream
//     /// closing without error.
//     Unsubscribe(String),
// }

// type PubsubStreamInner = mpsc::UnboundedReceiver<Result<resp::RespValue, error::Error>>;
// type PubsubSink = mpsc::UnboundedSender<Result<resp::RespValue, error::Error>>;

// /// A spawned future that handles a Pub/Sub connection and routes messages to streams for
// /// downstream consumption
// struct PubsubConnectionInner {
//     /// The actual Redis connection
//     connection: RespConnection,
//     /// A stream onto which subscription/unsubscription requests are read
//     out_rx: Fuse<mpsc::UnboundedReceiver<PubsubEvent>>,
//     /// Current subscriptions
//     subscriptions: BTreeMap<String, PubsubSink>,
//     /// Subscriptions that have not yet been confirmed
//     pending_subs: BTreeMap<String, (PubsubSink, oneshot::Sender<()>)>,
//     /// Any incomplete messages to-be-sent
//     send_pending: Option<resp::RespValue>,
// }

// impl PubsubConnectionInner {
//     fn new(con: RespConnection, out_rx: mpsc::UnboundedReceiver<PubsubEvent>) -> Self {
//         PubsubConnectionInner {
//             connection: con,
//             out_rx: out_rx.fuse(),
//             subscriptions: BTreeMap::new(),
//             pending_subs: BTreeMap::new(),
//             send_pending: None,
//         }
//     }

//     /// Returns true = OK, more can be sent, or false = sink is full, needs flushing
//     fn do_send(&mut self, msg: resp::RespValue) -> Result<bool, error::Error> {
//         match self.connection.start_send(msg)? {
//             AsyncSink::Ready => Ok(true),
//             AsyncSink::NotReady(msg) => {
//                 self.send_pending = Some(msg);
//                 Ok(false)
//             }
//         }
//     }

//     fn do_flush(&mut self) -> Result<(), error::Error> {
//         self.connection
//             .poll_complete()
//             .map(|_| ())
//             .map_err(|e| e.into())
//     }

//     // Returns true = flushing required.  false = no flushing required
//     fn handle_new_subs(&mut self) -> Result<bool, error::Error> {
//         let mut flushing_req = false;
//         if let Some(msg) = self.send_pending.take() {
//             flushing_req = true;
//             if !self.do_send(msg)? {
//                 return Ok(flushing_req);
//             }
//         }
//         loop {
//             match self.out_rx.poll() {
//                 Ok(Async::Ready(None)) => {
//                     return Ok(flushing_req);
//                 }
//                 Ok(Async::Ready(Some(pubsub_event))) => {
//                     let message = match pubsub_event {
//                         PubsubEvent::Subscribe(topic, sender, signal) => {
//                             self.pending_subs.insert(topic.clone(), (sender, signal));
//                             resp_array!["SUBSCRIBE", topic]
//                         }
//                         PubsubEvent::Unsubscribe(topic) => resp_array!["UNSUBSCRIBE", topic],
//                     };
//                     flushing_req = true;
//                     if !self.do_send(message)? {
//                         return Ok(flushing_req);
//                     }
//                 }
//                 Ok(Async::NotReady) => {
//                     return Ok(flushing_req);
//                 }
//                 Err(()) => {
//                     return Err(error::internal(
//                         "Unexpected error polling outgoing unbuffered channel",
//                     ));
//                 }
//             }
//         }
//     }

//     fn handle_message(&mut self, msg: resp::RespValue) -> Result<bool, error::Error> {
//         let (message_type, topic, msg) = match msg {
//             resp::RespValue::Array(mut messages) => match (
//                 messages.pop(),
//                 messages.pop(),
//                 messages.pop(),
//                 messages.pop(),
//             ) {
//                 (Some(msg), Some(topic), Some(message_type), None) => {
//                     match (msg, String::from_resp(topic), message_type) {
//                         (msg, Ok(topic), resp::RespValue::BulkString(bytes)) => (bytes, topic, msg),
//                         _ => return Err(error::unexpected("Incorrect format of a PUBSUB message")),
//                     }
//                 }
//                 _ => {
//                     return Err(error::unexpected(
//                         "Wrong number of parts for a PUBSUB message",
//                     ));
//                 }
//             },
//             _ => {
//                 return Err(error::unexpected(
//                     "PUBSUB message should be encoded as an array",
//                 ));
//             }
//         };

//         match message_type.as_slice() {
//             b"subscribe" => match self.pending_subs.remove(&topic) {
//                 Some((sender, signal)) => {
//                     self.subscriptions.insert(topic, sender);
//                     signal
//                         .send(())
//                         .map_err(|()| error::internal("Error confirming subscription"))?
//                 }
//                 None => {
//                     return Err(error::internal(format!(
//                         "Received unexpected subscribe notification for topic: {}",
//                         topic
//                     )));
//                 }
//             },
//             b"unsubscribe" => {
//                 match self.subscriptions.entry(topic) {
//                     Entry::Occupied(entry) => {
//                         entry.remove_entry();
//                     }
//                     Entry::Vacant(vacant) => {
//                         return Err(error::internal(format!(
//                             "Unexpected unsubscribe message: {}",
//                             vacant.key()
//                         )));
//                     }
//                 }
//                 if self.subscriptions.is_empty() {
//                     return Ok(false);
//                 }
//             }
//             b"message" => match self.subscriptions.get(&topic) {
//                 Some(sender) => sender.unbounded_send(Ok(msg)).expect("Cannot send message"),
//                 None => {
//                     return Err(error::internal(format!(
//                         "Unexpected message on topic: {}",
//                         topic
//                     )));
//                 }
//             },
//             t => {
//                 return Err(error::internal(format!(
//                     "Unexpected data on Pub/Sub connection: {}",
//                     String::from_utf8_lossy(t)
//                 )));
//             }
//         }

//         Ok(true)
//     }

//     /// Returns true, if there are still valid subscriptions at the end, or false if not, i.e. the whole thing can be dropped.
//     fn handle_messages(&mut self) -> Result<bool, error::Error> {
//         loop {
//             match self.connection.poll() {
//                 Ok(Async::Ready(None)) => {
//                     if self.subscriptions.is_empty() {
//                         return Ok(false);
//                     } else {
//                         // This can only happen if the connection is closed server-side
//                         for sub in self.subscriptions.values() {
//                             sub.unbounded_send(Err(error::Error::Connection(
//                                 ConnectionReason::NotConnected,
//                             )))
//                             .unwrap();
//                         }
//                         return Err(error::Error::Connection(ConnectionReason::NotConnected));
//                     }
//                 }
//                 Ok(Async::Ready(Some(message))) => {
//                     let message_result = self.handle_message(message)?;
//                     if !message_result {
//                         return Ok(false);
//                     }
//                 }
//                 Ok(Async::NotReady) => return Ok(true),
//                 Err(e) => {
//                     for sub in self.subscriptions.values() {
//                         sub.unbounded_send(Err(error::unexpected(format!(
//                             "Connection is in the process of failing due to: {}",
//                             e
//                         ))))
//                         .unwrap();
//                     }
//                     return Err(e);
//                 }
//             }
//         }
//     }
// }

// impl Future for PubsubConnectionInner {
//     type Item = ();
//     type Error = error::Error;

//     fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
//         let flush_req = self.handle_new_subs()?;
//         if flush_req {
//             self.do_flush()?;
//         }
//         let cont = self.handle_messages()?;
//         if cont {
//             Ok(Async::NotReady)
//         } else {
//             Ok(Async::Ready(()))
//         }
//     }
// }

// /// A shareable reference to subscribe to PUBSUB topics
// #[derive(Debug, Clone)]
// pub struct PubsubConnection {
//     out_tx_c: Reconnect<PubsubEvent, mpsc::UnboundedSender<PubsubEvent>>,
// }

// /// Used for Redis's PUBSUB functionality.
// ///
// /// Returns a future that resolves to a `PubsubConnection`. The future will only resolve once the
// /// connection is established; after the intial establishment, if the connection drops for any
// /// reason (e.g. Redis server being restarted), the connection will attempt re-connect, however
// /// any subscriptions will need to be re-subscribed.
// pub fn pubsub_connect(
//     addr: &SocketAddr,
// ) -> impl Future<Item = PubsubConnection, Error = error::Error> + Send {
//     let addr = *addr;
//     future::lazy(move || {
//         reconnect(
//             |con: &mpsc::UnboundedSender<PubsubEvent>, act| {
//                 Box::new(future::result(con.unbounded_send(act)).map_err(|e| e.into()))
//             },
//             move || {
//                 let con_f = connect(&addr).map_err(|e| e.into()).and_then(|connection| {
//                     let (out_tx, out_rx) = mpsc::unbounded();
//                     let pubsub_connection_inner = Box::new(
//                         PubsubConnectionInner::new(connection, out_rx)
//                             .map_err(|e| log::error!("Pub/Sub error: {:?}", e)),
//                     );
//                     let mut default_executor = DefaultExecutor::current();
//                     match default_executor.spawn(pubsub_connection_inner) {
//                         Ok(_) => Ok(out_tx),
//                         Err(e) => Err(error::Error::Internal(format!(
//                             "Cannot spawn a pubsub connection: {:?}",
//                             e
//                         ))),
//                     }
//                 });
//                 Box::new(con_f)
//             },
//         )
//     })
//     .map(|out_tx_c| PubsubConnection { out_tx_c })
// }

// impl PubsubConnection {
//     /// Subscribes to a particular PUBSUB topic.
//     ///
//     /// Returns a future that resolves to a `Stream` that contains all the messages published on
//     /// that particular topic.
//     ///
//     /// The resolved stream will end with `redis_async::error::Error::EndOfStream` if the
//     /// underlying connection is lost for unexpected reasons. In this situation, clients should
//     /// `subscribe` to re-subscribe; the underlying connect will automatically reconnect. However,
//     /// clients should be aware that resubscriptions will only succeed if the underlying connection
//     /// has re-established, so multiple calls to `subscribe` may be required.
//     pub fn subscribe(&self, topic: &str) -> impl Future<Item = PubsubStream, Error = error::Error> {
//         let (tx, rx) = mpsc::unbounded();
//         let (signal_t, signal_r) = oneshot::channel();
//         let do_work_f =
//             self.out_tx_c
//                 .do_work(PubsubEvent::Subscribe(topic.to_owned(), tx, signal_t));

//         let stream = PubsubStream {
//             topic: topic.to_owned(),
//             underlying: rx,
//             con: self.clone(),
//         };
//         do_work_f.and_then(|()| {
//             signal_r
//                 .map(|_| stream)
//                 .map_err(|_| error::internal("Subscription failed, try again later..."))
//         })
//     }

//     /// Tells the client to unsubscribe from a particular topic. This will return immediately, the
//     /// actual unsubscription will be confirmed when the stream returned from `subscribe` ends.
//     pub fn unsubscribe<T: Into<String>>(&self, topic: T) {
//         // Ignoring any results, as any errors communicating with Redis would de-facto unsubscribe
//         // anyway, and would be reported/logged elsewhere
//         let _ = self
//             .out_tx_c
//             .do_work(PubsubEvent::Unsubscribe(topic.into()));
//     }
// }

// pub struct PubsubStream {
//     topic: String,
//     underlying: PubsubStreamInner,
//     con: PubsubConnection,
// }

// impl Stream for PubsubStream {
//     type Item = resp::RespValue;
//     type Error = error::Error;

//     fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
//         match self.underlying.poll() {
//             Ok(Async::Ready(Some(Ok(v)))) => Ok(Async::Ready(Some(v))),
//             Ok(Async::Ready(Some(Err(e)))) => Err(e),
//             Ok(Async::Ready(None)) => Ok(Async::Ready(None)),
//             Ok(Async::NotReady) => Ok(Async::NotReady),
//             Err(()) => Err(error::internal(
//                 "Unexpected error from underlying PubsubStream",
//             )),
//         }
//     }
// }

// impl Drop for PubsubStream {
//     fn drop(&mut self) {
//         let topic: &str = self.topic.as_ref();
//         self.con.unsubscribe(topic);
//     }
// }
<commit_msg>Checkpoint<commit_after>/*
 * Copyright 2017-2019 Ben Ashford
 *
 * Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
 * http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
 * <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
 * option. This file may not be copied, modified, or distributed
 * except according to those terms.
 */

use std::{
    collections::{btree_map::Entry, BTreeMap},
    future::Future,
    net::SocketAddr,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
};

use futures_channel::{mpsc, oneshot};
use futures_util::{
    future::{FutureExt, TryFutureExt},
    pin_mut,
    sink::Sink,
    stream::{Fuse, StreamExt},
};

use super::connect::{connect, RespConnection};

use crate::{
    error::{self, ConnectionReason},
    reconnect::{reconnect, Reconnect},
    resp::{self, FromResp},
};

#[derive(Debug)]
enum PubsubEvent {
    /// The: topic, sink to send messages through, and a oneshot to signal subscription has
    /// occurred.
    Subscribe(String, PubsubSink, oneshot::Sender<()>),
    /// The name of the topic to unsubscribe from. Unsubscription will be signaled by the stream
    /// closing without error.
    Unsubscribe(String),
}

type PubsubStreamInner = mpsc::UnboundedReceiver<Result<resp::RespValue, error::Error>>;
type PubsubSink = mpsc::UnboundedSender<Result<resp::RespValue, error::Error>>;

/// A spawned future that handles a Pub/Sub connection and routes messages to streams for
/// downstream consumption
struct PubsubConnectionInner {
    /// The actual Redis connection
    connection: RespConnection,
    /// A stream onto which subscription/unsubscription requests are read
    out_rx: Fuse<mpsc::UnboundedReceiver<PubsubEvent>>,
    /// Current subscriptions
    subscriptions: BTreeMap<String, PubsubSink>,
    /// Subscriptions that have not yet been confirmed
    pending_subs: BTreeMap<String, (PubsubSink, oneshot::Sender<()>)>,
    /// Any incomplete messages to be sent...
    send_pending: Option<resp::RespValue>,
}

impl PubsubConnectionInner {
    fn new(con: RespConnection, out_rx: mpsc::UnboundedReceiver<PubsubEvent>) -> Self {
        PubsubConnectionInner {
            connection: con,
            out_rx: out_rx.fuse(),
            subscriptions: BTreeMap::new(),
            pending_subs: BTreeMap::new(),
            send_pending: None,
        }
    }

    /// Returns `true` if data sent, or `false` if stream not ready...
    fn do_send(&mut self, cx: &mut Context, msg: resp::RespValue) -> Result<bool, error::Error> {
        match Pin::new(&mut self.connection).poll_ready(cx) {
            Poll::Ready(_) => {
                Pin::new(&mut self.connection).start_send(msg)?;
                Ok(true)
            }
            Poll::Pending => {
                self.send_pending = Some(msg);
                Ok(false)
            }
        }
    }

    fn do_flush(&mut self, cx: &mut Context) -> Result<(), error::Error> {
        match Pin::new(&mut self.connection).poll_flush(cx) {
            Poll::Ready(r) => r.map_err(|e| e.into()),
            Poll::Pending => Ok(()),
        }
    }

    // Returns true = flushing required.  false = no flushing required
    fn handle_new_subs(&mut self, cx: &mut Context) -> Result<(), error::Error> {
        if let Some(msg) = self.send_pending.take() {
            if !self.do_send(cx, msg)? {
                return Ok(());
            }
        }
        loop {
            match self.out_rx.poll_next_unpin(cx) {
                Poll::Pending => return Ok(()),
                Poll::Ready(None) => return Ok(()),
                Poll::Ready(Some(pubsub_event)) => {
                    let message = match pubsub_event {
                        PubsubEvent::Subscribe(topic, sender, signal) => {
                            self.pending_subs.insert(topic.clone(), (sender, signal));
                            resp_array!["SUBSCRIBE", topic]
                        }
                        PubsubEvent::Unsubscribe(topic) => resp_array!["UNSUBSCRIBE", topic],
                    };
                    if !self.do_send(cx, message)? {
                        return Ok(());
                    }
                }
            }
        }
    }

    fn handle_message(&mut self, msg: resp::RespValue) -> Result<bool, error::Error> {
        let (message_type, topic, msg) = match msg {
            resp::RespValue::Array(mut messages) => match (
                messages.pop(),
                messages.pop(),
                messages.pop(),
                messages.pop(),
            ) {
                (Some(msg), Some(topic), Some(message_type), None) => {
                    match (msg, String::from_resp(topic), message_type) {
                        (msg, Ok(topic), resp::RespValue::BulkString(bytes)) => (bytes, topic, msg),
                        _ => return Err(error::unexpected("Incorrect format of a PUBSUB message")),
                    }
                }
                _ => {
                    return Err(error::unexpected(
                        "Wrong number of parts for a PUBSUB message",
                    ));
                }
            },
            _ => {
                return Err(error::unexpected(
                    "PUBSUB message should be encoded as an array",
                ));
            }
        };

        match message_type.as_slice() {
            b"subscribe" => match self.pending_subs.remove(&topic) {
                Some((sender, signal)) => {
                    self.subscriptions.insert(topic, sender);
                    signal
                        .send(())
                        .map_err(|()| error::internal("Error confirming subscription"))?
                }
                None => {
                    return Err(error::internal(format!(
                        "Received unexpected subscribe notification for topic: {}",
                        topic
                    )));
                }
            },
            b"unsubscribe" => {
                match self.subscriptions.entry(topic) {
                    Entry::Occupied(entry) => {
                        entry.remove_entry();
                    }
                    Entry::Vacant(vacant) => {
                        return Err(error::internal(format!(
                            "Unexpected unsubscribe message: {}",
                            vacant.key()
                        )));
                    }
                }
                if self.subscriptions.is_empty() {
                    return Ok(false);
                }
            }
            b"message" => match self.subscriptions.get(&topic) {
                Some(sender) => sender.unbounded_send(Ok(msg)).expect("Cannot send message"),
                None => {
                    return Err(error::internal(format!(
                        "Unexpected message on topic: {}",
                        topic
                    )));
                }
            },
            t => {
                return Err(error::internal(format!(
                    "Unexpected data on Pub/Sub connection: {}",
                    String::from_utf8_lossy(t)
                )));
            }
        }

        Ok(true)
    }

    /// Returns true, if there are still valid subscriptions at the end, or false if not, i.e. the whole thing can be dropped.
    fn handle_messages(&mut self, cx: &mut Context) -> Result<bool, error::Error> {
        loop {
            match self.connection.poll_next_unpin(cx) {
                Poll::Pending => return Ok(true),
                Poll::Ready(None) => {
                    if self.subscriptions.is_empty() {
                        return Ok(false);
                    } else {
                        // This can only happen if the connection is closed server-side
                        for sub in self.subscriptions.values() {
                            sub.unbounded_send(Err(error::Error::Connection(
                                ConnectionReason::NotConnected,
                            )))
                            .unwrap();
                        }
                        return Err(error::Error::Connection(ConnectionReason::NotConnected));
                    }
                }
                Poll::Ready(Some(Ok(message))) => {
                    let message_result = self.handle_message(message)?;
                    if !message_result {
                        return Ok(false);
                    }
                }
                Poll::Ready(Some(Err(e))) => {
                    for sub in self.subscriptions.values() {
                        sub.unbounded_send(Err(error::unexpected(format!(
                            "Connection is in the process of failing due to: {}",
                            e
                        ))))
                        .unwrap();
                    }
                    return Err(e);
                }
            }
        }
    }
}

impl Future for PubsubConnectionInner {
    type Output = Result<(), error::Error>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        let this_self = self.get_mut();
        let flush_req = this_self.handle_new_subs(cx)?;
        this_self.do_flush(cx)?;
        let cont = this_self.handle_messages(cx)?;
        if cont {
            Poll::Pending
        } else {
            Poll::Ready(Ok(()))
        }
    }
}

/// A shareable reference to subscribe to PUBSUB topics
#[derive(Debug, Clone)]
pub struct PubsubConnection {
    out_tx_c: Arc<Reconnect<PubsubEvent, mpsc::UnboundedSender<PubsubEvent>>>,
}

async fn inner_conn_fn(
    addr: &SocketAddr,
) -> Result<mpsc::UnboundedSender<PubsubEvent>, error::Error> {
    let connection = connect(addr).await?;
    let (out_tx, out_rx) = mpsc::unbounded();
    tokio::spawn(async {
        match PubsubConnectionInner::new(connection, out_rx).await {
            Ok(_) => (),
            Err(e) => log::error!("Pub/Sub error: {:?}", e),
        }
    });
    Ok(out_tx)
}

// /// Used for Redis's PUBSUB functionality.
// ///
// /// Returns a future that resolves to a `PubsubConnection`. The future will only resolve once the
// /// connection is established; after the intial establishment, if the connection drops for any
// /// reason (e.g. Redis server being restarted), the connection will attempt re-connect, however
// /// any subscriptions will need to be re-subscribed.
// pub async fn pubsub_connect(addr: &SocketAddr) -> Result<PubsubConnection, error::Error> {
//     reconnect(
//         |con: &mpsc::UnboundedSender<PubsubEvent>, act| {
//             con.unbounded_send(act).map_err(|e| e.into())
//         },
//         move || {
//             let con_f = inner_conn_fn(addr);
//             Box::new(Box::pin(con_f)) // TODO - do I need to box and pin? In this order?
//         },
//     )
//     .map(|out_tx_c| PubsubConnection { out_tx_c })
// }

// impl PubsubConnection {
//     /// Subscribes to a particular PUBSUB topic.
//     ///
//     /// Returns a future that resolves to a `Stream` that contains all the messages published on
//     /// that particular topic.
//     ///
//     /// The resolved stream will end with `redis_async::error::Error::EndOfStream` if the
//     /// underlying connection is lost for unexpected reasons. In this situation, clients should
//     /// `subscribe` to re-subscribe; the underlying connect will automatically reconnect. However,
//     /// clients should be aware that resubscriptions will only succeed if the underlying connection
//     /// has re-established, so multiple calls to `subscribe` may be required.
//     pub fn subscribe(&self, topic: &str) -> impl Future<Item = PubsubStream, Error = error::Error> {
//         let (tx, rx) = mpsc::unbounded();
//         let (signal_t, signal_r) = oneshot::channel();
//         let do_work_f =
//             self.out_tx_c
//                 .do_work(PubsubEvent::Subscribe(topic.to_owned(), tx, signal_t));

//         let stream = PubsubStream {
//             topic: topic.to_owned(),
//             underlying: rx,
//             con: self.clone(),
//         };
//         do_work_f.and_then(|()| {
//             signal_r
//                 .map(|_| stream)
//                 .map_err(|_| error::internal("Subscription failed, try again later..."))
//         })
//     }

//     /// Tells the client to unsubscribe from a particular topic. This will return immediately, the
//     /// actual unsubscription will be confirmed when the stream returned from `subscribe` ends.
//     pub fn unsubscribe<T: Into<String>>(&self, topic: T) {
//         // Ignoring any results, as any errors communicating with Redis would de-facto unsubscribe
//         // anyway, and would be reported/logged elsewhere
//         let _ = self
//             .out_tx_c
//             .do_work(PubsubEvent::Unsubscribe(topic.into()));
//     }
// }

// pub struct PubsubStream {
//     topic: String,
//     underlying: PubsubStreamInner,
//     con: PubsubConnection,
// }

// impl Stream for PubsubStream {
//     type Item = resp::RespValue;
//     type Error = error::Error;

//     fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
//         match self.underlying.poll() {
//             Ok(Async::Ready(Some(Ok(v)))) => Ok(Async::Ready(Some(v))),
//             Ok(Async::Ready(Some(Err(e)))) => Err(e),
//             Ok(Async::Ready(None)) => Ok(Async::Ready(None)),
//             Ok(Async::NotReady) => Ok(Async::NotReady),
//             Err(()) => Err(error::internal(
//                 "Unexpected error from underlying PubsubStream",
//             )),
//         }
//     }
// }

// impl Drop for PubsubStream {
//     fn drop(&mut self) {
//         let topic: &str = self.topic.as_ref();
//         self.con.unsubscribe(topic);
//     }
// }
<|endoftext|>


<commit_before>use std::ops::Deref;

use qvariant::*;
use types::*;
use qurl::*;
use qmeta::*;

extern "C" {
    fn dos_qapplication_create();
    fn dos_qapplication_exec();
    fn dos_qapplication_quit();
    fn dos_qapplication_delete();

    fn dos_qqmlapplicationengine_create() -> DosQmlApplicationEngine;
    fn dos_qqmlapplicationengine_load(vptr: DosQmlApplicationEngine, filename: DosCStr);
    fn dos_qqmlapplicationengine_load_url(vptr: DosQmlApplicationEngine, url: DosQUrl);
    fn dos_qqmlapplicationengine_load_data(vptr: DosQmlApplicationEngine, data: DosCStr);
    // fn dos_qqmlapplicationengine_add_import_path(vptr: *mut DosQQmlApplicationEngine, const char *path);
    fn dos_qqmlapplicationengine_context(vptr: DosQmlApplicationEngine) -> DosQQmlContext;
    fn dos_qqmlapplicationengine_delete(vptr: DosQmlApplicationEngine);

    fn dos_qqmlcontext_setcontextproperty(vptr: DosQQmlContext,
                                          name: DosCStr,
                                          value: DosQVariant);

}

/// Provides an entry point for building QML applications from Rust
pub struct QmlEngine {
    ptr: DosQmlApplicationEngine,
    stored: Vec<QVariant>,
}

impl QmlEngine {
    /// Creates a QML context of a non-headless application
    pub fn new() -> Self {
        unsafe {
            dos_qapplication_create();
            QmlEngine {
                ptr: dos_qqmlapplicationengine_create(),
                stored: Vec::new(),
            }
        }
    }

    pub fn registered_type<T: Deref<Target = QObjectMacro>>(&mut self, id: i32, registered: T) {}
    /// Loads a file as a qml file
    pub fn load_file(&self, path: &str) {
        let path_raw = ::std::env::current_dir().unwrap().join(path);
        let path = if cfg!(windows) {
            format!("file:///{}", path_raw.display())
        } else {
            format!("file://{}", path_raw.display())
        };
        unsafe { dos_qqmlapplicationengine_load_url(self.ptr, construct_qurl(&path)) }
    }

    /// Loads a string as a qml file
    pub fn load_data(&self, data: &str) {
        unsafe { dos_qqmlapplicationengine_load_data(self.ptr, stoptr(data)) }
    }

    /// Launches the application
    pub fn exec(&self) {
        unsafe {
            dos_qapplication_exec();
        }
    }
    /// Closes the application
    pub fn quit(&self) {
        unsafe {
            dos_qapplication_quit();
        }
    }

    /// Sets a property for this QML context
    ///
    /// This variant stores qvariant, so it is removed, only when this QmlEngine is removed.
    pub fn set_and_store_property<T: Into<QVariant>>(&mut self, name: &str, value: T) {
        let val = value.into();
        unsafe {
            let context = dos_qqmlapplicationengine_context(self.ptr);
            dos_qqmlcontext_setcontextproperty(context, stoptr(name), get_private_variant(&val));
        }
        self.stored.push(val);
    }

    /// Sets a property for this QML context
    pub fn set_property(&self, name: &str, value: &QVariant) {
        unsafe {
            let context = dos_qqmlapplicationengine_context(self.ptr);
            dos_qqmlcontext_setcontextproperty(context, stoptr(name), get_private_variant(value));
        }
    }
}

use utils::*;

impl Default for QmlEngine {
    fn default() -> Self {
        Self::new()
    }
}

impl Drop for QmlEngine {
    fn drop(&mut self) {
        unsafe {
            dos_qapplication_quit();
            dos_qqmlapplicationengine_delete(self.ptr);
            dos_qapplication_delete();
        }
    }
}
<commit_msg>Added a load_url function to qmlengine<commit_after>use std::ops::Deref;

use qvariant::*;
use types::*;
use qurl::*;
use qmeta::*;

extern "C" {
    fn dos_qapplication_create();
    fn dos_qapplication_exec();
    fn dos_qapplication_quit();
    fn dos_qapplication_delete();

    fn dos_qqmlapplicationengine_create() -> DosQmlApplicationEngine;
    fn dos_qqmlapplicationengine_load(vptr: DosQmlApplicationEngine, filename: DosCStr);
    fn dos_qqmlapplicationengine_load_url(vptr: DosQmlApplicationEngine, url: DosQUrl);
    fn dos_qqmlapplicationengine_load_data(vptr: DosQmlApplicationEngine, data: DosCStr);
    // fn dos_qqmlapplicationengine_add_import_path(vptr: *mut DosQQmlApplicationEngine, const char *path);
    fn dos_qqmlapplicationengine_context(vptr: DosQmlApplicationEngine) -> DosQQmlContext;
    fn dos_qqmlapplicationengine_delete(vptr: DosQmlApplicationEngine);

    fn dos_qqmlcontext_setcontextproperty(vptr: DosQQmlContext,
                                          name: DosCStr,
                                          value: DosQVariant);

}

/// Provides an entry point for building QML applications from Rust
pub struct QmlEngine {
    ptr: DosQmlApplicationEngine,
    stored: Vec<QVariant>,
}

impl QmlEngine {
    /// Creates a QML context of a non-headless application
    pub fn new() -> Self {
        unsafe {
            dos_qapplication_create();
            QmlEngine {
                ptr: dos_qqmlapplicationengine_create(),
                stored: Vec::new(),
            }
        }
    }

    pub fn registered_type<T: Deref<Target = QObjectMacro>>(&mut self, id: i32, registered: T) {}
    /// Loads a file as a qml file
    pub fn load_file(&self, path: &str) {
        let path_raw = ::std::env::current_dir().unwrap().join(path);
        let path = if cfg!(windows) {
            format!("file:///{}", path_raw.display())
        } else {
            format!("file://{}", path_raw.display())
        };
        unsafe { dos_qqmlapplicationengine_load_url(self.ptr, construct_qurl(&path)) }
    }

    /// Loads qml from a specified url (`file://`, `qrc://`, `http://`)
    pub fn load_url(&self, uri: &str) {
        unsafe { dos_qqmlapplicationengine_load_url(self.ptr, construct_qurl(uri)) }
    }

    /// Loads a string as a qml file
    pub fn load_data(&self, data: &str) {
        unsafe { dos_qqmlapplicationengine_load_data(self.ptr, stoptr(data)) }
    }

    /// Launches the application
    pub fn exec(&self) {
        unsafe {
            dos_qapplication_exec();
        }
    }
    /// Closes the application
    pub fn quit(&self) {
        unsafe {
            dos_qapplication_quit();
        }
    }

    /// Sets a property for this QML context
    ///
    /// This variant stores qvariant, so it is removed, only when this QmlEngine is removed.
    pub fn set_and_store_property<T: Into<QVariant>>(&mut self, name: &str, value: T) {
        let val = value.into();
        unsafe {
            let context = dos_qqmlapplicationengine_context(self.ptr);
            dos_qqmlcontext_setcontextproperty(context, stoptr(name), get_private_variant(&val));
        }
        self.stored.push(val);
    }

    /// Sets a property for this QML context
    pub fn set_property(&self, name: &str, value: &QVariant) {
        unsafe {
            let context = dos_qqmlapplicationengine_context(self.ptr);
            dos_qqmlcontext_setcontextproperty(context, stoptr(name), get_private_variant(value));
        }
    }
}

use utils::*;

impl Default for QmlEngine {
    fn default() -> Self {
        Self::new()
    }
}

impl Drop for QmlEngine {
    fn drop(&mut self) {
        unsafe {
            dos_qapplication_quit();
            dos_qqmlapplicationengine_delete(self.ptr);
            dos_qapplication_delete();
        }
    }
}
<|endoftext|>


<commit_before>//! Non-persistant pairwise geometric queries.

pub use self::closest_points::*;
pub use self::contact::*;
pub use self::distance::*;
pub use self::point::*;
pub use self::proximity::*;
pub use self::ray::*;
pub use self::time_of_impact::*;

pub mod algorithms;
mod closest_points;
mod contact;
mod distance;
mod point;
mod proximity;
mod ray;
mod time_of_impact;
pub mod visitors;
<commit_msg>Improve documentation of the query module.<commit_after>//! Non-persistent geometric queries.
//!
//! # General cases
//! The most general methods provided by this module are:
//!
//! * [`closest_points`](query::closest_points()) to compute the closest points between two shapes.
//! * [`distance`](query::distance()) to compute the distance between two shapes.
//! * [`contact`](query::contact()) to compute one pair of contact points between two shapes, including penetrating contact.
//! * [`proximity`](query::proximity()) to determine if two shapes are intersecting or not.
//! * [`time_of_impact`](query::time_of_impact()) to determine when two shapes undergoing translational motions hit for the first time.
//!
//! Ray-casting and point-projection can be achieved by importing traits:
//!
//! * [`RayCast`](query::RayCast) for ray-casting.
//! * [`PointQuery`](query::PointQuery) for point projection.
//!
//! # Specific cases
//! All the other functions exported by this module are more specific versions of the ones described above.
//! For example `distance_ball_ball` computes the distance between two shapes known at compile-time to be balls.
//! They are less convenient to use than the most generic version but will be slightly faster due to the lack of dynamic dispatch.
//! Generally, the specific functions have the form `[operation]_[shape1]_[shape2]()` where:
//!
//! * `[operation]` can be `closest_points`, `distance`, `contact`, `proximity` or `time_of_impact`.
//! * `[shape1]` is the type of the first shape passed to the function, e.g., `ball`, or `plane`.
//! * `[shape2]` is the type of the second shape passed to the function, e.g., `ball`, or `plane`.

pub use self::closest_points::*;
pub use self::contact::*;
pub use self::distance::*;
pub use self::point::*;
pub use self::proximity::*;
pub use self::ray::*;
pub use self::time_of_impact::*;

pub mod algorithms;
mod closest_points;
mod contact;
mod distance;
mod point;
mod proximity;
mod ray;
mod time_of_impact;
pub mod visitors;
<|endoftext|>


<commit_before>use std::cmp;
use std::collections::BinaryHeap;
use std::io;

use raw::build::Builder;
use raw::{Fst, FstStream, Output};
use Result;

pub fn union_ignore_outputs<'f, W, I>(
    bfst: &mut Builder<W>,
    fsts: I,
) -> Result<()>
where W: io::Write,
      I: IntoIterator<Item=&'f Fst> {
    let mut union = Union::new(fsts);
    while let Some(slot) = union.pop() {
        try!(bfst.add(slot.input()));
        union.refill(slot);
    }
    Ok(())
}

#[derive(Copy, Clone, Debug)]
pub struct UnionOutput {
    pub index: usize,
    pub output: u64,
}

pub fn union_with_outputs<'f, W, I, F>(
    bfst: &mut Builder<W>,
    fsts: I,
    mut merge: F,
) -> Result<()>
where W: io::Write,
      I: IntoIterator<Item=&'f Fst>,
      F: FnMut(&[u8], &[UnionOutput]) -> u64 {
    let mut union = Union::new(fsts);
    let mut outs = vec![];
    while let Some(slot) = union.pop() {
        unsafe { outs.set_len(0); }
        outs.push(slot.fst_output());
        while union.peek_is_duplicate(slot.input()) {
            let slot2 = union.pop().unwrap();
            outs.push(slot2.fst_output());
            union.refill(slot2);
        }
        try!(bfst.insert(slot.input(), merge(slot.input(), &outs)));
        union.refill(slot);
    }
    Ok(())
}

struct Union<'f> {
    rdrs: Vec<FstStream<'f>>,
    heap: BinaryHeap<Slot>,
}

impl<'f> Union<'f> {
    fn new<I: IntoIterator<Item=&'f Fst>>(fsts: I) -> Union<'f> {
        let mut u = Union {
            rdrs: fsts.into_iter().map(Fst::stream).collect(),
            heap: BinaryHeap::new(),
        };
        for i in 0..u.rdrs.len() {
            u.refill(Slot::new(i));
        }
        u
    }

    fn pop(&mut self) -> Option<Slot> {
        self.heap.pop()
    }

    fn peek_is_duplicate(&self, key: &[u8]) -> bool {
        self.heap.peek().map(|s| s.input() == key).unwrap_or(false)
    }

    fn refill(&mut self, mut slot: Slot) {
        if let Some((input, output)) = self.rdrs[slot.idx].next() {
            slot.set_input(input);
            slot.set_output(output);
            self.heap.push(slot);
        }
    }
}

#[derive(Debug, Eq, PartialEq)]
struct Slot {
    idx: usize,
    input: Vec<u8>,
    output: Output,
}

impl Slot {
    fn new(rdr_idx: usize) -> Slot {
        Slot {
            idx: rdr_idx,
            input: Vec::with_capacity(64),
            output: Output::zero(),
        }
    }

    fn fst_output(&self) -> UnionOutput {
        UnionOutput { index: self.idx, output: self.output.value() }
    }

    fn input(&self) -> &[u8] {
        &self.input
    }

    fn output(&self) -> Output {
        self.output
    }

    fn set_input(&mut self, input: &[u8]) {
        let addcap = input.len().checked_sub(self.input.len()).unwrap_or(0);
        self.input.reserve(addcap);
        unsafe { self.input.set_len(input.len()); }
        for (dst, &byte) in self.input.iter_mut().zip(input.iter()) {
            *dst = byte;
        }
    }

    fn set_output(&mut self, output: Output) {
        self.output = output;
    }
}

impl PartialOrd for Slot {
    fn partial_cmp(&self, other: &Slot) -> Option<cmp::Ordering> {
        (&self.input, self.output)
        .partial_cmp(&(&other.input, other.output))
        .map(|ord| ord.reverse())
    }
}

impl Ord for Slot {
    fn cmp(&self, other: &Slot) -> cmp::Ordering {
        self.partial_cmp(other).unwrap()
    }
}

#[cfg(test)]
mod tests {
    use raw::build::Builder;
    use raw::tests::{fst_map, fst_set, fst_inputstrs_outputs, fst_input_strs};
    use raw::Fst;
    use super::{union_ignore_outputs, union_with_outputs};

    fn s(string: &str) -> String { string.to_owned() }

    #[test]
    fn basic_set() {
        let set1 = fst_set(&["a", "b", "c"]);
        let set2 = fst_set(&["x", "y", "z"]);

        let mut bfst = Builder::memory();
        union_ignore_outputs(&mut bfst, &[set1, set2]).unwrap();
        let fst = Fst::from_bytes(bfst.into_inner().unwrap()).unwrap();
        assert_eq!(fst_input_strs(&fst), vec!["a", "b", "c", "x", "y", "z"]);
    }

    #[test]
    fn basic_set_dupes() {
        let set1 = fst_set(&["aa", "b", "cc"]);
        let set2 = fst_set(&["b", "cc", "z"]);

        let mut bfst = Builder::memory();
        union_ignore_outputs(&mut bfst, &[set1, set2]).unwrap();
        let fst = Fst::from_bytes(bfst.into_inner().unwrap()).unwrap();
        assert_eq!(fst_input_strs(&fst), vec!["aa", "b", "cc", "z"]);
    }

    #[test]
    fn basic_map() {
        let map1 = fst_map(vec![("a", 1), ("b", 2), ("c", 3)]);
        let map2 = fst_map(vec![("x", 1), ("y", 2), ("z", 3)]);

        let mut bfst = Builder::memory();
        union_with_outputs(
            &mut bfst,
            &[map1, map2],
            |_, os| os.iter().fold(0, |a, b| a + b.output),
        ).unwrap();
        let fst = Fst::from_bytes(bfst.into_inner().unwrap()).unwrap();
        assert_eq!(
            fst_inputstrs_outputs(&fst),
            vec![
                (s("a"), 1), (s("b"), 2), (s("c"), 3),
                (s("x"), 1), (s("y"), 2), (s("z"), 3),
            ]);
    }

    #[test]
    fn basic_map_dupes() {
        let map1 = fst_map(vec![("aa", 1), ("b", 2), ("cc", 3)]);
        let map2 = fst_map(vec![("b", 1), ("cc", 2), ("z", 3)]);
        let map3 = fst_map(vec![("b", 1)]);

        let mut bfst = Builder::memory();
        union_with_outputs(
            &mut bfst,
            &[map1, map2, map3],
            |_, os| os.iter().fold(0, |a, b| a + b.output),
        ).unwrap();
        let fst = Fst::from_bytes(bfst.into_inner().unwrap()).unwrap();
        assert_eq!(
            fst_inputstrs_outputs(&fst),
            vec![
                (s("aa"), 1), (s("b"), 4), (s("cc"), 5), (s("z"), 3),
            ]);
    }
}
<commit_msg>intersection<commit_after>use std::cmp;
use std::collections::BinaryHeap;
use std::io;

use raw::build::Builder;
use raw::{Fst, FstStream, Output};
use Result;

#[derive(Copy, Clone, Debug)]
pub struct FstOutput {
    pub index: usize,
    pub output: u64,
}

pub fn union<'f, W, I>(
    bfst: &mut Builder<W>,
    fsts: I,
) -> Result<()>
where W: io::Write,
      I: IntoIterator<Item=&'f Fst> {
    let mut heap = SlotHeap::new(fsts);
    while let Some(slot) = heap.pop() {
        try!(bfst.add(slot.input()));
        heap.refill(slot);
    }
    Ok(())
}

pub fn union_with<'f, W, I, F>(
    bfst: &mut Builder<W>,
    fsts: I,
    mut merge: F,
) -> Result<()>
where W: io::Write,
      I: IntoIterator<Item=&'f Fst>,
      F: FnMut(&[u8], &[FstOutput]) -> u64 {
    let mut heap = SlotHeap::new(fsts);
    let mut outs = vec![];
    while let Some(slot) = heap.pop() {
        outs.clear();
        outs.push(slot.fst_output());
        while let Some(slot2) = heap.pop_if_equal(slot.input()) {
            outs.push(slot2.fst_output());
            heap.refill(slot2);
        }
        try!(bfst.insert(slot.input(), merge(slot.input(), &outs)));
        heap.refill(slot);
    }
    Ok(())
}

pub fn intersect<'f, W, I>(
    bfst: &mut Builder<W>,
    fsts: I,
) -> Result<()>
where W: io::Write,
      I: IntoIterator<Item=&'f Fst> {
    let mut heap = SlotHeap::new(fsts);
    while let Some(slot) = heap.pop() {
        let mut popped: usize = 1;
        while let Some(slot2) = heap.pop_if_equal(slot.input()) {
            heap.refill(slot2);
            popped += 1;
        }
        // This term is only in the intersection if we found it in
        // every FST given.
        if popped == heap.num_slots() {
            try!(bfst.add(slot.input()));
        }
        heap.refill(slot);
    }
    Ok(())
}

pub fn intersect_with<'f, W, I, F>(
    bfst: &mut Builder<W>,
    fsts: I,
    mut merge: F,
) -> Result<()>
where W: io::Write,
      I: IntoIterator<Item=&'f Fst>,
      F: FnMut(&[u8], &[FstOutput]) -> u64 {
    let mut heap = SlotHeap::new(fsts);
    let mut outs = vec![];
    while let Some(slot) = heap.pop() {
        outs.clear();
        outs.push(slot.fst_output());
        let mut popped: usize = 1;
        while let Some(slot2) = heap.pop_if_equal(slot.input()) {
            outs.push(slot2.fst_output());
            heap.refill(slot2);
            popped += 1;
        }
        // This term is only in the intersection if we found it in
        // every FST given.
        if popped == heap.num_slots() {
            try!(bfst.insert(slot.input(), merge(slot.input(), &outs)));
        }
        heap.refill(slot);
    }
    Ok(())
}

struct SlotHeap<'f> {
    rdrs: Vec<FstStream<'f>>,
    heap: BinaryHeap<Slot>,
}

impl<'f> SlotHeap<'f> {
    fn new<I: IntoIterator<Item=&'f Fst>>(fsts: I) -> SlotHeap<'f> {
        let mut u = SlotHeap {
            rdrs: fsts.into_iter().map(Fst::stream).collect(),
            heap: BinaryHeap::new(),
        };
        for i in 0..u.rdrs.len() {
            u.refill(Slot::new(i));
        }
        u
    }

    fn pop(&mut self) -> Option<Slot> {
        self.heap.pop()
    }

    fn peek_is_duplicate(&self, key: &[u8]) -> bool {
        self.heap.peek().map(|s| s.input() == key).unwrap_or(false)
    }

    fn pop_if_equal(&mut self, key: &[u8]) -> Option<Slot> {
        if self.peek_is_duplicate(key) {
            self.pop()
        } else {
            None
        }
    }

    fn num_slots(&self) -> usize {
        self.rdrs.len()
    }

    fn refill(&mut self, mut slot: Slot) {
        if let Some((input, output)) = self.rdrs[slot.idx].next() {
            slot.set_input(input);
            slot.set_output(output);
            self.heap.push(slot);
        }
    }
}

#[derive(Debug, Eq, PartialEq)]
struct Slot {
    idx: usize,
    input: Vec<u8>,
    output: Output,
}

impl Slot {
    fn new(rdr_idx: usize) -> Slot {
        Slot {
            idx: rdr_idx,
            input: Vec::with_capacity(64),
            output: Output::zero(),
        }
    }

    fn fst_output(&self) -> FstOutput {
        FstOutput { index: self.idx, output: self.output.value() }
    }

    fn input(&self) -> &[u8] {
        &self.input
    }

    fn output(&self) -> Output {
        self.output
    }

    fn set_input(&mut self, input: &[u8]) {
        let addcap = input.len().checked_sub(self.input.len()).unwrap_or(0);
        self.input.clear();
        self.input.extend(input);
    }

    fn set_output(&mut self, output: Output) {
        self.output = output;
    }
}

impl PartialOrd for Slot {
    fn partial_cmp(&self, other: &Slot) -> Option<cmp::Ordering> {
        (&self.input, self.output)
        .partial_cmp(&(&other.input, other.output))
        .map(|ord| ord.reverse())
    }
}

impl Ord for Slot {
    fn cmp(&self, other: &Slot) -> cmp::Ordering {
        self.partial_cmp(other).unwrap()
    }
}

#[cfg(test)]
mod tests {
    use raw::build::Builder;
    use raw::tests::{fst_map, fst_set, fst_inputstrs_outputs, fst_input_strs};
    use raw::Fst;
    use super::{intersect, intersect_with, union, union_with};

    fn s(string: &str) -> String { string.to_owned() }

    #[test]
    fn union_set() {
        let set1 = fst_set(&["a", "b", "c"]);
        let set2 = fst_set(&["x", "y", "z"]);

        let mut bfst = Builder::memory();
        union(&mut bfst, &[set1, set2]).unwrap();
        let fst = Fst::from_bytes(bfst.into_inner().unwrap()).unwrap();
        assert_eq!(fst_input_strs(&fst), vec!["a", "b", "c", "x", "y", "z"]);
    }

    #[test]
    fn union_set_dupes() {
        let set1 = fst_set(&["aa", "b", "cc"]);
        let set2 = fst_set(&["b", "cc", "z"]);

        let mut bfst = Builder::memory();
        union(&mut bfst, &[set1, set2]).unwrap();
        let fst = Fst::from_bytes(bfst.into_inner().unwrap()).unwrap();
        assert_eq!(fst_input_strs(&fst), vec!["aa", "b", "cc", "z"]);
    }

    #[test]
    fn union_map() {
        let map1 = fst_map(vec![("a", 1), ("b", 2), ("c", 3)]);
        let map2 = fst_map(vec![("x", 1), ("y", 2), ("z", 3)]);

        let mut bfst = Builder::memory();
        union_with(
            &mut bfst,
            &[map1, map2],
            |_, os| os.iter().fold(0, |a, b| a + b.output),
        ).unwrap();
        let fst = Fst::from_bytes(bfst.into_inner().unwrap()).unwrap();
        assert_eq!(
            fst_inputstrs_outputs(&fst),
            vec![
                (s("a"), 1), (s("b"), 2), (s("c"), 3),
                (s("x"), 1), (s("y"), 2), (s("z"), 3),
            ]);
    }

    #[test]
    fn union_map_dupes() {
        let map1 = fst_map(vec![("aa", 1), ("b", 2), ("cc", 3)]);
        let map2 = fst_map(vec![("b", 1), ("cc", 2), ("z", 3)]);
        let map3 = fst_map(vec![("b", 1)]);

        let mut bfst = Builder::memory();
        union_with(
            &mut bfst,
            &[map1, map2, map3],
            |_, os| os.iter().fold(0, |a, b| a + b.output),
        ).unwrap();
        let fst = Fst::from_bytes(bfst.into_inner().unwrap()).unwrap();
        assert_eq!(
            fst_inputstrs_outputs(&fst),
            vec![
                (s("aa"), 1), (s("b"), 4), (s("cc"), 5), (s("z"), 3),
            ]);
    }

    #[test]
    fn intersect_set() {
        let set1 = fst_set(&["a", "b", "c"]);
        let set2 = fst_set(&["x", "y", "z"]);

        let mut bfst = Builder::memory();
        intersect(&mut bfst, &[set1, set2]).unwrap();
        let fst = Fst::from_bytes(bfst.into_inner().unwrap()).unwrap();
        assert_eq!(fst_input_strs(&fst), Vec::<&str>::new());
    }

    #[test]
    fn intersect_set_dupes() {
        let set1 = fst_set(&["aa", "b", "cc"]);
        let set2 = fst_set(&["b", "cc", "z"]);

        let mut bfst = Builder::memory();
        intersect(&mut bfst, &[set1, set2]).unwrap();
        let fst = Fst::from_bytes(bfst.into_inner().unwrap()).unwrap();
        assert_eq!(fst_input_strs(&fst), vec!["b", "cc"]);
    }

    #[test]
    fn intersect_map() {
        let map1 = fst_map(vec![("a", 1), ("b", 2), ("c", 3)]);
        let map2 = fst_map(vec![("x", 1), ("y", 2), ("z", 3)]);

        let mut bfst = Builder::memory();
        intersect_with(
            &mut bfst,
            &[map1, map2],
            |_, os| os.iter().fold(0, |a, b| a + b.output),
        ).unwrap();
        let fst = Fst::from_bytes(bfst.into_inner().unwrap()).unwrap();
        assert_eq!(fst_inputstrs_outputs(&fst), Vec::<(String, u64)>::new());
    }

    #[test]
    fn intersect_map_dupes() {
        let map1 = fst_map(vec![("aa", 1), ("b", 2), ("cc", 3)]);
        let map2 = fst_map(vec![("b", 1), ("cc", 2), ("z", 3)]);
        let map3 = fst_map(vec![("b", 1)]);

        let mut bfst = Builder::memory();
        intersect_with(
            &mut bfst,
            &[map1, map2, map3],
            |_, os| os.iter().fold(0, |a, b| a + b.output),
        ).unwrap();
        let fst = Fst::from_bytes(bfst.into_inner().unwrap()).unwrap();
        assert_eq!(fst_inputstrs_outputs(&fst), vec![(s("b"), 4)]);
    }
}
<|endoftext|>


<commit_before>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// https://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! A wrapper around another PRNG that reseeds it after it
//! generates a certain number of random bytes.

use {RngCore, SeedableRng, Error, ErrorKind};

/// A wrapper around any PRNG which reseeds the underlying PRNG after it has
/// generated a certain number of random bytes.
///
/// Reseeding is never strictly *necessary*. Cryptographic PRNGs don't have a
/// limited number of bytes they can output, or at least not a limit reachable
/// in any practical way. There is no such thing as 'running out of entropy'.
///
/// Some small non-cryptographic PRNGs can have very small periods, for
/// example less than 264. Would reseeding help to ensure that you do
/// not wrap around at the end of the period? A period of 264 still
/// takes several centuries of CPU-years on current hardware. Reseeding will
/// actually make things worse, because the reseeded PRNG will just continue
/// somewhere else *in the same period*, with a high chance of overlapping with
/// previously used parts of it.
///
/// # When should you use `ReseedingRng`?
///
/// - Reseeding can be seen as some form of 'security in depth'. Even if in the
///   future a cryptographic weakness is found in the CSPRNG being used,
///   occasionally reseeding should make exploiting it much more difficult or
///   even impossible.
/// - It can be used as a poor man's cryptography (not recommended, just use a
///   good CSPRNG). Previous implementations of `thread_rng` for example used
///   `ReseedingRng` with the ISAAC RNG. That algorithm, although apparently
///   strong and with no known attack, does not come with any proof of security
///   and does not meet the current standards for a cryptographically secure
///   PRNG. By reseeding it frequently (every 32 MiB) it seems safe to assume
///   there is no attack that can operate on the tiny window between reseeds.
///
/// # Error handling
///
/// If reseeding fails, `try_fill_bytes` is the only `Rng` method to report it.
/// For all other `Rng` methods, `ReseedingRng` will not panic but try to
/// handle the error intelligently; if handling the source error fails these
/// methods will continue generating data from the wrapped PRNG without
/// reseeding.
///
/// It is usually best to use the infallible methods `next_u32`, `next_u64` and
/// `fill_bytes` because they can make use of this error handling strategy.
/// Use `try_fill_bytes` and possibly `try_reseed` if you want to handle
/// reseeding errors explicitly.
#[derive(Debug)]
pub struct ReseedingRng<R, Rsdr> {
    rng: R,
    reseeder: Rsdr,
    threshold: i64,
    bytes_until_reseed: i64,
}

impl<R: RngCore + SeedableRng, Rsdr: RngCore> ReseedingRng<R, Rsdr> {
    /// Create a new `ReseedingRng` with the given parameters.
    ///
    /// # Arguments
    ///
    /// * `rng`: the random number generator to use.
    /// * `threshold`: the number of generated bytes after which to reseed the RNG.
    /// * `reseeder`: the RNG to use for reseeding.
    pub fn new(rng: R, threshold: u64, reseeder: Rsdr) -> ReseedingRng<R,Rsdr> {
        assert!(threshold <= ::core::i64::MAX as u64);
        ReseedingRng {
            rng: rng,
            reseeder: reseeder,
            threshold: threshold as i64,
            bytes_until_reseed: threshold as i64,
        }
    }

    /// Reseed the internal PRNG.
    ///
    /// This will try to work around errors in the RNG used for reseeding
    /// intelligently. If the error kind indicates retrying might help, it will
    /// immediately retry a couple of times. If the error kind indicates the
    /// seeding RNG is not ready, it will retry later, after `threshold / 256`
    /// generated bytes. On other errors in the source RNG, this will skip
    /// reseeding and continue using the internal PRNG, until another
    /// `threshold` bytes have been generated (at which point it will try
    /// reseeding again).
    #[inline(never)]
    pub fn reseed(&mut self) {
        trace!("Reseeding RNG after generating {} bytes",
               self.threshold - self.bytes_until_reseed);
        self.bytes_until_reseed = self.threshold;
        let mut err_count = 0;
        loop {
            if let Err(e) = R::from_rng(&mut self.reseeder)
                            .map(|result| self.rng = result) {
                let kind = e.kind();
                if kind.should_wait() {
                    self.bytes_until_reseed = self.threshold >> 8;
                    warn!("Reseeding RNG delayed for {} bytes",
                           self.bytes_until_reseed);
                } else if kind.should_retry() {
                    err_count += 1;
                    // Retry immediately for 5 times (arbitrary limit)
                    if err_count <= 5 { continue; }
                }
                warn!("Reseeding RNG failed; continuing without reseeding. Error: {}", e);
            }
            break; // Successfully reseeded, delayed, or given up.
        }
    }

    /// Reseed the internal RNG if the number of bytes that have been
    /// generated exceed the threshold.
    ///
    /// If reseeding fails, return an error with the original cause. Note that
    /// if the cause has a permanent failure, we report a transient error and
    /// skip reseeding; this means that only two error kinds can be reported
    /// from this method: `ErrorKind::Transient` and `ErrorKind::NotReady`.
    #[inline(never)]
    pub fn try_reseed(&mut self) -> Result<(), Error> {
        trace!("Reseeding RNG after {} generated bytes",
               self.threshold - self.bytes_until_reseed);
        if let Err(err) = R::from_rng(&mut self.reseeder)
                          .map(|result| self.rng = result) {
            let newkind = match err.kind() {
                a @ ErrorKind::NotReady => a,
                b @ ErrorKind::Transient => b,
                _ => {
                    self.bytes_until_reseed = self.threshold; // skip reseeding
                    ErrorKind::Transient
                }
            };
            return Err(Error::with_cause(newkind, "reseeding failed", err));
        }
        self.bytes_until_reseed = self.threshold;
        Ok(())
    }
}

impl<R: RngCore + SeedableRng, Rsdr: RngCore> RngCore for ReseedingRng<R, Rsdr> {
    fn next_u32(&mut self) -> u32 {
        let value = self.rng.next_u32();
        self.bytes_until_reseed -= 4;
        if self.bytes_until_reseed <= 0 {
            self.reseed();
        }
        value
    }

    fn next_u64(&mut self) -> u64 {
        let value = self.rng.next_u64();
        self.bytes_until_reseed -= 8;
        if self.bytes_until_reseed <= 0 {
            self.reseed();
        }
        value
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        self.rng.fill_bytes(dest);
        self.bytes_until_reseed -= dest.len() as i64;
        if self.bytes_until_reseed <= 0 {
            self.reseed();
        }
    }

    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        self.rng.try_fill_bytes(dest)?;
        self.bytes_until_reseed -= dest.len() as i64;
        if self.bytes_until_reseed <= 0 {
            self.try_reseed()?;
        }
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use {impls, le};
    use super::{ReseedingRng};
    use {SeedableRng, RngCore, Error};

    struct Counter {
        i: u32
    }

    impl RngCore for Counter {
        fn next_u32(&mut self) -> u32 {
            self.i += 1;
            // very random
            self.i - 1
        }
        fn next_u64(&mut self) -> u64 {
            impls::next_u64_via_u32(self)
        }

        fn fill_bytes(&mut self, dest: &mut [u8]) {
            impls::fill_bytes_via_u64(self, dest)
        }
    }
    impl SeedableRng for Counter {
        type Seed = [u8; 4];
        fn from_seed(seed: Self::Seed) -> Self {
            let mut seed_u32 = [0u32; 1];
            le::read_u32_into(&seed, &mut seed_u32);
            Counter { i: seed_u32[0] }
        }
    }

    #[derive(Debug, Clone)]
    struct ResetCounter;
    impl RngCore for ResetCounter {
        fn next_u32(&mut self) -> u32 { unimplemented!() }
        fn next_u64(&mut self) -> u64 { unimplemented!() }
        fn fill_bytes(&mut self, _dest: &mut [u8]) { unimplemented!() }
        fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
            for i in dest.iter_mut() { *i = 0; }
            Ok(())
        }
    }

    #[test]
    fn test_reseeding() {
        let mut rs = ReseedingRng::new(Counter {i:0}, 400, ResetCounter);

        let mut i = 0;
        for _ in 0..1000 {
            assert_eq!(rs.next_u32(), i % 100);
            i += 1;
        }
    }
}
<commit_msg>Reseeding test: replace custom RNGs with standard ones<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// https://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! A wrapper around another PRNG that reseeds it after it
//! generates a certain number of random bytes.

use {RngCore, SeedableRng, Error, ErrorKind};

/// A wrapper around any PRNG which reseeds the underlying PRNG after it has
/// generated a certain number of random bytes.
///
/// Reseeding is never strictly *necessary*. Cryptographic PRNGs don't have a
/// limited number of bytes they can output, or at least not a limit reachable
/// in any practical way. There is no such thing as 'running out of entropy'.
///
/// Some small non-cryptographic PRNGs can have very small periods, for
/// example less than 264. Would reseeding help to ensure that you do
/// not wrap around at the end of the period? A period of 264 still
/// takes several centuries of CPU-years on current hardware. Reseeding will
/// actually make things worse, because the reseeded PRNG will just continue
/// somewhere else *in the same period*, with a high chance of overlapping with
/// previously used parts of it.
///
/// # When should you use `ReseedingRng`?
///
/// - Reseeding can be seen as some form of 'security in depth'. Even if in the
///   future a cryptographic weakness is found in the CSPRNG being used,
///   occasionally reseeding should make exploiting it much more difficult or
///   even impossible.
/// - It can be used as a poor man's cryptography (not recommended, just use a
///   good CSPRNG). Previous implementations of `thread_rng` for example used
///   `ReseedingRng` with the ISAAC RNG. That algorithm, although apparently
///   strong and with no known attack, does not come with any proof of security
///   and does not meet the current standards for a cryptographically secure
///   PRNG. By reseeding it frequently (every 32 MiB) it seems safe to assume
///   there is no attack that can operate on the tiny window between reseeds.
///
/// # Error handling
///
/// If reseeding fails, `try_fill_bytes` is the only `Rng` method to report it.
/// For all other `Rng` methods, `ReseedingRng` will not panic but try to
/// handle the error intelligently; if handling the source error fails these
/// methods will continue generating data from the wrapped PRNG without
/// reseeding.
///
/// It is usually best to use the infallible methods `next_u32`, `next_u64` and
/// `fill_bytes` because they can make use of this error handling strategy.
/// Use `try_fill_bytes` and possibly `try_reseed` if you want to handle
/// reseeding errors explicitly.
#[derive(Debug)]
pub struct ReseedingRng<R, Rsdr> {
    rng: R,
    reseeder: Rsdr,
    threshold: i64,
    bytes_until_reseed: i64,
}

impl<R: RngCore + SeedableRng, Rsdr: RngCore> ReseedingRng<R, Rsdr> {
    /// Create a new `ReseedingRng` with the given parameters.
    ///
    /// # Arguments
    ///
    /// * `rng`: the random number generator to use.
    /// * `threshold`: the number of generated bytes after which to reseed the RNG.
    /// * `reseeder`: the RNG to use for reseeding.
    pub fn new(rng: R, threshold: u64, reseeder: Rsdr) -> ReseedingRng<R,Rsdr> {
        assert!(threshold <= ::core::i64::MAX as u64);
        ReseedingRng {
            rng: rng,
            reseeder: reseeder,
            threshold: threshold as i64,
            bytes_until_reseed: threshold as i64,
        }
    }

    /// Reseed the internal PRNG.
    ///
    /// This will try to work around errors in the RNG used for reseeding
    /// intelligently. If the error kind indicates retrying might help, it will
    /// immediately retry a couple of times. If the error kind indicates the
    /// seeding RNG is not ready, it will retry later, after `threshold / 256`
    /// generated bytes. On other errors in the source RNG, this will skip
    /// reseeding and continue using the internal PRNG, until another
    /// `threshold` bytes have been generated (at which point it will try
    /// reseeding again).
    #[inline(never)]
    pub fn reseed(&mut self) {
        trace!("Reseeding RNG after generating {} bytes",
               self.threshold - self.bytes_until_reseed);
        self.bytes_until_reseed = self.threshold;
        let mut err_count = 0;
        loop {
            if let Err(e) = R::from_rng(&mut self.reseeder)
                            .map(|result| self.rng = result) {
                let kind = e.kind();
                if kind.should_wait() {
                    self.bytes_until_reseed = self.threshold >> 8;
                    warn!("Reseeding RNG delayed for {} bytes",
                           self.bytes_until_reseed);
                } else if kind.should_retry() {
                    err_count += 1;
                    // Retry immediately for 5 times (arbitrary limit)
                    if err_count <= 5 { continue; }
                }
                warn!("Reseeding RNG failed; continuing without reseeding. Error: {}", e);
            }
            break; // Successfully reseeded, delayed, or given up.
        }
    }

    /// Reseed the internal RNG if the number of bytes that have been
    /// generated exceed the threshold.
    ///
    /// If reseeding fails, return an error with the original cause. Note that
    /// if the cause has a permanent failure, we report a transient error and
    /// skip reseeding; this means that only two error kinds can be reported
    /// from this method: `ErrorKind::Transient` and `ErrorKind::NotReady`.
    #[inline(never)]
    pub fn try_reseed(&mut self) -> Result<(), Error> {
        trace!("Reseeding RNG after {} generated bytes",
               self.threshold - self.bytes_until_reseed);
        if let Err(err) = R::from_rng(&mut self.reseeder)
                          .map(|result| self.rng = result) {
            let newkind = match err.kind() {
                a @ ErrorKind::NotReady => a,
                b @ ErrorKind::Transient => b,
                _ => {
                    self.bytes_until_reseed = self.threshold; // skip reseeding
                    ErrorKind::Transient
                }
            };
            return Err(Error::with_cause(newkind, "reseeding failed", err));
        }
        self.bytes_until_reseed = self.threshold;
        Ok(())
    }
}

impl<R: RngCore + SeedableRng, Rsdr: RngCore> RngCore for ReseedingRng<R, Rsdr> {
    fn next_u32(&mut self) -> u32 {
        let value = self.rng.next_u32();
        self.bytes_until_reseed -= 4;
        if self.bytes_until_reseed <= 0 {
            self.reseed();
        }
        value
    }

    fn next_u64(&mut self) -> u64 {
        let value = self.rng.next_u64();
        self.bytes_until_reseed -= 8;
        if self.bytes_until_reseed <= 0 {
            self.reseed();
        }
        value
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        self.rng.fill_bytes(dest);
        self.bytes_until_reseed -= dest.len() as i64;
        if self.bytes_until_reseed <= 0 {
            self.reseed();
        }
    }

    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        self.rng.try_fill_bytes(dest)?;
        self.bytes_until_reseed -= dest.len() as i64;
        if self.bytes_until_reseed <= 0 {
            self.try_reseed()?;
        }
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use {Rng, SeedableRng, StdRng};
    use mock::StepRng;
    use super::ReseedingRng;

    #[test]
    fn test_reseeding() {
        let mut zero = StepRng::new(0, 0);
        let rng = StdRng::from_rng(&mut zero).unwrap();
        let mut reseeding = ReseedingRng::new(rng, 32, zero);

        // Currently we only support for arrays up to length 32.
        // TODO: cannot generate seq via Rng::gen because it uses different alg
        let mut buf = [0u8; 32];
        reseeding.fill(&mut buf);
        let seq = buf;
        for _ in 0..10 {
            reseeding.fill(&mut buf);
            assert_eq!(buf, seq);
        }
    }
}
<|endoftext|>