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

use crate::{json_rpc::JsonRpcClientWrapper, TransactionContext};
use libra_crypto::{ed25519::Ed25519PublicKey, x25519};
use libra_global_constants::{
    CONSENSUS_KEY, FULLNODE_NETWORK_KEY, OWNER_ACCOUNT, VALIDATOR_NETWORK_KEY,
};
use libra_management::{
    error::Error,
    storage::{to_x25519, StorageWrapper},
};
use libra_network_address::{
    encrypted::{EncNetworkAddress, Key, RawEncNetworkAddress, TEST_SHARED_VAL_NETADDR_KEY},
    NetworkAddress, Protocol, RawNetworkAddress,
};
use libra_types::account_address::AccountAddress;
use std::convert::TryFrom;
use structopt::StructOpt;

// TODO: Load all chain IDs from the host
#[derive(Debug, StructOpt)]
pub struct SetValidatorConfig {
    #[structopt(long, help = "JSON-RPC Endpoint (e.g. http://localhost:8080")]
    host: String,
    #[structopt(flatten)]
    validator_config: libra_management::validator_config::ValidatorConfig,
    #[structopt(long, help = "Validator Network Address")]
    validator_address: Option<NetworkAddress>,
    #[structopt(long, help = "Full Node Network Address")]
    fullnode_address: Option<NetworkAddress>,
}

impl SetValidatorConfig {
    pub fn execute(self) -> Result<TransactionContext, Error> {
        let storage = StorageWrapper::new(
            self.validator_config.validator_backend.name(),
            &self.validator_config.validator_backend.validator_backend,
        )?;
        let owner_account = storage.account_address(OWNER_ACCOUNT)?;

        let client = JsonRpcClientWrapper::new(self.host);
        let sequence_number = client.sequence_number(owner_account)?;

        // Retrieve the current validator / fullnode addresses and update accordingly
        let validator_config = client.validator_config(owner_account)?;
        let validator_address = if let Some(validator_address) = self.validator_address {
            validator_address
        } else {
            decode_validator_address(
                validator_config.validator_network_address,
                &owner_account,
                &TEST_SHARED_VAL_NETADDR_KEY,
                0, // addr_idx
            )?
        };

        let fullnode_address = if let Some(fullnode_address) = self.fullnode_address {
            fullnode_address
        } else {
            decode_address(validator_config.full_node_network_address)?
        };

        let txn = self.validator_config.build_transaction(
            sequence_number,
            fullnode_address,
            validator_address,
            true,
        )?;
        client.submit_transaction(txn.as_signed_user_txn().unwrap().clone())
    }
}

#[derive(Debug, StructOpt)]
pub struct RotateKey {
    #[structopt(long, help = "JSON-RPC Endpoint (e.g. http://localhost:8080)")]
    pub(crate) host: String,
    #[structopt(flatten)]
    pub(crate) validator_config: libra_management::validator_config::ValidatorConfig,
}

impl RotateKey {
    pub fn execute(
        self,
        key_name: &'static str,
    ) -> Result<(TransactionContext, Ed25519PublicKey), Error> {
        let mut storage = StorageWrapper::new(
            self.validator_config.validator_backend.name(),
            &self.validator_config.validator_backend.validator_backend,
        )?;
        let key = storage.rotate_key(key_name)?;

        let set_validator_config = SetValidatorConfig {
            host: self.host,
            validator_config: self.validator_config,
            validator_address: None,
            fullnode_address: None,
        };

        set_validator_config.execute().map(|txn_ctx| (txn_ctx, key))
    }
}

#[derive(Debug, StructOpt)]
pub struct RotateConsensusKey {
    #[structopt(flatten)]
    rotate_key: RotateKey,
}

impl RotateConsensusKey {
    pub fn execute(self) -> Result<(TransactionContext, Ed25519PublicKey), Error> {
        self.rotate_key.execute(CONSENSUS_KEY)
    }
}

#[derive(Debug, StructOpt)]
pub struct RotateValidatorNetworkKey {
    #[structopt(flatten)]
    rotate_key: RotateKey,
}

impl RotateValidatorNetworkKey {
    pub fn execute(self) -> Result<(TransactionContext, x25519::PublicKey), Error> {
        let (txn_ctx, key) = self.rotate_key.execute(VALIDATOR_NETWORK_KEY)?;
        Ok((txn_ctx, to_x25519(key)?))
    }
}

#[derive(Debug, StructOpt)]
pub struct RotateFullNodeNetworkKey {
    #[structopt(flatten)]
    rotate_key: RotateKey,
}

impl RotateFullNodeNetworkKey {
    pub fn execute(self) -> Result<(TransactionContext, x25519::PublicKey), Error> {
        let (txn_ctx, key) = self.rotate_key.execute(FULLNODE_NETWORK_KEY)?;
        Ok((txn_ctx, to_x25519(key)?))
    }
}

fn decode_validator_address(
    address: RawEncNetworkAddress,
    account: &AccountAddress,
    key: &Key,
    addr_idx: u32,
) -> Result<NetworkAddress, Error> {
    let enc_addr = EncNetworkAddress::try_from(&address).map_err(|e| {
        Error::UnexpectedError(format!(
            "Failed to decode network address {}",
            e.to_string()
        ))
    })?;
    let raw_addr = enc_addr.decrypt(key, account, addr_idx).map_err(|e| {
        Error::UnexpectedError(format!(
            "Failed to decrypt network address {}",
            e.to_string()
        ))
    })?;
    decode_address(raw_addr)
}

fn decode_address(raw_address: RawNetworkAddress) -> Result<NetworkAddress, Error> {
    let network_address = NetworkAddress::try_from(&raw_address).map_err(|e| {
        Error::UnexpectedError(format!(
            "Failed to decode network address {}",
            e.to_string()
        ))
    })?;
    let protocols = network_address
        .as_slice()
        .iter()
        .filter(|protocol| match protocol {
            Protocol::Dns(_)
            | Protocol::Dns4(_)
            | Protocol::Dns6(_)
            | Protocol::Ip4(_)
            | Protocol::Ip6(_)
            | Protocol::Memory(_)
            | Protocol::Tcp(_) => true,
            _ => false,
        })
        .cloned()
        .collect::<Vec<_>>();
    Ok(NetworkAddress::try_from(protocols).unwrap())
}

#[derive(Debug, StructOpt)]
pub struct ValidatorConfig {
    #[structopt(long, help = "JSON-RPC Endpoint (e.g. http://localhost:8080")]
    json_server: String,
    #[structopt(long, help = "Validator account address to display the config")]
    account_address: AccountAddress,
}

impl ValidatorConfig {
    pub fn execute(self) -> Result<libra_types::validator_config::ValidatorConfig, Error> {
        let client = JsonRpcClientWrapper::new(self.json_server);
        client.validator_config(self.account_address)
    }
}
<commit_msg>[Operational Management] Fix bug in set_validator: fetch seq number from operator account.<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{json_rpc::JsonRpcClientWrapper, TransactionContext};
use libra_crypto::{ed25519::Ed25519PublicKey, x25519};
use libra_global_constants::{
    CONSENSUS_KEY, FULLNODE_NETWORK_KEY, OPERATOR_ACCOUNT, OWNER_ACCOUNT, VALIDATOR_NETWORK_KEY,
};
use libra_management::{
    error::Error,
    storage::{to_x25519, StorageWrapper},
};
use libra_network_address::{
    encrypted::{EncNetworkAddress, Key, RawEncNetworkAddress, TEST_SHARED_VAL_NETADDR_KEY},
    NetworkAddress, Protocol, RawNetworkAddress,
};
use libra_types::account_address::AccountAddress;
use std::convert::TryFrom;
use structopt::StructOpt;

// TODO: Load all chain IDs from the host
#[derive(Debug, StructOpt)]
pub struct SetValidatorConfig {
    #[structopt(long, help = "JSON-RPC Endpoint (e.g. http://localhost:8080")]
    host: String,
    #[structopt(flatten)]
    validator_config: libra_management::validator_config::ValidatorConfig,
    #[structopt(long, help = "Validator Network Address")]
    validator_address: Option<NetworkAddress>,
    #[structopt(long, help = "Full Node Network Address")]
    fullnode_address: Option<NetworkAddress>,
}

impl SetValidatorConfig {
    pub fn execute(self) -> Result<TransactionContext, Error> {
        let storage = StorageWrapper::new(
            self.validator_config.validator_backend.name(),
            &self.validator_config.validator_backend.validator_backend,
        )?;
        let operator_account = storage.account_address(OPERATOR_ACCOUNT)?;
        let owner_account = storage.account_address(OWNER_ACCOUNT)?;

        let client = JsonRpcClientWrapper::new(self.host);
        let sequence_number = client.sequence_number(operator_account)?;

        // Retrieve the current validator / fullnode addresses and update accordingly
        let validator_config = client.validator_config(owner_account)?;
        let validator_address = if let Some(validator_address) = self.validator_address {
            validator_address
        } else {
            decode_validator_address(
                validator_config.validator_network_address,
                &owner_account,
                &TEST_SHARED_VAL_NETADDR_KEY,
                0, // addr_idx
            )?
        };

        let fullnode_address = if let Some(fullnode_address) = self.fullnode_address {
            fullnode_address
        } else {
            decode_address(validator_config.full_node_network_address)?
        };

        let txn = self.validator_config.build_transaction(
            sequence_number,
            fullnode_address,
            validator_address,
            true,
        )?;
        client.submit_transaction(txn.as_signed_user_txn().unwrap().clone())
    }
}

#[derive(Debug, StructOpt)]
pub struct RotateKey {
    #[structopt(long, help = "JSON-RPC Endpoint (e.g. http://localhost:8080)")]
    pub(crate) host: String,
    #[structopt(flatten)]
    pub(crate) validator_config: libra_management::validator_config::ValidatorConfig,
}

