<commit_before>use super::placement::sm::client::SMClient;
use bifrost::rpc::*;
use bifrost_plugins::hash_ident;
use client::AsyncClient;
use dovahkiin::types::custom_types::id::Id;
use index::btree::storage::store_changed_nodes;
use index::lsmtree::service::inner::LSMTreeIns;
use index::lsmtree::split::check_and_split;
use index::lsmtree::tree::{LSMTree, LSMTreeResult};
use index::{EntryKey, Ordering};
use itertools::Itertools;
use parking_lot::RwLock;
use rayon::iter::IntoParallelRefIterator;
use rayon::iter::ParallelIterator;
use server::NebServer;
use smallvec::SmallVec;
use std::collections::HashMap;
use std::path::Component::CurDir;
use std::sync::atomic;
use std::sync::atomic::AtomicU64;
use std::thread;

mod inner;
#[cfg(test)]
mod test;

pub static DEFAULT_SERVICE_ID: u64 = hash_ident!(LSM_TREE_RPC_SERVICE) as u64;

#[derive(Serialize, Deserialize, Debug)]
pub enum LSMTreeSvrError {
    TreeNotFound,
    SMError,
}

#[derive(Serialize, Deserialize)]
pub struct LSMTreeSummary {
    id: Id,
    count: u64,
    epoch: u64,
    range: (Vec<u8>, Vec<u8>),
}

service! {
    rpc seek(tree_id: Id, key: Vec<u8>, ordering: Ordering, epoch: u64) -> LSMTreeResult<u64> | LSMTreeSvrError;
    rpc next(tree_id: Id, cursor_id: u64) -> Option<bool> | LSMTreeSvrError;
    rpc current(tree_id: Id, cursor_id: u64) -> Option<Option<Vec<u8>>> | LSMTreeSvrError;
    rpc complete(tree_id: Id, cursor_id: u64) -> bool | LSMTreeSvrError;
    rpc new_tree(start: Vec<u8>, end: Vec<u8>, id: Id) -> bool;
    rpc summary() -> Vec<LSMTreeSummary>;

    rpc insert(tree_id: Id, key: Vec<u8>, epoch: u64) -> LSMTreeResult<bool> | LSMTreeSvrError;
    rpc merge(tree_id: Id, keys: Vec<Vec<u8>>, epoch: u64) -> LSMTreeResult<()> | LSMTreeSvrError;
    rpc set_epoch(tree_id: Id, epoch: u64) -> u64 | LSMTreeSvrError;
}

pub struct LSMTreeService {
    neb_server: Arc<NebServer>,
    sm: Arc<SMClient>,
    counter: AtomicU64,
    trees: Arc<RwLock<HashMap<Id, Arc<LSMTreeIns>>>>,
}

dispatch_rpc_service_functions!(LSMTreeService);

impl Service for LSMTreeService {
    fn seek(
        &self,
        tree_id: Id,
        key: Vec<u8>,
        ordering: Ordering,
        epoch: u64,
    ) -> Box<Future<Item = LSMTreeResult<u64>, Error = LSMTreeSvrError>> {
        let trees = self.trees.read();
        box future::result(
            trees
                .get(&tree_id)
                .ok_or(LSMTreeSvrError::TreeNotFound)
                .map(|tree| {
                    tree.with_epoch_check(epoch, || tree.seek(&SmallVec::from(key), ordering))
                }),
        )
    }

    fn next(
        &self,
        tree_id: Id,
        cursor_id: u64,
    ) -> Box<Future<Item = Option<bool>, Error = LSMTreeSvrError>> {
        let trees = self.trees.read();
        box future::result(
            trees
                .get(&tree_id)
                .ok_or(LSMTreeSvrError::TreeNotFound)
                .map(|tree| tree.next(&cursor_id)),
        )
    }

    fn current(
        &self,
        tree_id: Id,
        cursor_id: u64,
    ) -> Box<Future<Item = Option<Option<Vec<u8>>>, Error = LSMTreeSvrError>> {
        let trees = self.trees.read();
        box future::result(
            trees
                .get(&tree_id)
                .ok_or(LSMTreeSvrError::TreeNotFound)
                .map(|tree| tree.current(&cursor_id)),
        )
    }

    fn complete(
        &self,
        tree_id: Id,
        cursor_id: u64,
    ) -> Box<Future<Item = bool, Error = LSMTreeSvrError>> {
        let trees = self.trees.read();
        box future::result(
            trees
                .get(&tree_id)
                .ok_or(LSMTreeSvrError::TreeNotFound)
                .map(|tree| tree.complete(&cursor_id)),
        )
    }

    fn new_tree(
        &self,
        start: Vec<u8>,
        end: Vec<u8>,
        id: Id,
    ) -> Box<Future<Item = bool, Error = ()>> {
        let mut trees = self.trees.write();
        let succeed = if trees.contains_key(&id) {
            false
        } else {
            trees.insert(
                id,
                Arc::new(LSMTreeIns::new(
                    (EntryKey::from(start), EntryKey::from(end)),
                    id,
                )),
            );
            true
        };
        box persist(self.neb_server.clone(), succeed)
    }

    fn summary(&self) -> Box<Future<Item = Vec<LSMTreeSummary>, Error = ()>> {
        let trees = self.trees.read();
        box future::ok(
            trees
                .iter()
                .map(|(id, tree)| LSMTreeSummary {
                    id: *id,
                    count: tree.count(),
                    epoch: tree.epoch(),
                    range: tree.range(),
                })
                .sorted_by_key(|tree| tree.range.0.clone())
                .collect(),
        )
    }

    fn insert(
        &self,
        tree_id: Id,
        key: Vec<u8>,
        epoch: u64,
    ) -> Box<Future<Item = LSMTreeResult<bool>, Error = LSMTreeSvrError>> {
        let trees = self.trees.read();
        let neb = self.neb_server.clone();
        box future::result(
            trees
                .get(&tree_id)
                .ok_or(LSMTreeSvrError::TreeNotFound)
                .map(|tree| tree.with_epoch_check(epoch, || tree.insert(SmallVec::from(key)))),
        )
        .and_then(|o| persist(neb, o))
    }

    fn merge(
        &self,
        tree_id: Id,
        keys: Vec<Vec<u8>>,
        epoch: u64,
    ) -> Box<Future<Item = LSMTreeResult<()>, Error = LSMTreeSvrError>> {
        let trees = self.trees.read();
        let neb = self.neb_server.clone();
        box future::result(
            trees
                .get(&tree_id)
                .ok_or(LSMTreeSvrError::TreeNotFound)
                .map(|tree| {
                    tree.with_epoch_check(epoch, || {
                        tree.merge(box keys.into_iter().map(|key| SmallVec::from(key)).collect())
                    })
                }),
        )
        .and_then(|o| persist(neb, o))
    }

    fn set_epoch(
        &self,
        tree_id: Id,
        epoch: u64,
    ) -> Box<Future<Item = u64, Error = LSMTreeSvrError>> {
        let trees = self.trees.read();
        let sm = self.sm.clone();
        box future::result(
            trees
                .get(&tree_id)
                .ok_or(LSMTreeSvrError::TreeNotFound)
                .map(|tree| tree.set_epoch(epoch)),
        )
        .and_then(move |_| {
            sm.update_epoch(&tree_id, &epoch)
                .map_err(|_| LSMTreeSvrError::SMError)
                .and_then(|r| r.map_err(|_| LSMTreeSvrError::SMError))
        })
    }
}

impl LSMTreeService {
    pub fn new(neb_server: &Arc<NebServer>, sm: &Arc<SMClient>) -> Arc<Self> {
        Arc::new(Self {
            neb_server: neb_server.clone(),
            counter: AtomicU64::new(0),
            trees: Arc::new(RwLock::new(HashMap::new())),
            sm: sm.clone(),
        })
    }
}

fn persist<T, E>(neb: Arc<NebServer>, val: T) -> impl Future<Item = T, Error = E> {
    store_changed_nodes(neb).then(move |_| future::ok(val))
}
<commit_msg>lsm tree service sentinel<commit_after>use super::placement::sm::client::SMClient;
use bifrost::rpc::*;
use bifrost_plugins::hash_ident;
use client::AsyncClient;
use dovahkiin::types::custom_types::id::Id;
use index::btree::storage::store_changed_nodes;
use index::lsmtree::service::inner::LSMTreeIns;
use index::lsmtree::split::check_and_split;
use index::lsmtree::tree::{LSMTree, LSMTreeResult};
use index::{EntryKey, Ordering};
use itertools::Itertools;
use parking_lot::RwLock;
use rayon::iter::IntoParallelRefIterator;
use rayon::iter::ParallelIterator;
use server::NebServer;
use smallvec::SmallVec;
use std::collections::HashMap;
use std::path::Component::CurDir;
use std::sync::atomic;
use std::sync::atomic::AtomicU64;
use std::thread;

mod inner;
#[cfg(test)]
mod test;

pub static DEFAULT_SERVICE_ID: u64 = hash_ident!(LSM_TREE_RPC_SERVICE) as u64;

#[derive(Serialize, Deserialize, Debug)]
pub enum LSMTreeSvrError {
    TreeNotFound,
    SMError,
}

#[derive(Serialize, Deserialize)]
pub struct LSMTreeSummary {
    id: Id,
    count: u64,
    epoch: u64,
    range: (Vec<u8>, Vec<u8>),
}

service! {
    rpc seek(tree_id: Id, key: Vec<u8>, ordering: Ordering, epoch: u64) -> LSMTreeResult<u64> | LSMTreeSvrError;
    rpc next(tree_id: Id, cursor_id: u64) -> Option<bool> | LSMTreeSvrError;
    rpc current(tree_id: Id, cursor_id: u64) -> Option<Option<Vec<u8>>> | LSMTreeSvrError;
    rpc complete(tree_id: Id, cursor_id: u64) -> bool | LSMTreeSvrError;
    rpc new_tree(start: Vec<u8>, end: Vec<u8>, id: Id) -> bool;
    rpc summary() -> Vec<LSMTreeSummary>;

    rpc insert(tree_id: Id, key: Vec<u8>, epoch: u64) -> LSMTreeResult<bool> | LSMTreeSvrError;
    rpc merge(tree_id: Id, keys: Vec<Vec<u8>>, epoch: u64) -> LSMTreeResult<()> | LSMTreeSvrError;
    rpc set_epoch(tree_id: Id, epoch: u64) -> u64 | LSMTreeSvrError;
}

pub struct LSMTreeService {
    neb_server: Arc<NebServer>,
    sm: Arc<SMClient>,
    counter: AtomicU64,
    trees: Arc<RwLock<HashMap<Id, Arc<LSMTreeIns>>>>,
}

dispatch_rpc_service_functions!(LSMTreeService);

impl Service for LSMTreeService {
    fn seek(
        &self,
        tree_id: Id,
        key: Vec<u8>,
        ordering: Ordering,
        epoch: u64,
    ) -> Box<Future<Item = LSMTreeResult<u64>, Error = LSMTreeSvrError>> {
        let trees = self.trees.read();
        box future::result(
            trees
                .get(&tree_id)
                .ok_or(LSMTreeSvrError::TreeNotFound)
                .map(|tree| {
                    tree.with_epoch_check(epoch, || tree.seek(&SmallVec::from(key), ordering))
                }),
        )
    }

    fn next(
        &self,
        tree_id: Id,
        cursor_id: u64,
    ) -> Box<Future<Item = Option<bool>, Error = LSMTreeSvrError>> {
        let trees = self.trees.read();
        box future::result(
            trees
                .get(&tree_id)
                .ok_or(LSMTreeSvrError::TreeNotFound)
                .map(|tree| tree.next(&cursor_id)),
        )
    }

    fn current(
        &self,
        tree_id: Id,
        cursor_id: u64,
    ) -> Box<Future<Item = Option<Option<Vec<u8>>>, Error = LSMTreeSvrError>> {
        let trees = self.trees.read();
        box future::result(
            trees
                .get(&tree_id)
                .ok_or(LSMTreeSvrError::TreeNotFound)
                .map(|tree| tree.current(&cursor_id)),
        )
    }

