<commit_before>//! # The Rust core allocation and collections library
//!
//! This library provides smart pointers and collections for managing
//! heap-allocated values.
//!
//! This library, like libcore, normally doesn’t need to be used directly
//! since its contents are re-exported in the [`std` crate](../std/index.html).
//! Crates that use the `#![no_std]` attribute however will typically
//! not depend on `std`, so they’d use this crate instead.
//!
//! ## Boxed values
//!
//! The [`Box`] type is a smart pointer type. There can only be one owner of a
//! [`Box`], and the owner can decide to mutate the contents, which live on the
//! heap.
//!
//! This type can be sent among threads efficiently as the size of a `Box` value
//! is the same as that of a pointer. Tree-like data structures are often built
//! with boxes because each node often has only one owner, the parent.
//!
//! ## Reference counted pointers
//!
//! The [`Rc`] type is a non-threadsafe reference-counted pointer type intended
//! for sharing memory within a thread. An [`Rc`] pointer wraps a type, `T`, and
//! only allows access to `&T`, a shared reference.
//!
//! This type is useful when inherited mutability (such as using [`Box`]) is too
//! constraining for an application, and is often paired with the [`Cell`] or
//! [`RefCell`] types in order to allow mutation.
//!
//! ## Atomically reference counted pointers
//!
//! The [`Arc`] type is the threadsafe equivalent of the [`Rc`] type. It
//! provides all the same functionality of [`Rc`], except it requires that the
//! contained type `T` is shareable. Additionally, [`Arc<T>`][`Arc`] is itself
//! sendable while [`Rc<T>`][`Rc`] is not.
//!
//! This type allows for shared access to the contained data, and is often
//! paired with synchronization primitives such as mutexes to allow mutation of
//! shared resources.
//!
//! ## Collections
//!
//! Implementations of the most common general purpose data structures are
//! defined in this library. They are re-exported through the
//! [standard collections library](../std/collections/index.html).
//!
//! ## Heap interfaces
//!
//! The [`alloc`](alloc/index.html) module defines the low-level interface to the
//! default global allocator. It is not compatible with the libc allocator API.
//!
//! [`Arc`]: sync
//! [`Box`]: boxed
//! [`Cell`]: core::cell
//! [`Rc`]: rc
//! [`RefCell`]: core::cell

#![allow(unused_attributes)]
#![stable(feature = "alloc", since = "1.36.0")]
#![doc(
    html_root_url = "https://doc.rust-lang.org/nightly/",
    html_playground_url = "https://play.rust-lang.org/",
    issue_tracker_base_url = "https://github.com/rust-lang/rust/issues/",
    test(no_crate_inject, attr(allow(unused_variables), deny(warnings)))
)]
#![no_std]
#![needs_allocator]
#![warn(deprecated_in_future)]
#![warn(missing_docs)]
#![warn(missing_debug_implementations)]
#![allow(explicit_outlives_requirements)]
#![deny(unsafe_op_in_unsafe_fn)]
#![feature(rustc_allow_const_fn_unstable)]
#![cfg_attr(not(test), feature(generator_trait))]
#![cfg_attr(test, feature(test))]
#![cfg_attr(test, feature(new_uninit))]
#![feature(allocator_api)]
#![feature(array_chunks)]
#![feature(array_methods)]
#![feature(array_windows)]
#![feature(allow_internal_unstable)]
#![feature(arbitrary_self_types)]
#![feature(async_stream)]
#![feature(box_patterns)]
#![feature(box_syntax)]
#![feature(cfg_sanitize)]
#![feature(cfg_target_has_atomic)]
#![feature(coerce_unsized)]
#![cfg_attr(not(no_global_oom_handling), feature(const_btree_new))]
#![feature(const_fn_trait_bound)]
#![feature(cow_is_borrowed)]
#![feature(const_cow_is_borrowed)]
#![feature(destructuring_assignment)]
#![feature(dispatch_from_dyn)]
#![feature(core_intrinsics)]
#![feature(dropck_eyepatch)]
#![feature(exact_size_is_empty)]
#![feature(exclusive_range_pattern)]
#![feature(extend_one)]
#![feature(fmt_internals)]
#![feature(fn_traits)]
#![feature(fundamental)]
#![feature(inplace_iteration)]
// Technically, this is a bug in rustdoc: rustdoc sees the documentation on `#[lang = slice_alloc]`
// blocks is for `&[T]`, which also has documentation using this feature in `core`, and gets mad
// that the feature-gate isn't enabled. Ideally, it wouldn't check for the feature gate for docs
// from other crates, but since this can only appear for lang items, it doesn't seem worth fixing.
#![feature(intra_doc_pointers)]
#![feature(iter_zip)]
#![feature(lang_items)]
#![feature(layout_for_ptr)]
#![feature(maybe_uninit_ref)]
#![feature(negative_impls)]
#![feature(never_type)]
#![feature(nll)]
#![feature(nonnull_slice_from_raw_parts)]
#![feature(auto_traits)]
#![feature(option_result_unwrap_unchecked)]
#![feature(pattern)]
#![feature(ptr_internals)]
#![feature(rustc_attrs)]
#![feature(receiver_trait)]
#![feature(min_specialization)]
#![feature(set_ptr_value)]
#![feature(slice_ptr_get)]
#![feature(slice_ptr_len)]
#![feature(slice_range)]
#![feature(staged_api)]
#![feature(str_internals)]
#![feature(trusted_len)]
#![feature(unboxed_closures)]
#![feature(unicode_internals)]
#![feature(unsize)]
#![feature(unsized_fn_params)]
#![feature(allocator_internals)]
#![feature(slice_partition_dedup)]
#![feature(maybe_uninit_extra, maybe_uninit_slice, maybe_uninit_uninit_array)]
#![feature(alloc_layout_extra)]
#![feature(trusted_random_access)]
#![cfg_attr(bootstrap, feature(try_trait))]
#![cfg_attr(not(bootstrap), feature(try_trait_v2))]
#![feature(min_type_alias_impl_trait)]
#![feature(associated_type_bounds)]
#![feature(slice_group_by)]
#![feature(decl_macro)]
// Allow testing this library

#[cfg(test)]
#[macro_use]
extern crate std;
#[cfg(test)]
extern crate test;

// Module with internal macros used by other modules (needs to be included before other modules).
#[macro_use]
mod macros;

// Heaps provided for low-level allocation strategies

pub mod alloc;

// Primitive types using the heaps above

// Need to conditionally define the mod from `boxed.rs` to avoid
// duplicating the lang-items when building in test cfg; but also need
// to allow code to have `use boxed::Box;` declarations.
#[cfg(not(test))]
pub mod boxed;
#[cfg(test)]
mod boxed {
    pub use std::boxed::Box;
}
pub mod borrow;
pub mod collections;
pub mod fmt;
pub mod prelude;
pub mod raw_vec;
pub mod rc;
pub mod slice;
pub mod str;
pub mod string;
#[cfg(target_has_atomic = "ptr")]
pub mod sync;
#[cfg(all(not(no_global_oom_handling), target_has_atomic = "ptr"))]
pub mod task;
#[cfg(test)]
mod tests;
pub mod vec;

#[doc(hidden)]
#[unstable(feature = "liballoc_internals", issue = "none", reason = "implementation detail")]
pub mod __export {
    pub use core::format_args;
}
<commit_msg>added back bindings_after_at as a cfg_attr<commit_after>//! # The Rust core allocation and collections library
//!
//! This library provides smart pointers and collections for managing
//! heap-allocated values.
//!
//! This library, like libcore, normally doesn’t need to be used directly
//! since its contents are re-exported in the [`std` crate](../std/index.html).
//! Crates that use the `#![no_std]` attribute however will typically
//! not depend on `std`, so they’d use this crate instead.
//!
//! ## Boxed values
//!
//! The [`Box`] type is a smart pointer type. There can only be one owner of a
//! [`Box`], and the owner can decide to mutate the contents, which live on the
//! heap.
//!
//! This type can be sent among threads efficiently as the size of a `Box` value
//! is the same as that of a pointer. Tree-like data structures are often built
//! with boxes because each node often has only one owner, the parent.
//!
//! ## Reference counted pointers
//!
//! The [`Rc`] type is a non-threadsafe reference-counted pointer type intended
//! for sharing memory within a thread. An [`Rc`] pointer wraps a type, `T`, and
//! only allows access to `&T`, a shared reference.
//!
//! This type is useful when inherited mutability (such as using [`Box`]) is too
//! constraining for an application, and is often paired with the [`Cell`] or
//! [`RefCell`] types in order to allow mutation.
//!
//! ## Atomically reference counted pointers
//!
//! The [`Arc`] type is the threadsafe equivalent of the [`Rc`] type. It
//! provides all the same functionality of [`Rc`], except it requires that the
//! contained type `T` is shareable. Additionally, [`Arc<T>`][`Arc`] is itself
//! sendable while [`Rc<T>`][`Rc`] is not.
//!
//! This type allows for shared access to the contained data, and is often
//! paired with synchronization primitives such as mutexes to allow mutation of
//! shared resources.
//!
//! ## Collections
//!
//! Implementations of the most common general purpose data structures are
//! defined in this library. They are re-exported through the
//! [standard collections library](../std/collections/index.html).
//!
//! ## Heap interfaces
//!
//! The [`alloc`](alloc/index.html) module defines the low-level interface to the
//! default global allocator. It is not compatible with the libc allocator API.
//!
//! [`Arc`]: sync
//! [`Box`]: boxed
//! [`Cell`]: core::cell
//! [`Rc`]: rc
//! [`RefCell`]: core::cell

#![allow(unused_attributes)]
#![stable(feature = "alloc", since = "1.36.0")]
#![doc(
    html_root_url = "https://doc.rust-lang.org/nightly/",
    html_playground_url = "https://play.rust-lang.org/",
    issue_tracker_base_url = "https://github.com/rust-lang/rust/issues/",
    test(no_crate_inject, attr(allow(unused_variables), deny(warnings)))
)]
#![no_std]
#![needs_allocator]
#![warn(deprecated_in_future)]
#![warn(missing_docs)]
#![warn(missing_debug_implementations)]
#![allow(explicit_outlives_requirements)]
#![deny(unsafe_op_in_unsafe_fn)]
#![feature(rustc_allow_const_fn_unstable)]
#![cfg_attr(not(test), feature(generator_trait))]
#![cfg_attr(test, feature(test))]
#![cfg_attr(test, feature(new_uninit))]
#![feature(allocator_api)]
#![feature(array_chunks)]
#![feature(array_methods)]
#![feature(array_windows)]
#![feature(allow_internal_unstable)]
#![feature(arbitrary_self_types)]
#![feature(async_stream)]
#![cfg_attr(bootstrap, feature(bindings_after_at))]
#![feature(box_patterns)]
#![feature(box_syntax)]
#![feature(cfg_sanitize)]
#![feature(cfg_target_has_atomic)]
#![feature(coerce_unsized)]
#![cfg_attr(not(no_global_oom_handling), feature(const_btree_new))]
#![feature(const_fn_trait_bound)]
#![feature(cow_is_borrowed)]
#![feature(const_cow_is_borrowed)]
#![feature(destructuring_assignment)]
#![feature(dispatch_from_dyn)]
#![feature(core_intrinsics)]
#![feature(dropck_eyepatch)]
#![feature(exact_size_is_empty)]
#![feature(exclusive_range_pattern)]
#![feature(extend_one)]
#![feature(fmt_internals)]
#![feature(fn_traits)]
#![feature(fundamental)]
#![feature(inplace_iteration)]
// Technically, this is a bug in rustdoc: rustdoc sees the documentation on `#[lang = slice_alloc]`
// blocks is for `&[T]`, which also has documentation using this feature in `core`, and gets mad
// that the feature-gate isn't enabled. Ideally, it wouldn't check for the feature gate for docs
// from other crates, but since this can only appear for lang items, it doesn't seem worth fixing.
#![feature(intra_doc_pointers)]
#![feature(iter_zip)]
#![feature(lang_items)]
#![feature(layout_for_ptr)]
#![feature(maybe_uninit_ref)]
#![feature(negative_impls)]
#![feature(never_type)]
#![feature(nll)]
#![feature(nonnull_slice_from_raw_parts)]
#![feature(auto_traits)]
#![feature(option_result_unwrap_unchecked)]
#![feature(pattern)]
#![feature(ptr_internals)]
#![feature(rustc_attrs)]
#![feature(receiver_trait)]
#![feature(min_specialization)]
#![feature(set_ptr_value)]
#![feature(slice_ptr_get)]
#![feature(slice_ptr_len)]
#![feature(slice_range)]
#![feature(staged_api)]
#![feature(str_internals)]
#![feature(trusted_len)]
#![feature(unboxed_closures)]
#![feature(unicode_internals)]
#![feature(unsize)]
#![feature(unsized_fn_params)]
#![feature(allocator_internals)]
#![feature(slice_partition_dedup)]
#![feature(maybe_uninit_extra, maybe_uninit_slice, maybe_uninit_uninit_array)]
#![feature(alloc_layout_extra)]
#![feature(trusted_random_access)]
#![cfg_attr(bootstrap, feature(try_trait))]
#![cfg_attr(not(bootstrap), feature(try_trait_v2))]
#![feature(min_type_alias_impl_trait)]
#![feature(associated_type_bounds)]
#![feature(slice_group_by)]
#![feature(decl_macro)]
// Allow testing this library