impl RotateKey {
    pub fn execute(
        self,
        key_name: &'static str,
    ) -> Result<(TransactionContext, Ed25519PublicKey), Error> {
        let mut storage = StorageWrapper::new(
            self.validator_config.validator_backend.name(),
            &self.validator_config.validator_backend.validator_backend,
        )?;
        let key = storage.rotate_key(key_name)?;

        let set_validator_config = SetValidatorConfig {
            host: self.host,
            validator_config: self.validator_config,
            validator_address: None,
            fullnode_address: None,
        };

        set_validator_config.execute().map(|txn_ctx| (txn_ctx, key))
    }
}

#[derive(Debug, StructOpt)]
pub struct RotateConsensusKey {
    #[structopt(flatten)]
    rotate_key: RotateKey,
}

impl RotateConsensusKey {
    pub fn execute(self) -> Result<(TransactionContext, Ed25519PublicKey), Error> {
        self.rotate_key.execute(CONSENSUS_KEY)
    }
}

#[derive(Debug, StructOpt)]
pub struct RotateValidatorNetworkKey {
    #[structopt(flatten)]
    rotate_key: RotateKey,
}

impl RotateValidatorNetworkKey {
    pub fn execute(self) -> Result<(TransactionContext, x25519::PublicKey), Error> {
        let (txn_ctx, key) = self.rotate_key.execute(VALIDATOR_NETWORK_KEY)?;
        Ok((txn_ctx, to_x25519(key)?))
    }
}

#[derive(Debug, StructOpt)]
pub struct RotateFullNodeNetworkKey {
    #[structopt(flatten)]
    rotate_key: RotateKey,
}

impl RotateFullNodeNetworkKey {
    pub fn execute(self) -> Result<(TransactionContext, x25519::PublicKey), Error> {
        let (txn_ctx, key) = self.rotate_key.execute(FULLNODE_NETWORK_KEY)?;
        Ok((txn_ctx, to_x25519(key)?))
    }
}

fn decode_validator_address(
    address: RawEncNetworkAddress,
    account: &AccountAddress,
    key: &Key,
    addr_idx: u32,
) -> Result<NetworkAddress, Error> {
    let enc_addr = EncNetworkAddress::try_from(&address).map_err(|e| {
        Error::UnexpectedError(format!(
            "Failed to decode network address {}",
            e.to_string()
        ))
    })?;
    let raw_addr = enc_addr.decrypt(key, account, addr_idx).map_err(|e| {
        Error::UnexpectedError(format!(
            "Failed to decrypt network address {}",
            e.to_string()
        ))
    })?;
    decode_address(raw_addr)
}

fn decode_address(raw_address: RawNetworkAddress) -> Result<NetworkAddress, Error> {
    let network_address = NetworkAddress::try_from(&raw_address).map_err(|e| {
        Error::UnexpectedError(format!(
            "Failed to decode network address {}",
            e.to_string()
        ))
    })?;
    let protocols = network_address
        .as_slice()
        .iter()
        .filter(|protocol| match protocol {
            Protocol::Dns(_)
            | Protocol::Dns4(_)
            | Protocol::Dns6(_)
            | Protocol::Ip4(_)
            | Protocol::Ip6(_)
            | Protocol::Memory(_)
            | Protocol::Tcp(_) => true,
            _ => false,
        })
        .cloned()
        .collect::<Vec<_>>();
    Ok(NetworkAddress::try_from(protocols).unwrap())
}

#[derive(Debug, StructOpt)]
pub struct ValidatorConfig {
    #[structopt(long, help = "JSON-RPC Endpoint (e.g. http://localhost:8080")]
    json_server: String,
    #[structopt(long, help = "Validator account address to display the config")]
    account_address: AccountAddress,
}

impl ValidatorConfig {
    pub fn execute(self) -> Result<libra_types::validator_config::ValidatorConfig, Error> {
        let client = JsonRpcClientWrapper::new(self.json_server);
        client.validator_config(self.account_address)
    }
}
<|endoftext|>


<commit_before>use if_chain::if_chain;
use rustc_errors::Applicability;
use rustc_hir::{Expr, ExprKind, QPath};
use rustc_lint::{LateContext, LateLintPass};
use rustc_middle::ty;
use rustc_session::{declare_lint_pass, declare_tool_lint};

use crate::utils::{any_parent_is_automatically_derived, match_def_path, paths, span_lint_and_sugg};

declare_clippy_lint! {
    /// **What it does:** Checks for literal calls to `Default::default()`.
    ///
    /// **Why is this bad?** It's more clear to the reader to use the name of the type whose default is
    /// being gotten than the generic `Default`.
    ///
    /// **Known problems:** None.
    ///
    /// **Example:**
    /// ```rust
    /// // Bad
    /// let s: String = Default::default();
    ///
    /// // Good
    /// let s = String::default();
    /// ```
    pub DEFAULT_TRAIT_ACCESS,
    pedantic,
    "checks for literal calls to `Default::default()`"
}

declare_lint_pass!(DefaultTraitAccess => [DEFAULT_TRAIT_ACCESS]);

impl<'tcx> LateLintPass<'tcx> for DefaultTraitAccess {
    fn check_expr(&mut self, cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) {
        if_chain! {
            if let ExprKind::Call(ref path, ..) = expr.kind;
            if !any_parent_is_automatically_derived(cx.tcx, expr.hir_id);
            if let ExprKind::Path(ref qpath) = path.kind;
            if let Some(def_id) = cx.qpath_res(qpath, path.hir_id).opt_def_id();
            if match_def_path(cx, def_id, &paths::DEFAULT_TRAIT_METHOD);
            // Detect and ignore <Foo as Default>::default() because these calls do explicitly name the type.
            if let QPath::Resolved(None, _path) = qpath;
            then {
                let expr_ty = cx.typeck_results().expr_ty(expr);
                if let ty::Adt(def, ..) = expr_ty.kind {
                    // TODO: Work out a way to put "whatever the imported way of referencing
                    // this type in this file" rather than a fully-qualified type.
                    let replacement = format!("{}::default()", cx.tcx.def_path_str(def.did));
                    span_lint_and_sugg(
                        cx,
                        DEFAULT_TRAIT_ACCESS,
                        expr.span,
                        &format!("calling `{}` is more clear than this expression", replacement),
                        "try",
                        replacement,
                        Applicability::Unspecified, // First resolve the TODO above
                    );
                }
            }
        }
    }
}
<commit_msg>Fix fallout from rustup<commit_after>use if_chain::if_chain;
use rustc_errors::Applicability;
use rustc_hir::{Expr, ExprKind, QPath};
use rustc_lint::{LateContext, LateLintPass};
use rustc_middle::ty;
use rustc_session::{declare_lint_pass, declare_tool_lint};

use crate::utils::{any_parent_is_automatically_derived, match_def_path, paths, span_lint_and_sugg};

declare_clippy_lint! {
    /// **What it does:** Checks for literal calls to `Default::default()`.
    ///
    /// **Why is this bad?** It's more clear to the reader to use the name of the type whose default is
    /// being gotten than the generic `Default`.
    ///
    /// **Known problems:** None.
    ///
    /// **Example:**
    /// ```rust
    /// // Bad
    /// let s: String = Default::default();
    ///
    /// // Good
    /// let s = String::default();
    /// ```
    pub DEFAULT_TRAIT_ACCESS,
    pedantic,
    "checks for literal calls to `Default::default()`"
}

declare_lint_pass!(DefaultTraitAccess => [DEFAULT_TRAIT_ACCESS]);