    fn complete(
        &self,
        tree_id: Id,
        cursor_id: u64,
    ) -> Box<Future<Item = bool, Error = LSMTreeSvrError>> {
        let trees = self.trees.read();
        box future::result(
            trees
                .get(&tree_id)
                .ok_or(LSMTreeSvrError::TreeNotFound)
                .map(|tree| tree.complete(&cursor_id)),
        )
    }

    fn new_tree(
        &self,
        start: Vec<u8>,
        end: Vec<u8>,
        id: Id,
    ) -> Box<Future<Item = bool, Error = ()>> {
        let mut trees = self.trees.write();
        let succeed = if trees.contains_key(&id) {
            false
        } else {
            trees.insert(
                id,
                Arc::new(LSMTreeIns::new(
                    (EntryKey::from(start), EntryKey::from(end)),
                    id,
                )),
            );
            true
        };
        box persist(self.neb_server.clone(), succeed)
    }

    fn summary(&self) -> Box<Future<Item = Vec<LSMTreeSummary>, Error = ()>> {
        let trees = self.trees.read();
        box future::ok(
            trees
                .iter()
                .map(|(id, tree)| LSMTreeSummary {
                    id: *id,
                    count: tree.count(),
                    epoch: tree.epoch(),
                    range: tree.range(),
                })
                .sorted_by_key(|tree| tree.range.0.clone())
                .collect(),
        )
    }

    fn insert(
        &self,
        tree_id: Id,
        key: Vec<u8>,
        epoch: u64,
    ) -> Box<Future<Item = LSMTreeResult<bool>, Error = LSMTreeSvrError>> {
        let trees = self.trees.read();
        let neb = self.neb_server.clone();
        box future::result(
            trees
                .get(&tree_id)
                .ok_or(LSMTreeSvrError::TreeNotFound)
                .map(|tree| tree.with_epoch_check(epoch, || tree.insert(SmallVec::from(key)))),
        )
        .and_then(|o| persist(neb, o))
    }

    fn merge(
        &self,
        tree_id: Id,
        keys: Vec<Vec<u8>>,
        epoch: u64,
    ) -> Box<Future<Item = LSMTreeResult<()>, Error = LSMTreeSvrError>> {
        let trees = self.trees.read();
        let neb = self.neb_server.clone();
        box future::result(
            trees
                .get(&tree_id)
                .ok_or(LSMTreeSvrError::TreeNotFound)
                .map(|tree| {
                    tree.with_epoch_check(epoch, || {
                        tree.merge(box keys.into_iter().map(|key| SmallVec::from(key)).collect())
                    })
                }),
        )
        .and_then(|o| persist(neb, o))
    }

    fn set_epoch(
        &self,
        tree_id: Id,
        epoch: u64,
    ) -> Box<Future<Item = u64, Error = LSMTreeSvrError>> {
        let trees = self.trees.read();
        let sm = self.sm.clone();
        box future::result(
            trees
                .get(&tree_id)
                .ok_or(LSMTreeSvrError::TreeNotFound)
                .map(|tree| tree.set_epoch(epoch)),
        )
        .and_then(move |_| {
            sm.update_epoch(&tree_id, &epoch)
                .map_err(|_| LSMTreeSvrError::SMError)
                .and_then(|r| r.map_err(|_| LSMTreeSvrError::SMError))
        })
    }
}

impl LSMTreeService {
    pub fn new(neb_server: &Arc<NebServer>, sm: &Arc<SMClient>) -> Arc<Self> {
        let trees = Arc::new(RwLock::new(HashMap::new()));
        let trees_clone = trees.clone();
        let neb = neb_server.clone();
        thread::Builder::new()
            .name("LSM-tree service sentinel".to_string())
            .spawn(move || {
                let tree_map = trees_clone.read();
                tree_map.values().cloned().collect_vec().par_iter().for_each(|tree: &Arc<LSMTreeIns>| {
                    tree.check_and_merge();
                });
                persist::<_, ()>(neb.clone(), ()).wait().unwrap()
            });
        Arc::new(Self {
            neb_server: neb_server.clone(),
            counter: AtomicU64::new(0),
            sm: sm.clone(),
            trees,
        })
    }
}

fn persist<T, E>(neb: Arc<NebServer>, val: T) -> impl Future<Item = T, Error = E> {
    store_changed_nodes(neb).then(move |_| future::ok(val))
}
<|endoftext|>


<commit_before>//! Reject certain requests as instance load reaches capacity.
//!
//! The primary goal of this middleware is to avoid starving the download endpoint of resources.
//! When bots send many parallel requests that run slow database queries, download requests may
//! block and eventually timeout waiting for a database connection.
//!
//! Bots must continue to respect our crawler policy, but until we can manually block bad clients
//! we should avoid dropping download requests even if that means rejecting some legitimate
//! requests to other endpoints.

use std::sync::atomic::{AtomicUsize, Ordering};

use super::prelude::*;
use conduit::{RequestExt, StatusCode};

#[derive(Default)]
pub(super) struct BalanceCapacity {
    handler: Option<Box<dyn Handler>>,
    capacity: usize,
    in_flight_requests: AtomicUsize,
}

impl BalanceCapacity {
    pub fn new(capacity: usize) -> Self {
        Self {
            handler: None,
            capacity,
            in_flight_requests: AtomicUsize::new(0),
        }
    }
}

impl AroundMiddleware for BalanceCapacity {
    fn with_handler(&mut self, handler: Box<dyn Handler>) {
        self.handler = Some(handler);
    }
}

impl Handler for BalanceCapacity {
    fn call(&self, request: &mut dyn RequestExt) -> AfterResult {
        // The _drop_on_exit ensures the counter is decremented for all exit paths (including panics)
        let (_drop_on_exit, count) = RequestCounter::add_one(&self.in_flight_requests);
        let handler = self.handler.as_ref().unwrap();
        let load = 100 * count / self.capacity;

        // Begin logging request count at 20% capacity
        if load >= 20 {
            super::log_request::add_custom_metadata(request, "in_flight_requests", count);
        }

        // Download requests are always accepted
        if request.path().starts_with("/api/v1/crates/") && request.path().ends_with("/download") {
            return handler.call(request);
        }

        // Reject read-only requests after reaching 70% load. Bots are likely to send only safe
        // requests and this helps prioritize requests that users may be reluctant to retry.
        if load >= 70 && request.method().is_safe() {
            return over_capcity_response();
        }

        // At 80% load, all non-download requests are rejected
        if load >= 80 {
            return over_capcity_response();
        }

        handler.call(request)
    }
}

fn over_capcity_response() -> AfterResult {
    // TODO: Generate an alert so we can investigate
    let body = "Service temporarily unavailable";
    Response::builder()
        .status(StatusCode::SERVICE_UNAVAILABLE)
        .header(header::CONTENT_LENGTH, body.len())
        .body(Body::from_static(body.as_bytes()))
        .map_err(box_error)
}

// FIXME(JTG): I've copied the following from my `conduit-hyper` crate.  Once we transition from
// `civet`, we could pass the in_flight_request count from `condut-hyper` via a request extension.

/// A struct that stores a reference to an atomic counter so it can be decremented when dropped
struct RequestCounter<'a> {
    counter: &'a AtomicUsize,
}

impl<'a> RequestCounter<'a> {
    fn add_one(counter: &'a AtomicUsize) -> (Self, usize) {
        let previous = counter.fetch_add(1, Ordering::SeqCst);
        (Self { counter }, previous + 1)
    }
}

impl<'a> Drop for RequestCounter<'a> {
    fn drop(&mut self) {
        self.counter.fetch_sub(1, Ordering::SeqCst);
    }
}
<commit_msg>Allow load limits to be configured from environment<commit_after>//! Reject certain requests as instance load reaches capacity.
//!
//! The primary goal of this middleware is to avoid starving the download endpoint of resources.
//! When bots send many parallel requests that run slow database queries, download requests may
//! block and eventually timeout waiting for a database connection.
//!
//! Bots must continue to respect our crawler policy, but until we can manually block bad clients
//! we should avoid dropping download requests even if that means rejecting some legitimate
//! requests to other endpoints.

use std::sync::atomic::{AtomicUsize, Ordering};

use super::prelude::*;
use conduit::{RequestExt, StatusCode};

#[derive(Default)]
pub(super) struct BalanceCapacity {
    handler: Option<Box<dyn Handler>>,
    capacity: usize,
    in_flight_requests: AtomicUsize,
    log_at_percentage: usize,
    throttle_at_percentage: usize,
    dl_only_at_percentage: usize,
}

impl BalanceCapacity {
    pub fn new(capacity: usize) -> Self {
        Self {
            handler: None,
            capacity,
            in_flight_requests: AtomicUsize::new(0),
            log_at_percentage: read_env_percentage("WEB_CAPACITY_LOG_PCT", 20),
            throttle_at_percentage: read_env_percentage("WEB_CAPACITY_THROTTLE_PCT", 70),
            dl_only_at_percentage: read_env_percentage("WEB_CAPACITY_DL_ONLY_PCT", 80),
        }
    }
}

impl AroundMiddleware for BalanceCapacity {
    fn with_handler(&mut self, handler: Box<dyn Handler>) {
        self.handler = Some(handler);
    }
}

impl Handler for BalanceCapacity {
    fn call(&self, request: &mut dyn RequestExt) -> AfterResult {
        // The _drop_on_exit ensures the counter is decremented for all exit paths (including panics)
        let (_drop_on_exit, count) = RequestCounter::add_one(&self.in_flight_requests);
        let handler = self.handler.as_ref().unwrap();
        let load = 100 * count / self.capacity;

        // Begin logging request count so early stages of load increase can be located
        if load >= self.log_at_percentage {
            super::log_request::add_custom_metadata(request, "in_flight_requests", count);
        }

        // Download requests are always accepted
        if request.path().starts_with("/api/v1/crates/") && request.path().ends_with("/download") {
            return handler.call(request);
        }

        // Reject read-only requests as load nears capacity. Bots are likely to send only safe
        // requests and this helps prioritize requests that users may be reluctant to retry.
        if load >= self.throttle_at_percentage && request.method().is_safe() {
            return over_capacity_response(request);
        }

        // As load reaches capacity, all non-download requests are rejected
        if load >= self.dl_only_at_percentage {
            return over_capacity_response(request);
        }

        handler.call(request)
    }
}

fn over_capacity_response(request: &mut dyn RequestExt) -> AfterResult {
    // TODO: Generate an alert so we can investigate
    super::log_request::add_custom_metadata(request, "cause", "over capacity");
    let body = "Service temporarily unavailable";
    Response::builder()
        .status(StatusCode::SERVICE_UNAVAILABLE)
        .header(header::CONTENT_LENGTH, body.len())
        .body(Body::from_static(body.as_bytes()))
        .map_err(box_error)
}

fn read_env_percentage(name: &str, default: usize) -> usize {
    if let Ok(value) = std::env::var(name) {
        value.parse().unwrap_or(default)
    } else {
        default
    }
}

// FIXME(JTG): I've copied the following from my `conduit-hyper` crate.  Once we transition from
// `civet`, we could pass the in_flight_request count from `condut-hyper` via a request extension.

/// A struct that stores a reference to an atomic counter so it can be decremented when dropped
struct RequestCounter<'a> {
    counter: &'a AtomicUsize,
}

impl<'a> RequestCounter<'a> {
    fn add_one(counter: &'a AtomicUsize) -> (Self, usize) {
        let previous = counter.fetch_add(1, Ordering::SeqCst);
        (Self { counter }, previous + 1)
    }
}

impl<'a> Drop for RequestCounter<'a> {
    fn drop(&mut self) {
        self.counter.fetch_sub(1, Ordering::SeqCst);
    }
}
<|endoftext|>


<commit_before>use std::pin::Pin;
use std::task::Context;
use std::task::Poll;