#[cfg(test)]
#[macro_use]
extern crate std;
#[cfg(test)]
extern crate test;

// Module with internal macros used by other modules (needs to be included before other modules).
#[macro_use]
mod macros;

// Heaps provided for low-level allocation strategies

pub mod alloc;

// Primitive types using the heaps above

// Need to conditionally define the mod from `boxed.rs` to avoid
// duplicating the lang-items when building in test cfg; but also need
// to allow code to have `use boxed::Box;` declarations.
#[cfg(not(test))]
pub mod boxed;
#[cfg(test)]
mod boxed {
    pub use std::boxed::Box;
}
pub mod borrow;
pub mod collections;
pub mod fmt;
pub mod prelude;
pub mod raw_vec;
pub mod rc;
pub mod slice;
pub mod str;
pub mod string;
#[cfg(target_has_atomic = "ptr")]
pub mod sync;
#[cfg(all(not(no_global_oom_handling), target_has_atomic = "ptr"))]
pub mod task;
#[cfg(test)]
mod tests;
pub mod vec;

#[doc(hidden)]
#[unstable(feature = "liballoc_internals", issue = "none", reason = "implementation detail")]
pub mod __export {
    pub use core::format_args;
}
<|endoftext|>


<commit_before>#![allow(bad_style)]

extern crate libc;
use libc::*;

pub type rs_magic_number = c_int;
pub const RS_DELTA_MAGIC: c_int = 0x72730236;
pub const RS_MD4_SIG_MAGIC: c_int = 0x72730136;
pub const RS_BLAKE2_SIG_MAGIC: c_int = 0x72730137;

pub type rs_result = c_int;
/// Completed successfully.
pub const RS_DONE: c_int = 0;
/// Blocked waiting for more data.
pub const RS_BLOCKED: c_int = 1;
/// The job is still running; and not yet finished or blocked.
///
/// (This value should never be seen by the application.)
pub const RS_RUNNING: c_int = 2;
/// Test neither passed or failed.
pub const RS_TEST_SKIPPED: c_int = 77;
/// Error in file or network IO.
pub const RS_IO_ERROR: c_int = 100;
/// Command line syntax error.
pub const RS_SYNTAX_ERROR: c_int = 101;
/// Out of memory.
pub const RS_MEM_ERROR: c_int = 102;
/// Unexpected end of input file.
///
/// Perhaps due to a truncated file or dropped network connection.
pub const RS_INPUT_ENDED: c_int = 103;
/// Bad magic number at start of stream.
///
/// Probably not a librsync file; or possibly the wrong kind of file or from an incompatible
/// library version.
pub const RS_BAD_MAGIC: c_int = 104;
/// Author is lazy.
pub const RS_UNIMPLEMENTED: c_int = 105;
/// Unbelievable value in stream.
pub const RS_CORRUPT: c_int = 106;
/// Probably a library bug.
pub const RS_INTERNAL_ERROR: c_int = 107;
/// Bad value passed in to library.
///
/// Probably an application bug.
pub const RS_PARAM_ERROR: c_int = 108;

#[repr(C)]
pub struct rs_job_t {
    pub dogtag: c_int,
    pub job_name: *const c_char,
    pub stream: *mut rs_buffers_t,
    pub statefn: Option<extern fn(*mut rs_job_t) -> rs_result>,
// TODO
//    rs_result final_result;
//    int                 block_len;
//    int                 strong_sum_len;
//    rs_signature_t      *signature;
//    unsigned char       op;
//    rs_weak_sum_t       weak_sig;
//    Rollsum             weak_sum;
//    rs_long_t           param1, param2;
//    struct rs_prototab_ent const *cmd;
//    rs_mdfour_t      output_md4;
//    rs_stats_t          stats;
//    rs_byte_t   *scoop_buf;          /* the allocation pointer */
//    rs_byte_t   *scoop_next;         /* the data pointer */
//    size_t      scoop_alloc;           /* the allocation size */
//    size_t      scoop_avail;           /* the data size */
//    size_t      scoop_pos;             /* the scan position */
//    rs_byte_t   write_buf[36];
//    int         write_len;
//    rs_long_t   copy_len;
//    rs_long_t       basis_pos, basis_len;
//    rs_copy_cb      *copy_cb;
//    void            *copy_arg;
//    int             magic;
}

pub struct rs_buffers_t {
    pub next_in: *mut c_char,
    pub avail_in: size_t,
    pub eof_in: c_int,
    pub next_out: *mut c_char,
    pub avail_out: size_t,
}

extern {
    pub fn rs_sig_begin(new_block_len: size_t,
                        strong_sum_len: size_t,
                        sig_magic: rs_magic_number)
                        -> *mut rs_job_t;
}
<commit_msg>Added job apis<commit_after>#![allow(bad_style)]

extern crate libc;
use libc::*;

pub type rs_magic_number = c_int;
pub const RS_DELTA_MAGIC: c_int = 0x72730236;
pub const RS_MD4_SIG_MAGIC: c_int = 0x72730136;
pub const RS_BLAKE2_SIG_MAGIC: c_int = 0x72730137;

pub type rs_result = c_int;
pub const RS_DONE: c_int = 0;
pub const RS_BLOCKED: c_int = 1;
pub const RS_RUNNING: c_int = 2;
pub const RS_TEST_SKIPPED: c_int = 77;
pub const RS_IO_ERROR: c_int = 100;
pub const RS_SYNTAX_ERROR: c_int = 101;
pub const RS_MEM_ERROR: c_int = 102;
pub const RS_INPUT_ENDED: c_int = 103;
pub const RS_BAD_MAGIC: c_int = 104;
pub const RS_UNIMPLEMENTED: c_int = 105;
pub const RS_CORRUPT: c_int = 106;
pub const RS_INTERNAL_ERROR: c_int = 107;
pub const RS_PARAM_ERROR: c_int = 108;

pub const RS_DEFAULT_BLOCK_LEN: size_t = 2048;

pub enum rs_job_t {}

#[repr(C)]
pub struct rs_buffers_t {
    pub next_in: *mut c_char,
    pub avail_in: size_t,
    pub eof_in: c_int,
    pub next_out: *mut c_char,
    pub avail_out: size_t,
}

pub type rs_driven_cb = extern "C" fn(*mut rs_job_t, *mut rs_buffers_t, *mut c_void) -> rs_result;

extern "C" {
    pub fn rs_sig_begin(new_block_len: size_t,
                        strong_sum_len: size_t,
                        sig_magic: rs_magic_number)
                        -> *mut rs_job_t;

    pub fn rs_job_iter(job: *mut rs_job_t, buffers: *mut rs_buffers_t) -> rs_result;
    pub fn fs_job_drive(job: *mut rs_job_t,
                        buf: *mut rs_buffers_t,
                        in_cb: rs_driven_cb,
                        in_opaque: *mut c_void,
                        out_cb: rs_driven_cb,
                        out_opaque: *mut c_void)
                        -> rs_result;
}
<|endoftext|>


<commit_before>/// This module parses a raw `tt::TokenStream` into macro-by-example token
/// stream. This is a *mostly* identify function, expect for handling of
/// `$var:tt_kind` and `$(repeat),*` constructs.
use crate::ParseError;
use crate::tt_cursor::TtCursor;

pub(crate) fn parse(tt: &tt::Subtree) -> Result<crate::MacroRules, ParseError> {
    let mut parser = TtCursor::new(tt);
    let mut rules = Vec::new();
    while !parser.is_eof() {
        rules.push(parse_rule(&mut parser)?);
        if let Err(e) = parser.expect_char(';') {
            if !parser.is_eof() {
                return Err(e);
            }
            break;
        }
    }
    Ok(crate::MacroRules { rules })
}

fn parse_rule(p: &mut TtCursor) -> Result<crate::Rule, ParseError> {
    let lhs = parse_subtree(p.eat_subtree()?)?;
    p.expect_char('=')?;
    p.expect_char('>')?;
    let mut rhs = parse_subtree(p.eat_subtree()?)?;
    rhs.delimiter = crate::Delimiter::None;
    Ok(crate::Rule { lhs, rhs })
}

fn parse_subtree(tt: &tt::Subtree) -> Result<crate::Subtree, ParseError> {
    let mut token_trees = Vec::new();
    let mut p = TtCursor::new(tt);
    while let Some(tt) = p.eat() {
        let child: crate::TokenTree = match tt {
            tt::TokenTree::Leaf(leaf) => match leaf {
                tt::Leaf::Punct(tt::Punct { char: '$', .. }) => {
                    if p.at_ident().is_some() {
                        crate::Leaf::from(parse_var(&mut p)?).into()
                    } else {
                        parse_repeat(&mut p)?.into()
                    }
                }
                tt::Leaf::Punct(punct) => crate::Leaf::from(*punct).into(),
                tt::Leaf::Ident(tt::Ident { text, id: _ }) => {
                    crate::Leaf::from(crate::Ident { text: text.clone() }).into()
                }
                tt::Leaf::Literal(tt::Literal { text }) => {
                    crate::Leaf::from(crate::Literal { text: text.clone() }).into()
                }
            },
            tt::TokenTree::Subtree(subtree) => parse_subtree(&subtree)?.into(),
        };
        token_trees.push(child);
    }
    Ok(crate::Subtree { token_trees, delimiter: tt.delimiter })
}

fn parse_var(p: &mut TtCursor) -> Result<crate::Var, ParseError> {
    let ident = p.eat_ident().unwrap();
    let text = ident.text.clone();
    let kind = if p.at_char(':') {
        p.bump();
        if let Some(ident) = p.eat_ident() {
            Some(ident.text.clone())
        } else {
            p.rev_bump();
            None
        }
    } else {
        None
    };
    Ok(crate::Var { text, kind })
}