impl<'tcx> LateLintPass<'tcx> for DefaultTraitAccess {
    fn check_expr(&mut self, cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) {
        if_chain! {
            if let ExprKind::Call(ref path, ..) = expr.kind;
            if !any_parent_is_automatically_derived(cx.tcx, expr.hir_id);
            if let ExprKind::Path(ref qpath) = path.kind;
            if let Some(def_id) = cx.qpath_res(qpath, path.hir_id).opt_def_id();
            if match_def_path(cx, def_id, &paths::DEFAULT_TRAIT_METHOD);
            // Detect and ignore <Foo as Default>::default() because these calls do explicitly name the type.
            if let QPath::Resolved(None, _path) = qpath;
            then {
                let expr_ty = cx.typeck_results().expr_ty(expr);
                if let ty::Adt(def, ..) = expr_ty.kind() {
                    // TODO: Work out a way to put "whatever the imported way of referencing
                    // this type in this file" rather than a fully-qualified type.
                    let replacement = format!("{}::default()", cx.tcx.def_path_str(def.did));
                    span_lint_and_sugg(
                        cx,
                        DEFAULT_TRAIT_ACCESS,
                        expr.span,
                        &format!("calling `{}` is more clear than this expression", replacement),
                        "try",
                        replacement,
                        Applicability::Unspecified, // First resolve the TODO above
                    );
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use if_chain::if_chain;
use rustc_errors::Applicability;
use rustc_hir::{Expr, ExprKind, QPath};
use rustc_lint::{LateContext, LateLintPass};
use rustc_middle::ty;
use rustc_session::{declare_lint_pass, declare_tool_lint};

use crate::utils::{any_parent_is_automatically_derived, match_def_path, paths, span_lint_and_sugg};

declare_clippy_lint! {
    /// **What it does:** Checks for literal calls to `Default::default()`.
    ///
    /// **Why is this bad?** It's more clear to the reader to use the name of the type whose default is
    /// being gotten than the generic `Default`.
    ///
    /// **Known problems:** None.
    ///
    /// **Example:**
    /// ```rust
    /// // Bad
    /// let s: String = Default::default();
    ///
    /// // Good
    /// let s = String::default();
    /// ```
    pub DEFAULT_TRAIT_ACCESS,
    pedantic,
    "checks for literal calls to `Default::default()`"
}

declare_lint_pass!(DefaultTraitAccess => [DEFAULT_TRAIT_ACCESS]);

impl<'tcx> LateLintPass<'tcx> for DefaultTraitAccess {
    fn check_expr(&mut self, cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) {
        if_chain! {
            if let ExprKind::Call(ref path, ..) = expr.kind;
            if !any_parent_is_automatically_derived(cx.tcx, expr.hir_id);
            if let ExprKind::Path(ref qpath) = path.kind;
            if let Some(def_id) = cx.qpath_res(qpath, path.hir_id).opt_def_id();
            if match_def_path(cx, def_id, &paths::DEFAULT_TRAIT_METHOD);
            then {
                match qpath {
                    QPath::Resolved(..) => {
                        if_chain! {
                            // Detect and ignore <Foo as Default>::default() because these calls do
                            // explicitly name the type.
                            if let ExprKind::Call(ref method, ref _args) = expr.kind;
                            if let ExprKind::Path(ref p) = method.kind;
                            if let QPath::Resolved(Some(_ty), _path) = p;
                            then {
                                return;
                            }
                        }

                        // TODO: Work out a way to put "whatever the imported way of referencing
                        // this type in this file" rather than a fully-qualified type.
                        let expr_ty = cx.typeck_results().expr_ty(expr);
                        if let ty::Adt(def, ..) = expr_ty.kind {
                            let replacement = format!("{}::default()", cx.tcx.def_path_str(def.did));
                            span_lint_and_sugg(
                                cx,
                                DEFAULT_TRAIT_ACCESS,
                                expr.span,
                                &format!("calling `{}` is more clear than this expression", replacement),
                                "try",
                                replacement,
                                Applicability::Unspecified, // First resolve the TODO above
                            );
                         }
                    },
                    QPath::TypeRelative(..) | QPath::LangItem(..) => {},
                }
            }
        }
    }
}
<commit_msg>Simplify `clippy::default_trait_access`<commit_after>use if_chain::if_chain;
use rustc_errors::Applicability;
use rustc_hir::{Expr, ExprKind, QPath};
use rustc_lint::{LateContext, LateLintPass};
use rustc_middle::ty;
use rustc_session::{declare_lint_pass, declare_tool_lint};

use crate::utils::{any_parent_is_automatically_derived, match_def_path, paths, span_lint_and_sugg};

declare_clippy_lint! {
    /// **What it does:** Checks for literal calls to `Default::default()`.
    ///
    /// **Why is this bad?** It's more clear to the reader to use the name of the type whose default is
    /// being gotten than the generic `Default`.
    ///
    /// **Known problems:** None.
    ///
    /// **Example:**
    /// ```rust
    /// // Bad
    /// let s: String = Default::default();
    ///
    /// // Good
    /// let s = String::default();
    /// ```
    pub DEFAULT_TRAIT_ACCESS,
    pedantic,
    "checks for literal calls to `Default::default()`"
}

declare_lint_pass!(DefaultTraitAccess => [DEFAULT_TRAIT_ACCESS]);

impl<'tcx> LateLintPass<'tcx> for DefaultTraitAccess {
    fn check_expr(&mut self, cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) {
        if_chain! {
            if let ExprKind::Call(ref path, ..) = expr.kind;
            if !any_parent_is_automatically_derived(cx.tcx, expr.hir_id);
            if let ExprKind::Path(ref qpath) = path.kind;
            if let Some(def_id) = cx.qpath_res(qpath, path.hir_id).opt_def_id();
            if match_def_path(cx, def_id, &paths::DEFAULT_TRAIT_METHOD);
            // Detect and ignore <Foo as Default>::default() because these calls do explicitly name the type.
            if let QPath::Resolved(None, _path) = qpath;
            then {
                let expr_ty = cx.typeck_results().expr_ty(expr);
                if let ty::Adt(def, ..) = expr_ty.kind {
                    // TODO: Work out a way to put "whatever the imported way of referencing
                    // this type in this file" rather than a fully-qualified type.
                    let replacement = format!("{}::default()", cx.tcx.def_path_str(def.did));
                    span_lint_and_sugg(
                        cx,
                        DEFAULT_TRAIT_ACCESS,
                        expr.span,
                        &format!("calling `{}` is more clear than this expression", replacement),
                        "try",
                        replacement,
                        Applicability::Unspecified, // First resolve the TODO above
                    );
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use story_builder::story_builder::*;
use story_builder::article_provider::*;
use std::collections::HashMap;
use std::rc::Rc;

static EXPECTED_SUGGESTION: &'static str = "Cannot find wikipedia article for <not-found>, try one of the following suggestions:\r\n\
- Suggestion 1\r\n\
- Suggestion 2\r\n\
- Suggestion 3\r\n";

#[test]
/// For: build_suggestions_msg
fn build_suggestions_msg_is_working() {
    struct TestProvider {}
    impl ArticleProvider for TestProvider {
        #[allow(unused_variables)]
        fn get(&self, topic: &str) -> Option<Box<Article>> {
            panic!("get() should never be called in this test.");
        }
        #[allow(unused_variables)]
        fn search(&self, topic: &str) -> Vec<String> {
            vec!["Suggestion 1".to_owned(),
            "Suggestion 2".to_owned(),
            "Suggestion 3".to_owned()]
        }
    }
    let provider = TestProvider {};
    let story_builder = StoryBuilder::new(&provider);

    assert_eq!(story_builder.build_suggestions_msg("not-found"), EXPECTED_SUGGESTION);
}

#[test]
/// For: build_story
fn build_story_cannot_find_first_article_suggest() {
    struct TestArticle {
        topics: Vec<Paragraph>
    }
    impl Article for TestArticle {
        fn get_paragraphs(&self) -> &Vec<Paragraph> {&self.topics}
    }
    struct TestProvider {}
    impl ArticleProvider for TestProvider {
        fn get(&self, topic: &str) -> Option<Box<Article>> {
            if topic == "found" {
                Some(Box::new(TestArticle {topics: vec![]}))
            }
            else {
                None
            }
        }
        #[allow(unused_variables)]
        fn search(&self, topic: &str) -> Vec<String> {
            vec!["Suggestion 1".to_owned(),
            "Suggestion 2".to_owned(),
            "Suggestion 3".to_owned()]
        }
    }

    let provider = TestProvider {};
    let story_builder = StoryBuilder::new(&provider);
    assert_eq!(story_builder.build_story("not-found", "found"), Err(EXPECTED_SUGGESTION.to_owned()));
}

#[test]
/// For: build_story
fn build_story_cannot_find_second_article_suggest() {
    struct TestArticle {
        topics: Vec<Paragraph>
    }
    impl Article for TestArticle {
        fn get_paragraphs(&self) -> &Vec<Paragraph> {&self.topics}
    }
    struct TestProvider {}
    impl ArticleProvider for TestProvider {
        fn get(&self, topic: &str) -> Option<Box<Article>> {
            if topic == "found"
                {Some(Box::new(TestArticle {topics: vec![]}))}
            else
                {None}
        }
        #[allow(unused_variables)]
        fn search(&self, topic: &str) -> Vec<String> {
            vec!["Suggestion 1".to_owned(),
            "Suggestion 2".to_owned(),
            "Suggestion 3".to_owned()]
        }
    }

    let provider = TestProvider {};
    let story_builder = StoryBuilder::new(&provider);
    assert_eq!(story_builder.build_story("found", "not-found"), Err(EXPECTED_SUGGESTION.to_owned()));
}

#[test]
/// For: build_story
fn build_story_from_same_start_and_end_should_err() {
    struct TestProvider {}
    impl ArticleProvider for TestProvider {
        #[allow(unused_variables)]
        fn get(&self, topic: &str) -> Option<Box<Article>> {panic!("get() should not get called in this test.")}
        #[allow(unused_variables)]
        fn search(&self, topic: &str) -> Vec<String> {panic!("get() should not get called in this test.")}
    }
    let provider = TestProvider {};
    let story_builder = StoryBuilder::new(&provider);
    assert_eq!(story_builder.build_story("similar topic", "similar topic"),
               Err("No story to build; same start and end topics.".to_owned()));
}

#[test]
/// For: build_story
fn build_story_empty_start_topic_should_err() {
    struct TestProvider {}
    impl ArticleProvider for TestProvider {
        #[allow(unused_variables)]
        fn get(&self, topic: &str) -> Option<Box<Article>> {panic!("get() should not get called in this test.")}
        #[allow(unused_variables)]
        fn search(&self, topic: &str) -> Vec<String> {panic!("get() should not get called in this test.")}
    }
    let provider = TestProvider {};
    let story_builder = StoryBuilder::new(&provider);
    assert_eq!(story_builder.build_story("", "Other topic"),
               Err("Missing start topic.".to_owned()));
}

#[test]
/// For: build_story
fn build_story_empty_end_topic_should_err() {
    struct TestProvider {}
    impl ArticleProvider for TestProvider {
        #[allow(unused_variables)]
        fn get(&self, topic: &str) -> Option<Box<Article>> {panic!("get() should not get called in this test.")}
        #[allow(unused_variables)]
        fn search(&self, topic: &str) -> Vec<String> {panic!("get() should not get called in this test.")}
    }
    let provider = TestProvider {};
    let story_builder = StoryBuilder::new(&provider);
    assert_eq!(story_builder.build_story("First topic", ""),
               Err("Missing end topic.".to_owned()));
}

#[test]
/// For: build_story
fn build_story_end_topic_found_in_start_article() {
    let mut prebuilt_rels = Rc::new(HashMap::new());
    Rc::get_mut(&mut prebuilt_rels).unwrap().insert("start", vec![Paragraph {text: "Paragraph 1".to_owned(), topics: vec!["topic 1".to_owned(), "topic 2".to_owned(), "topic 3".to_owned()]},
                                                                  Paragraph {text: "Paragraph 2".to_owned(), topics: vec!["topic 3".to_owned(), "end".to_owned(), "topic 5".to_owned()]},
                                                                  Paragraph {text: "Paragraph 3".to_owned(), topics: vec!["topic 3".to_owned(), "topic 1".to_owned(), "topic 5".to_owned()]}]);
    struct TestArticle {
        topic: String,
        prebuilt_rels: Rc<HashMap<&'static str, Vec<Paragraph>>>,
    }

    impl TestArticle {
        fn new(topic: String, prebuilt_rels: Rc<HashMap<&'static str, Vec<Paragraph>>>) -> TestArticle {
            TestArticle {topic, prebuilt_rels}
        }
    }

    impl Article for TestArticle {
            fn get_paragraphs(&self) -> &Vec<Paragraph> {
            self.prebuilt_rels.get::<str>(&self.topic).expect("Tried to access a node that doesn't exist!")
        }
    }
    struct TestProvider {
        prebuilt_rels: Rc<HashMap<&'static str, Vec<Paragraph>>>
    }
    impl TestProvider {
        fn new(prebuilt_rels: Rc<HashMap<&'static str, Vec<Paragraph>>>) -> TestProvider {
            TestProvider {prebuilt_rels}
        }
    }
    impl ArticleProvider for TestProvider {
        fn get(&self, topic: &str) -> Option<Box<Article>> {
                let new_rels: Rc<HashMap<&'static str, Vec<Paragraph>>> = self.prebuilt_rels.clone();
                Some(Box::new(TestArticle::new(topic.to_owned(), new_rels)))
        }
        fn search(&self, topic: &str) -> Vec<String> {panic!("search() should not be called in this test.");}
    }
    let provider = TestProvider::new(prebuilt_rels.clone());
    let story_builder = StoryBuilder::new(&provider);
    assert_eq!(story_builder.build_story("start", "end"), Ok("Paragraph 2\r\n".to_owned()));
}
<commit_msg>Added test for 2 level deep solution.<commit_after>use story_builder::story_builder::*;
use story_builder::article_provider::*;
use std::collections::HashMap;
use std::rc::Rc;

static EXPECTED_SUGGESTION: &'static str = "Cannot find wikipedia article for <not-found>, try one of the following suggestions:\r\n\
- Suggestion 1\r\n\
- Suggestion 2\r\n\
- Suggestion 3\r\n";

#[test]
/// For: build_suggestions_msg
fn build_suggestions_msg_is_working() {
    struct TestProvider {}
    impl ArticleProvider for TestProvider {
        #[allow(unused_variables)]
        fn get(&self, topic: &str) -> Option<Box<Article>> {
            panic!("get() should never be called in this test.");
        }
        #[allow(unused_variables)]
        fn search(&self, topic: &str) -> Vec<String> {
            vec!["Suggestion 1".to_owned(),
            "Suggestion 2".to_owned(),
            "Suggestion 3".to_owned()]
        }
    }
    let provider = TestProvider {};
    let story_builder = StoryBuilder::new(&provider);

    assert_eq!(story_builder.build_suggestions_msg("not-found"), EXPECTED_SUGGESTION);
}

#[test]
/// For: build_story
fn build_story_cannot_find_first_article_suggest() {
    struct TestArticle {
        topics: Vec<Paragraph>
    }
    impl Article for TestArticle {
        fn get_paragraphs(&self) -> &Vec<Paragraph> {&self.topics}
    }
    struct TestProvider {}
    impl ArticleProvider for TestProvider {
        fn get(&self, topic: &str) -> Option<Box<Article>> {
            if topic == "found" {
                Some(Box::new(TestArticle {topics: vec![]}))
            }
            else {
                None
            }
        }
        #[allow(unused_variables)]
        fn search(&self, topic: &str) -> Vec<String> {
            vec!["Suggestion 1".to_owned(),
            "Suggestion 2".to_owned(),
            "Suggestion 3".to_owned()]
        }
    }

    let provider = TestProvider {};
    let story_builder = StoryBuilder::new(&provider);
    assert_eq!(story_builder.build_story("not-found", "found"), Err(EXPECTED_SUGGESTION.to_owned()));
}

#[test]
/// For: build_story
fn build_story_cannot_find_second_article_suggest() {
    struct TestArticle {
        topics: Vec<Paragraph>
    }
    impl Article for TestArticle {
        fn get_paragraphs(&self) -> &Vec<Paragraph> {&self.topics}
    }
    struct TestProvider {}
    impl ArticleProvider for TestProvider {
        fn get(&self, topic: &str) -> Option<Box<Article>> {
            if topic == "found"
                {Some(Box::new(TestArticle {topics: vec![]}))}
            else
                {None}
        }
        #[allow(unused_variables)]
        fn search(&self, topic: &str) -> Vec<String> {
            vec!["Suggestion 1".to_owned(),
            "Suggestion 2".to_owned(),
            "Suggestion 3".to_owned()]
        }
    }

    let provider = TestProvider {};
    let story_builder = StoryBuilder::new(&provider);
    assert_eq!(story_builder.build_story("found", "not-found"), Err(EXPECTED_SUGGESTION.to_owned()));
}

#[test]
/// For: build_story
fn build_story_from_same_start_and_end_should_err() {
    struct TestProvider {}
    impl ArticleProvider for TestProvider {
        #[allow(unused_variables)]
        fn get(&self, topic: &str) -> Option<Box<Article>> {panic!("get() should not get called in this test.")}
        #[allow(unused_variables)]
        fn search(&self, topic: &str) -> Vec<String> {panic!("get() should not get called in this test.")}
    }
    let provider = TestProvider {};
    let story_builder = StoryBuilder::new(&provider);
    assert_eq!(story_builder.build_story("similar topic", "similar topic"),
               Err("No story to build; same start and end topics.".to_owned()));
}

#[test]
/// For: build_story
fn build_story_empty_start_topic_should_err() {
    struct TestProvider {}
    impl ArticleProvider for TestProvider {
        #[allow(unused_variables)]
        fn get(&self, topic: &str) -> Option<Box<Article>> {panic!("get() should not get called in this test.")}
        #[allow(unused_variables)]
        fn search(&self, topic: &str) -> Vec<String> {panic!("get() should not get called in this test.")}
    }
    let provider = TestProvider {};
    let story_builder = StoryBuilder::new(&provider);
    assert_eq!(story_builder.build_story("", "Other topic"),
               Err("Missing start topic.".to_owned()));
}

#[test]
/// For: build_story
fn build_story_empty_end_topic_should_err() {
    struct TestProvider {}
    impl ArticleProvider for TestProvider {
        #[allow(unused_variables)]
        fn get(&self, topic: &str) -> Option<Box<Article>> {panic!("get() should not get called in this test.")}
        #[allow(unused_variables)]
        fn search(&self, topic: &str) -> Vec<String> {panic!("get() should not get called in this test.")}
    }
    let provider = TestProvider {};
    let story_builder = StoryBuilder::new(&provider);
    assert_eq!(story_builder.build_story("First topic", ""),
               Err("Missing end topic.".to_owned()));
}

#[test]
/// For: build_story
fn build_story_end_topic_found_in_start_article() {
    let mut prebuilt_rels = Rc::new(HashMap::new());
    Rc::get_mut(&mut prebuilt_rels).unwrap().insert("start", vec![Paragraph {text: "Paragraph 1".to_owned(), topics: vec!["topic 1".to_owned(), "topic 2".to_owned(), "topic 3".to_owned()]},
                                                                  Paragraph {text: "Paragraph 2".to_owned(), topics: vec!["topic 3".to_owned(), "end".to_owned(), "topic 5".to_owned()]},
                                                                  Paragraph {text: "Paragraph 3".to_owned(), topics: vec!["topic 3".to_owned(), "topic 1".to_owned(), "topic 5".to_owned()]}]);
    struct TestArticle {
        topic: String,
        prebuilt_rels: Rc<HashMap<&'static str, Vec<Paragraph>>>,
    }

    impl TestArticle {
        fn new(topic: String, prebuilt_rels: Rc<HashMap<&'static str, Vec<Paragraph>>>) -> TestArticle {
            TestArticle {topic, prebuilt_rels}
        }
    }

    impl Article for TestArticle {
            fn get_paragraphs(&self) -> &Vec<Paragraph> {
            self.prebuilt_rels.get::<str>(&self.topic).expect("Tried to access a node that doesn't exist!")
        }
    }
    struct TestProvider {
        prebuilt_rels: Rc<HashMap<&'static str, Vec<Paragraph>>>
    }
    impl TestProvider {
        fn new(prebuilt_rels: Rc<HashMap<&'static str, Vec<Paragraph>>>) -> TestProvider {
            TestProvider {prebuilt_rels}
        }
    }
    impl ArticleProvider for TestProvider {
        fn get(&self, topic: &str) -> Option<Box<Article>> {
                let new_rels: Rc<HashMap<&'static str, Vec<Paragraph>>> = self.prebuilt_rels.clone();
                Some(Box::new(TestArticle::new(topic.to_owned(), new_rels)))
        }
        fn search(&self, topic: &str) -> Vec<String> {panic!("search() should not be called in this test.");}
    }
    let provider = TestProvider::new(prebuilt_rels.clone());
    let story_builder = StoryBuilder::new(&provider);
    assert_eq!(story_builder.build_story("start", "end"), Ok("Paragraph 2\r\n".to_owned()));
}

#[test]
/// For: build_story
fn build_story_end_topic_found_in_second_level() {
    let mut prebuilt_rels = Rc::new(HashMap::new());
    Rc::get_mut(&mut prebuilt_rels).unwrap().insert("start", vec![Paragraph {text: "Paragraph 1".to_owned(), topics: vec!["topic 1".to_owned(), "topic 2".to_owned(), "topic 3".to_owned()]},
                                                                  Paragraph {text: "Paragraph 2".to_owned(), topics: vec!["topic 4".to_owned(), "topic 2".to_owned(), "topic 4".to_owned()]},
                                                                  Paragraph {text: "Paragraph 3".to_owned(), topics: vec!["topic 1".to_owned(), "topic 1".to_owned(), "topic 2".to_owned()]}]);

    Rc::get_mut(&mut prebuilt_rels).unwrap().insert("topic 1", vec![Paragraph {text: "Paragraph 1".to_owned(), topics: vec!["topic 1".to_owned(), "topic 2".to_owned(), "topic 3".to_owned()]},
                                                                    Paragraph {text: "Paragraph 2".to_owned(), topics: vec!["end".to_owned(), "topic 1".to_owned(), "topic 2".to_owned()]}]);
    struct TestArticle {
        topic: String,
        prebuilt_rels: Rc<HashMap<&'static str, Vec<Paragraph>>>,
    }

    impl TestArticle {
        fn new(topic: String, prebuilt_rels: Rc<HashMap<&'static str, Vec<Paragraph>>>) -> TestArticle {
            TestArticle {topic, prebuilt_rels}
        }
    }

    impl Article for TestArticle {
            fn get_paragraphs(&self) -> &Vec<Paragraph> {
            self.prebuilt_rels.get::<str>(&self.topic).expect("Tried to access a node that doesn't exist!")
        }
    }
    struct TestProvider {
        prebuilt_rels: Rc<HashMap<&'static str, Vec<Paragraph>>>
    }
    impl TestProvider {
        fn new(prebuilt_rels: Rc<HashMap<&'static str, Vec<Paragraph>>>) -> TestProvider {
            TestProvider {prebuilt_rels}
        }
    }
    impl ArticleProvider for TestProvider {
        fn get(&self, topic: &str) -> Option<Box<Article>> {
                let new_rels: Rc<HashMap<&'static str, Vec<Paragraph>>> = self.prebuilt_rels.clone();
                Some(Box::new(TestArticle::new(topic.to_owned(), new_rels)))
        }
        fn search(&self, topic: &str) -> Vec<String> {panic!("search() should not be called in this test.");}
    }
    let provider = TestProvider::new(prebuilt_rels.clone());
    let story_builder = StoryBuilder::new(&provider);
    assert_eq!(story_builder.build_story("start", "end"), Ok("Paragraph 1\r\nParagraph 2\r\n".to_owned()));
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#[macro_use]
extern crate criterion;
use criterion::Criterion;

extern crate arrow;

use arrow::array::*;
use arrow::compute::kernels::boolean as boolean_kernels;

///  Helper function to create arrays
fn create_boolean_array(size: usize) -> BooleanArray {
    let mut builder = BooleanBuilder::new(size);
    for i in 0..size {
        if i % 2 == 0 {
            builder.append_value(true).unwrap();
        } else {
            builder.append_value(false).unwrap();
        }
    }
    builder.finish()
}

/// Benchmark for `AND`
fn bench_and(size: usize) {
    let buffer_a = create_boolean_array(size);
    let buffer_b = create_boolean_array(size);
    criterion::black_box(boolean_kernels::and(&buffer_a, &buffer_b).unwrap());
}

/// Benchmark for `OR`
fn bench_or(size: usize) {
    let buffer_a = create_boolean_array(size);
    let buffer_b = create_boolean_array(size);
    criterion::black_box(boolean_kernels::or(&buffer_a, &buffer_b).unwrap());
}

/// Benchmark for `NOT`
fn bench_not(size: usize) {
    let buffer = create_boolean_array(size);
    criterion::black_box(boolean_kernels::not(&buffer).unwrap());
}

fn add_benchmark(c: &mut Criterion) {
    c.bench_function("and", |b| b.iter(|| bench_and(512)));
    c.bench_function("or", |b| b.iter(|| bench_or(512)));
    c.bench_function("not", |b| b.iter(|| bench_not(512)));
}

criterion_group!(benches, add_benchmark);
criterion_main!(benches);
<commit_msg>ARROW-11025: [Rust] Fixed bench for binary boolean kernels<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#[macro_use]
extern crate criterion;
use criterion::Criterion;

use rand::distributions::{Distribution, Standard};
use rand::Rng;

use arrow::util::test_util::seedable_rng;

extern crate arrow;

use arrow::array::*;
use arrow::compute::kernels::boolean as boolean_kernels;

fn create_boolean(size: usize) -> BooleanArray
where
    Standard: Distribution<bool>,
{
    seedable_rng()
        .sample_iter(&Standard)
        .take(size)
        .map(Some)
        .collect()
}

fn bench_and(lhs: &BooleanArray, rhs: &BooleanArray) {
    criterion::black_box(boolean_kernels::and(lhs, rhs).unwrap());
}

fn bench_or(lhs: &BooleanArray, rhs: &BooleanArray) {
    criterion::black_box(boolean_kernels::or(lhs, rhs).unwrap());
}

fn bench_not(array: &BooleanArray) {
    criterion::black_box(boolean_kernels::not(&array).unwrap());
}

fn add_benchmark(c: &mut Criterion) {
    let size = 2usize.pow(15);
    let array1 = create_boolean(size);
    let array2 = create_boolean(size);
    c.bench_function("and", |b| b.iter(|| bench_and(&array1, &array2)));
    c.bench_function("or", |b| b.iter(|| bench_or(&array1, &array2)));
    c.bench_function("not", |b| b.iter(|| bench_not(&array1)));
}

criterion_group!(benches, add_benchmark);
criterion_main!(benches);
<|endoftext|>


<commit_before>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Alex Lyon <arcterus@mail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

// spell-checker:ignore (ToDO) cmdline evec seps rvec fdata

#[macro_use]
extern crate uucore;

use clap::{crate_version, App, Arg};
use rand::Rng;
use std::fs::File;
use std::io::{stdin, stdout, BufReader, BufWriter, Read, Write};
use uucore::InvalidEncodingHandling;

enum Mode {
    Default(String),
    Echo(Vec<String>),
    InputRange((usize, usize)),
}

static NAME: &str = "shuf";
static USAGE: &str = r#"shuf [OPTION]... [FILE]
  or:  shuf -e [OPTION]... [ARG]...
  or:  shuf -i LO-HI [OPTION]...
Write a random permutation of the input lines to standard output.

With no FILE, or when FILE is -, read standard input.
"#;
static TEMPLATE: &str = "Usage: {usage}\nMandatory arguments to long options are mandatory for short options too.\n{unified}";

struct Options {
    head_count: usize,
    output: Option<String>,
    random_source: Option<String>,
    repeat: bool,
    sep: u8,
}

mod options {
    pub static ECHO: &str = "echo";
    pub static INPUT_RANGE: &str = "input-range";
    pub static HEAD_COUNT: &str = "head-count";
    pub static OUTPUT: &str = "output";
    pub static RANDOM_SOURCE: &str = "random-source";
    pub static REPEAT: &str = "repeat";
    pub static ZERO_TERMINATED: &str = "zero-terminated";
    pub static FILE: &str = "file";
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let args = args
        .collect_str(InvalidEncodingHandling::ConvertLossy)
        .accept_any();

    let matches = App::new(executable!())
        .name(NAME)
        .version(crate_version!())
        .template(TEMPLATE)
        .usage(USAGE)
        .arg(
            Arg::with_name(options::ECHO)
                .short("e")
                .long(options::ECHO)
                .takes_value(true)
                .value_name("ARG")
                .help("treat each ARG as an input line")
                .multiple(true)
                .use_delimiter(false)
                .min_values(0)
                .conflicts_with(options::INPUT_RANGE),
        )
        .arg(
            Arg::with_name(options::INPUT_RANGE)
                .short("i")
                .long(options::INPUT_RANGE)
                .takes_value(true)
                .value_name("LO-HI")
                .help("treat each number LO through HI as an input line")
                .conflicts_with(options::FILE),
        )
        .arg(
            Arg::with_name(options::HEAD_COUNT)
                .short("n")
                .long(options::HEAD_COUNT)
                .takes_value(true)
                .value_name("COUNT")
                .help("output at most COUNT lines"),
        )
        .arg(
            Arg::with_name(options::OUTPUT)
                .short("o")
                .long(options::OUTPUT)
                .takes_value(true)
                .value_name("FILE")
                .help("write result to FILE instead of standard output"),
        )
        .arg(
            Arg::with_name(options::RANDOM_SOURCE)
                .long(options::RANDOM_SOURCE)
                .takes_value(true)
                .value_name("FILE")
                .help("get random bytes from FILE"),
        )
        .arg(
            Arg::with_name(options::REPEAT)
                .short("r")
                .long(options::REPEAT)
                .help("output lines can be repeated"),
        )
        .arg(
            Arg::with_name(options::ZERO_TERMINATED)
                .short("z")
                .long(options::ZERO_TERMINATED)
                .help("line delimiter is NUL, not newline"),
        )
        .arg(Arg::with_name(options::FILE).takes_value(true))
        .get_matches_from(args);

    let mode = if let Some(args) = matches.values_of(options::ECHO) {
        Mode::Echo(args.map(String::from).collect())
    } else if let Some(range) = matches.value_of(options::INPUT_RANGE) {
        match parse_range(range) {
            Ok(m) => Mode::InputRange(m),
            Err(msg) => {
                crash!(1, "{}", msg);
            }
        }
    } else {
        Mode::Default(matches.value_of(options::FILE).unwrap_or("-").to_string())
    };

    let options = Options {
        head_count: match matches.value_of(options::HEAD_COUNT) {
            Some(count) => match count.parse::<usize>() {
                Ok(val) => val,
                Err(_) => {
                    show_error!("invalid line count: '{}'", count);
                    return 1;
                }
            },
            None => std::usize::MAX,
        },
        output: matches.value_of(options::OUTPUT).map(String::from),
        random_source: matches.value_of(options::RANDOM_SOURCE).map(String::from),
        repeat: matches.is_present(options::REPEAT),
        sep: if matches.is_present(options::ZERO_TERMINATED) {
            0x00_u8
        } else {
            0x0a_u8
        },
    };

    match mode {
        Mode::Echo(args) => {
            let mut evec = args.iter().map(String::as_bytes).collect::<Vec<_>>();
            find_seps(&mut evec, options.sep);
            shuf_bytes(&mut evec, options);
        }
        Mode::InputRange((b, e)) => {
            let rvec = (b..e).map(|x| format!("{}", x)).collect::<Vec<String>>();
            let mut rvec = rvec.iter().map(String::as_bytes).collect::<Vec<&[u8]>>();
            shuf_bytes(&mut rvec, options);
        }
        Mode::Default(filename) => {
            let fdata = read_input_file(&filename);
            let mut fdata = vec![&fdata[..]];
            find_seps(&mut fdata, options.sep);
            shuf_bytes(&mut fdata, options);
        }
    }

    0
}

fn read_input_file(filename: &str) -> Vec<u8> {
    let mut file = BufReader::new(if filename == "-" {
        Box::new(stdin()) as Box<dyn Read>
    } else {
        match File::open(filename) {
            Ok(f) => Box::new(f) as Box<dyn Read>,
            Err(e) => crash!(1, "failed to open '{}': {}", filename, e),
        }
    });

    let mut data = Vec::new();
    if let Err(e) = file.read_to_end(&mut data) {
        crash!(1, "failed reading '{}': {}", filename, e)
    };

    data
}

fn find_seps(data: &mut Vec<&[u8]>, sep: u8) {
    // need to use for loop so we don't borrow the vector as we modify it in place
    // basic idea:
    // * We don't care about the order of the result. This lets us slice the slices
    //   without making a new vector.
    // * Starting from the end of the vector, we examine each element.
    // * If that element contains the separator, we remove it from the vector,
    //   and then sub-slice it into slices that do not contain the separator.
    // * We maintain the invariant throughout that each element in the vector past
    //   the ith element does not have any separators remaining.
    for i in (0..data.len()).rev() {
        if data[i].contains(&sep) {
            let this = data.swap_remove(i);
            let mut p = 0;
            let mut i = 1;
            loop {
                if i == this.len() {
                    break;
                }

                if this[i] == sep {
                    data.push(&this[p..i]);
                    p = i + 1;
                }
                i += 1;
            }
            if p < this.len() {
                data.push(&this[p..i]);
            }
        }
    }
}

fn shuf_bytes(input: &mut Vec<&[u8]>, opts: Options) {
    let mut output = BufWriter::new(match opts.output {
        None => Box::new(stdout()) as Box<dyn Write>,
        Some(s) => match File::create(&s[..]) {
            Ok(f) => Box::new(f) as Box<dyn Write>,
            Err(e) => crash!(1, "failed to open '{}' for writing: {}", &s[..], e),
        },
    });

    let mut rng = match opts.random_source {
        Some(r) => WrappedRng::RngFile(rand::read::ReadRng::new(match File::open(&r[..]) {
            Ok(f) => f,
            Err(e) => crash!(1, "failed to open random source '{}': {}", &r[..], e),
        })),
        None => WrappedRng::RngDefault(rand::thread_rng()),
    };

    // we're generating a random usize. To keep things fair, we take this number mod ceil(log2(length+1))
    let mut len_mod = 1;
    let mut len = input.len();
    while len > 0 {
        len >>= 1;
        len_mod <<= 1;
    }

    let mut count = opts.head_count;
    while count > 0 && !input.is_empty() {
        let mut r = input.len();
        while r >= input.len() {
            r = rng.next_usize() % len_mod;
        }

        // write the randomly chosen value and the separator
        output
            .write_all(input[r])
            .unwrap_or_else(|e| crash!(1, "write failed: {}", e));
        output
            .write_all(&[opts.sep])
            .unwrap_or_else(|e| crash!(1, "write failed: {}", e));

        // if we do not allow repeats, remove the chosen value from the input vector
        if !opts.repeat {
            // shrink the mask if we will drop below a power of 2
            if input.len() % 2 == 0 && len_mod > 2 {
                len_mod >>= 1;
            }
            input.swap_remove(r);
        }

        count -= 1;
    }
}

fn parse_range(input_range: &str) -> Result<(usize, usize), String> {
    let split: Vec<&str> = input_range.split('-').collect();
    if split.len() != 2 {
        Err(format!("invalid input range: '{}'", input_range))
    } else {
        let begin = match split[0].parse::<usize>() {
            Ok(m) => m,
            Err(_) => return Err(format!("invalid input range: '{}'", split[0])),
        };
        let end = match split[1].parse::<usize>() {
            Ok(m) => m,
            Err(_) => return Err(format!("invalid input range: '{}'", split[1])),
        };
        Ok((begin, end + 1))
    }
}

enum WrappedRng {
    RngFile(rand::read::ReadRng<File>),
    RngDefault(rand::ThreadRng),
}

impl WrappedRng {
    fn next_usize(&mut self) -> usize {
        match *self {
            WrappedRng::RngFile(ref mut r) => r.gen(),
            WrappedRng::RngDefault(ref mut r) => r.gen(),
        }
    }
}
<commit_msg>shuf: clean up returning Err<commit_after>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Alex Lyon <arcterus@mail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

// spell-checker:ignore (ToDO) cmdline evec seps rvec fdata

#[macro_use]
extern crate uucore;

use clap::{crate_version, App, Arg};
use rand::Rng;
use std::fs::File;
use std::io::{stdin, stdout, BufReader, BufWriter, Read, Write};
use uucore::InvalidEncodingHandling;

enum Mode {
    Default(String),
    Echo(Vec<String>),
    InputRange((usize, usize)),
}

static NAME: &str = "shuf";
static USAGE: &str = r#"shuf [OPTION]... [FILE]
  or:  shuf -e [OPTION]... [ARG]...
  or:  shuf -i LO-HI [OPTION]...
Write a random permutation of the input lines to standard output.

With no FILE, or when FILE is -, read standard input.
"#;
static TEMPLATE: &str = "Usage: {usage}\nMandatory arguments to long options are mandatory for short options too.\n{unified}";

struct Options {
    head_count: usize,
    output: Option<String>,
    random_source: Option<String>,
    repeat: bool,
    sep: u8,
}

mod options {
    pub static ECHO: &str = "echo";
    pub static INPUT_RANGE: &str = "input-range";
    pub static HEAD_COUNT: &str = "head-count";
    pub static OUTPUT: &str = "output";
    pub static RANDOM_SOURCE: &str = "random-source";
    pub static REPEAT: &str = "repeat";
    pub static ZERO_TERMINATED: &str = "zero-terminated";
    pub static FILE: &str = "file";
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let args = args
        .collect_str(InvalidEncodingHandling::ConvertLossy)
        .accept_any();

    let matches = App::new(executable!())
        .name(NAME)
        .version(crate_version!())
        .template(TEMPLATE)
        .usage(USAGE)
        .arg(
            Arg::with_name(options::ECHO)
                .short("e")
                .long(options::ECHO)
                .takes_value(true)
                .value_name("ARG")
                .help("treat each ARG as an input line")
                .multiple(true)
                .use_delimiter(false)
                .min_values(0)
                .conflicts_with(options::INPUT_RANGE),
        )
        .arg(
            Arg::with_name(options::INPUT_RANGE)
                .short("i")
                .long(options::INPUT_RANGE)
                .takes_value(true)
                .value_name("LO-HI")
                .help("treat each number LO through HI as an input line")
                .conflicts_with(options::FILE),
        )
        .arg(
            Arg::with_name(options::HEAD_COUNT)
                .short("n")
                .long(options::HEAD_COUNT)
                .takes_value(true)
                .value_name("COUNT")
                .help("output at most COUNT lines"),
        )
        .arg(
            Arg::with_name(options::OUTPUT)
                .short("o")
                .long(options::OUTPUT)
                .takes_value(true)
                .value_name("FILE")
                .help("write result to FILE instead of standard output"),
        )
        .arg(
            Arg::with_name(options::RANDOM_SOURCE)
                .long(options::RANDOM_SOURCE)
                .takes_value(true)
                .value_name("FILE")
                .help("get random bytes from FILE"),
        )
        .arg(
            Arg::with_name(options::REPEAT)
                .short("r")
                .long(options::REPEAT)
                .help("output lines can be repeated"),
        )
        .arg(
            Arg::with_name(options::ZERO_TERMINATED)
                .short("z")
                .long(options::ZERO_TERMINATED)
                .help("line delimiter is NUL, not newline"),
        )
        .arg(Arg::with_name(options::FILE).takes_value(true))
        .get_matches_from(args);

    let mode = if let Some(args) = matches.values_of(options::ECHO) {
        Mode::Echo(args.map(String::from).collect())
    } else if let Some(range) = matches.value_of(options::INPUT_RANGE) {
        match parse_range(range) {
            Ok(m) => Mode::InputRange(m),
            Err(msg) => {
                crash!(1, "{}", msg);
            }
        }
    } else {
        Mode::Default(matches.value_of(options::FILE).unwrap_or("-").to_string())
    };

    let options = Options {
        head_count: match matches.value_of(options::HEAD_COUNT) {
            Some(count) => match count.parse::<usize>() {
                Ok(val) => val,
                Err(_) => {
                    show_error!("invalid line count: '{}'", count);
                    return 1;
                }
            },
            None => std::usize::MAX,
        },
        output: matches.value_of(options::OUTPUT).map(String::from),
        random_source: matches.value_of(options::RANDOM_SOURCE).map(String::from),
        repeat: matches.is_present(options::REPEAT),
        sep: if matches.is_present(options::ZERO_TERMINATED) {
            0x00_u8
        } else {
            0x0a_u8
        },
    };

    match mode {
        Mode::Echo(args) => {
            let mut evec = args.iter().map(String::as_bytes).collect::<Vec<_>>();
            find_seps(&mut evec, options.sep);
            shuf_bytes(&mut evec, options);
        }
        Mode::InputRange((b, e)) => {
            let rvec = (b..e).map(|x| format!("{}", x)).collect::<Vec<String>>();
            let mut rvec = rvec.iter().map(String::as_bytes).collect::<Vec<&[u8]>>();
            shuf_bytes(&mut rvec, options);
        }
        Mode::Default(filename) => {
            let fdata = read_input_file(&filename);
            let mut fdata = vec![&fdata[..]];
            find_seps(&mut fdata, options.sep);
            shuf_bytes(&mut fdata, options);
        }
    }

    0
}

fn read_input_file(filename: &str) -> Vec<u8> {
    let mut file = BufReader::new(if filename == "-" {
        Box::new(stdin()) as Box<dyn Read>
    } else {
        match File::open(filename) {
            Ok(f) => Box::new(f) as Box<dyn Read>,
            Err(e) => crash!(1, "failed to open '{}': {}", filename, e),
        }
    });

    let mut data = Vec::new();
    if let Err(e) = file.read_to_end(&mut data) {
        crash!(1, "failed reading '{}': {}", filename, e)
    };

    data
}

fn find_seps(data: &mut Vec<&[u8]>, sep: u8) {
    // need to use for loop so we don't borrow the vector as we modify it in place
    // basic idea:
    // * We don't care about the order of the result. This lets us slice the slices
    //   without making a new vector.
    // * Starting from the end of the vector, we examine each element.
    // * If that element contains the separator, we remove it from the vector,
    //   and then sub-slice it into slices that do not contain the separator.
    // * We maintain the invariant throughout that each element in the vector past
    //   the ith element does not have any separators remaining.
    for i in (0..data.len()).rev() {
        if data[i].contains(&sep) {
            let this = data.swap_remove(i);
            let mut p = 0;
            let mut i = 1;
            loop {
                if i == this.len() {
                    break;
                }

                if this[i] == sep {
                    data.push(&this[p..i]);
                    p = i + 1;
                }
                i += 1;
            }
            if p < this.len() {
                data.push(&this[p..i]);
            }
        }
    }
}

fn shuf_bytes(input: &mut Vec<&[u8]>, opts: Options) {
    let mut output = BufWriter::new(match opts.output {
        None => Box::new(stdout()) as Box<dyn Write>,
        Some(s) => match File::create(&s[..]) {
            Ok(f) => Box::new(f) as Box<dyn Write>,
            Err(e) => crash!(1, "failed to open '{}' for writing: {}", &s[..], e),
        },
    });

    let mut rng = match opts.random_source {
        Some(r) => WrappedRng::RngFile(rand::read::ReadRng::new(match File::open(&r[..]) {
            Ok(f) => f,
            Err(e) => crash!(1, "failed to open random source '{}': {}", &r[..], e),
        })),
        None => WrappedRng::RngDefault(rand::thread_rng()),
    };

    // we're generating a random usize. To keep things fair, we take this number mod ceil(log2(length+1))
    let mut len_mod = 1;
    let mut len = input.len();
    while len > 0 {
        len >>= 1;
        len_mod <<= 1;
    }

    let mut count = opts.head_count;
    while count > 0 && !input.is_empty() {
        let mut r = input.len();
        while r >= input.len() {
            r = rng.next_usize() % len_mod;
        }

        // write the randomly chosen value and the separator
        output
            .write_all(input[r])
            .unwrap_or_else(|e| crash!(1, "write failed: {}", e));
        output
            .write_all(&[opts.sep])
            .unwrap_or_else(|e| crash!(1, "write failed: {}", e));

        // if we do not allow repeats, remove the chosen value from the input vector
        if !opts.repeat {
            // shrink the mask if we will drop below a power of 2
            if input.len() % 2 == 0 && len_mod > 2 {
                len_mod >>= 1;
            }
            input.swap_remove(r);
        }

        count -= 1;
    }
}

fn parse_range(input_range: &str) -> Result<(usize, usize), String> {
    let split: Vec<&str> = input_range.split('-').collect();
    if split.len() != 2 {
        Err(format!("invalid input range: '{}'", input_range))
    } else {
        let begin = split[0]
            .parse::<usize>()
            .map_err(|_| format!("invalid input range: '{}'", split[0]))?;
        let end = split[1]
            .parse::<usize>()
            .map_err(|_| format!("invalid input range: '{}'", split[1]))?;
        Ok((begin, end + 1))
    }
}

enum WrappedRng {
    RngFile(rand::read::ReadRng<File>),
    RngDefault(rand::ThreadRng),
}

impl WrappedRng {
    fn next_usize(&mut self) -> usize {
        match *self {
            WrappedRng::RngFile(ref mut r) => r.gen(),
            WrappedRng::RngDefault(ref mut r) => r.gen(),
        }
    }
}
<|endoftext|>


<commit_before>use pagecache::{Measure, M};

use super::*;

/// An iterator over keys and values in a `Tree`.
pub struct Iter<'a> {
    pub(super) id: PageId,
    pub(super) inner: &'a Tree,
    pub(super) last_key: Key,
    pub(super) inclusive: bool,
    pub(super) broken: Option<Error<()>>,
    pub(super) done: bool,
    pub(super) guard: Guard,
    // TODO we have to refactor this in light of pages being deleted
}

impl<'a> Iterator for Iter<'a> {
    type Item = Result<(Vec<u8>, PinnedValue), ()>;

    fn next(&mut self) -> Option<Self::Item> {
        let _measure = Measure::new(&M.tree_scan);

        if self.done {
            return None;
        } else if let Some(broken) = self.broken.take() {
            self.done = true;
            return Some(Err(broken));
        };

        loop {
            let pin_guard = self.guard.clone();
            let get_guard = self.guard.clone();

            let res = self
                .inner
                .pages
                .get(self.id, &get_guard)
                .map(|page_get| page_get.unwrap());

            if let Err(e) = res {
                error!("iteration failed: {:?}", e);
                self.done = true;
                return Some(Err(e.danger_cast()));
            }

            // TODO (when implementing merge support) this could
            // be None if the node was removed since the last
            // iteration, and we need to just get the inner
            // node again...
            let (frag, _ptr) = res.unwrap();
            let node = frag.unwrap_base();
            let leaf =
                node.data.leaf_ref().expect("node should be a leaf");
            let prefix = node.lo.inner();

            let search = if self.inclusive {
                self.inclusive = false;
                leaf.binary_search_by(|&(ref k, ref _v)| {
                    prefix_cmp_encoded(k, &self.last_key, prefix)
                }).ok()
                .or_else(|| {
                    binary_search_gt(leaf, |&(ref k, ref _v)| {
                        prefix_cmp_encoded(k, &self.last_key, prefix)
                    })
                })
            } else {
                binary_search_gt(leaf, |&(ref k, ref _v)| {
                    prefix_cmp_encoded(k, &self.last_key, prefix)
                })
            };

            if let Some(idx) = search {
                let (k, v) = &leaf[idx];
                let decoded_k = prefix_decode(prefix, &k);
                self.last_key = decoded_k.to_vec();
                let ret =
                    Ok((decoded_k, PinnedValue::new(&*v, pin_guard)));
                return Some(ret);
            }

            match node.next {
                Some(id) => self.id = id,
                None => {
                    assert_eq!(
                        node.hi,
                        Bound::Inf,
                        "if a node has no right sibling, \
                         it must be the upper-bound node"
                    );
                    return None;
                }
            }
        }
    }
}

impl<'a> DoubleEndedIterator for Iter<'a> {
    fn next_back(&mut self) -> Option<Self::Item> {
        let _measure = Measure::new(&M.tree_scan);

        if self.done {
            return None;
        } else if let Some(broken) = self.broken.take() {
            self.done = true;
            return Some(Err(broken));
        };

        loop {
            let pin_guard = self.guard.clone();
            let get_guard = self.guard.clone();

            let res = self
                .inner
                .pages
                .get(self.id, &get_guard)
                .map(|page_get| page_get.unwrap());

            if let Err(e) = res {
                error!("iteration failed: {:?}", e);
                self.done = true;
                return Some(Err(e.danger_cast()));
            }

            // TODO (when implementing merge support) this could
            // be None if the node was removed since the last
            // iteration, and we need to just get the inner
            // node again...
            let (frag, _ptr) = res.unwrap();
            let node = frag.unwrap_base();
            let leaf =
                node.data.leaf_ref().expect("node should be a leaf");
            let prefix = node.lo.inner();

            let search = if self.inclusive {
                self.inclusive = false;
                leaf.binary_search_by(|&(ref k, ref _v)| {
                    prefix_cmp_encoded(k, &self.last_key, prefix)
                }).ok()
                .or_else(|| {
                    binary_search_lt(leaf, |&(ref k, ref _v)| {
                        prefix_cmp_encoded(k, &self.last_key, prefix)
                    })
                })
            } else {
                binary_search_lt(leaf, |&(ref k, ref _v)| {
                    prefix_cmp_encoded(k, &self.last_key, prefix)
                })
            };

            if let Some(idx) = search {
                let (k, v) = &leaf[idx];
                let decoded_k = prefix_decode(prefix, &k);
                self.last_key = decoded_k.to_vec();
                let ret =
                    Ok((decoded_k, PinnedValue::new(&*v, pin_guard)));
                return Some(ret);
            }

            // we need to get the node to the right of ours by
            // guessing a key that might land on it, and then
            // fast-forwarding through the right child pointers
            // if we went too far to the left.
            let pred = possible_predecessor(prefix)?;
            let mut next_node =
                match self.inner.path_for_key(pred, &get_guard) {
                    Err(e) => {
                        error!("next_back iteration failed: {:?}", e);
                        self.done = true;
                        return Some(Err(e.danger_cast()));
                    }
                    Ok(path) => path.last().unwrap().0.unwrap_base(),
                };
            while next_node.next != Some(node.id)
                && next_node.lo < node.lo
            {
                let res = self
                    .inner
                    .pages
                    .get(next_node.next?, &get_guard)
                    .map(|page_get| page_get.unwrap());

                if let Err(e) = res {
                    error!("iteration failed: {:?}", e);
                    self.done = true;
                    return Some(Err(e.danger_cast()));
                }
                let (frag, _ptr) = res.unwrap();
                next_node = frag.unwrap_base();
            }

            self.id = next_node.id;
        }
    }
}

fn possible_predecessor(s: &[u8]) -> Option<Vec<u8>> {
    let mut ret = s.to_vec();
    match ret.pop() {
        None => None,
        Some(i) if i == 0 => Some(ret),
        Some(i) => {
            ret.push(i - 1);
            Some(ret)
        }
    }
}

#[test]
fn test_possible_predecessor() {
    assert_eq!(possible_predecessor(b""), None);
    assert_eq!(possible_predecessor(&[0]), Some(vec![]));
    assert_eq!(possible_predecessor(&[0, 0]), Some(vec![0]));
    assert_eq!(
        possible_predecessor(&[0, 1]),
        Some(vec![0, 0, 255, 255, 255, 255])
    );
    assert_eq!(
        possible_predecessor(&[0, 2]),
        Some(vec![0, 1, 255, 255, 255, 255])
    );
    assert_eq!(possible_predecessor(&[1, 0]), Some(vec![1]));
    assert_eq!(
        possible_predecessor(&[1, 1]),
        Some(vec![1, 0, 255, 255, 255, 255])
    );
    assert_eq!(
        possible_predecessor(&[155]),
        Some(vec![154, 255, 255, 255, 255])
    );
}
<commit_msg>Add additional high bytes to predecessor guesses<commit_after>use pagecache::{Measure, M};

use super::*;

/// An iterator over keys and values in a `Tree`.
pub struct Iter<'a> {
    pub(super) id: PageId,
    pub(super) inner: &'a Tree,
    pub(super) last_key: Key,
    pub(super) inclusive: bool,
    pub(super) broken: Option<Error<()>>,
    pub(super) done: bool,
    pub(super) guard: Guard,
    // TODO we have to refactor this in light of pages being deleted
}

impl<'a> Iterator for Iter<'a> {
    type Item = Result<(Vec<u8>, PinnedValue), ()>;

    fn next(&mut self) -> Option<Self::Item> {
        let _measure = Measure::new(&M.tree_scan);

        if self.done {
            return None;
        } else if let Some(broken) = self.broken.take() {
            self.done = true;
            return Some(Err(broken));
        };

        loop {
            let pin_guard = self.guard.clone();
            let get_guard = self.guard.clone();

            let res = self
                .inner
                .pages
                .get(self.id, &get_guard)
                .map(|page_get| page_get.unwrap());

            if let Err(e) = res {
                error!("iteration failed: {:?}", e);
                self.done = true;
                return Some(Err(e.danger_cast()));
            }

            // TODO (when implementing merge support) this could
            // be None if the node was removed since the last
            // iteration, and we need to just get the inner
            // node again...
            let (frag, _ptr) = res.unwrap();
            let node = frag.unwrap_base();
            let leaf =
                node.data.leaf_ref().expect("node should be a leaf");
            let prefix = node.lo.inner();

            let search = if self.inclusive {
                self.inclusive = false;
                leaf.binary_search_by(|&(ref k, ref _v)| {
                    prefix_cmp_encoded(k, &self.last_key, prefix)
                }).ok()
                .or_else(|| {
                    binary_search_gt(leaf, |&(ref k, ref _v)| {
                        prefix_cmp_encoded(k, &self.last_key, prefix)
                    })
                })
            } else {
                binary_search_gt(leaf, |&(ref k, ref _v)| {
                    prefix_cmp_encoded(k, &self.last_key, prefix)
                })
            };

            if let Some(idx) = search {
                let (k, v) = &leaf[idx];
                let decoded_k = prefix_decode(prefix, &k);
                self.last_key = decoded_k.to_vec();
                let ret =
                    Ok((decoded_k, PinnedValue::new(&*v, pin_guard)));
                return Some(ret);
            }

            match node.next {
                Some(id) => self.id = id,
                None => {
                    assert_eq!(
                        node.hi,
                        Bound::Inf,
                        "if a node has no right sibling, \
                         it must be the upper-bound node"
                    );
                    return None;
                }
            }
        }
    }
}

impl<'a> DoubleEndedIterator for Iter<'a> {
    fn next_back(&mut self) -> Option<Self::Item> {
        let _measure = Measure::new(&M.tree_scan);

        if self.done {
            return None;
        } else if let Some(broken) = self.broken.take() {
            self.done = true;
            return Some(Err(broken));
        };

        loop {
            let pin_guard = self.guard.clone();
            let get_guard = self.guard.clone();

            let res = self
                .inner
                .pages
                .get(self.id, &get_guard)
                .map(|page_get| page_get.unwrap());

            if let Err(e) = res {
                error!("iteration failed: {:?}", e);
                self.done = true;
                return Some(Err(e.danger_cast()));
            }

            // TODO (when implementing merge support) this could
            // be None if the node was removed since the last
            // iteration, and we need to just get the inner
            // node again...
            let (frag, _ptr) = res.unwrap();
            let node = frag.unwrap_base();
            let leaf =
                node.data.leaf_ref().expect("node should be a leaf");
            let prefix = node.lo.inner();

            let search = if self.inclusive {
                self.inclusive = false;
                leaf.binary_search_by(|&(ref k, ref _v)| {
                    prefix_cmp_encoded(k, &self.last_key, prefix)
                }).ok()
                .or_else(|| {
                    binary_search_lt(leaf, |&(ref k, ref _v)| {
                        prefix_cmp_encoded(k, &self.last_key, prefix)
                    })
                })
            } else {
                binary_search_lt(leaf, |&(ref k, ref _v)| {
                    prefix_cmp_encoded(k, &self.last_key, prefix)
                })
            };

            if let Some(idx) = search {
                let (k, v) = &leaf[idx];
                let decoded_k = prefix_decode(prefix, &k);
                self.last_key = decoded_k.to_vec();
                let ret =
                    Ok((decoded_k, PinnedValue::new(&*v, pin_guard)));
                return Some(ret);
            }

            // we need to get the node to the right of ours by
            // guessing a key that might land on it, and then
            // fast-forwarding through the right child pointers
            // if we went too far to the left.
            let pred = possible_predecessor(prefix)?;
            let mut next_node =
                match self.inner.path_for_key(pred, &get_guard) {
                    Err(e) => {
                        error!("next_back iteration failed: {:?}", e);
                        self.done = true;
                        return Some(Err(e.danger_cast()));
                    }
                    Ok(path) => path.last().unwrap().0.unwrap_base(),
                };
            while next_node.next != Some(node.id)
                && next_node.lo < node.lo
            {
                let res = self
                    .inner
                    .pages
                    .get(next_node.next?, &get_guard)
                    .map(|page_get| page_get.unwrap());

                if let Err(e) = res {
                    error!("iteration failed: {:?}", e);
                    self.done = true;
                    return Some(Err(e.danger_cast()));
                }
                let (frag, _ptr) = res.unwrap();
                next_node = frag.unwrap_base();
            }

            self.id = next_node.id;
        }
    }
}

fn possible_predecessor(s: &[u8]) -> Option<Vec<u8>> {
    let mut ret = s.to_vec();
    match ret.pop() {
        None => None,
        Some(i) if i == 0 => Some(ret),
        Some(i) => {
            ret.push(i - 1);
            for _ in 0..4 {
                ret.push(255);
            }
            Some(ret)
        }
    }
}

#[test]
fn test_possible_predecessor() {
    assert_eq!(possible_predecessor(b""), None);
    assert_eq!(possible_predecessor(&[0]), Some(vec![]));
    assert_eq!(possible_predecessor(&[0, 0]), Some(vec![0]));
    assert_eq!(
        possible_predecessor(&[0, 1]),
        Some(vec![0, 0, 255, 255, 255, 255])
    );
    assert_eq!(
        possible_predecessor(&[0, 2]),
        Some(vec![0, 1, 255, 255, 255, 255])
    );
    assert_eq!(possible_predecessor(&[1, 0]), Some(vec![1]));
    assert_eq!(
        possible_predecessor(&[1, 1]),
        Some(vec![1, 0, 255, 255, 255, 255])
    );
    assert_eq!(
        possible_predecessor(&[155]),
        Some(vec![154, 255, 255, 255, 255])
    );
}
<|endoftext|>