use futures::stream::{Fuse, Stream};

pub trait SelectWithWeakExt: Stream {
	fn select_with_weak<S>(self, other: S) -> SelectWithWeak<Self, S>
	where
		S: Stream<Item = Self::Item>,
		Self: Sized;
}

impl<T> SelectWithWeakExt for T
where
	T: Stream,
{
	fn select_with_weak<S>(self, other: S) -> SelectWithWeak<Self, S>
	where
		S: Stream<Item = Self::Item>,
		Self: Sized,
	{
		new(self, other)
	}
}

/// An adapter for merging the output of two streams.
///
/// The merged stream produces items from either of the underlying streams as
/// they become available, and the streams are polled in a round-robin fashion.
/// Errors, however, are not merged: you get at most one error at a time.
///
/// Finishes when strong stream finishes
#[derive(Debug)]
#[must_use = "streams do nothing unless polled"]
pub struct SelectWithWeak<S1, S2> {
	strong: Fuse<S1>,
	weak: Fuse<S2>,
	use_strong: bool,
}

fn new<S1, S2>(stream1: S1, stream2: S2) -> SelectWithWeak<S1, S2>
where
	S1: Stream,
	S2: Stream<Item = S1::Item>,
{
	use futures::StreamExt;
	SelectWithWeak {
		strong: stream1.fuse(),
		weak: stream2.fuse(),
		use_strong: false,
	}
}