fn parse_repeat(p: &mut TtCursor) -> Result<crate::Repeat, ParseError> {
    let subtree = p.eat_subtree().unwrap();
    let mut subtree = parse_subtree(subtree)?;
    subtree.delimiter = crate::Delimiter::None;
    let sep = p.eat_punct().ok_or(ParseError::Expected(String::from("separator")))?;
    let (separator, rep) = match sep.char {
        '*' | '+' | '?' => (None, sep.char),
        char => {
            (Some(char), p.eat_punct().ok_or(ParseError::Expected(String::from("separator")))?.char)
        }
    };

    let kind = match rep {
        '*' => crate::RepeatKind::ZeroOrMore,
        '+' => crate::RepeatKind::OneOrMore,
        '?' => crate::RepeatKind::ZeroOrOne,
        _ => return Err(ParseError::Expected(String::from("repeat"))),
    };
    Ok(crate::Repeat { subtree, kind, separator })
}

#[cfg(test)]
mod tests {
    use ra_syntax::{ast, AstNode};

    use super::*;
    use crate::ast_to_token_tree;

    #[test]
    fn test_invalid_parse() {
        expect_err("invalid", "subtree");

        is_valid("($i:ident) => ()");
        expect_err("$i:ident => ()", "subtree");
        expect_err("($i:ident) ()", "`=`");
        expect_err("($($i:ident)_) => ()", "separator");
    }

    fn expect_err(macro_body: &str, expected: &str) {
        assert_eq!(
            create_rules(&format_macro(macro_body)),
            Err(ParseError::Expected(String::from(expected)))
        );
    }

    fn is_valid(macro_body: &str) {
        assert!(create_rules(&format_macro(macro_body)).is_ok());
    }

    fn format_macro(macro_body: &str) -> String {
        format!(
            "
        macro_rules! foo {{
            {}
        }}
",
            macro_body
        )
    }

    fn create_rules(macro_definition: &str) -> Result<crate::MacroRules, ParseError> {
        let source_file = ast::SourceFile::parse(macro_definition);
        let macro_definition =
            source_file.syntax().descendants().find_map(ast::MacroCall::cast).unwrap();

        let (definition_tt, _) = ast_to_token_tree(macro_definition.token_tree().unwrap()).unwrap();
        parse(&definition_tt)
    }

}
<commit_msg>fix mbe_parser rhs colon parsing<commit_after>/// This module parses a raw `tt::TokenStream` into macro-by-example token
/// stream. This is a *mostly* identify function, expect for handling of
/// `$var:tt_kind` and `$(repeat),*` constructs.
use crate::ParseError;
use crate::tt_cursor::TtCursor;

pub(crate) fn parse(tt: &tt::Subtree) -> Result<crate::MacroRules, ParseError> {
    let mut parser = TtCursor::new(tt);
    let mut rules = Vec::new();
    while !parser.is_eof() {
        rules.push(parse_rule(&mut parser)?);
        if let Err(e) = parser.expect_char(';') {
            if !parser.is_eof() {
                return Err(e);
            }
            break;
        }
    }
    Ok(crate::MacroRules { rules })
}

fn parse_rule(p: &mut TtCursor) -> Result<crate::Rule, ParseError> {
    let lhs = parse_subtree(p.eat_subtree()?, false)?;
    p.expect_char('=')?;
    p.expect_char('>')?;
    let mut rhs = parse_subtree(p.eat_subtree()?, true)?;
    rhs.delimiter = crate::Delimiter::None;
    Ok(crate::Rule { lhs, rhs })
}

fn parse_subtree(tt: &tt::Subtree, transcriber: bool) -> Result<crate::Subtree, ParseError> {
    let mut token_trees = Vec::new();
    let mut p = TtCursor::new(tt);
    while let Some(tt) = p.eat() {
        let child: crate::TokenTree = match tt {
            tt::TokenTree::Leaf(leaf) => match leaf {
                tt::Leaf::Punct(tt::Punct { char: '$', .. }) => {
                    if p.at_ident().is_some() {
                        crate::Leaf::from(parse_var(&mut p, transcriber)?).into()
                    } else {
                        parse_repeat(&mut p, transcriber)?.into()
                    }
                }
                tt::Leaf::Punct(punct) => crate::Leaf::from(*punct).into(),
                tt::Leaf::Ident(tt::Ident { text, id: _ }) => {
                    crate::Leaf::from(crate::Ident { text: text.clone() }).into()
                }
                tt::Leaf::Literal(tt::Literal { text }) => {
                    crate::Leaf::from(crate::Literal { text: text.clone() }).into()
                }
            },
            tt::TokenTree::Subtree(subtree) => parse_subtree(&subtree, transcriber)?.into(),
        };
        token_trees.push(child);
    }
    Ok(crate::Subtree { token_trees, delimiter: tt.delimiter })
}

fn parse_var(p: &mut TtCursor, transcriber: bool) -> Result<crate::Var, ParseError> {
    let ident = p.eat_ident().unwrap();
    let text = ident.text.clone();
    let kind = if !transcriber && p.at_char(':') {
        p.bump();
        if let Some(ident) = p.eat_ident() {
            Some(ident.text.clone())
        } else {
            p.rev_bump();
            None
        }
    } else {
        None
    };

    Ok(crate::Var { text, kind })
}

fn parse_repeat(p: &mut TtCursor, transcriber: bool) -> Result<crate::Repeat, ParseError> {
    let subtree = p.eat_subtree().unwrap();
    let mut subtree = parse_subtree(subtree, transcriber)?;
    subtree.delimiter = crate::Delimiter::None;
    let sep = p.eat_punct().ok_or(ParseError::Expected(String::from("separator")))?;
    let (separator, rep) = match sep.char {
        '*' | '+' | '?' => (None, sep.char),
        char => {
            (Some(char), p.eat_punct().ok_or(ParseError::Expected(String::from("separator")))?.char)
        }
    };

    let kind = match rep {
        '*' => crate::RepeatKind::ZeroOrMore,
        '+' => crate::RepeatKind::OneOrMore,
        '?' => crate::RepeatKind::ZeroOrOne,
        _ => return Err(ParseError::Expected(String::from("repeat"))),
    };
    Ok(crate::Repeat { subtree, kind, separator })
}

#[cfg(test)]
mod tests {
    use ra_syntax::{ast, AstNode};

    use super::*;
    use crate::ast_to_token_tree;

    #[test]
    fn test_invalid_parse() {
        expect_err("invalid", "subtree");

        is_valid("($i:ident) => ()");
        expect_err("$i:ident => ()", "subtree");
        expect_err("($i:ident) ()", "`=`");
        expect_err("($($i:ident)_) => ()", "separator");
    }

    fn expect_err(macro_body: &str, expected: &str) {
        assert_eq!(
            create_rules(&format_macro(macro_body)),
            Err(ParseError::Expected(String::from(expected)))
        );
    }

    fn is_valid(macro_body: &str) {
        assert!(create_rules(&format_macro(macro_body)).is_ok());
    }

    fn format_macro(macro_body: &str) -> String {
        format!(
            "
        macro_rules! foo {{
            {}
        }}
",
            macro_body
        )
    }

    fn create_rules(macro_definition: &str) -> Result<crate::MacroRules, ParseError> {
        let source_file = ast::SourceFile::parse(macro_definition);
        let macro_definition =
            source_file.syntax().descendants().find_map(ast::MacroCall::cast).unwrap();

        let (definition_tt, _) = ast_to_token_tree(macro_definition.token_tree().unwrap()).unwrap();
        parse(&definition_tt)
    }

}
<|endoftext|>


<commit_before>use regex::Regex;
use std::fs::File;
use std::io::{self, BufRead};

#[macro_use]
extern crate lazy_static;

struct State<T>(Box<dyn Iterator<Item = T>>, Option<T>);
pub struct Merge {
    count: usize,
    states: Vec<State<u32>>,
}

impl Merge {
    pub fn new(filename: &str) -> Self {
        let count = get_number(filename).unwrap();
        let states: Vec<State<u32>> = (0..count).map(|i| {
            let collection = Box::new(get_collection(filename, i).unwrap());
            State(collection, None)
        }).collect();
        Merge {
            count: count,
            states: states,
        }
    }
}

impl Iterator for Merge {
    type Item = u32;
    fn next(&mut self) -> Option<Self::Item> {
        let count = self.count;
        let states = &mut self.states;
        loop {
            for i in 0..count {
                let mut state = &mut states[i];
                if state.1.is_none() {
                    state.1 = state.0.next();
                }
            }
            let not_done = states.iter().any(|s| s.1.is_some());
            if !not_done {
                return None;
            }

            let mut min = u32::MAX;
            let mut min_index = usize::MIN;
            for i in 0..count {
                let state = &mut states[i];
                if let Some(val) = state.1 {
                    if val < min {
                        min = val;
                        min_index = i;
                    }
                }
            }
            states[min_index].1 = None;
            return Some(min);
        }
    }
}

pub fn get_number(filename: &str) -> io::Result<usize> {
    lazy_static! {
        static ref RE: Regex = Regex::new(r"NUMBER=(\d+)").unwrap();
    }
    let mut matches = match_lines(filename, &*RE)?;
    return Ok(matches.next().unwrap().parse().unwrap());
}

pub fn get_output(filename: &str) -> io::Result<impl Iterator<Item = u32>> {
    lazy_static! {
        static ref RE: Regex = Regex::new(r"OUTPUT:(.*)").unwrap();
    }
    get_numbers(filename, &*RE)
}

pub fn get_collection(filename: &str, idx: usize) -> io::Result<impl Iterator<Item = u32>> {
    let regex: Regex = {
        let r = format!("COLLECTION{}:(.*)", idx);
        Regex::new(&r).unwrap()
    };
    get_numbers(filename, &regex)
}

fn get_numbers(filename: &str, regex: &Regex) -> io::Result<impl Iterator<Item = u32>> {
    let matches = match_lines(filename, regex)?;
    let to_int = |v: &str| -> Option<u32> {
        if let Ok(i) = v.parse::<u32>() {
            Some(i)
        } else {
            None
        }
    };
    let values = matches
        .map(move |s| {
            let numbers = s.split(',').filter_map(|v| to_int(&v));
            numbers.collect::<Vec<_>>()
        })
        .flatten();
    Ok(values)
}

fn match_lines(filename: &str, regex: &Regex) -> io::Result<impl Iterator<Item = String>> {
    let read_lines = move || -> io::Result<_> {
        let file = File::open(filename)?;
        Ok(io::BufReader::new(file).lines().map(|line| line.unwrap()))
    };
    let lines = read_lines()?;
    let r = regex.clone(); // TODO
    let matches =
        lines.filter_map(move |line| r.captures(&line).map(|captures| captures[1].to_string()));
    Ok(matches)
}

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

    #[test]
    fn test_number() {
        assert_eq!(get_number("../data/test.txt").unwrap(), 3);
    }

    #[test]
    fn test_output() {
        let output = [1, 1, 2, 2, 3, 4, 4, 6, 7, 9, 9, 20, 21];
        assert_eq!(
            get_output("../data/test.txt").unwrap().collect::<Vec<_>>(),
            output
        );
    }

    #[test]
    fn test_collection() {
        let output = [2, 4, 6, 9, 20];
        assert_eq!(
            get_collection("../data/test.txt", 1)
                .unwrap()
                .collect::<Vec<_>>(),
            output
        );
    }

    fn merge(filename: &str) -> (impl Iterator<Item = u32>, impl Iterator<Item = u32>) {
        let result = Merge::new(filename);
        let output = get_output(filename).unwrap();
        (result, output)
    }

    fn assert_iterator(
        iter1: &mut impl Iterator<Item = u32>,
        iter2: &mut impl Iterator<Item = u32>,
    ) -> bool {
        loop {
            match (iter1.next(), iter2.next()) {
                (None, None) => return true,
                (Some(v1), Some(v2)) if v1 == v2 => continue,
                _ => return false,
            }
        }
    }

    #[test]
    fn test_merge() {
        let filename = "../data/test.txt";
        let (mut result, mut output) = merge(filename);
        assert!(assert_iterator(&mut result, &mut output));
    }

    #[test]
    #[should_panic]
    fn test_bad_merge() {
        let filename = "../data/bad.txt";
        let (_result, _output) = merge(filename);
    }

    #[test]
    fn test_zero_merge() {
        let filename = "../data/zero.txt";
        let (mut result, mut output) = merge(filename);
        assert!(result.next().is_none());
        assert!(output.next().is_none());
    }

    #[test]
    fn test_err() {
        let filename = "../data/err.txt";
        let (mut result, mut output) = merge(filename);
        assert!(!assert_iterator(&mut result, &mut output));
    }
}
<commit_msg>reduce mut<commit_after>use regex::Regex;
use std::fs::File;
use std::io::{self, BufRead};

#[macro_use]
extern crate lazy_static;

struct State<T>(Box<dyn Iterator<Item = T>>, Option<T>);
pub struct Merge {
    count: usize,
    states: Vec<State<u32>>,
}

impl Merge {
    pub fn new(filename: &str) -> Self {
        let count = get_number(filename).unwrap();
        let states: Vec<State<u32>> = (0..count).map(|i| {
            let collection = Box::new(get_collection(filename, i).unwrap());
            State(collection, None)
        }).collect();
        Merge {
            count: count,
            states: states,
        }
    }
}

impl Iterator for Merge {
    type Item = u32;
    fn next(&mut self) -> Option<Self::Item> {
        let count = self.count;
        let states = &mut self.states;
        loop {
            for i in 0..count {
                let state = &mut states[i];
                if state.1.is_none() {
                    state.1 = state.0.next();
                }
            }
            let not_done = states.iter().any(|s| s.1.is_some());
            if !not_done {
                return None;
            }

            let (mut min, mut min_index) = (u32::MAX, usize::MIN);
            for i in 0..count {
                if let Some(val) = states[i].1 {
                    if val < min {
                        min = val;
                        min_index = i;
                    }
                }
            }
            states[min_index].1 = None;
            return Some(min);
        }
    }
}

pub fn get_number(filename: &str) -> io::Result<usize> {
    lazy_static! {
        static ref RE: Regex = Regex::new(r"NUMBER=(\d+)").unwrap();
    }
    let mut matches = match_lines(filename, &*RE)?;
    return Ok(matches.next().unwrap().parse().unwrap());
}

pub fn get_output(filename: &str) -> io::Result<impl Iterator<Item = u32>> {
    lazy_static! {
        static ref RE: Regex = Regex::new(r"OUTPUT:(.*)").unwrap();
    }
    get_numbers(filename, &*RE)
}

pub fn get_collection(filename: &str, idx: usize) -> io::Result<impl Iterator<Item = u32>> {
    let regex: Regex = {
        let r = format!("COLLECTION{}:(.*)", idx);
        Regex::new(&r).unwrap()
    };
    get_numbers(filename, &regex)
}

fn get_numbers(filename: &str, regex: &Regex) -> io::Result<impl Iterator<Item = u32>> {
    let matches = match_lines(filename, regex)?;
    let to_int = |v: &str| -> Option<u32> {
        if let Ok(i) = v.parse::<u32>() {
            Some(i)
        } else {
            None
        }
    };
    let values = matches
        .map(move |s| {
            let numbers = s.split(',').filter_map(|v| to_int(&v));
            numbers.collect::<Vec<_>>()
        })
        .flatten();
    Ok(values)
}

fn match_lines(filename: &str, regex: &Regex) -> io::Result<impl Iterator<Item = String>> {
    let read_lines = move || -> io::Result<_> {
        let file = File::open(filename)?;
        Ok(io::BufReader::new(file).lines().map(|line| line.unwrap()))
    };
    let lines = read_lines()?;
    let r = regex.clone(); // TODO
    let matches =
        lines.filter_map(move |line| r.captures(&line).map(|captures| captures[1].to_string()));
    Ok(matches)
}

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

    #[test]
    fn test_number() {
        assert_eq!(get_number("../data/test.txt").unwrap(), 3);
    }

    #[test]
    fn test_output() {
        let output = [1, 1, 2, 2, 3, 4, 4, 6, 7, 9, 9, 20, 21];
        assert_eq!(
            get_output("../data/test.txt").unwrap().collect::<Vec<_>>(),
            output
        );
    }

    #[test]
    fn test_collection() {
        let output = [2, 4, 6, 9, 20];
        assert_eq!(
            get_collection("../data/test.txt", 1)
                .unwrap()
                .collect::<Vec<_>>(),
            output
        );
    }

    fn merge(filename: &str) -> (impl Iterator<Item = u32>, impl Iterator<Item = u32>) {
        let result = Merge::new(filename);
        let output = get_output(filename).unwrap();
        (result, output)
    }

    fn assert_iterator(
        iter1: &mut impl Iterator<Item = u32>,
        iter2: &mut impl Iterator<Item = u32>,
    ) -> bool {
        loop {
            match (iter1.next(), iter2.next()) {
                (None, None) => return true,
                (Some(v1), Some(v2)) if v1 == v2 => continue,
                _ => return false,
            }
        }
    }

    #[test]
    fn test_merge() {
        let filename = "../data/test.txt";
        let (mut result, mut output) = merge(filename);
        assert!(assert_iterator(&mut result, &mut output));
    }

    #[test]
    #[should_panic]
    fn test_bad_merge() {
        let filename = "../data/bad.txt";
        let (_result, _output) = merge(filename);
    }

    #[test]
    fn test_zero_merge() {
        let filename = "../data/zero.txt";
        let (mut result, mut output) = merge(filename);
        assert!(result.next().is_none());
        assert!(output.next().is_none());
    }

    #[test]
    fn test_err() {
        let filename = "../data/err.txt";
        let (mut result, mut output) = merge(filename);
        assert!(!assert_iterator(&mut result, &mut output));
    }
}
<|endoftext|>


<commit_before>use encoding;

import io::writer_util;
import encoding::codec;

fn main(args: ~[~str]) {
    let binary = copy args[0];
    let stdout = io::stdout();
    let stderr = io::stderr();

    if args.len() < 2u {
        stderr.write_line(fmt!{"Usage: %s <mode> <filename>", binary});
        return;
    }

    match io::read_whole_file(args[2]) {
        result::ok(data) => {
            let res = match args[1] {
                ~"encode" => data.encode(encoding::base64)
              , ~"decode" => data.decode(encoding::base64)
              , _         => return
            };
            stdout.write(res);
        }
        result::err(msg) => {
            stderr.write_line(#fmt["Error: %s", msg]);
        }
    }
}
<commit_msg>match expr as rhs<commit_after>use encoding;

import io::writer_util;
import encoding::codec;

fn main(args: ~[~str]) {
    let binary = copy args[0];
    let stdout = io::stdout();
    let stderr = io::stderr();

    if args.len() < 2u {
        stderr.write_line(fmt!("Usage: %s <mode> <filename>", binary));
        return;
    }

    match io::read_whole_file(args[2]) {
        result::ok(data) => match args[1] {
            ~"encode" => stdout.write(data.encode(encoding::base64))
          , ~"decode" => stdout.write(data.decode(encoding::base64))
          , _         => return
        }
        result::err(msg) => {
            stderr.write_line(fmt!("Error: %s", msg));
        }
    }
}
<|endoftext|>


<commit_before>// xfail-stage0

fn start(chan[int] c, int n) {
    let int i = n;

    while(i > 0) {
        c <| 0;
        i = i - 1;
    }
}

fn main() {
    let port[int] p = port();
    // Spawn a task that sends us back messages. The parent task
    // is likely to terminate before the child completes, so from
    // the child's point of view the receiver may die. We should
    // drop messages on the floor in this case, and not crash!
    auto child = spawn start(chan(p), 10);
    auto c; p |> c;
}
<commit_msg>Disable task-comm-15. Fails with multiple threads.<commit_after>// xfail-stage1
// xfail-stage2
// xfail-stage3
// This test fails when run with multiple threads

fn start(chan[int] c, int n) {
    let int i = n;

    while(i > 0) {
        c <| 0;
        i = i - 1;
    }
}

fn main() {
    let port[int] p = port();
    // Spawn a task that sends us back messages. The parent task
    // is likely to terminate before the child completes, so from
    // the child's point of view the receiver may die. We should
    // drop messages on the floor in this case, and not crash!
    auto child = spawn start(chan(p), 10);
    auto c; p |> c;
}
<|endoftext|>


<commit_before>use syntax::ast::{ItemKind, Mutability, Stmt, Ty, TyKind, Unsafety};
use syntax::ast::{self, Arg, Attribute, Expr, FnHeader, Generics, Ident};
use syntax::attr::check_builtin_macro_attribute;
use syntax::ext::allocator::{AllocatorKind, AllocatorMethod, AllocatorTy, ALLOCATOR_METHODS};
use syntax::ext::base::{Annotatable, ExtCtxt};
use syntax::ext::build::AstBuilder;
use syntax::ext::hygiene::SyntaxContext;
use syntax::ptr::P;
use syntax::symbol::{kw, sym, Symbol};
use syntax_pos::Span;

pub fn expand(
    ecx: &mut ExtCtxt<'_>,
    _span: Span,
    meta_item: &ast::MetaItem,
    item: Annotatable,
) -> Vec<Annotatable> {
    check_builtin_macro_attribute(ecx, meta_item, sym::global_allocator);

    let not_static = |item: Annotatable| {
        ecx.parse_sess.span_diagnostic.span_err(item.span(), "allocators must be statics");
        vec![item]
    };
    let item = match item {
        Annotatable::Item(item) => match item.node {
            ItemKind::Static(..) => item,
            _ => return not_static(Annotatable::Item(item)),
        }
        _ => return not_static(item),
    };

    // Generate a bunch of new items using the AllocFnFactory
    let span = item.span.with_ctxt(SyntaxContext::empty().apply_mark(ecx.current_expansion.mark));
    let f = AllocFnFactory {
        span,
        kind: AllocatorKind::Global,
        global: item.ident,
        cx: ecx,
    };

    // Generate item statements for the allocator methods.
    let stmts = ALLOCATOR_METHODS.iter().map(|method| f.allocator_fn(method)).collect();

    // Generate anonymous constant serving as container for the allocator methods.
    let const_ty = ecx.ty(span, TyKind::Tup(Vec::new()));
    let const_body = ecx.expr_block(ecx.block(span, stmts));
    let const_item =
        ecx.item_const(span, Ident::with_empty_ctxt(kw::Underscore), const_ty, const_body);

    // Return the original item and the new methods.
    vec![Annotatable::Item(item), Annotatable::Item(const_item)]
}

struct AllocFnFactory<'a, 'b> {
    span: Span,
    kind: AllocatorKind,
    global: Ident,
    cx: &'b ExtCtxt<'a>,
}

impl AllocFnFactory<'_, '_> {
    fn allocator_fn(&self, method: &AllocatorMethod) -> Stmt {
        let mut abi_args = Vec::new();
        let mut i = 0;
        let ref mut mk = || {
            let name = Ident::from_str(&format!("arg{}", i));
            i += 1;
            name
        };
        let args = method
            .inputs
            .iter()
            .map(|ty| self.arg_ty(ty, &mut abi_args, mk))
            .collect();
        let result = self.call_allocator(method.name, args);
        let (output_ty, output_expr) = self.ret_ty(&method.output, result);
        let kind = ItemKind::Fn(
            self.cx.fn_decl(abi_args, ast::FunctionRetTy::Ty(output_ty)),
            FnHeader {
                unsafety: Unsafety::Unsafe,
                ..FnHeader::default()
            },
            Generics::default(),
            self.cx.block_expr(output_expr),
        );
        let item = self.cx.item(
            self.span,
            Ident::from_str(&self.kind.fn_name(method.name)),
            self.attrs(),
            kind,
        );
        self.cx.stmt_item(self.span, item)
    }

    fn call_allocator(&self, method: &str, mut args: Vec<P<Expr>>) -> P<Expr> {
        let method = self.cx.std_path(&[
            Symbol::intern("alloc"),
            Symbol::intern("GlobalAlloc"),
            Symbol::intern(method),
        ]);
        let method = self.cx.expr_path(self.cx.path(self.span, method));
        let allocator = self.cx.path_ident(self.span, self.global);
        let allocator = self.cx.expr_path(allocator);
        let allocator = self.cx.expr_addr_of(self.span, allocator);
        args.insert(0, allocator);

        self.cx.expr_call(self.span, method, args)
    }

    fn attrs(&self) -> Vec<Attribute> {
        let special = sym::rustc_std_internal_symbol;
        let special = self.cx.meta_word(self.span, special);
        vec![self.cx.attribute(self.span, special)]
    }

    fn arg_ty(
        &self,
        ty: &AllocatorTy,
        args: &mut Vec<Arg>,
        ident: &mut dyn FnMut() -> Ident,
    ) -> P<Expr> {
        match *ty {
            AllocatorTy::Layout => {
                let usize = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::usize));
                let ty_usize = self.cx.ty_path(usize);
                let size = ident();
                let align = ident();
                args.push(self.cx.arg(self.span, size, ty_usize.clone()));
                args.push(self.cx.arg(self.span, align, ty_usize));

                let layout_new = self.cx.std_path(&[
                    Symbol::intern("alloc"),
                    Symbol::intern("Layout"),
                    Symbol::intern("from_size_align_unchecked"),
                ]);
                let layout_new = self.cx.expr_path(self.cx.path(self.span, layout_new));
                let size = self.cx.expr_ident(self.span, size);
                let align = self.cx.expr_ident(self.span, align);
                let layout = self.cx.expr_call(self.span, layout_new, vec![size, align]);
                layout
            }

            AllocatorTy::Ptr => {
                let ident = ident();
                args.push(self.cx.arg(self.span, ident, self.ptr_u8()));
                let arg = self.cx.expr_ident(self.span, ident);
                self.cx.expr_cast(self.span, arg, self.ptr_u8())
            }

            AllocatorTy::Usize => {
                let ident = ident();
                args.push(self.cx.arg(self.span, ident, self.usize()));
                self.cx.expr_ident(self.span, ident)
            }

            AllocatorTy::ResultPtr | AllocatorTy::Unit => {
                panic!("can't convert AllocatorTy to an argument")
            }
        }
    }

    fn ret_ty(&self, ty: &AllocatorTy, expr: P<Expr>) -> (P<Ty>, P<Expr>) {
        match *ty {
            AllocatorTy::ResultPtr => {
                // We're creating:
                //
                //      #expr as *mut u8

                let expr = self.cx.expr_cast(self.span, expr, self.ptr_u8());
                (self.ptr_u8(), expr)
            }

            AllocatorTy::Unit => (self.cx.ty(self.span, TyKind::Tup(Vec::new())), expr),

            AllocatorTy::Layout | AllocatorTy::Usize | AllocatorTy::Ptr => {
                panic!("can't convert `AllocatorTy` to an output")
            }
        }
    }

    fn usize(&self) -> P<Ty> {
        let usize = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::usize));
        self.cx.ty_path(usize)
    }

    fn ptr_u8(&self) -> P<Ty> {
        let u8 = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::u8));
        let ty_u8 = self.cx.ty_path(u8);
        self.cx.ty_ptr(self.span, ty_u8, Mutability::Mutable)
    }
}
<commit_msg>Fix rebase<commit_after>use syntax::ast::{ItemKind, Mutability, Stmt, Ty, TyKind, Unsafety};
use syntax::ast::{self, Arg, Attribute, Expr, FnHeader, Generics, Ident};
use syntax::attr::check_builtin_macro_attribute;
use syntax::ext::allocator::{AllocatorKind, AllocatorMethod, AllocatorTy, ALLOCATOR_METHODS};
use syntax::ext::base::{Annotatable, ExtCtxt};
use syntax::ext::build::AstBuilder;
use syntax::ext::hygiene::SyntaxContext;
use syntax::ptr::P;
use syntax::symbol::{kw, sym, Symbol};
use syntax_pos::Span;