impl<S1, S2> Stream for SelectWithWeak<S1, S2>
where
	S1: Stream + Unpin,
	S2: Stream<Item = S1::Item> + Unpin,
{
	type Item = S1::Item;

	fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
		let mut this = Pin::into_inner(self);
		let mut checked_strong = false;
		loop {
			if this.use_strong {
				match Pin::new(&mut this.strong).poll_next(cx) {
					Poll::Ready(Some(item)) => {
						this.use_strong = false;
						return Poll::Ready(Some(item));
					}
					Poll::Ready(None) => return Poll::Ready(None),
					Poll::Pending => {
						if !checked_strong {
							this.use_strong = false;
						} else {
							return Poll::Pending;
						}
					}
				}
				checked_strong = true;
			} else {
				this.use_strong = true;
				match Pin::new(&mut this.strong).poll_next(cx) {
					Poll::Ready(Some(item)) => return Poll::Ready(Some(item)),
					Poll::Ready(None) | Poll::Pending => (),
				}
			}
		}
	}
}
<commit_msg>fix typo in select_with_weak future (#610)<commit_after>use std::pin::Pin;
use std::task::Context;
use std::task::Poll;

use futures::stream::{Fuse, Stream};

pub trait SelectWithWeakExt: Stream {
	fn select_with_weak<S>(self, other: S) -> SelectWithWeak<Self, S>
	where
		S: Stream<Item = Self::Item>,
		Self: Sized;
}

impl<T> SelectWithWeakExt for T
where
	T: Stream,
{
	fn select_with_weak<S>(self, other: S) -> SelectWithWeak<Self, S>
	where
		S: Stream<Item = Self::Item>,
		Self: Sized,
	{
		new(self, other)
	}
}

/// An adapter for merging the output of two streams.
///
/// The merged stream produces items from either of the underlying streams as
/// they become available, and the streams are polled in a round-robin fashion.
/// Errors, however, are not merged: you get at most one error at a time.
///
/// Finishes when strong stream finishes
#[derive(Debug)]
#[must_use = "streams do nothing unless polled"]
pub struct SelectWithWeak<S1, S2> {
	strong: Fuse<S1>,
	weak: Fuse<S2>,
	use_strong: bool,
}

fn new<S1, S2>(stream1: S1, stream2: S2) -> SelectWithWeak<S1, S2>
where
	S1: Stream,
	S2: Stream<Item = S1::Item>,
{
	use futures::StreamExt;
	SelectWithWeak {
		strong: stream1.fuse(),
		weak: stream2.fuse(),
		use_strong: false,
	}
}

impl<S1, S2> Stream for SelectWithWeak<S1, S2>
where
	S1: Stream + Unpin,
	S2: Stream<Item = S1::Item> + Unpin,
{
	type Item = S1::Item;

	fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
		let mut this = Pin::into_inner(self);
		let mut checked_strong = false;
		loop {
			if this.use_strong {
				match Pin::new(&mut this.strong).poll_next(cx) {
					Poll::Ready(Some(item)) => {
						this.use_strong = false;
						return Poll::Ready(Some(item));
					}
					Poll::Ready(None) => return Poll::Ready(None),
					Poll::Pending => {
						if !checked_strong {
							this.use_strong = false;
						} else {
							return Poll::Pending;
						}
					}
				}
				checked_strong = true;
			} else {
				this.use_strong = true;
				match Pin::new(&mut this.weak).poll_next(cx) {
					Poll::Ready(Some(item)) => return Poll::Ready(Some(item)),
					Poll::Ready(None) | Poll::Pending => (),
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>use std::io;
use byteorder::WriteBytesExt;
use super::{MessageType, WriteTo};
use super::messages::{AddEntry, EntryExpired, RemoveEntry};

pub enum Message {
    Ping,
    Pong,
    AddEntry(AddEntry),
    RemoveEntry(RemoveEntry),
    EntryExpired(EntryExpired),
    #[doc(hidden)]
    __NonExhaustive,
}

impl Message {
    pub fn message_type(&self) -> MessageType {
        match self {
            &Message::Ping => MessageType::Ping,
            &Message::Pong => MessageType::Pong,
            &Message::AddEntry(_) => MessageType::AddEntry,
            &Message::RemoveEntry(_) => MessageType::RemoveEntry,
            &Message::EntryExpired(_) => MessageType::EntryExpired,
            _ => panic!("invalid Message")
        }
    }

    pub fn is_command(&self) -> bool {
        self.message_type().is_command()
    }
}

impl WriteTo for Message {
    fn write_to<W: io::Write>(&self, target: &mut W) -> io::Result<()> {
        target.write_u8(self.message_type().into())?;

        match self {
            &Message::Ping => Ok(()),
            &Message::Pong => Ok(()),
            &Message::RemoveEntry(ref msg) => msg.write_to(target),
            &Message::AddEntry(ref msg) => msg.write_to(target),
            &Message::EntryExpired(ref msg) => msg.write_to(target),
            _ => panic!("invalid Message")
        }
    }
}

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

    macro_rules! test_message {
        ($test:ident, $ty:ident) => {
            test_message!($test, Message::$ty, MessageType::$ty);
        };
        ($test:ident, $msg:expr, $ty:expr) => {
            #[test]
            fn $test() {
                let msg = $msg;

                assert_eq!($ty, msg.message_type());

                let mut vec = Vec::new();
                assert!(msg.write_to(&mut vec).is_ok());
                assert_eq!(msg.message_type().to_u8(), vec[0]);
            }
        }
    }

    test_message!(test_ping, Ping);
    test_message!(test_pong, Pong);

    test_message!(test_add_entry,
                  Message::AddEntry(AddEntry::new(0, vec![])),
                  MessageType::AddEntry);

    test_message!(test_remove_entry,
                  Message::RemoveEntry(RemoveEntry::new(0, 0)),
                  MessageType::RemoveEntry);

    test_message!(test_entry_expired,
                  Message::EntryExpired(EntryExpired::new(0, vec![])),
                  MessageType::EntryExpired);
}<commit_msg>Implement ReadFrom for Message<commit_after>use std::io;
use byteorder::WriteBytesExt;
use super::{MessageType, ReadFrom, ReadError, WriteTo};
use super::messages::{AddEntry, EntryExpired, RemoveEntry};

pub enum Message {
    Ping,
    Pong,
    AddEntry(AddEntry),
    RemoveEntry(RemoveEntry),
    EntryExpired(EntryExpired),
    #[doc(hidden)]
    __NonExhaustive,
}

impl Message {
    pub fn message_type(&self) -> MessageType {
        match self {
            &Message::Ping => MessageType::Ping,
            &Message::Pong => MessageType::Pong,
            &Message::AddEntry(_) => MessageType::AddEntry,
            &Message::RemoveEntry(_) => MessageType::RemoveEntry,
            &Message::EntryExpired(_) => MessageType::EntryExpired,
            _ => panic!("invalid message")
        }
    }

    pub fn is_command(&self) -> bool {
        self.message_type().is_command()
    }
}

impl ReadFrom for Message {
    fn read_from<R: io::Read>(source: &mut R) -> Result<Self, ReadError> {
        let msg_type = MessageType::read_from(source)?;

        match msg_type {
            MessageType::Ping => Ok(Message::Ping),
            MessageType::Pong => Ok(Message::Pong),
            MessageType::AddEntry => Ok(Message::AddEntry(AddEntry::read_from(source)?)),
            MessageType::RemoveEntry => Ok(Message::RemoveEntry(RemoveEntry::read_from(source)?)),
            MessageType::EntryExpired => Ok(Message::EntryExpired(EntryExpired::read_from(source)?)),
            _ => panic!("invalid message type")
        }
    }
}

impl WriteTo for Message {
    fn write_to<W: io::Write>(&self, target: &mut W) -> io::Result<()> {
        target.write_u8(self.message_type().into())?;

        match self {
            &Message::Ping => Ok(()),
            &Message::Pong => Ok(()),
            &Message::RemoveEntry(ref msg) => msg.write_to(target),
            &Message::AddEntry(ref msg) => msg.write_to(target),
            &Message::EntryExpired(ref msg) => msg.write_to(target),
            _ => panic!("invalid message")
        }
    }
}

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

    macro_rules! test_message {
        ($test:ident, $ty:ident) => {
            test_message!($test, Message::$ty, MessageType::$ty);
        };
        ($test:ident, $msg:expr, $ty:expr) => {
            #[test]
            fn $test() {
                let msg = $msg;

                assert_eq!($ty, msg.message_type());

                let mut vec = Vec::new();
                assert!(msg.write_to(&mut vec).is_ok());
                assert_eq!(msg.message_type().to_u8(), vec[0]);
            }
        }
    }

    test_message!(test_ping, Ping);
    test_message!(test_pong, Pong);

    test_message!(test_add_entry,
                  Message::AddEntry(AddEntry::new(0, vec![])),
                  MessageType::AddEntry);

    test_message!(test_remove_entry,
                  Message::RemoveEntry(RemoveEntry::new(0, 0)),
                  MessageType::RemoveEntry);

    test_message!(test_entry_expired,
                  Message::EntryExpired(EntryExpired::new(0, vec![])),
                  MessageType::EntryExpired);
}<|endoftext|>


<commit_before>use leveldb::management::*;
use leveldb::options::*;
use utils::{open_database,tmpdir};

#[test]
fn test_destroy_database() {
    let tmp = tmpdir("destroy");
    let database = open_database::<i32>(tmp.path(), true);
    drop(database);
    let options = Options::new();
    let res = destroy(tmp.path(), options);
    assert!(res.is_ok());
}

#[test]
fn test_repair_database() {
    let tmp = tmpdir("repair");
    let database = open_database::<i32>(tmp.path(), true);
    drop(database);
    let options = Options::new();
    let res = repair(tmp.path(), options);
    assert!(res.is_ok());
}

#[test]
fn test_destroy_open_database() {
    let tmp = tmpdir("destroy_open");
    let database = open_database::<i32>(tmp.path(), true);
    let options = Options::new();
    let res = destroy(tmp.path(), options);
    assert!(res.is_err());
    drop(database);
}
<commit_msg>Remove a test due to library version dependence<commit_after>use leveldb::management::*;
use leveldb::options::*;
use utils::{open_database,tmpdir};

#[test]
fn test_destroy_database() {
    let tmp = tmpdir("destroy");
    let database = open_database::<i32>(tmp.path(), true);
    drop(database);
    let options = Options::new();
    let res = destroy(tmp.path(), options);
    assert!(res.is_ok());
}

#[test]
fn test_repair_database() {
    let tmp = tmpdir("repair");
    let database = open_database::<i32>(tmp.path(), true);
    drop(database);
    let options = Options::new();
    let res = repair(tmp.path(), options);
    assert!(res.is_ok());
}

// Deactivated due do library version dependence
//#[test]
//fn test_destroy_open_database() {
//    let tmp = tmpdir("destroy_open");
//    let database = open_database::<i32>(tmp.path(), true);
//    let options = Options::new();
//    let res = destroy(tmp.path(), options);
//    assert!(res.is_err());
//    drop(database);
//}
<|endoftext|>


<commit_before>extern crate flac;

fn run_metaflac(fname: &Path) -> String {
    use std::io::Command;

    // Run metaflac on the specified file and print all streaminfo data.
    let mut child = Command::new("metaflac")
                        .arg("--show-min-blocksize")
                        .arg("--show-max-blocksize")
                        .arg("--show-min-framesize")
                        .arg("--show-max-framesize")
                        .arg("--show-sample-rate")
                        .arg("--show-channels")
                        .arg("--show-bps")
                        .arg("--show-total-samples")
                        .arg("--show-md5sum")
                        .arg(fname.as_str().expect("unsupported filename"))
                        .spawn().ok().expect("failed to run metaflac");

    assert!(child.wait().unwrap().success());

    let output = match child.stdout {
        Some(ref mut stdout) => stdout.read_to_string().ok()
                                      .expect("failed to read metaflac stdout"),
        None => panic!("failed to open metaflac stdout")
    };

    output
}

fn print_hex(seq: &[u8]) -> String {
    let vec: Vec<String> = seq.iter().map(|x| format!("{:x}", *x)).collect();
    vec.concat()
}

fn read_streaminfo(fname: &Path) -> String {
    use std::io::fs::File;
    use flac::{read_stream_header, read_metadata_block_header, read_streaminfo_block};

    // Open the file and extract its streaminfo block.
    // TODO: avoid repetition in the tests and library;
    // there might be some kind of metadata block iterator in the future.
    let mut file = File::open(fname).ok().expect("failed to open flac file");
    read_stream_header(&mut file).ok().expect("invalid flac header");
    read_metadata_block_header(&mut file).ok().expect("invalid metadata block header");
    let streaminfo = read_streaminfo_block(&mut file).ok().expect("failed to read streaminfo");

    // Format the streaminfo in the same way that metaflac prints it.
    format!("{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n",
            streaminfo.min_block_size,
            streaminfo.max_block_size,
            streaminfo.min_frame_size.unwrap_or(0),
            streaminfo.max_frame_size.unwrap_or(0),
            streaminfo.sample_rate,
            streaminfo.n_channels,
            streaminfo.bits_per_sample,
            streaminfo.n_samples.unwrap_or(0),
            print_hex(&streaminfo.md5sum)) // TODO implement LowerHex for &[u8] and submit a PR.
}

fn compare_metaflac(fname: &Path) {
    let metaflac = run_metaflac(fname);
    let streaminfo = read_streaminfo(fname);
    let mut mf_lines = metaflac.lines_any();
    let mut si_lines = streaminfo.lines_any();
    while let (Some(mf), Some(si)) = (mf_lines.next(), si_lines.next()) {
        if mf != si {
            println!("metaflac\n--------\n{}", metaflac);
            println!("streaminfo\n----------\n{}", streaminfo);
            panic!("metaflac disagrees on parsed streaminfo");
        }
    };
}

#[test]
fn verify_streaminfo() {
    use std::io::fs::{readdir, PathExtensions};

    // Compare our streaminfo parsing with metaflac on all flac files in the
    // current directory.
    let dir = readdir(&Path::new(".")).ok().expect("failed to enumerate flac files");
    for path in dir.iter() {
        if path.is_file() && path.extension_str() == Some("flac") {
            print!("    comparing {} ...", path.as_str().expect("unsupported filename"));
            compare_metaflac(path);
            println!(" ok");
        }
    }
}
<commit_msg>fix bug in formatting byte arrays as hex<commit_after>extern crate flac;

fn run_metaflac(fname: &Path) -> String {
    use std::io::Command;

    // Run metaflac on the specified file and print all streaminfo data.
    let mut child = Command::new("metaflac")
                        .arg("--show-min-blocksize")
                        .arg("--show-max-blocksize")
                        .arg("--show-min-framesize")
                        .arg("--show-max-framesize")
                        .arg("--show-sample-rate")
                        .arg("--show-channels")
                        .arg("--show-bps")
                        .arg("--show-total-samples")
                        .arg("--show-md5sum")
                        .arg(fname.as_str().expect("unsupported filename"))
                        .spawn().ok().expect("failed to run metaflac");

    assert!(child.wait().unwrap().success());

    let output = match child.stdout {
        Some(ref mut stdout) => stdout.read_to_string().ok()
                                      .expect("failed to read metaflac stdout"),
        None => panic!("failed to open metaflac stdout")
    };

    output
}

fn print_hex(seq: &[u8]) -> String {
    let vec: Vec<String> = seq.iter().map(|x| format!("{:0>2x}", *x)).collect();
    vec.concat()
}

fn read_streaminfo(fname: &Path) -> String {
    use std::io::fs::File;
    use flac::{read_stream_header, read_metadata_block_header, read_streaminfo_block};

    // Open the file and extract its streaminfo block.
    // TODO: avoid repetition in the tests and library;
    // there might be some kind of metadata block iterator in the future.
    let mut file = File::open(fname).ok().expect("failed to open flac file");
    read_stream_header(&mut file).ok().expect("invalid flac header");
    read_metadata_block_header(&mut file).ok().expect("invalid metadata block header");
    let streaminfo = read_streaminfo_block(&mut file).ok().expect("failed to read streaminfo");

    // Format the streaminfo in the same way that metaflac prints it.
    format!("{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n",
            streaminfo.min_block_size,
            streaminfo.max_block_size,
            streaminfo.min_frame_size.unwrap_or(0),
            streaminfo.max_frame_size.unwrap_or(0),
            streaminfo.sample_rate,
            streaminfo.n_channels,
            streaminfo.bits_per_sample,
            streaminfo.n_samples.unwrap_or(0),
            print_hex(&streaminfo.md5sum)) // TODO implement LowerHex for &[u8] and submit a PR.
}

fn compare_metaflac(fname: &Path) {
    let metaflac = run_metaflac(fname);
    let streaminfo = read_streaminfo(fname);
    let mut mf_lines = metaflac.lines_any();
    let mut si_lines = streaminfo.lines_any();
    while let (Some(mf), Some(si)) = (mf_lines.next(), si_lines.next()) {
        if mf != si {
            println!("metaflac\n--------\n{}", metaflac);
            println!("streaminfo\n----------\n{}", streaminfo);
            panic!("metaflac disagrees on parsed streaminfo");
        }
    };
}

#[test]
fn verify_streaminfo() {
    use std::io::fs::{readdir, PathExtensions};

    // Compare our streaminfo parsing with metaflac on all flac files in the
    // current directory.
    let dir = readdir(&Path::new(".")).ok().expect("failed to enumerate flac files");
    for path in dir.iter() {
        if path.is_file() && path.extension_str() == Some("flac") {
            print!("    comparing {} ...", path.as_str().expect("unsupported filename"));
            compare_metaflac(path);
            println!(" ok");
        }
    }
}
<|endoftext|>


<commit_before>use zookeeper::CreateMode::*;
use zookeeper::{WatchedEvent, ZooKeeper};
use zookeeper::acls::*;
use zookeeper::recipes::cache::PathChildrenCache;

use ZkCluster;

use std::sync::Arc;
use std::time::Duration;
use std::thread;
use env_logger;

#[test]
fn path_children_cache_test() {
    let _ = env_logger::init();

    // Create a test cluster
    //let cluster = ZkCluster::start(3);

    // Connect to the test cluster
    let zk = Arc::new(ZooKeeper::connect("localhost:2181", //&cluster.connect_string,
                                         Duration::from_secs(30),
                                         |event: &WatchedEvent| info!("{:?}", event)).unwrap());

    zk.create("/cache/a", vec![1,4], OPEN_ACL_UNSAFE.clone(), Ephemeral).unwrap();
    zk.create("/cache/b", vec![2,4], OPEN_ACL_UNSAFE.clone(), Ephemeral).unwrap();
    zk.create("/cache/c", vec![3,4], OPEN_ACL_UNSAFE.clone(), Ephemeral).unwrap();

    let path_children_cache = Arc::new(PathChildrenCache::new(zk, "/cache").unwrap());

    let data = path_children_cache.get_current_data();

    info!("Data: {:?}", data);

    thread::sleep_ms(9999999);
}<commit_msg>Update test_cache.rs<commit_after>use zookeeper::CreateMode::*;
use zookeeper::{WatchedEvent, ZooKeeper, ZooKeeperExt};
use zookeeper::acls::*;
use zookeeper::recipes::cache::PathChildrenCache;

use ZkCluster;

use std::sync::Arc;
use std::time::Duration;
use std::thread;
use env_logger;

#[test]
fn path_children_cache_test() {
    let _ = env_logger::init();

    // Create a test cluster
    let cluster = ZkCluster::start(1);

    // Connect to the test cluster
    let zk = Arc::new(ZooKeeper::connect(&cluster.connect_string,
                                         Duration::from_secs(30),
                                         |event: &WatchedEvent| info!("{:?}", event)).unwrap());

    zk.ensure_path("/cache").unwrap();
    zk.create("/cache/a", vec![1,4], OPEN_ACL_UNSAFE.clone(), Ephemeral).unwrap();
    zk.create("/cache/b", vec![2,4], OPEN_ACL_UNSAFE.clone(), Ephemeral).unwrap();
    zk.create("/cache/c", vec![3,4], OPEN_ACL_UNSAFE.clone(), Ephemeral).unwrap();

    let path_children_cache = Arc::new(PathChildrenCache::new(zk, "/cache").unwrap());

    let data = path_children_cache.get_current_data();

    info!("Data: {:?}", data);

    thread::sleep_ms(9999999);
}
<|endoftext|>


<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use std::path::{Path, PathBuf};

use codespan_reporting::term::termcolor::Buffer;

use anyhow::anyhow;
use diem_temppath::TempPath;
use itertools::Itertools;
use move_prover::{cli::Options, run_move_prover};
use move_prover_test_utils::{
    baseline_test::verify_or_update_baseline, extract_test_directives, read_env_var,
};

#[allow(unused_imports)]
use log::{debug, warn};
use std::sync::atomic::{AtomicBool, Ordering};

const ENV_FLAGS: &str = "MVP_TEST_FLAGS";
const STDLIB_FLAGS: &[&str] = &["--dependency=../stdlib/modules"];

static NOT_CONFIGURED_WARNED: AtomicBool = AtomicBool::new(false);

fn test_runner(path: &Path) -> datatest_stable::Result<()> {
    let no_boogie = read_env_var("BOOGIE_EXE").is_empty() || read_env_var("Z3_EXE").is_empty();
    let baseline_valid =
        !no_boogie || !extract_test_directives(path, "// no-boogie-test")?.is_empty();

    let temp_dir = TempPath::new();
    std::fs::create_dir_all(temp_dir.path())?;
    let (flags, baseline_path) = get_flags(temp_dir.path(), path)?;

    let mut args = vec!["mvp_test".to_string()];
    args.extend(flags);
    args.push("--verbose=warn".to_owned());
    args.push("--num-instances=2".to_owned()); // run two Boogie instances with different seeds
    args.push("--sequential".to_owned());
    args.push(path.to_string_lossy().to_string());

    args.extend(shell_words::split(&read_env_var(ENV_FLAGS))?);

    let mut options = Options::create_from_args(&args)?;
    options.setup_logging_for_test();
    if no_boogie {
        options.prover.generate_only = true;
        if !NOT_CONFIGURED_WARNED.compare_and_swap(false, true, Ordering::Relaxed) {
            warn!(
                "Prover tools are not configured, verification tests will be skipped. \
        See https://github.com/diem/diem/tree/master/language/move-prover/doc/user/install.md \
        for instructions."
            );
        }
    }
    options.prover.stable_test_output = true;

    let mut error_writer = Buffer::no_color();
    let mut diags = match run_move_prover(&mut error_writer, options) {
        Ok(()) => "".to_string(),
        Err(err) => format!("Move prover returns: {}\n", err),
    };
    if baseline_valid {
        if let Some(path) = baseline_path {
            diags += &String::from_utf8_lossy(&error_writer.into_inner()).to_string();
            verify_or_update_baseline(path.as_path(), &diags)?
        } else if !diags.is_empty() {
            return Err(anyhow!(
                "Unexpected prover output (expected none): {}{}",
                diags,
                String::from_utf8_lossy(&error_writer.into_inner())
            )
            .into());
        }
    }
    Ok(())
}

fn test_runner_stdlib(path: &Path) -> datatest_stable::Result<()> {
    // Gives the standard test runner a different name in test output, which is useful because
    // the datatest infrastructure drops `..` in test file paths.
    test_runner(path)
}

fn get_flags(temp_dir: &Path, path: &Path) -> anyhow::Result<(Vec<String>, Option<PathBuf>)> {
    // Determine the way how to configure tests based on directory of the path.
    let path_str = path.to_string_lossy();
    let (base_flags, baseline_path, modifier) = if path_str.contains("../stdlib/") {
        (STDLIB_FLAGS, None, "std_")
    } else if path_str.contains("tests/sources/functional/")
        || path_str.contains("tests/sources/regression/")
    {
        (STDLIB_FLAGS, Some(path.with_extension("exp")), "func_")
    } else {
        return Err(anyhow!(
            "do not know how to run tests for `{}` because it's directory is not configured",
            path_str
        ));
    };
    let mut flags = base_flags.iter().map(|s| (*s).to_string()).collect_vec();
    // Add any flags specified in the source.
    flags.extend(extract_test_directives(path, "// flag:")?);

    // Create a temporary file for output. We inject the modifier to potentially prevent
    // any races between similar named files in different directories, as it appears TempPath
    // isn't working always.
    let base_name = format!(
        "{}{}.bpl",
        modifier,
        path.file_stem().unwrap().to_str().unwrap()
    );
    let output = temp_dir.join(base_name).to_str().unwrap().to_string();
    flags.push(format!("--output={}", output));
    Ok((flags, baseline_path))
}

datatest_stable::harness!(
    // Run tests for the content of our tests directory.
    test_runner,
    "tests/sources",
    r".*\.move$",
    // Run tests for the content of the stdlib directory.
    test_runner_stdlib,
    "../stdlib",
    r".*\.move$"
);
<commit_msg>Adds cvc4 to tests<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use std::path::{Path, PathBuf};

use codespan_reporting::term::termcolor::Buffer;

use anyhow::anyhow;
use diem_temppath::TempPath;
use itertools::Itertools;
use move_prover::{cli::Options, run_move_prover};
use move_prover_test_utils::{
    baseline_test::verify_or_update_baseline, extract_test_directives, read_env_var,
};

#[allow(unused_imports)]
use log::{debug, warn};
use std::sync::atomic::{AtomicBool, Ordering};

const ENV_FLAGS: &str = "MVP_TEST_FLAGS";
const STDLIB_FLAGS: &[&str] = &["--dependency=../stdlib/modules"];

static NOT_CONFIGURED_WARNED: AtomicBool = AtomicBool::new(false);

fn test_runner(path: &Path) -> datatest_stable::Result<()> {
    let no_boogie = read_env_var("BOOGIE_EXE").is_empty() || read_env_var("Z3_EXE").is_empty();
    let baseline_valid =
        !no_boogie || !extract_test_directives(path, "// no-boogie-test")?.is_empty();

    let temp_dir = TempPath::new();
    std::fs::create_dir_all(temp_dir.path())?;
    let (flags, baseline_path) = get_flags(temp_dir.path(), path)?;

    let mut args = vec!["mvp_test".to_string()];
    args.extend(flags);
    args.push("--verbose=warn".to_owned());
    args.push("--num-instances=2".to_owned()); // run two Boogie instances with different seeds
    args.push("--sequential".to_owned());
    args.push(path.to_string_lossy().to_string());

    args.extend(shell_words::split(&read_env_var(ENV_FLAGS))?);

    let mut options = Options::create_from_args(&args)?;
    options.setup_logging_for_test();
    if no_boogie {
        options.prover.generate_only = true;
        if !NOT_CONFIGURED_WARNED.compare_and_swap(false, true, Ordering::Relaxed) {
            warn!(
                "Prover tools are not configured, verification tests will be skipped. \
        See https://github.com/diem/diem/tree/master/language/move-prover/doc/user/install.md \
        for instructions."
            );
        }
    }
    options.prover.stable_test_output = true;

    let mut error_writer = Buffer::no_color();
    let mut diags = match run_move_prover(&mut error_writer, options) {
        Ok(()) => "".to_string(),
        Err(err) => format!("Move prover returns: {}\n", err),
    };
    if baseline_valid {
        if let Some(ref path) = baseline_path {
            diags += &String::from_utf8_lossy(&error_writer.into_inner()).to_string();
            verify_or_update_baseline(path.as_path(), &diags)?
        } else if !diags.is_empty() {
            return Err(anyhow!(
                "Unexpected prover output (expected none): {}{}",
                diags,
                String::from_utf8_lossy(&error_writer.into_inner())
            )
            .into());
        }
    }

    // Run again with cvc4
    if !no_boogie && !read_env_var("CVC4_EXE").is_empty() && !cvc4blacklisted(path) {
        args.push("--use-cvc4".to_owned());
        options = Options::create_from_args(&args)?;
        options.setup_logging_for_test();
        options.prover.stable_test_output = true;
        error_writer = Buffer::no_color();
        diags = match run_move_prover(&mut error_writer, options) {
            Ok(()) => "".to_string(),
            Err(err) => format!("Move prover returns: {}\n", err),
        };
        if let Some(path) = baseline_path {
            diags += &String::from_utf8_lossy(&error_writer.into_inner()).to_string();
            verify_or_update_baseline(path.as_path(), &diags)?
        } else if !diags.is_empty() {
            return Err(anyhow!(
                "Unexpected prover output (expected none): {}{}",
                diags,
                String::from_utf8_lossy(&error_writer.into_inner())
            )
            .into());
        }
    }

    Ok(())
}

fn cvc4blacklisted(path: &Path) -> bool {
    let path_str = path.to_str().unwrap();
    if path_str == "../stdlib/modules/ValidatorOperatorConfig.move" ||
        path_str == "../stdlib/modules/Option.move" ||
        path_str == "../stdlib/modules/RegisteredCurrencies.move" ||
        path_str == "../stdlib/modules/AccountFreezing.move" ||
        path_str == "../stdlib/modules/DiemTransactionPublishingOption.move" ||
        path_str == "../stdlib/modules/VASP.move" ||
        path_str == "../stdlib/modules/ValidatorConfig.move" ||
        path_str == "../stdlib/modules/DiemConfig.move" ||
        path_str == "../stdlib/modules/DiemSystem.move" ||
        path_str == "../stdlib/modules/XUS.move" ||
        path_str == "../stdlib/modules/DualAttestation.move" ||
        path_str == "../stdlib/modules/XDX.move" ||
        path_str == "tests/sources/functional/cast.move" ||
        path_str == "../stdlib/modules/RecoveryAddress.move" ||
        path_str == "tests/sources/functional/loops.move" ||
        path_str == "../stdlib/transaction_scripts/add_validator_and_reconfigure.move" ||
        path_str == "../stdlib/transaction_scripts/rotate_authentication_key.move" ||
        path_str == "tests/sources/functional/aborts_if_assume_assert.move" ||
        path_str == "../stdlib/transaction_scripts/remove_validator_and_reconfigure.move" ||
        path_str == "../stdlib/modules/DesignatedDealer.move" ||
        path_str == "tests/sources/functional/marketcap.move" ||
        path_str == "tests/sources/functional/invariants.move" ||
        path_str == "tests/sources/functional/module_invariants.move" ||
        path_str == "tests/sources/functional/ModifiesSchemaTest.move" ||
        path_str == "tests/sources/functional/resources.move" ||
        path_str == "tests/sources/functional/schema_exp.move" ||
        path_str == "tests/sources/functional/marketcap_generic.move" ||
        path_str == "tests/sources/functional/aborts_if_with_code.move" ||
        path_str == "tests/sources/functional/address_serialization_constant_size.move" ||
        path_str == "../stdlib/transaction_scripts/set_validator_config_and_reconfigure.move" ||
        path_str == "tests/sources/functional/global_invariants.move" ||
        path_str == "tests/sources/functional/nested_invariants.move" ||
        path_str == "../stdlib/transaction_scripts/rotate_authentication_key_with_recovery_address.move" ||
        path_str == "tests/sources/functional/marketcap_as_schema_apply.move" ||
        path_str == "tests/sources/functional/global_vars.move" ||
        path_str == "../stdlib/transaction_scripts/rotate_authentication_key_with_nonce.move" ||
        path_str == "tests/sources/functional/specs_in_fun_ref.move" ||
        path_str == "tests/sources/functional/references.move" ||
        path_str == "tests/sources/functional/mut_ref_unpack.move" ||
        path_str == "tests/sources/functional/hash_model.move" ||
        path_str == "tests/sources/functional/ModifiesErrorTest.move" ||
        path_str == "tests/sources/functional/consts.move" ||
        path_str == "tests/sources/functional/type_values.move" ||
        path_str == "tests/sources/functional/pragma.move" ||
        path_str == "tests/sources/functional/exists_in_vector.move" ||
        path_str == "tests/sources/functional/aborts_with_check.move" ||
        path_str == "tests/sources/functional/aborts_with_negative_check.move" ||
        path_str == "tests/sources/functional/opaque.move" ||
        path_str == "tests/sources/functional/marketcap_as_schema.move" ||
        path_str == "tests/sources/functional/aborts_if.move" ||
        path_str == "tests/sources/functional/address_quant.move" ||
        path_str == "tests/sources/functional/hash_model_invalid.move" ||
        path_str == "tests/sources/functional/serialize_model.move" ||
        path_str == "tests/sources/functional/return_values.move" ||
        path_str == "tests/sources/functional/pack_unpack.move" ||
        path_str == "tests/sources/functional/specs_in_fun.move" ||
        path_str == "tests/sources/functional/arithm.move" ||
        path_str == "tests/sources/regression/Escape.move" ||
        path_str == "tests/sources/functional/mut_ref_accross_modules.move" ||
        path_str == "tests/sources/regression/type_param_bug_200228.move" ||
        path_str == "tests/sources/regression/trace200527.move" ||
        path_str == "tests/sources/regression/generic_invariant200518.move" ||
        path_str == "../stdlib/transaction_scripts/rotate_authentication_key_with_nonce_admin.move" ||
        path_str == "tests/sources/functional/simple_vector_client.move" ||
        path_str == "../stdlib/transaction_scripts/publish_shared_ed25519_public_key.move" ||
        path_str == "tests/sources/functional/verify_vector.move" ||
        path_str == "../stdlib/transaction_scripts/create_designated_dealer.move" ||
        path_str == "../stdlib/transaction_scripts/create_parent_vasp_account.move" ||
        path_str == "../stdlib/modules/Diem.move" ||
        path_str == "../stdlib/transaction_scripts/create_child_vasp_account.move" ||
        path_str == "../stdlib/modules/FixedPoint32.move" ||
        path_str == "../stdlib/modules/Genesis.move" ||
        path_str == "../stdlib/modules/DiemAccount.move" ||

        path_str == "../stdlib/modules/AccountLimits.move" || // This one takes over a minute

        true
    {
        return true;
    }
    false
}

fn test_runner_stdlib(path: &Path) -> datatest_stable::Result<()> {
    // Gives the standard test runner a different name in test output, which is useful because
    // the datatest infrastructure drops `..` in test file paths.
    test_runner(path)
}

fn get_flags(temp_dir: &Path, path: &Path) -> anyhow::Result<(Vec<String>, Option<PathBuf>)> {
    // Determine the way how to configure tests based on directory of the path.
    let path_str = path.to_string_lossy();
    let (base_flags, baseline_path, modifier) = if path_str.contains("../stdlib/") {
        (STDLIB_FLAGS, None, "std_")
    } else if path_str.contains("tests/sources/functional/")
        || path_str.contains("tests/sources/regression/")
    {
        (STDLIB_FLAGS, Some(path.with_extension("exp")), "func_")
    } else {
        return Err(anyhow!(
            "do not know how to run tests for `{}` because it's directory is not configured",
            path_str
        ));
    };
    let mut flags = base_flags.iter().map(|s| (*s).to_string()).collect_vec();
    // Add any flags specified in the source.
    flags.extend(extract_test_directives(path, "// flag:")?);

    // Create a temporary file for output. We inject the modifier to potentially prevent
    // any races between similar named files in different directories, as it appears TempPath
    // isn't working always.
    let base_name = format!(
        "{}{}.bpl",
        modifier,
        path.file_stem().unwrap().to_str().unwrap()
    );
    let output = temp_dir.join(base_name).to_str().unwrap().to_string();
    flags.push(format!("--output={}", output));
    Ok((flags, baseline_path))
}

datatest_stable::harness!(
    // Run tests for the content of our tests directory.
    test_runner,
    "tests/sources",
    r".*\.move$",
    // Run tests for the content of the stdlib directory.
    test_runner_stdlib,
    "../stdlib",
    r".*\.move$"
);
<|endoftext|>


<commit_before>extern crate portmidi as pm;

extern crate rustc_serialize;
extern crate docopt;

use std::time::Duration;
use std::thread;

const USAGE: &'static str = r#"
portmidi-rs: monitor-device example

Usage:
    monitor-device <device-id>

Options:
    -h --help   Show this screen.

Omitting <device-id> will list the available devices.
"#;

#[derive(Debug, RustcDecodable)]
struct Args {
    arg_device_id: i32,
}

fn print_devices(pm: &pm::PortMidi) {
    for dev in pm.devices().unwrap() {
        println!("{}", dev);
    }
}

fn main() {
    // initialize the PortMidi context.
    let context = pm::PortMidi::new().unwrap();

    // setup the command line interface
    let args: Args = docopt::Docopt::new(USAGE).and_then(|d| d.decode()).unwrap_or_else(|err| {
        print_devices(&context);
        err.exit();
    });

    // get the device info for the given id
    let info = context.device(args.arg_device_id).unwrap();
    println!("Listening on: {}) {}", info.id(), info.name());

    // get the device's input port
    let mut in_port = context.input_port(info, 1024).unwrap();

    while let Ok(_) = in_port.poll() {
        if let Ok(Some(event)) = in_port.read_n(1024) {
            println!("{:?}", event);
            // there is no blocking receive method in PortMidi, therefore
            // we have to sleep some time to prevent a busy-wait loop
            thread::sleep(Duration::new(0,50))
        }
    }
}
<commit_msg>Remove mutability from `in_port`<commit_after>extern crate portmidi as pm;

extern crate rustc_serialize;
extern crate docopt;

use std::time::Duration;
use std::thread;

const USAGE: &'static str = r#"
portmidi-rs: monitor-device example

Usage:
    monitor-device <device-id>

Options:
    -h --help   Show this screen.

Omitting <device-id> will list the available devices.
"#;

#[derive(Debug, RustcDecodable)]
struct Args {
    arg_device_id: i32,
}

fn print_devices(pm: &pm::PortMidi) {
    for dev in pm.devices().unwrap() {
        println!("{}", dev);
    }
}

fn main() {
    // initialize the PortMidi context.
    let context = pm::PortMidi::new().unwrap();

    // setup the command line interface
    let args: Args = docopt::Docopt::new(USAGE).and_then(|d| d.decode()).unwrap_or_else(|err| {
        print_devices(&context);
        err.exit();
    });

    // get the device info for the given id
    let info = context.device(args.arg_device_id).unwrap();
    println!("Listening on: {}) {}", info.id(), info.name());

    // get the device's input port
    let in_port = context.input_port(info, 1024).unwrap();

    while let Ok(_) = in_port.poll() {
        if let Ok(Some(event)) = in_port.read_n(1024) {
            println!("{:?}", event);
            // there is no blocking receive method in PortMidi, therefore
            // we have to sleep some time to prevent a busy-wait loop
            thread::sleep(Duration::new(0,50))
        }
    }
}
<|endoftext|>


<commit_before>// examples/blinky.rs
//! Blinks an LED

#![feature(const_fn)]
#![feature(used)]
#![no_std]

// version = "0.2.0"
extern crate cortex_m_rt;

// version = "0.1.0"
#[macro_use]
extern crate cortex_m_rtfm as rtfm;

extern crate bluepill;

use bluepill::pin::{halPin, Pin, Mode};
use bluepill::frequency;
use bluepill::stm32f103xx::interrupt::Tim3;
use bluepill::stm32f103xx;
use bluepill::timer::{halTimer, Timer};
use rtfm::{Local, P0, P1, T0, T1, TMax};

// CONFIGURATION
const TICKS: u32 = 900_000; // one rotation a second
const ORDER:[[bool; 4]; 9] = [[false,false,false,true],
                [false,false,true,true],
                [false,false,true,false],
                [false,true,true,false],
                [false,true,false,false],
                [true,true,false,false],
                [true,false,false,false],
                [true,false,false,true],
                [false,false,false,false]];

// RESOURCES
peripherals!(stm32f103xx, {
    GPIOA: Peripheral {
        register_block: Gpioa,
        ceiling: C1,
    },
    RCC: Peripheral {
        register_block: Rcc,
        ceiling: C0,
    },
    TIM3: Peripheral {
        register_block: Tim3,
        ceiling: C1,
    },
    FLASH: Peripheral {
        register_block: Flash,
        ceiling: C0,
    },
});

// INITIALIZATION PHASE
fn init(ref priority: P0, threshold: &TMax) {
    let gpioa = GPIOA.access(priority, threshold);
    let rcc = RCC.access(priority, threshold);
    let tim3 = TIM3.access(priority, threshold);
    let flash = FLASH.access(priority, threshold);
    let timer = Timer::new(&**tim3);
    let in1 = Pin{pin: 1, port: &**gpioa};
    let in2 = Pin{pin: 2, port: &**gpioa};
    let in3 = Pin{pin: 3, port: &**gpioa};
    let in4 = Pin{pin: 4, port: &**gpioa};

    // set clock to 72Mhz
    frequency::init(&rcc, &flash, frequency::Speed::S72Mhz);

    // configure pins for output
    in1.init(&rcc, Mode::OUTPUT);
    in2.init(&rcc, Mode::OUTPUT);
    in3.init(&rcc, Mode::OUTPUT);
    in4.init(&rcc, Mode::OUTPUT);

    // Configure TIM2 for periodic update events
    timer.init(&rcc, TICKS);

    // Start the timer
    timer.resume();
}

// IDLE LOOP
fn idle(_priority: P0, _threshold: T0) -> ! {
    // Sleep
    loop {
        rtfm::wfi();
    }
}

// TASKS
tasks!(stm32f103xx, {
    periodic: Task {
        interrupt: Tim3,
        priority: P1,
        enabled: true,
    },
});

fn periodic(mut task: Tim3, ref priority: P1, ref threshold: T1) {
    // Task local data
    static STEP: Local<u16, Tim3> = Local::new(0);

    let tim3 = TIM3.access(priority, threshold);
    let timer = Timer{timer: &**tim3};
    let gpioa = GPIOA.access(priority, threshold);
    let in1 = Pin{pin: 1, port: &**gpioa};
    let in2 = Pin{pin: 2, port: &**gpioa};
    let in3 = Pin{pin: 3, port: &**gpioa};
    let in4 = Pin{pin: 4, port: &**gpioa};

    if timer.clear_update_flag().is_ok() {
        let step = STEP.borrow_mut(&mut task);
        let current = ORDER[*step as usize];

        match current[0]{
            true => in1.on(),
            false => in1.off(),
        }
        match current[1]{
            true => in2.on(),
            false => in2.off(),
        }
        match current[2]{
            true => in3.on(),
            false => in3.off(),
        }
        match current[3]{
            true => in4.on(),
            false => in4.off(),
        }

        if *step < 8 {
            *step += 1;
        } else {
            *step = 0;
        }
    } else {
        // Only reachable through `rtfm::request(periodic)`
        //#[cfg(debug_assertion)]
        //unreachable!()
    }
}
<commit_msg>stepper was going too slow<commit_after>// examples/blinky.rs
//! Blinks an LED

#![feature(const_fn)]
#![feature(used)]
#![no_std]

// version = "0.2.0"
extern crate cortex_m_rt;

// version = "0.1.0"
#[macro_use]
extern crate cortex_m_rtfm as rtfm;

extern crate bluepill;

use bluepill::pin::{halPin, Pin, Mode};
use bluepill::frequency;
use bluepill::stm32f103xx::interrupt::Tim3;
use bluepill::stm32f103xx;
use bluepill::timer::{halTimer, Timer};
use rtfm::{Local, P0, P1, T0, T1, TMax};

// CONFIGURATION
const TICKS: u32 = 400_000;
const ORDER:[[bool; 4]; 9] = [[false,false,false,true],
                [false,false,true,true],
                [false,false,true,false],
                [false,true,true,false],
                [false,true,false,false],
                [true,true,false,false],
                [true,false,false,false],
                [true,false,false,true],
                [false,false,false,false]];

// RESOURCES
peripherals!(stm32f103xx, {
    GPIOA: Peripheral {
        register_block: Gpioa,
        ceiling: C1,
    },
    RCC: Peripheral {
        register_block: Rcc,
        ceiling: C0,
    },
    TIM3: Peripheral {
        register_block: Tim3,
        ceiling: C1,
    },
    FLASH: Peripheral {
        register_block: Flash,
        ceiling: C0,
    },
});

// INITIALIZATION PHASE
fn init(ref priority: P0, threshold: &TMax) {
    let gpioa = GPIOA.access(priority, threshold);
    let rcc = RCC.access(priority, threshold);
    let tim3 = TIM3.access(priority, threshold);
    let flash = FLASH.access(priority, threshold);
    let timer = Timer::new(&**tim3);
    let in1 = Pin{pin: 1, port: &**gpioa};
    let in2 = Pin{pin: 2, port: &**gpioa};
    let in3 = Pin{pin: 3, port: &**gpioa};
    let in4 = Pin{pin: 4, port: &**gpioa};

    // set clock to 72Mhz
    frequency::init(&rcc, &flash, frequency::Speed::S72Mhz);

    // configure pins for output
    in1.init(&rcc, Mode::OUTPUT);
    in2.init(&rcc, Mode::OUTPUT);
    in3.init(&rcc, Mode::OUTPUT);
    in4.init(&rcc, Mode::OUTPUT);

    // Configure TIM2 for periodic update events
    timer.init(&rcc, TICKS);

    // Start the timer
    timer.resume();
}

// IDLE LOOP
fn idle(_priority: P0, _threshold: T0) -> ! {
    // Sleep
    loop {
        rtfm::wfi();
    }
}

// TASKS
tasks!(stm32f103xx, {
    periodic: Task {
        interrupt: Tim3,
        priority: P1,
        enabled: true,
    },
});

fn periodic(mut task: Tim3, ref priority: P1, ref threshold: T1) {
    // Task local data
    static STEP: Local<u16, Tim3> = Local::new(0);

    let tim3 = TIM3.access(priority, threshold);
    let timer = Timer{timer: &**tim3};
    let gpioa = GPIOA.access(priority, threshold);
    let in1 = Pin{pin: 1, port: &**gpioa};
    let in2 = Pin{pin: 2, port: &**gpioa};
    let in3 = Pin{pin: 3, port: &**gpioa};
    let in4 = Pin{pin: 4, port: &**gpioa};

    if timer.clear_update_flag().is_ok() {
        let step = STEP.borrow_mut(&mut task);
        let current = ORDER[*step as usize];

        match current[0]{
            true => in1.on(),
            false => in1.off(),
        }
        match current[1]{
            true => in2.on(),
            false => in2.off(),
        }
        match current[2]{
            true => in3.on(),
            false => in3.off(),
        }
        match current[3]{
            true => in4.on(),
            false => in4.off(),
        }

        if *step < 8 {
            *step += 1;
        } else {
            *step = 0;
        }
    } else {
        // Only reachable through `rtfm::request(periodic)`
        //#[cfg(debug_assertion)]
        //unreachable!()
    }
}
<|endoftext|>


<commit_before>//! Macros to generate the implementation of the serialization traits for Rust
//! structs, as ASN.1 sequences.
//!
//! You can either use the shortcut `ber_sequence!` macro, or each of
//! `asn1_sequence_info!`, `ber_sequence_serialize!` and `ber_sequence_deserialize!`.
//!
//! ```
//! #[macro_use] extern crate asn1_cereal; fn main() {
//!   struct ShortSequence {
//!     z: u64,
//!     y: u32,
//!   }
//!
//!   ber_sequence!(
//!     ShortSequence,
//!     "SHORT_SEQUENCE",
//!     z;
//!     y;
//!   );
//!
//!   // OR
//!
//!   struct SomeSequence {
//!     a: u64,
//!     b: u32,
//!     c: String,
//!   }
//!
//!   asn1_sequence_info!(SomeSequence, "SOME_SEQUENCE");
//!   ber_sequence_serialize!(SomeSequence, a; b; c;);
//!   ber_sequence_deserialize!(SomeSequence, a; b; c;);
//! }
//! ```
//!
//! Note that these macros won't handle SEQUENCE OF.

#[macro_export]
/// This macro is a compact way of defining all three of the
/// Asn1 traits - Asn1Info, BerSerialize and BerDeserialize -
/// for a rust struct, that represents an ASN.1 sequence.
///
/// Note that the order the fields are placed in will affect the order
/// that they are encoded to, and decoded from ASN.1. If some form of
/// procedural macros are eventually stabilised, listing the fields
/// in the macro might no longer be required.
macro_rules! ber_sequence {
  ($rs_type:ident, $asn1_ty:expr, $($args:tt)*) => (
    asn1_sequence_info!($rs_type, $asn1_ty);
    ber_sequence_serialize!($rs_type, $($args)*);
    ber_sequence_deserialize!($rs_type, $($args)*);
  )
}

#[macro_export]
/// This macro defines the Asn1Info trait for a rust struct. This allows the other
/// traits to get information about this type. If you need to provide a custom
/// class or tag, consider using the asn1_info! macro.
macro_rules! asn1_sequence_info {
  ($rs_type:ident, $asn1_ty:expr) => (
    impl $crate::Asn1Info for $rs_type {
      fn asn1_tag() -> Option<$crate::tag::Tag> {
        Some($crate::tag::Tag {
          class: $crate::tag::Class::Universal,
          tagnum: (0x10 as u8).into(),
          constructed: true,
        })
      }

      fn asn1_type() -> $crate::tag::Type {
        $asn1_ty.into()
      }
    }
  )
}

#[macro_export]
/// This macro defines the BerSerialize trait for a rust struct. The code generated
/// will serialize the specified fields in the order that they are given.
macro_rules! ber_sequence_serialize {
  ($rs_type:ty, $($args:tt)*) => (
    impl $crate::BerSerialize for $rs_type {
      fn serialize_value<E: $crate::BerEncRules, W: std::io::Write>
          (&self, e: E, writer: &mut W) -> Result<(), $crate::err::EncodeError> {
        let mut bytes = Vec::new();
        let mut _count: u64 = 0;
        // For each declared sequence member, serialize it onto the stream.
        ber_sequence_serialize!(_ { self e writer bytes _count } $($args)*);
        Ok(())
      }
    }
  );

  // Handle field creation.
  // Expand a field with no options.
  (_ { $this:ident $e:ident $writer:ident $bytes:ident $count:ident }
      $item:ident; $($args:tt)*) => (
    ber_sequence_serialize!(_ { $this $e $writer $bytes $count } $item (); $($args)*);
  );
  // Expand an OPTIONAL field with no tag.
  (_ { $this:ident $e:ident $writer:ident $bytes:ident $count:ident }
      $item:ident (OPTIONAL); $($args:tt)*) => (
    ber_sequence_serialize!(_ { $this $e $writer $bytes $count } $item ([] OPTIONAL); $($args)*);
  );
  // Expand an OPTIONAL field with a tag.
  (_ { $this:ident $e:ident $writer:ident $bytes:ident $count:ident }
      $item:ident ([$($opts:tt)*] OPTIONAL); $($args:tt)*) => (
    let tag = asn1_spec_tag!({ $count } [$($opts)*]);
    if let Some(ref val) = $this.$item {
      ber_sequence_serialize!(_ { $this $e $writer $bytes $count tag, val } $($args)*);
    } else {
      ber_sequence_serialize!(_ { $this $e $writer $bytes $count } $($args)*);
    }
  );
  // Expand a field with a DEFAULT.
  (_ { $this:ident $e:ident $writer:ident $bytes:ident $count:ident }
      $item:ident (DEFAULT $default:expr); $($args:tt)*) => (
    ber_sequence_serialize!(_ { $this $e $writer $bytes $count } $item ([] DEFAULT $default); $($args)*);
  );
  // Expand a field with a DEFAULT and a tag.
  (_ { $this:ident $e:ident $writer:ident $bytes:ident $count:ident }
      $item:ident ([$($opts:tt)*] DEFAULT $default:expr); $($args:tt)*) => (
    let tag = asn1_spec_tag!({ $count } [$($opts)*]);
    if &$this.$item != &$default {
      ber_sequence_serialize!(_ { $this $e $writer $bytes $count tag, &$this.$item } $($args)*);
    } else {
      ber_sequence_serialize!(_ { $this $e $writer $bytes $count } $($args)*);
    }
  );
  // Expand a field that has options (assumes its a tag here).
  (_ { $this:ident $e:ident $writer:ident $bytes:ident $count:ident }
      $item:ident ($($opts:tt)*); $($args:tt)*) => (
    let tag = asn1_spec_tag!({ $count } [$($opts)*]);
    ber_sequence_serialize!(_ { $this $e $writer $bytes $count tag, &$this.$item } $($args)*);
  );
  // Create the implementation for a field with options.
  (_ { $this:ident $e:ident $writer:ident $bytes:ident $count:ident $tag:expr, $value:expr }
      $($args:tt)*) => (
    // If encoding uses implicit tag, skip context-specific tag.
    if E::tag_rules() == $crate::ber::enc::TagEnc::Implicit {
      try!($crate::BerSerialize::serialize_enc($value, $e, $writer));
    // Otherwise encode the context-specific tag.
    } else {
      try!($crate::BerSerialize::serialize_enc($value, $e, &mut $bytes));
      let len: $crate::tag::Len = Some($bytes.len() as $crate::tag::LenNum).into();
      try!($crate::tag::write_taglen($tag, len, $writer));
      try!($writer.write_all(&mut $bytes));
      $bytes.clear();
    }

    ber_sequence_serialize!(_ { $this $e $writer $bytes $count } $($args)*);
  );
  (_ { $($args:tt)* } ) => ();
}

#[macro_export]
/// This macro defines the BerDeserialize trait for a rust struct. The code generated
/// will deserialize the specified fields in the order that they are given.
macro_rules! ber_sequence_deserialize {
  // Handle fields.
  // Expand a field with no options.
  (_ { $rs_type:ident $e:ident $reader:ident $count:ident $tag:ident [$($fields:ident)*] }
      $item:ident; $($args:tt)*) => (
    ber_sequence_deserialize!(
      _ { $rs_type $e $reader $count $tag [$($fields)*] } $item ([]); $($args)*
    );
  );
  // Expand a field with empty options.
  (_ { $rs_type:ident $e:ident $reader:ident $count:ident $tag:ident [$($fields:ident)*] }
      $item:ident (); $($args:tt)*) => (
    ber_sequence_deserialize!(
      _ { $rs_type $e $reader $count $tag [$($fields)*] } $item ([]); $($args)*
    );
  );
  // Expand a field with empty options and a tag.
  (_ { $rs_type:ident $e:ident $reader:ident $count:ident $tag:ident [$($fields:ident)*] }
      $field:ident ([$($opts:tt)*]); $($args:tt)*) => (
    let $field = {
      let tag = match $tag {
        Some(t) => t,
        None => {
          let t = try!($crate::tag::Tag::read_tag($reader));
          $tag = Some(t); t
        },
      };
      let our_tag = asn1_spec_tag!({ $count } [$($opts)*]);

      ber_sequence_deserialize!(_ { $rs_type $e $reader $count } tag, our_tag )
    };
    ber_sequence_deserialize!(_ { $rs_type $e $reader $count $tag [$($fields)* $field] } $($args)*);
  );
  // Expand an OPTIONAL field with no tag.
  (_ { $rs_type:ident $e:ident $reader:ident $count:ident $tag:ident [$($fields:ident)*] }
      $item:ident (OPTIONAL); $($args:tt)*) => (
    ber_sequence_deserialize!(
      _ { $rs_type $e $reader $count $tag [$($fields)*] } $item ([] OPTIONAL); $($args)*
    );
  );
  // Expand an OPTIONAL field with a tag.
  (_ { $rs_type:ident $e:ident $reader:ident $count:ident $tag:ident [$($fields:ident)*] }
      $field:ident ([$($opts:tt)*] OPTIONAL); $($args:tt)*) => (
    let $field = {
      let tag = match $tag {
        Some(t) => t,
        None => {
          let t = try!($crate::tag::Tag::read_tag($reader));
          $tag = Some(t); t
        },
      };
      let our_tag = asn1_spec_tag!({ $count } [$($opts)*]);

      if tag == our_tag {
        let len = try!($crate::tag::Len::read_len($reader));
        Some(try!($crate::BerDeserialize::deserialize_with_tag($e, $reader, tag, len)));
      } else {
        None
      }
    };
    ber_sequence_deserialize!({ $rs_type $e $reader $count $tag [$($fields)* $field] } $($args)*);
  );
  // Create the implementation for a field.
  (_ { $rs_type:ident $e:ident $reader:ident $count:ident } $tag:expr, $our_tag:expr) => ({
    // If encoding uses implicit tagging, throw an error if this isn't an implicit tag.
    if E::tag_rules() == $crate::ber::enc::TagEnc::Implicit && $tag == $our_tag {
      return Err($crate::err::DecodeError::ExplicitTag);
    }

    let len = try!($crate::tag::Len::read_len($reader));

    // If the tag matches our tag, decode the len and call the normal deserialize function.
    if $tag == $our_tag {
      // We don't have anything to do with the len, technically we should use it to
      // check the length decoded.
      let _ = len;
      try!($crate::BerDeserialize::deserialize_enc($e, $reader))
    // Otherwise decode it as the inner type. (We give the tag that we
    // decoded, and the function will decode the length itself).
    } else {
      try!($crate::BerDeserialize::deserialize_with_tag($e, $reader, $tag, len))
    }
  });
  // When no fields are left, build the struct.
  (_ { $rs_type:ident $e:ident $reader:ident $count:ident $tag:ident [ $($field:ident)* ] }) => (
    return Ok($rs_type { $(
      $field: $field
    ),* })
  );

  ($rs_type:ident, $($args:tt)*) => (
    impl $crate::BerDeserialize for $rs_type {
      fn deserialize_value<E: $crate::BerEncRules, I: Iterator<Item=std::io::Result<u8>>>
          (e: E, reader: &mut I, _: $crate::tag::Len) -> Result<Self, $crate::err::DecodeError> {
        let mut _count: u64 = 0;
        let mut _tag: Option<$crate::tag::Tag> = None;
        ber_sequence_deserialize!(_ { $rs_type e reader _count _tag [ ] } $($args)*);
      }
    }
  );
}
<commit_msg>OPTIONAL encoding/decoding now works in sequences<commit_after>//! Macros to generate the implementation of the serialization traits for Rust
//! structs, as ASN.1 sequences.
//!
//! You can either use the shortcut `ber_sequence!` macro, or each of
//! `asn1_sequence_info!`, `ber_sequence_serialize!` and `ber_sequence_deserialize!`.
//!
//! ```
//! #[macro_use] extern crate asn1_cereal; fn main() {
//!   struct ShortSequence {
//!     z: u64,
//!     y: u32,
//!   }
//!
//!   ber_sequence!(
//!     ShortSequence,
//!     "SHORT_SEQUENCE",
//!     z;
//!     y;
//!   );
//!
//!   // OR
//!
//!   struct SomeSequence {
//!     a: u64,
//!     b: u32,
//!     c: String,
//!   }
//!
//!   asn1_sequence_info!(SomeSequence, "SOME_SEQUENCE");
//!   ber_sequence_serialize!(SomeSequence, a; b; c;);
//!   ber_sequence_deserialize!(SomeSequence, a; b; c;);
//! }
//! ```
//!
//! Note that these macros won't handle SEQUENCE OF.

#[macro_export]
/// This macro is a compact way of defining all three of the
/// Asn1 traits - Asn1Info, BerSerialize and BerDeserialize -
/// for a rust struct, that represents an ASN.1 sequence.
///
/// Note that the order the fields are placed in will affect the order
/// that they are encoded to, and decoded from ASN.1. If some form of
/// procedural macros are eventually stabilised, listing the fields
/// in the macro might no longer be required.
macro_rules! ber_sequence {
  ($rs_type:ident, $asn1_ty:expr, $($args:tt)*) => (
    asn1_sequence_info!($rs_type, $asn1_ty);
    ber_sequence_serialize!($rs_type, $($args)*);
    ber_sequence_deserialize!($rs_type, $($args)*);
  )
}

#[macro_export]
/// This macro defines the Asn1Info trait for a rust struct. This allows the other
/// traits to get information about this type. If you need to provide a custom
/// class or tag, consider using the asn1_info! macro.
macro_rules! asn1_sequence_info {
  ($rs_type:ident, $asn1_ty:expr) => (
    impl $crate::Asn1Info for $rs_type {
      fn asn1_tag() -> Option<$crate::tag::Tag> {
        Some($crate::tag::Tag {
          class: $crate::tag::Class::Universal,
          tagnum: (0x10 as u8).into(),
          constructed: true,
        })
      }

      fn asn1_type() -> $crate::tag::Type {
        $asn1_ty.into()
      }
    }
  )
}

#[macro_export]
/// This macro defines the BerSerialize trait for a rust struct. The code generated
/// will serialize the specified fields in the order that they are given.
macro_rules! ber_sequence_serialize {
  ($rs_type:ty, $($args:tt)*) => (
    impl $crate::BerSerialize for $rs_type {
      fn serialize_value<E: $crate::BerEncRules, W: std::io::Write>
          (&self, e: E, writer: &mut W) -> Result<(), $crate::err::EncodeError> {
        let mut bytes = Vec::new();
        let mut _count: u64 = 0;
        // For each declared sequence member, serialize it onto the stream.
        ber_sequence_serialize!(_ { self e writer bytes _count } $($args)*);
        Ok(())
      }
    }
  );

  // Handle field creation.
  // Expand a field with no options.
  (_ { $this:ident $e:ident $writer:ident $bytes:ident $count:ident }
      $item:ident; $($args:tt)*) => (
    ber_sequence_serialize!(_ { $this $e $writer $bytes $count } $item (); $($args)*);
  );
  // Expand an OPTIONAL field with no tag.
  (_ { $this:ident $e:ident $writer:ident $bytes:ident $count:ident }
      $item:ident (OPTIONAL); $($args:tt)*) => (
    ber_sequence_serialize!(_ { $this $e $writer $bytes $count } $item ([] OPTIONAL); $($args)*);
  );
  // Expand an OPTIONAL field with a tag.
  (_ { $this:ident $e:ident $writer:ident $bytes:ident $count:ident }
      $item:ident ([$($opts:tt)*] OPTIONAL); $($args:tt)*) => (
    let tag = asn1_spec_tag!({ $count } [$($opts)*]);
    if let Some(ref val) = $this.$item {
      ber_sequence_serialize!(_ { $this $e $writer $bytes $count tag, val } $($args)*);
    } else {
      ber_sequence_serialize!(_ { $this $e $writer $bytes $count } $($args)*);
    }
  );
  // Expand a field with a DEFAULT.
  (_ { $this:ident $e:ident $writer:ident $bytes:ident $count:ident }
      $item:ident (DEFAULT $default:expr); $($args:tt)*) => (
    ber_sequence_serialize!(_ { $this $e $writer $bytes $count } $item ([] DEFAULT $default); $($args)*);
  );
  // Expand a field with a DEFAULT and a tag.
  (_ { $this:ident $e:ident $writer:ident $bytes:ident $count:ident }
      $item:ident ([$($opts:tt)*] DEFAULT $default:expr); $($args:tt)*) => (
    let tag = asn1_spec_tag!({ $count } [$($opts)*]);
    if &$this.$item != &$default {
      ber_sequence_serialize!(_ { $this $e $writer $bytes $count tag, &$this.$item } $($args)*);
    } else {
      ber_sequence_serialize!(_ { $this $e $writer $bytes $count } $($args)*);
    }
  );
  // Expand a field that has options.
  (_ { $this:ident $e:ident $writer:ident $bytes:ident $count:ident }
      $item:ident ($($opts:tt)*); $($args:tt)*) => (
    let tag = asn1_spec_tag!({ $count } [$($opts)*]);

    // If encoding uses implicit tag, skip context-specific tag.
    if E::tag_rules() == $crate::ber::enc::TagEnc::Implicit {
      try!($crate::BerSerialize::serialize_enc(&$this.$item, $e, $writer));
      ber_sequence_serialize!(_ { $this $e $writer $bytes $count } $($args)*);
    // Otherwise encode the context-specific tag.
    } else {
      ber_sequence_serialize!(_ { $this $e $writer $bytes $count tag, &$this.$item } $($args)*);
    }
  );
  // Write a field with a tag.
  (_ { $this:ident $e:ident $writer:ident $bytes:ident $count:ident $tag:expr, $value:expr }
      $($args:tt)*) => (
    try!($crate::BerSerialize::serialize_enc($value, $e, &mut $bytes));
    let len: $crate::tag::Len = Some($bytes.len() as $crate::tag::LenNum).into();
    try!($crate::tag::write_taglen($tag, len, $writer));
    try!($writer.write_all(&mut $bytes));
    $bytes.clear();

    ber_sequence_serialize!(_ { $this $e $writer $bytes $count } $($args)*);
  );
  (_ { $($args:tt)* } ) => ();
}

#[macro_export]
/// This macro defines the BerDeserialize trait for a rust struct. The code generated
/// will deserialize the specified fields in the order that they are given.
macro_rules! ber_sequence_deserialize {
  // Handle fields.
  // Expand a field with no options.
  (_ { $rs_type:ident $e:ident $reader:ident $count:ident $tag:ident [$($fields:ident)*] }
      $item:ident; $($args:tt)*) => (
    ber_sequence_deserialize!(
      _ { $rs_type $e $reader $count $tag [$($fields)*] } $item ([]); $($args)*
    );
  );
  // Expand a field with empty options.
  (_ { $rs_type:ident $e:ident $reader:ident $count:ident $tag:ident [$($fields:ident)*] }
      $item:ident (); $($args:tt)*) => (
    ber_sequence_deserialize!(
      _ { $rs_type $e $reader $count $tag [$($fields)*] } $item ([]); $($args)*
    );
  );
  // Expand a field with empty options and a tag.
  (_ { $rs_type:ident $e:ident $reader:ident $count:ident $tag:ident [$($fields:ident)*] }
      $field:ident ([$($opts:tt)*]); $($args:tt)*) => (
    let $field = {
      let tag = $tag.unwrap_or(try!($crate::tag::Tag::read_tag($reader)));
      $tag = None;
      let our_tag = asn1_spec_tag!({ $count } [$($opts)*]);

      ber_sequence_deserialize!(_ { $rs_type $e $reader $count } tag, our_tag )
    };
    ber_sequence_deserialize!(_ { $rs_type $e $reader $count $tag [$($fields)* $field] } $($args)*);
  );
  // Expand an OPTIONAL field with no tag.
  (_ { $rs_type:ident $e:ident $reader:ident $count:ident $tag:ident [$($fields:ident)*] }
      $item:ident (OPTIONAL); $($args:tt)*) => (
    ber_sequence_deserialize!(
      _ { $rs_type $e $reader $count $tag [$($fields)*] } $item ([] OPTIONAL); $($args)*
    );
  );
  // Expand an OPTIONAL field with a tag.
  (_ { $rs_type:ident $e:ident $reader:ident $count:ident $tag:ident [$($fields:ident)*] }
      $field:ident ([$($opts:tt)*] OPTIONAL); $($args:tt)*) => (
    let $field = {
      let tag = match $tag {
        Some(t) => t,
        None => {
          let t = try!($crate::tag::Tag::read_tag($reader));
          $tag = Some(t); t
        },
      };

      let our_tag = asn1_spec_tag!({ $count } [$($opts)*]);

      if tag == our_tag {
        $tag = None;
        let _ = try!($crate::tag::Len::read_len($reader));
        Some(try!($crate::BerDeserialize::deserialize_enc($e, $reader)))
      } else {
        None
      }
    };
    ber_sequence_deserialize!(_ { $rs_type $e $reader $count $tag [$($fields)* $field] } $($args)*);
  );
  // Create the implementation for a field.
  (_ { $rs_type:ident $e:ident $reader:ident $count:ident } $tag:expr, $our_tag:expr) => ({
    // If encoding uses implicit tagging, throw an error if this isn't an implicit tag.
    if E::tag_rules() == $crate::ber::enc::TagEnc::Implicit && $tag == $our_tag {
      return Err($crate::err::DecodeError::ExplicitTag);
    }

    let len = try!($crate::tag::Len::read_len($reader));

    // If the tag matches our tag, decode the len and call the normal deserialize function.
    if $tag == $our_tag {
      // We don't have anything to do with the len, technically we should use it to
      // check the length decoded.
      let _ = len;
      try!($crate::BerDeserialize::deserialize_enc($e, $reader))
    // Otherwise decode it as the inner type. (We give the tag that we
    // decoded, and the function will decode the length itself).
    } else {
      try!($crate::BerDeserialize::deserialize_with_tag($e, $reader, $tag, len))
    }
  });
  // When no fields are left, build the struct.
  (_ { $rs_type:ident $e:ident $reader:ident $count:ident $tag:ident [ $($field:ident)* ] }) => (
    return Ok($rs_type { $(
      $field: $field
    ),* })
  );

  ($rs_type:ident, $($args:tt)*) => (
    impl $crate::BerDeserialize for $rs_type {
      fn deserialize_value<E: $crate::BerEncRules, I: Iterator<Item=std::io::Result<u8>>>
          (e: E, reader: &mut I, _: $crate::tag::Len) -> Result<Self, $crate::err::DecodeError> {
        let mut _count: u64 = 0;
        let mut _tag: Option<$crate::tag::Tag> = None;
        ber_sequence_deserialize!(_ { $rs_type e reader _count _tag [ ] } $($args)*);
      }
    }
  );
}
<|endoftext|>