pub fn expand(
    ecx: &mut ExtCtxt<'_>,
    _span: Span,
    meta_item: &ast::MetaItem,
    item: Annotatable,
) -> Vec<Annotatable> {
    check_builtin_macro_attribute(ecx, meta_item, sym::global_allocator);

    let not_static = |item: Annotatable| {
        ecx.parse_sess.span_diagnostic.span_err(item.span(), "allocators must be statics");
        vec![item]
    };
    let item = match item {
        Annotatable::Item(item) => match item.node {
            ItemKind::Static(..) => item,
            _ => return not_static(Annotatable::Item(item)),
        }
        _ => return not_static(item),
    };

    // Generate a bunch of new items using the AllocFnFactory
    let span = item.span.with_ctxt(SyntaxContext::empty().apply_mark(ecx.current_expansion.id));
    let f = AllocFnFactory {
        span,
        kind: AllocatorKind::Global,
        global: item.ident,
        cx: ecx,
    };

    // Generate item statements for the allocator methods.
    let stmts = ALLOCATOR_METHODS.iter().map(|method| f.allocator_fn(method)).collect();

    // Generate anonymous constant serving as container for the allocator methods.
    let const_ty = ecx.ty(span, TyKind::Tup(Vec::new()));
    let const_body = ecx.expr_block(ecx.block(span, stmts));
    let const_item =
        ecx.item_const(span, Ident::with_empty_ctxt(kw::Underscore), const_ty, const_body);

    // Return the original item and the new methods.
    vec![Annotatable::Item(item), Annotatable::Item(const_item)]
}

struct AllocFnFactory<'a, 'b> {
    span: Span,
    kind: AllocatorKind,
    global: Ident,
    cx: &'b ExtCtxt<'a>,
}

impl AllocFnFactory<'_, '_> {
    fn allocator_fn(&self, method: &AllocatorMethod) -> Stmt {
        let mut abi_args = Vec::new();
        let mut i = 0;
        let ref mut mk = || {
            let name = Ident::from_str(&format!("arg{}", i));
            i += 1;
            name
        };
        let args = method
            .inputs
            .iter()
            .map(|ty| self.arg_ty(ty, &mut abi_args, mk))
            .collect();
        let result = self.call_allocator(method.name, args);
        let (output_ty, output_expr) = self.ret_ty(&method.output, result);
        let kind = ItemKind::Fn(
            self.cx.fn_decl(abi_args, ast::FunctionRetTy::Ty(output_ty)),
            FnHeader {
                unsafety: Unsafety::Unsafe,
                ..FnHeader::default()
            },
            Generics::default(),
            self.cx.block_expr(output_expr),
        );
        let item = self.cx.item(
            self.span,
            Ident::from_str(&self.kind.fn_name(method.name)),
            self.attrs(),
            kind,
        );
        self.cx.stmt_item(self.span, item)
    }

    fn call_allocator(&self, method: &str, mut args: Vec<P<Expr>>) -> P<Expr> {
        let method = self.cx.std_path(&[
            Symbol::intern("alloc"),
            Symbol::intern("GlobalAlloc"),
            Symbol::intern(method),
        ]);
        let method = self.cx.expr_path(self.cx.path(self.span, method));
        let allocator = self.cx.path_ident(self.span, self.global);
        let allocator = self.cx.expr_path(allocator);
        let allocator = self.cx.expr_addr_of(self.span, allocator);
        args.insert(0, allocator);

        self.cx.expr_call(self.span, method, args)
    }

    fn attrs(&self) -> Vec<Attribute> {
        let special = sym::rustc_std_internal_symbol;
        let special = self.cx.meta_word(self.span, special);
        vec![self.cx.attribute(self.span, special)]
    }

    fn arg_ty(
        &self,
        ty: &AllocatorTy,
        args: &mut Vec<Arg>,
        ident: &mut dyn FnMut() -> Ident,
    ) -> P<Expr> {
        match *ty {
            AllocatorTy::Layout => {
                let usize = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::usize));
                let ty_usize = self.cx.ty_path(usize);
                let size = ident();
                let align = ident();
                args.push(self.cx.arg(self.span, size, ty_usize.clone()));
                args.push(self.cx.arg(self.span, align, ty_usize));

                let layout_new = self.cx.std_path(&[
                    Symbol::intern("alloc"),
                    Symbol::intern("Layout"),
                    Symbol::intern("from_size_align_unchecked"),
                ]);
                let layout_new = self.cx.expr_path(self.cx.path(self.span, layout_new));
                let size = self.cx.expr_ident(self.span, size);
                let align = self.cx.expr_ident(self.span, align);
                let layout = self.cx.expr_call(self.span, layout_new, vec![size, align]);
                layout
            }

            AllocatorTy::Ptr => {
                let ident = ident();
                args.push(self.cx.arg(self.span, ident, self.ptr_u8()));
                let arg = self.cx.expr_ident(self.span, ident);
                self.cx.expr_cast(self.span, arg, self.ptr_u8())
            }

            AllocatorTy::Usize => {
                let ident = ident();
                args.push(self.cx.arg(self.span, ident, self.usize()));
                self.cx.expr_ident(self.span, ident)
            }

            AllocatorTy::ResultPtr | AllocatorTy::Unit => {
                panic!("can't convert AllocatorTy to an argument")
            }
        }
    }

    fn ret_ty(&self, ty: &AllocatorTy, expr: P<Expr>) -> (P<Ty>, P<Expr>) {
        match *ty {
            AllocatorTy::ResultPtr => {
                // We're creating:
                //
                //      #expr as *mut u8

                let expr = self.cx.expr_cast(self.span, expr, self.ptr_u8());
                (self.ptr_u8(), expr)
            }

            AllocatorTy::Unit => (self.cx.ty(self.span, TyKind::Tup(Vec::new())), expr),

            AllocatorTy::Layout | AllocatorTy::Usize | AllocatorTy::Ptr => {
                panic!("can't convert `AllocatorTy` to an output")
            }
        }
    }

    fn usize(&self) -> P<Ty> {
        let usize = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::usize));
        self.cx.ty_path(usize)
    }

    fn ptr_u8(&self) -> P<Ty> {
        let u8 = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::u8));
        let ty_u8 = self.cx.ty_path(u8);
        self.cx.ty_ptr(self.span, ty_u8, Mutability::Mutable)
    }
}
<|endoftext|>


<commit_before>//! Various builders to generate/alter wasm components

mod invoke;
mod module;
mod code;
mod misc;
mod import;
mod memory;
mod table;
mod export;
mod global;
mod data;

pub use self::code::{
	signatures, signature, function, SignatureBuilder, SignaturesBuilder,
	FunctionBuilder, TypeRefBuilder, FuncBodyBuilder, FunctionDefinition,
};
pub use self::data::DataSegmentBuilder;
pub use self::export::{export, ExportBuilder, ExportInternalBuilder};
pub use self::global::{global, GlobalBuilder};
pub use self::import::{import, ImportBuilder};
pub use self::invoke::Identity;
pub use self::memory::MemoryBuilder;
pub use self::module::{module, from_module, ModuleBuilder};
pub use self::table::{TableBuilder, TableDefinition, TableEntryDefinition};
<commit_msg>Export CodeLocation (#301)<commit_after>//! Various builders to generate/alter wasm components

mod invoke;
mod module;
mod code;
mod misc;
mod import;
mod memory;
mod table;
mod export;
mod global;
mod data;

pub use self::code::{
	signatures, signature, function, SignatureBuilder, SignaturesBuilder,
	FunctionBuilder, TypeRefBuilder, FuncBodyBuilder, FunctionDefinition,
};
pub use self::data::DataSegmentBuilder;
pub use self::export::{export, ExportBuilder, ExportInternalBuilder};
pub use self::global::{global, GlobalBuilder};
pub use self::import::{import, ImportBuilder};
pub use self::invoke::Identity;
pub use self::memory::MemoryBuilder;
pub use self::module::{module, from_module, ModuleBuilder, CodeLocation};
pub use self::table::{TableBuilder, TableDefinition, TableEntryDefinition};
<|endoftext|>


<commit_before>//! A set of builders used to make using methods on certain structs simpler to
//! use.
//!
//! These are used when not all parameters are required, all parameters are
//! optional, and/or sane default values for required parameters can be applied
//! by a builder.

mod create_embed;
mod create_channel;
mod create_invite;
mod create_message;
mod edit_channel;
mod edit_guild;
mod edit_member;
mod edit_message;
mod edit_profile;
mod edit_role;
mod execute_webhook;
mod get_messages;

pub use self::{
    create_embed::{CreateEmbed, CreateEmbedAuthor, CreateEmbedFooter},
    create_channel::CreateChannel,
    create_invite::CreateInvite,
    create_message::CreateMessage,
    edit_channel::EditChannel,
    edit_guild::EditGuild,
    edit_member::EditMember,
    edit_message::EditMessage,
    edit_profile::EditProfile,
    edit_role::EditRole,
    execute_webhook::ExecuteWebhook,
    get_messages::GetMessages
};
<commit_msg>Reexport `Timestamp` from the builder's module.<commit_after>//! A set of builders used to make using methods on certain structs simpler to
//! use.
//!
//! These are used when not all parameters are required, all parameters are
//! optional, and/or sane default values for required parameters can be applied
//! by a builder.

mod create_embed;
mod create_channel;
mod create_invite;
mod create_message;
mod edit_channel;
mod edit_guild;
mod edit_member;
mod edit_message;
mod edit_profile;
mod edit_role;
mod execute_webhook;
mod get_messages;

pub use self::{
    create_embed::{CreateEmbed, CreateEmbedAuthor, CreateEmbedFooter, Timestamp},
    create_channel::CreateChannel,
    create_invite::CreateInvite,
    create_message::CreateMessage,
    edit_channel::EditChannel,
    edit_guild::EditGuild,
    edit_member::EditMember,
    edit_message::EditMessage,
    edit_profile::EditProfile,
    edit_role::EditRole,
    execute_webhook::ExecuteWebhook,
    get_messages::GetMessages
};
<|endoftext|>


<commit_before>// Bloom
//
// HTTP REST API caching middleware
// Copyright: 2017, Valerian Saliou <valerian@valeriansaliou.name>
// License: Mozilla Public License v2.0 (MPL v2.0)

use std::cmp;
use std::time::Duration;
use r2d2::Pool;
use r2d2::config::Config;
use r2d2_redis::{RedisConnectionManager, Error};
use redis::{self, Value, Connection, Commands, PipelineCommands};

use APP_CONF;

pub struct CacheStoreBuilder;

pub struct CacheStore {
    pool: Pool<RedisConnectionManager>,
}

#[derive(Debug)]
pub enum CacheStoreError {
    Disconnected,
    Failed,
    Invalid,
    Corrupted,
    TooLarge,
}

#[derive(Debug)]
pub enum CachePurgeVariant {
    Bucket,
    Auth,
}

type CacheResult = Result<Option<String>, CacheStoreError>;

impl CacheStoreBuilder {
    pub fn new() -> CacheStore {
        info!("binding to store backend at {}", APP_CONF.redis.inet);

        let tcp_addr_raw =
            format!(
            "redis://{}:{}/{}",
            APP_CONF.redis.inet.ip(),
            APP_CONF.redis.inet.port(),
            APP_CONF.redis.database,
        );

        match RedisConnectionManager::new(tcp_addr_raw.as_ref()) {
            Ok(manager) => {
                let config = Config::<Connection, Error>::builder()
                    .pool_size(APP_CONF.redis.pool_size)
                    .idle_timeout(Some(
                        Duration::from_secs(APP_CONF.redis.idle_timeout_seconds),
                    ))
                    .connection_timeout(Duration::from_secs(
                        APP_CONF.redis.connection_timeout_seconds,
                    ))
                    .build();

                match Pool::new(config, manager) {
                    Ok(pool) => {
                        info!("bound to store backend");

                        CacheStore { pool: pool }
                    }
                    Err(_) => panic!("could not spawn redis pool"),
                }
            }
            Err(_) => panic!("could not create redis connection manager"),
        }
    }
}

impl CacheStore {
    pub fn get(&self, key: &str) -> CacheResult {
        get_cache_store_client!(self, client {
            match (*client).get::<_, Value>(key) {
                Ok(value) => {
                    match value {
                        Value::Data(bytes) => {
                            // Decode raw bytes to string
                            if let Ok(string) = String::from_utf8(bytes) {
                                Ok(Some(string))
                            } else {
                                Err(CacheStoreError::Corrupted)
                            }
                        },
                        Value::Nil => Ok(None),
                        _ => Err(CacheStoreError::Invalid),
                    }
                },
                _ => Err(CacheStoreError::Failed),
            }
        })
    }

    pub fn set(
        &self,
        key: &str,
        value: &str,
        ttl: usize,
        key_buckets: Option<Vec<String>>,
    ) -> CacheResult {
        get_cache_store_client!(self, client {
            // Cap TTL to 'max_key_expiration'
            let ttl_cap = cmp::min(ttl, APP_CONF.redis.max_key_expiration);

            // Ensure value is not larger than 'max_key_size'
            if value.len() > APP_CONF.redis.max_key_size {
                Err(CacheStoreError::TooLarge)
            } else {
                match key_buckets {
                    Some(key_buckets_value) => {
                        // Bucket (MULTI operation for main data + bucket marker)
                        gen_cache_store_empty_result!(
                            redis::pipe()
                                .atomic()
                                .set_ex(key, value, ttl_cap).ignore()
                                .set_ex(key_buckets_value, "", ttl_cap).ignore()
                                .query::<()>(&*client)
                        )
                    },
                    None => {
                        gen_cache_store_empty_result!(
                            (*client).set_ex::<_, _, ()>(key, value, ttl_cap)
                        )
                    },
                }
            }
        })
    }

    pub fn purge_pattern(&self, variant: &CachePurgeVariant, key_pattern: &str) -> CacheResult {
        get_cache_store_client!(self, client {
            // Invoke keyspace cleanup script for key pattern
            gen_cache_store_empty_result!(
                redis::Script::new(variant.get_script())
                    .arg(key_pattern)
                    .invoke::<()>(&*client)
            )
        })
    }
}

impl CachePurgeVariant {
    fn get_script(&self) -> &'static str {
        match *self {
            CachePurgeVariant::Bucket => {
                r#"
                    local re = '^(.+):b:.+$'
                    local targets = {}

                    for _, bucket_key in pairs(redis.call('KEYS', ARGV[1])) do
                        local base_key = bucket_key:match(re)

                        if base_key then
                            table.insert(targets, base_key)
                        end

                        table.insert(targets, bucket_key)
                    end

                    if next(targets) then
                        redis.call('DEL', unpack(targets))
                    end
                "#
            }
            CachePurgeVariant::Auth => {
                r#"
                    local targets = redis.call('KEYS', ARGV[1])

                    if next(targets) then
                        redis.call('DEL', unpack(targets))
                    end
                "#
            }
        }
    }
}
<commit_msg>Support for multi buckets in Redis<commit_after>// Bloom
//
// HTTP REST API caching middleware
// Copyright: 2017, Valerian Saliou <valerian@valeriansaliou.name>
// License: Mozilla Public License v2.0 (MPL v2.0)

use std::cmp;
use std::time::Duration;
use r2d2::Pool;
use r2d2::config::Config;
use r2d2_redis::{RedisConnectionManager, Error};
use redis::{self, Value, Connection, Commands, PipelineCommands};

use APP_CONF;

pub struct CacheStoreBuilder;

pub struct CacheStore {
    pool: Pool<RedisConnectionManager>,
}

#[derive(Debug)]
pub enum CacheStoreError {
    Disconnected,
    Failed,
    Invalid,
    Corrupted,
    TooLarge,
}

#[derive(Debug)]
pub enum CachePurgeVariant {
    Bucket,
    Auth,
}

type CacheResult = Result<Option<String>, CacheStoreError>;

impl CacheStoreBuilder {
    pub fn new() -> CacheStore {
        info!("binding to store backend at {}", APP_CONF.redis.inet);

        let tcp_addr_raw =
            format!(
            "redis://{}:{}/{}",
            APP_CONF.redis.inet.ip(),
            APP_CONF.redis.inet.port(),
            APP_CONF.redis.database,
        );

        match RedisConnectionManager::new(tcp_addr_raw.as_ref()) {
            Ok(manager) => {
                let config = Config::<Connection, Error>::builder()
                    .pool_size(APP_CONF.redis.pool_size)
                    .idle_timeout(Some(
                        Duration::from_secs(APP_CONF.redis.idle_timeout_seconds),
                    ))
                    .connection_timeout(Duration::from_secs(
                        APP_CONF.redis.connection_timeout_seconds,
                    ))
                    .build();

                match Pool::new(config, manager) {
                    Ok(pool) => {
                        info!("bound to store backend");

                        CacheStore { pool: pool }
                    }
                    Err(_) => panic!("could not spawn redis pool"),
                }
            }
            Err(_) => panic!("could not create redis connection manager"),
        }
    }
}

impl CacheStore {
    pub fn get(&self, key: &str) -> CacheResult {
        get_cache_store_client!(self, client {
            match (*client).get::<_, Value>(key) {
                Ok(value) => {
                    match value {
                        Value::Data(bytes) => {
                            // Decode raw bytes to string
                            if let Ok(string) = String::from_utf8(bytes) {
                                Ok(Some(string))
                            } else {
                                Err(CacheStoreError::Corrupted)
                            }
                        },
                        Value::Nil => Ok(None),
                        _ => Err(CacheStoreError::Invalid),
                    }
                },
                _ => Err(CacheStoreError::Failed),
            }
        })
    }

    pub fn set(
        &self,
        key: &str,
        value: &str,
        ttl: usize,
        key_buckets: Option<Vec<String>>,
    ) -> CacheResult {
        get_cache_store_client!(self, client {
            // Cap TTL to 'max_key_expiration'
            let ttl_cap = cmp::min(ttl, APP_CONF.redis.max_key_expiration);

            // Ensure value is not larger than 'max_key_size'
            if value.len() > APP_CONF.redis.max_key_size {
                Err(CacheStoreError::TooLarge)
            } else {
                match key_buckets {
                    Some(key_buckets_value) => {
                        let mut pipeline = redis::pipe();

                        pipeline.set_ex(key, value, ttl_cap).ignore();

                        for key_bucket_value in key_buckets_value {
                            pipeline.set_ex(key_bucket_value, "", ttl_cap).ignore();
                        }

                        // Bucket (MULTI operation for main data + bucket marker)
                        gen_cache_store_empty_result!(
                            pipeline.query::<()>(&*client)
                        )
                    },
                    None => {
                        gen_cache_store_empty_result!(
                            (*client).set_ex::<_, _, ()>(key, value, ttl_cap)
                        )
                    },
                }
            }
        })
    }

    pub fn purge_pattern(&self, variant: &CachePurgeVariant, key_pattern: &str) -> CacheResult {
        get_cache_store_client!(self, client {
            // Invoke keyspace cleanup script for key pattern
            gen_cache_store_empty_result!(
                redis::Script::new(variant.get_script())
                    .arg(key_pattern)
                    .invoke::<()>(&*client)
            )
        })
    }
}

impl CachePurgeVariant {
    fn get_script(&self) -> &'static str {
        match *self {
            CachePurgeVariant::Bucket => {
                r#"
                    local re = '^(.+):b:.+$'
                    local targets = {}

                    for _, bucket_key in pairs(redis.call('KEYS', ARGV[1])) do
                        local base_key = bucket_key:match(re)

                        if base_key then
                            table.insert(targets, base_key)
                        end

                        table.insert(targets, bucket_key)
                    end

                    if next(targets) then
                        redis.call('DEL', unpack(targets))
                    end
                "#
            }
            CachePurgeVariant::Auth => {
                r#"
                    local targets = redis.call('KEYS', ARGV[1])

                    if next(targets) then
                        redis.call('DEL', unpack(targets))
                    end
                "#
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::{
    sync::{
        atomic::{AtomicBool, AtomicUsize, Ordering},
        Arc,
    },
    thread,
};

use rand::{thread_rng, Rng};

#[cfg_attr(
    // only enable jemalloc on linux and macos by default
    all(
        any(target_os = "linux", target_os = "macos"),
        feature = "jemalloc",
    ),
    global_allocator
)]
#[cfg(
    // only enable jemalloc on linux and macos by default
    all(
        any(target_os = "linux", target_os = "macos"),
        feature = "jemalloc",
    ),
)]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

static TOTAL: AtomicUsize = AtomicUsize::new(0);

const USAGE: &str = "
Usage: stress [--threads=<#>] [--burn-in] [--duration=<s>] \
    [--key-len=<l>] [--val-len=<l>] \
    [--get-prop=<p>] \
    [--set-prop=<p>] \
    [--del-prop=<p>] \
    [--cas-prop=<p>] \
    [--scan-prop=<p>] \
    [--merge-prop=<p>] \
    [--entries=<n>] \
    [--sequential] \
    [--total-ops=<n>] \
    [--flush-every=<ms>]

Options:
    --threads=<#>      Number of threads [default: 4].
    --burn-in          Don't halt until we receive a signal.
    --duration=<s>     Seconds to run for [default: 10].
    --key-len=<l>      The length of keys [default: 10].
    --val-len=<l>      The length of values [default: 100].
    --get-prop=<p>     The relative proportion of get requests [default: 94].
    --set-prop=<p>     The relative proportion of set requests [default: 2].
    --del-prop=<p>     The relative proportion of del requests [default: 1].
    --cas-prop=<p>     The relative proportion of cas requests [default: 1].
    --scan-prop=<p>    The relative proportion of scan requests [default: 1].
    --merge-prop=<p>   The relative proportion of merge requests [default: 1].
    --entries=<n>      The total keyspace [default: 100000].
    --sequential       Run the test in sequential mode instead of random.
    --total-ops=<n>    Stop test after executing a total number of operations.
    --flush-every=<m>  Flush and sync the database every ms [default: 200].
";

#[derive(Clone, Copy)]
struct Args {
    threads: usize,
    burn_in: bool,
    duration: u64,
    key_len: usize,
    val_len: usize,
    get_prop: usize,
    set_prop: usize,
    del_prop: usize,
    cas_prop: usize,
    scan_prop: usize,
    merge_prop: usize,
    entries: usize,
    sequential: bool,
    total_ops: Option<usize>,
    flush_every: u64,
}

impl Default for Args {
    fn default() -> Args {
        Args {
            threads: 4,
            burn_in: false,
            duration: 10,
            key_len: 10,
            val_len: 100,
            get_prop: 94,
            set_prop: 2,
            del_prop: 1,
            cas_prop: 1,
            scan_prop: 1,
            merge_prop: 1,
            entries: 100000,
            sequential: false,
            total_ops: None,
            flush_every: 200,
        }
    }
}

fn parse<'a, I, T>(mut iter: I) -> T
where
    I: Iterator<Item = &'a str>,
    T: std::str::FromStr,
    <T as std::str::FromStr>::Err: std::fmt::Debug,
{
    iter.next().expect(USAGE).parse().expect(USAGE)
}

impl Args {
    fn parse() -> Args {
        let mut args = Args::default();
        for raw_arg in std::env::args().skip(1) {
            let mut splits = raw_arg[2..].split('=');
            match splits.next().unwrap() {
                "threads" => args.threads = parse(&mut splits),
                "burn-in" => args.burn_in = true,
                "duration" => args.duration = parse(&mut splits),
                "key-len" => args.key_len = parse(&mut splits),
                "val-len" => args.val_len = parse(&mut splits),
                "get-prop" => args.get_prop = parse(&mut splits),
                "set-prop" => args.set_prop = parse(&mut splits),
                "del-prop" => args.del_prop = parse(&mut splits),
                "cas-prop" => args.cas_prop = parse(&mut splits),
                "scan-prop" => args.scan_prop = parse(&mut splits),
                "merge-prop" => args.merge_prop = parse(&mut splits),
                "entries" => args.entries = parse(&mut splits),
                "sequential" => args.sequential = true,
                "total-ops" => args.total_ops = Some(parse(&mut splits)),
                "flush-every" => args.flush_every = parse(&mut splits),
                other => panic!("unknown option: {}, {}", other, USAGE),
            }
        }
        args
    }
}

fn report(shutdown: Arc<AtomicBool>) {
    let mut last = 0;
    while !shutdown.load(Ordering::Relaxed) {
        thread::sleep(std::time::Duration::from_secs(1));
        let total = TOTAL.load(Ordering::Acquire);

        println!("did {} ops, {}mb RSS", total - last, rss() / (1024 * 1024));

        last = total;
    }
}

fn concatenate_merge(
    _key: &[u8],              // the key being merged
    old_value: Option<&[u8]>, // the previous value, if one existed
    merged_bytes: &[u8],      // the new bytes being merged in
) -> Option<Vec<u8>> {
    // set the new value, return None to delete
    let mut ret = old_value.map(|ov| ov.to_vec()).unwrap_or_else(Vec::new);

    ret.extend_from_slice(merged_bytes);

    Some(ret)
}

fn run(args: Args, tree: Arc<sled::Db>, shutdown: Arc<AtomicBool>) {
    let get_max = args.get_prop;
    let set_max = get_max + args.set_prop;
    let del_max = set_max + args.del_prop;
    let cas_max = del_max + args.cas_prop;
    let merge_max = cas_max + args.merge_prop;
    let scan_max = merge_max + args.scan_prop;

    let keygen = |len| -> sled::IVec {
        static SEQ: AtomicUsize = AtomicUsize::new(0);
        let i = if args.sequential {
            SEQ.fetch_add(1, Ordering::Relaxed)
        } else {
            thread_rng().gen::<usize>()
        } % args.entries;

        let i_keygen = i.to_be_bytes();

        i_keygen
            .iter()
            .skip_while(|v| **v == 0)
            .cycle()
            .take(len)
            .copied()
            .collect()
    };

    let valgen = |len| -> sled::IVec {
        if len == 0 {
            return vec![].into();
        }

        let i: usize = thread_rng().gen::<usize>() % (len * 8);

        let i_keygen = i.to_be_bytes();

        i_keygen
            .iter()
            .skip_while(|v| **v == 0)
            .cycle()
            .take(len)
            .copied()
            .collect()
    };

    let mut rng = thread_rng();

    while !shutdown.load(Ordering::Relaxed) {
        let op = TOTAL.fetch_add(1, Ordering::Release);
        let key = keygen(args.key_len);
        let choice = rng.gen_range(0, scan_max + 1);

        match choice {
            v if v <= get_max => {
                tree.get(&key).unwrap();
            }
            v if v > get_max && v <= set_max => {
                let value = valgen(args.val_len);
                tree.insert(&key, value).unwrap();
            }
            v if v > set_max && v <= del_max => {
                tree.remove(&key).unwrap();
            }
            v if v > del_max && v <= cas_max => {
                let old = if rng.gen::<bool>() {
                    let value = valgen(args.val_len);
                    Some(value)
                } else {
                    None
                };

                let new = if rng.gen::<bool>() {
                    let value = valgen(args.val_len);
                    Some(value)
                } else {
                    None
                };

                if let Err(e) = tree.compare_and_swap(&key, old, new) {
                    panic!("operational error: {:?}", e);
                }
            }
            v if v > cas_max && v <= merge_max => {
                let value = valgen(args.val_len);
                tree.merge(&key, value).unwrap();
            }
            _ => {
                let iter = tree.range(key..).map(|res| res.unwrap());

                if op % 2 == 0 {
                    let _ = iter.take(rng.gen_range(0, 15)).collect::<Vec<_>>();
                } else {
                    let _ = iter
                        .rev()
                        .take(rng.gen_range(0, 15))
                        .collect::<Vec<_>>();
                }
            }
        }
    }
}

fn rss() -> usize {
    #[cfg(target_os = "linux")]
    {
        use std::io::prelude::*;
        use std::io::BufReader;

        let mut buf = String::new();
        let mut f =
            BufReader::new(std::fs::File::open("/proc/self/statm").unwrap());
        f.read_line(&mut buf).unwrap();
        let mut parts = buf.split_whitespace();
        let rss_pages = parts.nth(1).unwrap().parse::<usize>().unwrap();
        rss_pages * 4096
    }
    #[cfg(not(target_os = "linux"))]
    {
        0
    }
}

fn main() {
    #[cfg(feature = "logging")]
    setup_logger();

    let args = Args::parse();

    let shutdown = Arc::new(AtomicBool::new(false));

    let config = sled::Config::new()
        .cache_capacity(256 * 1024 * 1024)
        .flush_every_ms(Some(args.flush_every))
        .print_profile_on_drop(true);

    let tree = Arc::new(config.open().unwrap());
    tree.set_merge_operator(concatenate_merge);

    let mut threads = vec![];

    let now = std::time::Instant::now();

    let n_threads = args.threads;

    for i in 0..=n_threads {
        let tree = tree.clone();
        let shutdown = shutdown.clone();

        let t = if i == 0 {
            thread::Builder::new()
                .name("reporter".into())
                .spawn(move || report(shutdown))
                .unwrap()
        } else {
            thread::spawn(move || run(args, tree, shutdown))
        };

        threads.push(t);
    }

    if let Some(ops) = args.total_ops {
        assert!(!args.burn_in, "don't set both --burn-in and --total-ops");
        while TOTAL.load(Ordering::Relaxed) < ops {
            thread::sleep(std::time::Duration::from_millis(50));
        }
        shutdown.store(true, Ordering::SeqCst);
    } else if !args.burn_in {
        thread::sleep(std::time::Duration::from_secs(args.duration));
        shutdown.store(true, Ordering::SeqCst);
    }

    for t in threads.into_iter() {
        t.join().unwrap();
    }

    let ops = TOTAL.load(Ordering::SeqCst);
    let time = now.elapsed().as_secs() as usize;

    println!(
        "did {} total ops in {} seconds. {} ops/s",
        ops,
        time,
        (ops * 1_000) / (time * 1_000)
    );
}

#[cfg(feature = "logging")]
pub fn setup_logger() {
    use std::io::Write;

    color_backtrace::install();

    fn tn() -> String {
        std::thread::current().name().unwrap_or("unknown").to_owned()
    }

    let mut builder = env_logger::Builder::new();
    builder
        .format(|buf, record| {
            writeln!(
                buf,
                "{:05} {:25} {:10} {}",
                record.level(),
                tn(),
                record.module_path().unwrap().split("::").last().unwrap(),
                record.args()
            )
        })
        .filter(None, log::LevelFilter::Info);

    if let Ok(env) = std::env::var("RUST_LOG") {
        builder.parse_filters(&env);
    }

    let _r = builder.try_init();
}
<commit_msg>Fix monotonicity of sequential key generator<commit_after>use std::{
    sync::{
        atomic::{AtomicBool, AtomicUsize, Ordering},
        Arc,
    },
    thread,
};

use rand::{thread_rng, Rng};

#[cfg_attr(
    // only enable jemalloc on linux and macos by default
    all(
        any(target_os = "linux", target_os = "macos"),
        feature = "jemalloc",
    ),
    global_allocator
)]
#[cfg(
    // only enable jemalloc on linux and macos by default
    all(
        any(target_os = "linux", target_os = "macos"),
        feature = "jemalloc",
    ),
)]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

static TOTAL: AtomicUsize = AtomicUsize::new(0);

const USAGE: &str = "
Usage: stress [--threads=<#>] [--burn-in] [--duration=<s>] \
    [--key-len=<l>] [--val-len=<l>] \
    [--get-prop=<p>] \
    [--set-prop=<p>] \
    [--del-prop=<p>] \
    [--cas-prop=<p>] \
    [--scan-prop=<p>] \
    [--merge-prop=<p>] \
    [--entries=<n>] \
    [--sequential] \
    [--total-ops=<n>] \
    [--flush-every=<ms>]

Options:
    --threads=<#>      Number of threads [default: 4].
    --burn-in          Don't halt until we receive a signal.
    --duration=<s>     Seconds to run for [default: 10].
    --key-len=<l>      The length of keys [default: 10].
    --val-len=<l>      The length of values [default: 100].
    --get-prop=<p>     The relative proportion of get requests [default: 94].
    --set-prop=<p>     The relative proportion of set requests [default: 2].
    --del-prop=<p>     The relative proportion of del requests [default: 1].
    --cas-prop=<p>     The relative proportion of cas requests [default: 1].
    --scan-prop=<p>    The relative proportion of scan requests [default: 1].
    --merge-prop=<p>   The relative proportion of merge requests [default: 1].
    --entries=<n>      The total keyspace [default: 100000].
    --sequential       Run the test in sequential mode instead of random.
    --total-ops=<n>    Stop test after executing a total number of operations.
    --flush-every=<m>  Flush and sync the database every ms [default: 200].
";

#[derive(Clone, Copy)]
struct Args {
    threads: usize,
    burn_in: bool,
    duration: u64,
    key_len: usize,
    val_len: usize,
    get_prop: usize,
    set_prop: usize,
    del_prop: usize,
    cas_prop: usize,
    scan_prop: usize,
    merge_prop: usize,
    entries: usize,
    sequential: bool,
    total_ops: Option<usize>,
    flush_every: u64,
}

impl Default for Args {
    fn default() -> Args {
        Args {
            threads: 4,
            burn_in: false,
            duration: 10,
            key_len: 10,
            val_len: 100,
            get_prop: 94,
            set_prop: 2,
            del_prop: 1,
            cas_prop: 1,
            scan_prop: 1,
            merge_prop: 1,
            entries: 100000,
            sequential: false,
            total_ops: None,
            flush_every: 200,
        }
    }
}

fn parse<'a, I, T>(mut iter: I) -> T
where
    I: Iterator<Item = &'a str>,
    T: std::str::FromStr,
    <T as std::str::FromStr>::Err: std::fmt::Debug,
{
    iter.next().expect(USAGE).parse().expect(USAGE)
}

impl Args {
    fn parse() -> Args {
        let mut args = Args::default();
        for raw_arg in std::env::args().skip(1) {
            let mut splits = raw_arg[2..].split('=');
            match splits.next().unwrap() {
                "threads" => args.threads = parse(&mut splits),
                "burn-in" => args.burn_in = true,
                "duration" => args.duration = parse(&mut splits),
                "key-len" => args.key_len = parse(&mut splits),
                "val-len" => args.val_len = parse(&mut splits),
                "get-prop" => args.get_prop = parse(&mut splits),
                "set-prop" => args.set_prop = parse(&mut splits),
                "del-prop" => args.del_prop = parse(&mut splits),
                "cas-prop" => args.cas_prop = parse(&mut splits),
                "scan-prop" => args.scan_prop = parse(&mut splits),
                "merge-prop" => args.merge_prop = parse(&mut splits),
                "entries" => args.entries = parse(&mut splits),
                "sequential" => args.sequential = true,
                "total-ops" => args.total_ops = Some(parse(&mut splits)),
                "flush-every" => args.flush_every = parse(&mut splits),
                other => panic!("unknown option: {}, {}", other, USAGE),
            }
        }
        args
    }
}

fn report(shutdown: Arc<AtomicBool>) {
    let mut last = 0;
    while !shutdown.load(Ordering::Relaxed) {
        thread::sleep(std::time::Duration::from_secs(1));
        let total = TOTAL.load(Ordering::Acquire);

        println!("did {} ops, {}mb RSS", total - last, rss() / (1024 * 1024));

        last = total;
    }
}

fn concatenate_merge(
    _key: &[u8],              // the key being merged
    old_value: Option<&[u8]>, // the previous value, if one existed
    merged_bytes: &[u8],      // the new bytes being merged in
) -> Option<Vec<u8>> {
    // set the new value, return None to delete
    let mut ret = old_value.map(|ov| ov.to_vec()).unwrap_or_else(Vec::new);

    ret.extend_from_slice(merged_bytes);

    Some(ret)
}

fn run(args: Args, tree: Arc<sled::Db>, shutdown: Arc<AtomicBool>) {
    let get_max = args.get_prop;
    let set_max = get_max + args.set_prop;
    let del_max = set_max + args.del_prop;
    let cas_max = del_max + args.cas_prop;
    let merge_max = cas_max + args.merge_prop;
    let scan_max = merge_max + args.scan_prop;

    let keygen = |len| -> sled::IVec {
        static SEQ: AtomicUsize = AtomicUsize::new(0);
        let i = if args.sequential {
            SEQ.fetch_add(1, Ordering::Relaxed)
        } else {
            thread_rng().gen::<usize>()
        } % args.entries;

        let start = if len < 8 { 8 - len } else { 0 };

        let i_keygen = &i.to_be_bytes()[start..];

        i_keygen.iter().cycle().take(len).copied().collect()
    };

    let valgen = |len| -> sled::IVec {
        if len == 0 {
            return vec![].into();
        }

        let i: usize = thread_rng().gen::<usize>() % (len * 8);

        let i_keygen = i.to_be_bytes();

        i_keygen
            .iter()
            .skip_while(|v| **v == 0)
            .cycle()
            .take(len)
            .copied()
            .collect()
    };

    let mut rng = thread_rng();

    while !shutdown.load(Ordering::Relaxed) {
        let op = TOTAL.fetch_add(1, Ordering::Release);
        let key = keygen(args.key_len);
        let choice = rng.gen_range(0, scan_max + 1);

        match choice {
            v if v <= get_max => {
                tree.get(&key).unwrap();
            }
            v if v > get_max && v <= set_max => {
                let value = valgen(args.val_len);
                tree.insert(&key, value).unwrap();
            }
            v if v > set_max && v <= del_max => {
                tree.remove(&key).unwrap();
            }
            v if v > del_max && v <= cas_max => {
                let old = if rng.gen::<bool>() {
                    let value = valgen(args.val_len);
                    Some(value)
                } else {
                    None
                };

                let new = if rng.gen::<bool>() {
                    let value = valgen(args.val_len);
                    Some(value)
                } else {
                    None
                };

                if let Err(e) = tree.compare_and_swap(&key, old, new) {
                    panic!("operational error: {:?}", e);
                }
            }
            v if v > cas_max && v <= merge_max => {
                let value = valgen(args.val_len);
                tree.merge(&key, value).unwrap();
            }
            _ => {
                let iter = tree.range(key..).map(|res| res.unwrap());

                if op % 2 == 0 {
                    let _ = iter.take(rng.gen_range(0, 15)).collect::<Vec<_>>();
                } else {
                    let _ = iter
                        .rev()
                        .take(rng.gen_range(0, 15))
                        .collect::<Vec<_>>();
                }
            }
        }
    }
}

fn rss() -> usize {
    #[cfg(target_os = "linux")]
    {
        use std::io::prelude::*;
        use std::io::BufReader;

        let mut buf = String::new();
        let mut f =
            BufReader::new(std::fs::File::open("/proc/self/statm").unwrap());
        f.read_line(&mut buf).unwrap();
        let mut parts = buf.split_whitespace();
        let rss_pages = parts.nth(1).unwrap().parse::<usize>().unwrap();
        rss_pages * 4096
    }
    #[cfg(not(target_os = "linux"))]
    {
        0
    }
}

fn main() {
    #[cfg(feature = "logging")]
    setup_logger();

    let args = Args::parse();

    let shutdown = Arc::new(AtomicBool::new(false));

    let config = sled::Config::new()
        .cache_capacity(256 * 1024 * 1024)
        .flush_every_ms(Some(args.flush_every))
        .print_profile_on_drop(true);

    let tree = Arc::new(config.open().unwrap());
    tree.set_merge_operator(concatenate_merge);

    let mut threads = vec![];

    let now = std::time::Instant::now();

    let n_threads = args.threads;

    for i in 0..=n_threads {
        let tree = tree.clone();
        let shutdown = shutdown.clone();

        let t = if i == 0 {
            thread::Builder::new()
                .name("reporter".into())
                .spawn(move || report(shutdown))
                .unwrap()
        } else {
            thread::spawn(move || run(args, tree, shutdown))
        };

        threads.push(t);
    }

    if let Some(ops) = args.total_ops {
        assert!(!args.burn_in, "don't set both --burn-in and --total-ops");
        while TOTAL.load(Ordering::Relaxed) < ops {
            thread::sleep(std::time::Duration::from_millis(50));
        }
        shutdown.store(true, Ordering::SeqCst);
    } else if !args.burn_in {
        thread::sleep(std::time::Duration::from_secs(args.duration));
        shutdown.store(true, Ordering::SeqCst);
    }

    for t in threads.into_iter() {
        t.join().unwrap();
    }

    let ops = TOTAL.load(Ordering::SeqCst);
    let time = now.elapsed().as_secs() as usize;

    println!(
        "did {} total ops in {} seconds. {} ops/s",
        ops,
        time,
        (ops * 1_000) / (time * 1_000)
    );
}

#[cfg(feature = "logging")]
pub fn setup_logger() {
    use std::io::Write;

    color_backtrace::install();

    fn tn() -> String {
        std::thread::current().name().unwrap_or("unknown").to_owned()
    }

    let mut builder = env_logger::Builder::new();
    builder
        .format(|buf, record| {
            writeln!(
                buf,
                "{:05} {:25} {:10} {}",
                record.level(),
                tn(),
                record.module_path().unwrap().split("::").last().unwrap(),
                record.args()
            )
        })
        .filter(None, log::LevelFilter::Info);

    if let Ok(env) = std::env::var("RUST_LOG") {
        builder.parse_filters(&env);
    }

    let _r = builder.try_init();
}
<|endoftext|>