<commit_before>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::{
    move_types::{account_address::AccountAddress, language_storage::TypeTag},
    types::{
        account_config::{xdx_type_tag, xus_tag, XDX_NAME, XUS_NAME},
        chain_id::ChainId,
        transaction::{authenticator::AuthenticationKey, RawTransaction, TransactionPayload},
    },
};
use aptos_crypto::ed25519::Ed25519PublicKey;
use serde::{Deserialize, Serialize};
use std::fmt;

pub use aptos_transaction_builder::{aptos_stdlib, stdlib};
use aptos_types::transaction::{
    authenticator::AuthenticationKeyPreimage, ChangeSet, ModuleBundle, Script, ScriptFunction,
    WriteSetPayload,
};

pub struct TransactionBuilder {
    sender: Option<AccountAddress>,
    sequence_number: Option<u64>,
    payload: TransactionPayload,
    max_gas_amount: u64,
    gas_unit_price: u64,
    gas_currency_code: String,
    expiration_timestamp_secs: u64,
    chain_id: ChainId,
}

impl TransactionBuilder {
    pub fn sender(mut self, sender: AccountAddress) -> Self {
        self.sender = Some(sender);
        self
    }

    pub fn sequence_number(mut self, sequence_number: u64) -> Self {
        self.sequence_number = Some(sequence_number);
        self
    }

    pub fn max_gas_amount(mut self, max_gas_amount: u64) -> Self {
        self.max_gas_amount = max_gas_amount;
        self
    }

    pub fn gas_unit_price(mut self, gas_unit_price: u64) -> Self {
        self.gas_unit_price = gas_unit_price;
        self
    }

    pub fn gas_currency_code<T: Into<String>>(mut self, gas_currency_code: T) -> Self {
        self.gas_currency_code = gas_currency_code.into();
        self
    }

    pub fn chain_id(mut self, chain_id: ChainId) -> Self {
        self.chain_id = chain_id;
        self
    }

    pub fn expiration_timestamp_secs(mut self, expiration_timestamp_secs: u64) -> Self {
        self.expiration_timestamp_secs = expiration_timestamp_secs;
        self
    }

    pub fn build(self) -> RawTransaction {
        RawTransaction::new(
            self.sender.expect("sender must have been set"),
            self.sequence_number
                .expect("sequence number must have been set"),
            self.payload,
            self.max_gas_amount,
            self.gas_unit_price,
            self.gas_currency_code,
            self.expiration_timestamp_secs,
            self.chain_id,
        )
    }
}

#[derive(Clone, Debug)]
pub struct TransactionFactory {
    max_gas_amount: u64,
    gas_unit_price: u64,
    gas_currency: Currency,
    transaction_expiration_time: u64,
    chain_id: ChainId,
}

impl TransactionFactory {
    pub fn new(chain_id: ChainId) -> Self {
        Self {
            max_gas_amount: 1_000,
            gas_unit_price: 0,
            gas_currency: Currency::XUS,
            transaction_expiration_time: 30,
            chain_id,
        }
    }

    pub fn with_max_gas_amount(mut self, max_gas_amount: u64) -> Self {
        self.max_gas_amount = max_gas_amount;
        self
    }

    pub fn with_gas_unit_price(mut self, gas_unit_price: u64) -> Self {
        self.gas_unit_price = gas_unit_price;
        self
    }

    pub fn with_transaction_expiration_time(mut self, transaction_expiration_time: u64) -> Self {
        self.transaction_expiration_time = transaction_expiration_time;
        self
    }

    pub fn with_chain_id(mut self, chain_id: ChainId) -> Self {
        self.chain_id = chain_id;
        self
    }

    pub fn payload(&self, payload: TransactionPayload) -> TransactionBuilder {
        self.transaction_builder(payload)
    }

    pub fn module(&self, code: Vec<u8>) -> TransactionBuilder {
        self.payload(TransactionPayload::ModuleBundle(ModuleBundle::singleton(
            code,
        )))
    }

    pub fn change_set(&self, change_set: ChangeSet) -> TransactionBuilder {
        self.payload(TransactionPayload::WriteSet(WriteSetPayload::Direct(
            change_set,
        )))
    }

    pub fn script_function(&self, func: ScriptFunction) -> TransactionBuilder {
        self.payload(TransactionPayload::ScriptFunction(func))
    }

    pub fn create_user_account(&self, public_key: &Ed25519PublicKey) -> TransactionBuilder {
        let preimage = AuthenticationKeyPreimage::ed25519(public_key);
        self.payload(aptos_stdlib::encode_create_account_script_function(
            AuthenticationKey::from_preimage(&preimage).derived_address(),
            preimage.into_vec(),
        ))
    }

    pub fn transfer(&self, to: AccountAddress, amount: u64) -> TransactionBuilder {
        self.payload(aptos_stdlib::encode_transfer_script_function(to, amount))
    }

    pub fn mint(&self, to: AccountAddress, amount: u64) -> TransactionBuilder {
        self.payload(aptos_stdlib::encode_mint_script_function(to, amount))
    }

    //
    // Internal Helpers
    //

    pub fn script(&self, script: Script) -> TransactionBuilder {
        self.payload(TransactionPayload::Script(script))
    }

    fn transaction_builder(&self, payload: TransactionPayload) -> TransactionBuilder {
        TransactionBuilder {
            sender: None,
            sequence_number: None,
            payload,
            max_gas_amount: self.max_gas_amount,
            gas_unit_price: self.gas_unit_price,
            gas_currency_code: self.gas_currency.into(),
            expiration_timestamp_secs: self.expiration_timestamp(),
            chain_id: self.chain_id,
        }
    }

    fn expiration_timestamp(&self) -> u64 {
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs()
            + self.transaction_expiration_time
    }
}

pub struct DualAttestationMessage {
    message: Box<[u8]>,
}

impl DualAttestationMessage {
    pub fn new<M: Into<Vec<u8>>>(metadata: M, reciever: AccountAddress, amount: u64) -> Self {
        let mut message = metadata.into();
        bcs::serialize_into(&mut message, &reciever).unwrap();
        bcs::serialize_into(&mut message, &amount).unwrap();
        message.extend(b"@@$$DIEM_ATTEST$$@@");

        Self {
            message: message.into(),
        }
    }

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

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(rename_all = "UPPERCASE")]
pub enum Currency {
    XDX,
    XUS,
}

impl Currency {
    pub fn as_str(&self) -> &str {
        match self {
            Currency::XDX => XDX_NAME,
            Currency::XUS => XUS_NAME,
        }
    }

    pub fn type_tag(&self) -> TypeTag {
        match self {
            Currency::XDX => xdx_type_tag(),
            Currency::XUS => xus_tag(),
        }
    }
}

impl PartialEq<str> for Currency {
    fn eq(&self, other: &str) -> bool {
        self.as_str().eq(other)
    }
}

impl PartialEq<Currency> for str {
    fn eq(&self, other: &Currency) -> bool {
        other.as_str().eq(self)
    }
}

impl PartialEq<String> for Currency {
    fn eq(&self, other: &String) -> bool {
        self.as_str().eq(other)
    }
}

impl PartialEq<Currency> for String {
    fn eq(&self, other: &Currency) -> bool {
        other.as_str().eq(self)
    }
}

impl From<Currency> for String {
    fn from(currency: Currency) -> Self {
        currency.as_str().to_owned()
    }
}

impl fmt::Display for Currency {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.as_str())
    }
}
<commit_msg>[Address32][2/3] replace script with script function<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::{
    move_types::{account_address::AccountAddress, language_storage::TypeTag},
    types::{
        account_config::{xdx_type_tag, xus_tag, XDX_NAME, XUS_NAME},
        chain_id::ChainId,
        transaction::{authenticator::AuthenticationKey, RawTransaction, TransactionPayload},
    },
};
use aptos_crypto::ed25519::Ed25519PublicKey;
use serde::{Deserialize, Serialize};
use std::fmt;

pub use aptos_transaction_builder::{aptos_stdlib, stdlib};
use aptos_types::transaction::{
    authenticator::AuthenticationKeyPreimage, ChangeSet, ModuleBundle, Script, ScriptFunction,
    WriteSetPayload,
};

pub struct TransactionBuilder {
    sender: Option<AccountAddress>,
    sequence_number: Option<u64>,
    payload: TransactionPayload,
    max_gas_amount: u64,
    gas_unit_price: u64,
    gas_currency_code: String,
    expiration_timestamp_secs: u64,
    chain_id: ChainId,
}

impl TransactionBuilder {
    pub fn sender(mut self, sender: AccountAddress) -> Self {
        self.sender = Some(sender);
        self
    }

    pub fn sequence_number(mut self, sequence_number: u64) -> Self {
        self.sequence_number = Some(sequence_number);
        self
    }

    pub fn max_gas_amount(mut self, max_gas_amount: u64) -> Self {
        self.max_gas_amount = max_gas_amount;
        self
    }

    pub fn gas_unit_price(mut self, gas_unit_price: u64) -> Self {
        self.gas_unit_price = gas_unit_price;
        self
    }

    pub fn gas_currency_code<T: Into<String>>(mut self, gas_currency_code: T) -> Self {
        self.gas_currency_code = gas_currency_code.into();
        self
    }

    pub fn chain_id(mut self, chain_id: ChainId) -> Self {
        self.chain_id = chain_id;
        self
    }

    pub fn expiration_timestamp_secs(mut self, expiration_timestamp_secs: u64) -> Self {
        self.expiration_timestamp_secs = expiration_timestamp_secs;
        self
    }

    pub fn build(self) -> RawTransaction {
        RawTransaction::new(
            self.sender.expect("sender must have been set"),
            self.sequence_number
                .expect("sequence number must have been set"),
            self.payload,
            self.max_gas_amount,
            self.gas_unit_price,
            self.gas_currency_code,
            self.expiration_timestamp_secs,
            self.chain_id,
        )
    }
}

#[derive(Clone, Debug)]
pub struct TransactionFactory {
    max_gas_amount: u64,
    gas_unit_price: u64,
    gas_currency: Currency,
    transaction_expiration_time: u64,
    chain_id: ChainId,
}

impl TransactionFactory {
    pub fn new(chain_id: ChainId) -> Self {
        Self {
            max_gas_amount: 1_000,
            gas_unit_price: 0,
            gas_currency: Currency::XUS,
            transaction_expiration_time: 30,
            chain_id,
        }
    }

    pub fn with_max_gas_amount(mut self, max_gas_amount: u64) -> Self {
        self.max_gas_amount = max_gas_amount;
        self
    }

    pub fn with_gas_unit_price(mut self, gas_unit_price: u64) -> Self {
        self.gas_unit_price = gas_unit_price;
        self
    }

    pub fn with_transaction_expiration_time(mut self, transaction_expiration_time: u64) -> Self {
        self.transaction_expiration_time = transaction_expiration_time;
        self
    }

    pub fn with_chain_id(mut self, chain_id: ChainId) -> Self {
        self.chain_id = chain_id;
        self
    }

    pub fn payload(&self, payload: TransactionPayload) -> TransactionBuilder {
        self.transaction_builder(payload)
    }

    pub fn module(&self, code: Vec<u8>) -> TransactionBuilder {
        self.payload(TransactionPayload::ModuleBundle(ModuleBundle::singleton(
            code,
        )))
    }

    pub fn change_set(&self, change_set: ChangeSet) -> TransactionBuilder {
        self.payload(TransactionPayload::WriteSet(WriteSetPayload::Direct(
            change_set,
        )))
    }

    pub fn script_function(&self, func: ScriptFunction) -> TransactionBuilder {
        self.payload(TransactionPayload::ScriptFunction(func))
    }

<<<<<<< HEAD
    pub fn create_user_account(&self, public_key: &Ed25519PublicKey) -> TransactionBuilder {
        let preimage = AuthenticationKeyPreimage::ed25519(public_key);
        self.payload(aptos_stdlib::encode_create_account_script_function(
            AuthenticationKey::from_preimage(&preimage).derived_address(),
            preimage.into_vec(),
        ))
    }

    pub fn transfer(&self, to: AccountAddress, amount: u64) -> TransactionBuilder {
        self.payload(aptos_stdlib::encode_transfer_script_function(to, amount))
    }

    pub fn mint(&self, to: AccountAddress, amount: u64) -> TransactionBuilder {
        self.payload(aptos_stdlib::encode_mint_script_function(to, amount))
    }

    //
    // Internal Helpers
    //

    pub fn script(&self, script: Script) -> TransactionBuilder {
        self.payload(TransactionPayload::Script(script))
    }

    fn transaction_builder(&self, payload: TransactionPayload) -> TransactionBuilder {
        TransactionBuilder {
            sender: None,
            sequence_number: None,
            payload,
            max_gas_amount: self.max_gas_amount,
            gas_unit_price: self.gas_unit_price,
            gas_currency_code: self.gas_currency.into(),
            expiration_timestamp_secs: self.expiration_timestamp(),
            chain_id: self.chain_id,
        }
    }

    fn expiration_timestamp(&self) -> u64 {
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs()
            + self.transaction_expiration_time
    }
}

pub struct DualAttestationMessage {
    message: Box<[u8]>,
}

impl DualAttestationMessage {
    pub fn new<M: Into<Vec<u8>>>(metadata: M, reciever: AccountAddress, amount: u64) -> Self {
        let mut message = metadata.into();
        bcs::serialize_into(&mut message, &reciever).unwrap();
        bcs::serialize_into(&mut message, &amount).unwrap();
        message.extend(b"@@$$DIEM_ATTEST$$@@");

        Self {
            message: message.into(),
        }
    }

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

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(rename_all = "UPPERCASE")]
pub enum Currency {
    XDX,
    XUS,
}

impl Currency {
    pub fn as_str(&self) -> &str {
        match self {
            Currency::XDX => XDX_NAME,
            Currency::XUS => XUS_NAME,
        }
    }

    pub fn type_tag(&self) -> TypeTag {
        match self {
            Currency::XDX => xdx_type_tag(),
            Currency::XUS => xus_tag(),
        }
    }
}

impl PartialEq<str> for Currency {
    fn eq(&self, other: &str) -> bool {
        self.as_str().eq(other)
    }
}

impl PartialEq<Currency> for str {
    fn eq(&self, other: &Currency) -> bool {
        other.as_str().eq(self)
    }
}

impl PartialEq<String> for Currency {
    fn eq(&self, other: &String) -> bool {
        self.as_str().eq(other)
    }
}

impl PartialEq<Currency> for String {
    fn eq(&self, other: &Currency) -> bool {
        other.as_str().eq(self)
    }
}

impl From<Currency> for String {
    fn from(currency: Currency) -> Self {
        currency.as_str().to_owned()
    }
}

impl fmt::Display for Currency {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.as_str())
    }
}
<|endoftext|>


<commit_before>use crate::build::scope::BreakableTarget;
use crate::build::{BlockAnd, BlockAndExtension, BlockFrame, Builder};
use crate::thir::*;
use rustc_middle::middle::region;
use rustc_middle::mir::*;

impl<'a, 'tcx> Builder<'a, 'tcx> {
    /// Builds a block of MIR statements to evaluate the THIR `expr`.
    /// If the original expression was an AST statement,
    /// (e.g., `some().code(&here());`) then `opt_stmt_span` is the
    /// span of that statement (including its semicolon, if any).
    /// The scope is used if a statement temporary must be dropped.
    crate fn stmt_expr(
        &mut self,
        mut block: BasicBlock,
        expr: &Expr<'_, 'tcx>,
        statement_scope: Option<region::Scope>,
    ) -> BlockAnd<()> {
        let this = self;
        let expr_span = expr.span;
        let source_info = this.source_info(expr.span);
        // Handle a number of expressions that don't need a destination at all. This
        // avoids needing a mountain of temporary `()` variables.
        match expr.kind {
            ExprKind::Scope { region_scope, lint_level, value } => {
                this.in_scope((region_scope, source_info), lint_level, |this| {
                    this.stmt_expr(block, value, statement_scope)
                })
            }
            ExprKind::Assign { lhs, rhs } => {
                let lhs_span = lhs.span;

                // Note: we evaluate assignments right-to-left. This
                // is better for borrowck interaction with overloaded
                // operators like x[j] = x[i].

                debug!("stmt_expr Assign block_context.push(SubExpr) : {:?}", expr);
                this.block_context.push(BlockFrame::SubExpr);

                // Generate better code for things that don't need to be
                // dropped.
                if lhs.ty.needs_drop(this.tcx, this.param_env) {
                    let rhs = unpack!(block = this.as_local_operand(block, rhs));
                    let lhs = unpack!(block = this.as_place(block, lhs));
                    unpack!(block = this.build_drop_and_replace(block, lhs_span, lhs, rhs));
                } else {
                    let rhs = unpack!(block = this.as_local_rvalue(block, rhs));
                    let lhs = unpack!(block = this.as_place(block, lhs));
                    this.cfg.push_assign(block, source_info, lhs, rhs);
                }

                this.block_context.pop();
                block.unit()
            }
            ExprKind::AssignOp { op, lhs, rhs } => {
                // FIXME(#28160) there is an interesting semantics
                // question raised here -- should we "freeze" the
                // value of the lhs here?  I'm inclined to think not,
                // since it seems closer to the semantics of the
                // overloaded version, which takes `&mut self`.  This
                // only affects weird things like `x += {x += 1; x}`
                // -- is that equal to `x + (x + 1)` or `2*(x+1)`?

                let lhs_ty = lhs.ty;

                debug!("stmt_expr AssignOp block_context.push(SubExpr) : {:?}", expr);
                this.block_context.push(BlockFrame::SubExpr);

                // As above, RTL.
                let rhs = unpack!(block = this.as_local_operand(block, rhs));
                let lhs = unpack!(block = this.as_place(block, lhs));

                // we don't have to drop prior contents or anything
                // because AssignOp is only legal for Copy types
                // (overloaded ops should be desugared into a call).
                let result = unpack!(
                    block =
                        this.build_binary_op(block, op, expr_span, lhs_ty, Operand::Copy(lhs), rhs)
                );
                this.cfg.push_assign(block, source_info, lhs, result);

                this.block_context.pop();
                block.unit()
            }
            ExprKind::Continue { label } => {
                this.break_scope(block, None, BreakableTarget::Continue(label), source_info)
            }
            ExprKind::Break { label, value } => this.break_scope(
                block,
                value.as_deref(),
                BreakableTarget::Break(label),
                source_info,
            ),
            ExprKind::Return { value } => {
                this.break_scope(block, value.as_deref(), BreakableTarget::Return, source_info)
            }
            ExprKind::LlvmInlineAsm { asm, outputs, inputs } => {
                debug!("stmt_expr LlvmInlineAsm block_context.push(SubExpr) : {:?}", expr);
                this.block_context.push(BlockFrame::SubExpr);
                let outputs = outputs
                    .into_iter()
                    .map(|output| unpack!(block = this.as_place(block, &output)))
                    .collect::<Vec<_>>()
                    .into_boxed_slice();
                let inputs = inputs
                    .into_iter()
                    .map(|input| {
                        (input.span, unpack!(block = this.as_local_operand(block, &input)))
                    })
                    .collect::<Vec<_>>()
                    .into_boxed_slice();
                this.cfg.push(
                    block,
                    Statement {
                        source_info,
                        kind: StatementKind::LlvmInlineAsm(box LlvmInlineAsm {
                            asm: asm.clone(),
                            outputs,
                            inputs,
                        }),
                    },
                );
                this.block_context.pop();
                block.unit()
            }
            _ => {
                assert!(
                    statement_scope.is_some(),
                    "Should not be calling `stmt_expr` on a general expression \
                     without a statement scope",
                );

                // Issue #54382: When creating temp for the value of
                // expression like:
                //
                // `{ side_effects(); { let l = stuff(); the_value } }`
                //
                // it is usually better to focus on `the_value` rather
                // than the entirety of block(s) surrounding it.
                let adjusted_span = (|| {
                    if let ExprKind::Block { body } = &expr.kind {
                        if let Some(tail_expr) = &body.expr {
                            let mut expr = &*tail_expr;
                            while let ExprKind::Block { body: subblock } = &expr.kind {
                                if let Some(subtail_expr) = &subblock.expr {
                                    expr = subtail_expr
                                } else {
                                    break;
                                }
                            }
                            this.block_context.push(BlockFrame::TailExpr {
                                tail_result_is_ignored: true,
                                span: expr.span,
                            });
                            return Some(expr.span);
                        }
                    }
                    None
                })();

                let temp =
                    unpack!(block = this.as_temp(block, statement_scope, expr, Mutability::Not));

                if let Some(span) = adjusted_span {
                    this.local_decls[temp].source_info.span = span;
                    this.block_context.pop();
                }

                block.unit()
            }
        }
    }
}
<commit_msg>Take into account `ExprKind::Scope` when calculating span of temporary<commit_after>use crate::build::scope::BreakableTarget;
use crate::build::{BlockAnd, BlockAndExtension, BlockFrame, Builder};
use crate::thir::*;
use rustc_middle::middle::region;
use rustc_middle::mir::*;

impl<'a, 'tcx> Builder<'a, 'tcx> {
    /// Builds a block of MIR statements to evaluate the THIR `expr`.
    /// If the original expression was an AST statement,
    /// (e.g., `some().code(&here());`) then `opt_stmt_span` is the
    /// span of that statement (including its semicolon, if any).
    /// The scope is used if a statement temporary must be dropped.
    crate fn stmt_expr(
        &mut self,
        mut block: BasicBlock,
        expr: &Expr<'_, 'tcx>,
        statement_scope: Option<region::Scope>,
    ) -> BlockAnd<()> {
        let this = self;
        let expr_span = expr.span;
        let source_info = this.source_info(expr.span);
        // Handle a number of expressions that don't need a destination at all. This
        // avoids needing a mountain of temporary `()` variables.
        match expr.kind {
            ExprKind::Scope { region_scope, lint_level, value } => {
                this.in_scope((region_scope, source_info), lint_level, |this| {
                    this.stmt_expr(block, value, statement_scope)
                })
            }
            ExprKind::Assign { lhs, rhs } => {
                let lhs_span = lhs.span;

                // Note: we evaluate assignments right-to-left. This
                // is better for borrowck interaction with overloaded
                // operators like x[j] = x[i].

                debug!("stmt_expr Assign block_context.push(SubExpr) : {:?}", expr);
                this.block_context.push(BlockFrame::SubExpr);

                // Generate better code for things that don't need to be
                // dropped.
                if lhs.ty.needs_drop(this.tcx, this.param_env) {
                    let rhs = unpack!(block = this.as_local_operand(block, rhs));
                    let lhs = unpack!(block = this.as_place(block, lhs));
                    unpack!(block = this.build_drop_and_replace(block, lhs_span, lhs, rhs));
                } else {
                    let rhs = unpack!(block = this.as_local_rvalue(block, rhs));
                    let lhs = unpack!(block = this.as_place(block, lhs));
                    this.cfg.push_assign(block, source_info, lhs, rhs);
                }

                this.block_context.pop();
                block.unit()
            }
            ExprKind::AssignOp { op, lhs, rhs } => {
                // FIXME(#28160) there is an interesting semantics
                // question raised here -- should we "freeze" the
                // value of the lhs here?  I'm inclined to think not,
                // since it seems closer to the semantics of the
                // overloaded version, which takes `&mut self`.  This
                // only affects weird things like `x += {x += 1; x}`
                // -- is that equal to `x + (x + 1)` or `2*(x+1)`?

                let lhs_ty = lhs.ty;

                debug!("stmt_expr AssignOp block_context.push(SubExpr) : {:?}", expr);
                this.block_context.push(BlockFrame::SubExpr);

                // As above, RTL.
                let rhs = unpack!(block = this.as_local_operand(block, rhs));
                let lhs = unpack!(block = this.as_place(block, lhs));

                // we don't have to drop prior contents or anything
                // because AssignOp is only legal for Copy types
                // (overloaded ops should be desugared into a call).
                let result = unpack!(
                    block =
                        this.build_binary_op(block, op, expr_span, lhs_ty, Operand::Copy(lhs), rhs)
                );
                this.cfg.push_assign(block, source_info, lhs, result);

                this.block_context.pop();
                block.unit()
            }
            ExprKind::Continue { label } => {
                this.break_scope(block, None, BreakableTarget::Continue(label), source_info)
            }
            ExprKind::Break { label, value } => this.break_scope(
                block,
                value.as_deref(),
                BreakableTarget::Break(label),
                source_info,
            ),
            ExprKind::Return { value } => {
                this.break_scope(block, value.as_deref(), BreakableTarget::Return, source_info)
            }
            ExprKind::LlvmInlineAsm { asm, outputs, inputs } => {
                debug!("stmt_expr LlvmInlineAsm block_context.push(SubExpr) : {:?}", expr);
                this.block_context.push(BlockFrame::SubExpr);
                let outputs = outputs
                    .into_iter()
                    .map(|output| unpack!(block = this.as_place(block, &output)))
                    .collect::<Vec<_>>()
                    .into_boxed_slice();
                let inputs = inputs
                    .into_iter()
                    .map(|input| {
                        (input.span, unpack!(block = this.as_local_operand(block, &input)))
                    })
                    .collect::<Vec<_>>()
                    .into_boxed_slice();
                this.cfg.push(
                    block,
                    Statement {
                        source_info,
                        kind: StatementKind::LlvmInlineAsm(box LlvmInlineAsm {
                            asm: asm.clone(),
                            outputs,
                            inputs,
                        }),
                    },
                );
                this.block_context.pop();
                block.unit()
            }
            _ => {
                assert!(
                    statement_scope.is_some(),
                    "Should not be calling `stmt_expr` on a general expression \
                     without a statement scope",
                );

                // Issue #54382: When creating temp for the value of
                // expression like:
                //
                // `{ side_effects(); { let l = stuff(); the_value } }`
                //
                // it is usually better to focus on `the_value` rather
                // than the entirety of block(s) surrounding it.
                let adjusted_span = (|| {
                    if let ExprKind::Block { body } = &expr.kind {
                        if let Some(tail_expr) = &body.expr {
                            let mut expr = &*tail_expr;
                            while let ExprKind::Block {
                                body: Block { expr: Some(nested_expr), .. },
                            }
                            | ExprKind::Scope { value: nested_expr, .. } = &expr.kind
                            {
                                expr = nested_expr;
                            }
                            this.block_context.push(BlockFrame::TailExpr {
                                tail_result_is_ignored: true,
                                span: expr.span,
                            });
                            return Some(expr.span);
                        }
                    }
                    None
                })();

                let temp =
                    unpack!(block = this.as_temp(block, statement_scope, expr, Mutability::Not));

                if let Some(span) = adjusted_span {
                    this.local_decls[temp].source_info.span = span;
                    this.block_context.pop();
                }

                block.unit()
            }
        }
    }
}
<|endoftext|>


<commit_before>use rand;
use rand::distributions::{IndependentSample, Range};

use std::{cmp, thread, time};
use std::fmt::Debug;

#[derive(Debug, PartialEq, Eq)]
pub enum RetryAction {
    OKAY,
    WAIT
}

pub trait RetryPolicy: Debug + PartialEq + Eq {
    fn max_tries(&self) -> usize;
    fn current_tries(&self) -> usize;

    fn fail(&mut self);
    fn succeed(&mut self);

    fn can_try(&self) -> Option<RetryAction> {
        if self.current_tries() >= self.max_tries() {
            None
        } else {
            Some(RetryAction::OKAY)
        }
    }

    fn is_down(&self) -> bool;
}

#[derive(Debug, PartialEq, Eq)]
pub enum RetryPolicyWrapper {
    ExponentialBackoff(ExponentialBackoffPolicy)
}

#[derive(Debug, PartialEq, Eq)]
pub struct ExponentialBackoffPolicy {
    max_tries: usize,
    current_tries: usize,
    last_try: time::Instant,
    wait: time::Duration
}

impl ExponentialBackoffPolicy {
    pub fn new(max_tries: usize) -> Self {
        ExponentialBackoffPolicy {
            max_tries,
            current_tries: 0,
            last_try: time::Instant::now(),
            wait: time::Duration::default()
        }
    }
}

impl RetryPolicy for ExponentialBackoffPolicy {
    fn max_tries(&self) -> usize {
        self.max_tries
    }

    fn current_tries(&self) -> usize {
        self.current_tries
    }

    fn fail(&mut self) {
        if self.last_try.elapsed().lt(&self.wait) {
          //we're already in back off
          return;
        }

        let max_secs = cmp::max(1, 1 << self.current_tries);
        let wait = if max_secs == 1 {
            1
        } else {
            let mut rng = rand::thread_rng();
            let range = Range::new(1, max_secs);
            range.ind_sample(&mut rng)
        };

        self.wait = time::Duration::from_secs(wait);
        self.last_try = time::Instant::now();
        self.current_tries += 1;
    }

    fn succeed(&mut self) {
        self.wait = time::Duration::default();
        self.last_try = time::Instant::now();
        self.current_tries = 0;
    }

    fn can_try(&self) -> Option<RetryAction> {
        if self.current_tries() >= self.max_tries() {
            return None;
        }

        let action = if self.last_try.elapsed().gt(&self.wait) {
            RetryAction::OKAY
        } else {
            RetryAction::WAIT
        };

        Some(action)
    }

    fn is_down(&self) -> bool {
      self.current_tries() >= self.max_tries()
    }
}

impl Into<RetryPolicyWrapper> for ExponentialBackoffPolicy {
    fn into(self) -> RetryPolicyWrapper {
        RetryPolicyWrapper::ExponentialBackoff(self)
    }
}

impl RetryPolicy for RetryPolicyWrapper {
    fn max_tries(&self) -> usize {
        match *self {
            RetryPolicyWrapper::ExponentialBackoff(ref policy) => policy
        }.max_tries()
    }

    fn current_tries(&self) -> usize {
        match *self {
            RetryPolicyWrapper::ExponentialBackoff(ref policy) => policy
        }.current_tries()
    }

    fn fail(&mut self) {
        match *self {
            RetryPolicyWrapper::ExponentialBackoff(ref mut policy) => policy
        }.fail()
    }

    fn succeed(&mut self) {
        match *self {
            RetryPolicyWrapper::ExponentialBackoff(ref mut policy) => policy
        }.succeed()
    }

    fn can_try(&self) -> Option<RetryAction> {
        match *self {
            RetryPolicyWrapper::ExponentialBackoff(ref policy) => policy
        }.can_try()
    }

    fn is_down(&self) -> bool {
        match *self {
            RetryPolicyWrapper::ExponentialBackoff(ref policy) => policy
        }.is_down()
    }
}

#[cfg(test)]
mod tests {
    use std::time;

    use super::{RetryAction, RetryPolicy, ExponentialBackoffPolicy};

    const MAX_FAILS: usize = 10;

    #[test]
    fn no_fail() {
        let policy = ExponentialBackoffPolicy::new(MAX_FAILS);
        let can_try = policy.can_try();

        assert_eq!(Some(RetryAction::OKAY), can_try)
    }

    #[test]
    fn single_fail() {
        let mut policy = ExponentialBackoffPolicy::new(MAX_FAILS);
        policy.fail();
        let can_try = policy.can_try();

        // The wait will be >= 1s, so we'll be WAIT by the time we do the assert
        assert_eq!(Some(RetryAction::WAIT), can_try)
    }

    #[test]
    fn max_fails() {
        let mut policy = ExponentialBackoffPolicy::new(MAX_FAILS);

        for _ in 0..MAX_FAILS {
            policy.fail();
        }

        let can_try = policy.can_try();

        assert_eq!(None, can_try)
    }

    #[test]
    fn recover_from_fail() {
        let mut policy = ExponentialBackoffPolicy::new(MAX_FAILS);

        // Stop just before total failure
        for _ in 0..(MAX_FAILS - 1) {
            policy.fail();
        }

        policy.succeed();
        policy.fail();
        policy.fail();
        policy.fail();

        let can_try = policy.can_try();

        assert_eq!(Some(RetryAction::WAIT), can_try)
    }
}
<commit_msg>avoid panic when overflowing with current_tries<commit_after>use rand;
use rand::distributions::{IndependentSample, Range};

use std::{cmp, thread, time};
use std::fmt::Debug;

#[derive(Debug, PartialEq, Eq)]
pub enum RetryAction {
    OKAY,
    WAIT
}

pub trait RetryPolicy: Debug + PartialEq + Eq {
    fn max_tries(&self) -> usize;
    fn current_tries(&self) -> usize;

    fn fail(&mut self);
    fn succeed(&mut self);

    fn can_try(&self) -> Option<RetryAction> {
        if self.current_tries() >= self.max_tries() {
            None
        } else {
            Some(RetryAction::OKAY)
        }
    }

    fn is_down(&self) -> bool;
}

#[derive(Debug, PartialEq, Eq)]
pub enum RetryPolicyWrapper {
    ExponentialBackoff(ExponentialBackoffPolicy)
}

#[derive(Debug, PartialEq, Eq)]
pub struct ExponentialBackoffPolicy {
    max_tries: usize,
    current_tries: usize,
    last_try: time::Instant,
    wait: time::Duration
}

impl ExponentialBackoffPolicy {
    pub fn new(max_tries: usize) -> Self {
        ExponentialBackoffPolicy {
            max_tries,
            current_tries: 0,
            last_try: time::Instant::now(),
            wait: time::Duration::default()
        }
    }
}

impl RetryPolicy for ExponentialBackoffPolicy {
    fn max_tries(&self) -> usize {
        self.max_tries
    }

    fn current_tries(&self) -> usize {
        self.current_tries
    }

    fn fail(&mut self) {
        if self.last_try.elapsed().lt(&self.wait) {
          //we're already in back off
          return;
        }

        let max_secs = cmp::max(1, 1u64.wrapping_shl(self.current_tries as u32));
        let wait = if max_secs == 1 {
            1
        } else {
            let mut rng = rand::thread_rng();
            let range = Range::new(1, max_secs);
            range.ind_sample(&mut rng)
        };

        self.wait = time::Duration::from_secs(wait);
        self.last_try = time::Instant::now();
        self.current_tries += 1;
    }

    fn succeed(&mut self) {
        self.wait = time::Duration::default();
        self.last_try = time::Instant::now();
        self.current_tries = 0;
    }

    fn can_try(&self) -> Option<RetryAction> {
        if self.current_tries() >= self.max_tries() {
            return None;
        }

        let action = if self.last_try.elapsed().gt(&self.wait) {
            RetryAction::OKAY
        } else {
            RetryAction::WAIT
        };

        Some(action)
    }

    fn is_down(&self) -> bool {
      self.current_tries() >= self.max_tries()
    }
}

impl Into<RetryPolicyWrapper> for ExponentialBackoffPolicy {
    fn into(self) -> RetryPolicyWrapper {
        RetryPolicyWrapper::ExponentialBackoff(self)
    }
}

impl RetryPolicy for RetryPolicyWrapper {
    fn max_tries(&self) -> usize {
        match *self {
            RetryPolicyWrapper::ExponentialBackoff(ref policy) => policy
        }.max_tries()
    }

    fn current_tries(&self) -> usize {
        match *self {
            RetryPolicyWrapper::ExponentialBackoff(ref policy) => policy
        }.current_tries()
    }

    fn fail(&mut self) {
        match *self {
            RetryPolicyWrapper::ExponentialBackoff(ref mut policy) => policy
        }.fail()
    }

    fn succeed(&mut self) {
        match *self {
            RetryPolicyWrapper::ExponentialBackoff(ref mut policy) => policy
        }.succeed()
    }

    fn can_try(&self) -> Option<RetryAction> {
        match *self {
            RetryPolicyWrapper::ExponentialBackoff(ref policy) => policy
        }.can_try()
    }

    fn is_down(&self) -> bool {
        match *self {
            RetryPolicyWrapper::ExponentialBackoff(ref policy) => policy
        }.is_down()
    }
}

#[cfg(test)]
mod tests {
    use std::time;

    use super::{RetryAction, RetryPolicy, ExponentialBackoffPolicy};

    const MAX_FAILS: usize = 10;

    #[test]
    fn no_fail() {
        let policy = ExponentialBackoffPolicy::new(MAX_FAILS);
        let can_try = policy.can_try();

        assert_eq!(Some(RetryAction::OKAY), can_try)
    }

    #[test]
    fn single_fail() {
        let mut policy = ExponentialBackoffPolicy::new(MAX_FAILS);
        policy.fail();
        let can_try = policy.can_try();

        // The wait will be >= 1s, so we'll be WAIT by the time we do the assert
        assert_eq!(Some(RetryAction::WAIT), can_try)
    }

    #[test]
    fn max_fails() {
        let mut policy = ExponentialBackoffPolicy::new(MAX_FAILS);

        for _ in 0..MAX_FAILS {
            policy.fail();
        }

        let can_try = policy.can_try();

        assert_eq!(None, can_try)
    }

    #[test]
    fn recover_from_fail() {
        let mut policy = ExponentialBackoffPolicy::new(MAX_FAILS);

        // Stop just before total failure
        for _ in 0..(MAX_FAILS - 1) {
            policy.fail();
        }

        policy.succeed();
        policy.fail();
        policy.fail();
        policy.fail();

        let can_try = policy.can_try();

        assert_eq!(Some(RetryAction::WAIT), can_try)
    }
}
<|endoftext|>


<commit_before>use std::boxed::{Box, FnBox};
use std::sync::{Arc, Mutex, Condvar};
use std::fmt;
use std::time::Duration;

use mio;

use raftserver::{Result, send_msg, Error};
use kvproto::raft_serverpb::RaftMessage;
use kvproto::raft_cmdpb::{RaftCmdRequest, RaftCmdResponse};

pub type Callback = Box<FnBox(RaftCmdResponse) -> Result<()> + Send>;

#[derive(Debug)]
pub enum Tick {
    Raft,
    RaftLogGc,
    SplitRegionCheck,
    ReplicaCheck,
}

pub enum Msg {
    Quit,

    // For notify.
    RaftMessage(RaftMessage),
    RaftCmd {
        request: RaftCmdRequest,
        callback: Callback,
    },

    // For split check
    SplitCheckResult {
        region_id: u64,
        split_key: Vec<u8>,
    },
}

impl fmt::Debug for Msg {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Msg::Quit => write!(fmt, "Quit"),
            Msg::RaftMessage(_) => write!(fmt, "Raft Message"),
            Msg::RaftCmd{..} => write!(fmt, "Raft Command"),
            Msg::SplitCheckResult{..} => write!(fmt, "Split Check Result"),
        }
    }
}

// Send the request and wait the response until timeout.
// Use Condvar to support call timeout. if timeout, return None.
// We should know that even timeout happens, the command may still
// be handled in store later.
pub fn call_command(sendch: &SendCh,
                    request: RaftCmdRequest,
                    timeout: Duration)
                    -> Result<RaftCmdResponse> {
    let resp: Option<RaftCmdResponse> = None;
    let pair = Arc::new((Mutex::new(resp), Condvar::new()));
    let pair2 = pair.clone();

    try!(sendch.send(Msg::RaftCmd {
        request: request,
        callback: Box::new(move |resp: RaftCmdResponse| -> Result<()> {
            let &(ref lock, ref cvar) = &*pair2;
            let mut v = lock.lock().unwrap();
            *v = Some(resp);
            cvar.notify_one();
            Ok(())
        }),
    }));

    let &(ref lock, ref cvar) = &*pair;
    let mut v = lock.lock().unwrap();
    while v.is_none() {
        let (resp, timeout_res) = cvar.wait_timeout(v, timeout).unwrap();
        if timeout_res.timed_out() {
            return Err(Error::Timeout(format!("request timeout for {:?}", timeout)));
        }

        v = resp
    }

    Ok(v.take().unwrap())
}

#[derive(Debug)]
pub struct SendCh {
    ch: mio::Sender<Msg>,
}

impl Clone for SendCh {
    fn clone(&self) -> SendCh {
        SendCh { ch: self.ch.clone() }
    }
}

impl SendCh {
    pub fn new(ch: mio::Sender<Msg>) -> SendCh {
        SendCh { ch: ch }
    }

    pub fn send(&self, msg: Msg) -> Result<()> {
        try!(send_msg(&self.ch, msg));
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::thread;
    use std::sync::mpsc::channel;
    use std::time::Duration;

    use mio::{EventLoop, Handler};

    use super::*;
    use kvproto::raft_cmdpb::{RaftCmdRequest, RaftCmdResponse};
    use raftserver::{Result, Error};

    struct TestHandler;

    impl Handler for TestHandler {
        type Timeout = ();
        type Message = Msg;

        fn notify(&mut self, event_loop: &mut EventLoop<Self>, msg: Self::Message) {
            match msg {
                Msg::Quit => event_loop.shutdown(),
                Msg::RaftCmd{callback, request} => {
                    // a trick for test timeout.
                    if request.get_header().get_region_id() == u64::max_value() {
                        thread::sleep(Duration::from_millis(100));
                    }
                    callback.call_box((RaftCmdResponse::new(),)).unwrap()
                }
                // we only test above message types, others panic.
                _ => unreachable!(),
            }
        }
    }

    #[test]
    fn test_sender() {
        let mut event_loop = EventLoop::new().unwrap();
        let sendch = &SendCh::new(event_loop.channel());

        let t = thread::spawn(move || {
            event_loop.run(&mut TestHandler).unwrap();
        });

        let (tx, rx) = channel();
        let cmd = Msg::RaftCmd {
            request: RaftCmdRequest::new(),
            callback: Box::new(move |_: RaftCmdResponse| -> Result<()> {
                tx.send(1).unwrap();
                Ok(())
            }),
        };
        sendch.send(cmd).unwrap();

        rx.recv().unwrap();

        let mut request = RaftCmdRequest::new();
        request.mut_header().set_region_id(u64::max_value());
        assert!(call_command(sendch, request.clone(), Duration::from_millis(500)).is_ok());
        match call_command(sendch, request.clone(), Duration::from_millis(10)) {
            Err(Error::Timeout(_)) => {}
            _ => panic!("should failed with timeout"),
        }

        sendch.send(Msg::Quit).unwrap();

        t.join().unwrap();
    }
}
<commit_msg>Tiny clean up<commit_after>use std::boxed::{Box, FnBox};
use std::sync::{Arc, Mutex, Condvar};
use std::fmt;
use std::time::Duration;

use mio;

use raftserver::{Result, send_msg, Error};
use kvproto::raft_serverpb::RaftMessage;
use kvproto::raft_cmdpb::{RaftCmdRequest, RaftCmdResponse};

pub type Callback = Box<FnBox(RaftCmdResponse) -> Result<()> + Send>;

#[derive(Debug)]
pub enum Tick {
    Raft,
    RaftLogGc,
    SplitRegionCheck,
    ReplicaCheck,
}

pub enum Msg {
    Quit,

    // For notify.
    RaftMessage(RaftMessage),
    RaftCmd {
        request: RaftCmdRequest,
        callback: Callback,
    },

    // For split check
    SplitCheckResult {
        region_id: u64,
        split_key: Vec<u8>,
    },
}

impl fmt::Debug for Msg {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Msg::Quit => write!(fmt, "Quit"),
            Msg::RaftMessage(_) => write!(fmt, "Raft Message"),
            Msg::RaftCmd{..} => write!(fmt, "Raft Command"),
            Msg::SplitCheckResult{..} => write!(fmt, "Split Check Result"),
        }
    }
}

// Send the request and wait the response until timeout.
// Use Condvar to support call timeout. if timeout, return None.
// We should know that even timeout happens, the command may still
// be handled in store later.
pub fn call_command(sendch: &SendCh,
                    request: RaftCmdRequest,
                    timeout: Duration)
                    -> Result<RaftCmdResponse> {
    let resp: Option<RaftCmdResponse> = None;
    let pair = Arc::new((Mutex::new(resp), Condvar::new()));
    let pair2 = pair.clone();

    try!(sendch.send(Msg::RaftCmd {
        request: request,
        callback: box move |resp| {
            let &(ref lock, ref cvar) = &*pair2;
            let mut v = lock.lock().unwrap();
            *v = Some(resp);
            cvar.notify_one();
            Ok(())
        },
    }));

    let &(ref lock, ref cvar) = &*pair;
    let mut v = lock.lock().unwrap();
    while v.is_none() {
        let (resp, timeout_res) = cvar.wait_timeout(v, timeout).unwrap();
        if timeout_res.timed_out() {
            return Err(Error::Timeout(format!("request timeout for {:?}", timeout)));
        }

        v = resp
    }

    Ok(v.take().unwrap())
}

#[derive(Debug)]
pub struct SendCh {
    ch: mio::Sender<Msg>,
}

impl Clone for SendCh {
    fn clone(&self) -> SendCh {
        SendCh { ch: self.ch.clone() }
    }
}

impl SendCh {
    pub fn new(ch: mio::Sender<Msg>) -> SendCh {
        SendCh { ch: ch }
    }

    pub fn send(&self, msg: Msg) -> Result<()> {
        try!(send_msg(&self.ch, msg));
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::thread;
    use std::sync::mpsc::channel;
    use std::time::Duration;

    use mio::{EventLoop, Handler};

    use super::*;
    use kvproto::raft_cmdpb::{RaftCmdRequest, RaftCmdResponse};
    use raftserver::Error;

    struct TestHandler;

    impl Handler for TestHandler {
        type Timeout = ();
        type Message = Msg;

        fn notify(&mut self, event_loop: &mut EventLoop<Self>, msg: Self::Message) {
            match msg {
                Msg::Quit => event_loop.shutdown(),
                Msg::RaftCmd{callback, request} => {
                    // a trick for test timeout.
                    if request.get_header().get_region_id() == u64::max_value() {
                        thread::sleep(Duration::from_millis(100));
                    }
                    callback.call_box((RaftCmdResponse::new(),)).unwrap()
                }
                // we only test above message types, others panic.
                _ => unreachable!(),
            }
        }
    }

    #[test]
    fn test_sender() {
        let mut event_loop = EventLoop::new().unwrap();
        let sendch = &SendCh::new(event_loop.channel());

        let t = thread::spawn(move || {
            event_loop.run(&mut TestHandler).unwrap();
        });

        let (tx, rx) = channel();
        let cmd = Msg::RaftCmd {
            request: RaftCmdRequest::new(),
            callback: box move |_| {
                tx.send(1).unwrap();
                Ok(())
            },
        };
        sendch.send(cmd).unwrap();

        rx.recv().unwrap();

        let mut request = RaftCmdRequest::new();
        request.mut_header().set_region_id(u64::max_value());
        assert!(call_command(sendch, request.clone(), Duration::from_millis(500)).is_ok());
        match call_command(sendch, request.clone(), Duration::from_millis(10)) {
            Err(Error::Timeout(_)) => {}
            _ => panic!("should failed with timeout"),
        }

        sendch.send(Msg::Quit).unwrap();

        t.join().unwrap();
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::sync::Arc;

use parking_lot::{Mutex, MutexGuard};

use cpu::irq;
use io::regs::IoReg;

#[derive(Clone, Copy, Debug)]
pub enum Prescaler {
    Div1 = 0,
    Div64 = 1,
    Div256 = 2,
    Div1024 = 3,
}

impl Prescaler {
    fn new(val: u16) -> Prescaler {
        match val {
            0 => Prescaler::Div1,
            1 => Prescaler::Div64,
            2 => Prescaler::Div256,
            3 => Prescaler::Div1024,
            _ => unreachable!()
        }
    }
}

bfdesc!(CntReg: u16, {
    prescaler: 0 => 1,
    count_up: 2 => 2,
    irq_enable: 6 => 6,
    started: 7 => 7
});

fn get_regs(dev: &mut TimerDevice, index: usize) -> (&mut IoReg<u16>, &mut IoReg<u16>) {
    match index {
        0 => (&mut dev.val0, &mut dev.cnt0),
        1 => (&mut dev.val1, &mut dev.cnt1),
        2 => (&mut dev.val2, &mut dev.cnt2),
        3 => (&mut dev.val3, &mut dev.cnt3),
        _ => unreachable!()
    }
}

fn reg_val_update(dev: &mut TimerDevice, index: usize) {
    let val = {
        let (val, _) = get_regs(dev, index);
        val.get()
    };
    let states = &dev._internal_state;
    for mut t in TimerIter::new(states).filter(|t| t.has_index(index)) {
        t.set_val_hword(index, val)
    }
}

fn reg_val_read(dev: &mut TimerDevice, index: usize) {
    let new_val = {
        let states = &dev._internal_state;
        TimerIter::new(states).filter_map(|t| t.val_hword(index))
                              .next().unwrap()
    };
    let (val, _) = get_regs(dev, index);
    val.set_unchecked(new_val);
}

fn reg_cnt_update(dev: &mut TimerDevice, index: usize) {
    let (val, cnt) = {
        let (val, cnt) = get_regs(dev, index);
        (val.get(), cnt.get())
    };
    let mut timer = {
        let states = &dev._internal_state;
        TimerIter::new(states).filter(|t| t.has_index(index))
                              .next().unwrap()
    };
    {
        let state = match timer {
            Timer::Unit(_, ref mut state) | Timer::Joined { lowstate: ref mut state, .. } => state
        };
        state.started = bf!(cnt @ CntReg::started) == 1;
        state.count_up = bf!(cnt @ CntReg::count_up) == 1;
        state.prescaler = Prescaler::new(bf!(cnt @ CntReg::prescaler));
    }

    timer.set_val_hword(index, val);
    trace!("Setting TIMER CNT{}: {:?}", index, timer);
}

enum Timer<'a> {
    Joined { bitset: u8, lowstate: MutexGuard<'a, TimerState> },
    Unit(u8, MutexGuard<'a, TimerState>)
}

impl<'a> fmt::Debug for Timer<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Timer::Unit(num, ref state) => write!(f, "Timer::Unit({}, {:?})", num, **state),
            Timer::Joined { bitset: bs, lowstate: ref state } => {
                write!(f, "Timer::Joined {{ bitset: 0b{:b}, lowstate: {:?} }}", bs, **state)
            }
        }
    }
}

impl<'a> Timer<'a> {
    fn started(&self) -> bool {
        match *self {
            Timer::Unit(_, ref state) | Timer::Joined { lowstate: ref state, .. } => state.started
        }
    }

    fn val(&self) -> u64 {
        let state = match *self {
            Timer::Unit(_, ref state) | Timer::Joined { lowstate: ref state, .. } => state
        };
        scale(state.val_cycles, state.prescaler)
    }

    fn incr_val(&mut self, clock_diff: u64) {
        let state = match *self {
            Timer::Unit(_, ref mut state) | Timer::Joined { lowstate: ref mut state, .. } => state
        };
        state.val_cycles += clock_diff;
    }

    fn val_diff(&self, clock_diff: u64) -> u64 {
        let state = match *self {
            Timer::Unit(_, ref state) | Timer::Joined { lowstate: ref state, .. } => state
        };

        let new_val = scale(state.val_cycles + clock_diff, state.prescaler);
        new_val - self.val()
    }

    fn will_overflow(&self, clock_diff: u64) -> Option<u8> {
        let diff = self.val_diff(clock_diff);
        match *self {
            Timer::Unit(num, _) => {
                let overflows = (self.val() as u16).checked_add(diff as u16).is_none();
                if diff >> 16 != 0 || overflows {
                    return Some(num)
                } else {
                    None
                }
            }
            Timer::Joined { .. } => unimplemented!()
        }
    }

    fn has_index(&self, t_index: usize) -> bool {
        let t_index = t_index as u8;

        match *self {
            Timer::Unit(num, ..) => num == t_index,
            Timer::Joined { bitset: bs, .. } => bs & (1 << t_index) != 0
        }
    }

    fn val_hword(&self, t_index: usize) -> Option<u16> {
        let t_index = t_index as u8;

        match *self {
            Timer::Unit(num, ..) if num == t_index => {
                Some(self.val() as u16)
            }
            Timer::Joined { bitset: bs, .. } if bs & (1 << t_index) != 0 => {
                Some((self.val() >> (16*t_index)) as u16)
            }
            _ => None
        }
    }

    fn set_val_hword(&mut self, t_index: usize, val: u16) {
        let t_index = t_index as u8;
        let new = match *self {
            Timer::Unit(num, ..) if num == t_index => {
                val as u64
            }
            Timer::Joined { bitset: bs, .. } if bs & (1 << t_index) != 0 => {
                self.val() & !(0xFFFF << (16*t_index)) | ((val as u64) << (16*t_index))
            }
            _ => return
        };

        let state = match *self {
            Timer::Unit(_, ref mut state) | Timer::Joined { lowstate: ref mut state, .. } => state
        };
        state.val_cycles = unscale(new, state.prescaler);
    }
}

struct TimerIter<'a> {
    next_timer: u8,
    timer_states: Vec<MutexGuard<'a, TimerState>>,
}

impl<'a> TimerIter<'a> {
    fn new<'b>(states: &'b TimerStates) -> TimerIter<'b> {
        TimerIter {
            next_timer: 3,
            timer_states: states.0.iter().map(|x| x.lock()).collect(),
        }
    }
}

impl<'a> Iterator for TimerIter<'a> {
    type Item = Timer<'a>;
    fn next(&mut self) -> Option<Self::Item> {
        if self.next_timer == 0xFF {
            return None
        }
        let mut item = self.timer_states.pop().unwrap();
        if !item.count_up {
            let ret = Timer::Unit(self.next_timer, item);
            self.next_timer -= 1;
            return Some(ret)
        }

        // assert!(bf!(cnt(0) @ CntReg::count_up) == 0);

        let mut joined_bitset = 0;
        loop {
            joined_bitset |= 1 << self.next_timer;
            self.next_timer = self.next_timer.wrapping_sub(1);
            if !item.count_up || self.next_timer == 0xFF {
                break
            }
            item = self.timer_states.pop().unwrap();
        }
        Some(Timer::Joined { bitset: joined_bitset, lowstate: item })
    }
}

fn scale(cycles: u64, prescaler: Prescaler) -> u64 {
    match prescaler {
        Prescaler::Div1 => cycles,
        Prescaler::Div64 => cycles >> 6,
        Prescaler::Div256 => cycles >> 8,
        Prescaler::Div1024 => cycles >> 10
    }
}

fn unscale(clock_ticks: u64, prescaler: Prescaler) -> u64 {
    match prescaler {
        Prescaler::Div1 => clock_ticks,
        Prescaler::Div64 => clock_ticks << 6,
        Prescaler::Div256 => clock_ticks << 8,
        Prescaler::Div1024 => clock_ticks << 10
    }
}

fn irq(t_index: usize) -> irq::IrqType {
    match t_index {
        0 => irq::IrqType::Timer0,
        1 => irq::IrqType::Timer1,
        2 => irq::IrqType::Timer2,
        3 => irq::IrqType::Timer3,
        _ => unreachable!()
    }
}

#[derive(Clone)]
pub struct TimerStates(Arc<[Mutex<TimerState>; 4]>);
impl TimerStates {
    pub fn new() -> TimerStates {
        TimerStates(Arc::new([
            Mutex::new(TimerState::new()), Mutex::new(TimerState::new()),
            Mutex::new(TimerState::new()), Mutex::new(TimerState::new()),
        ]))
    }
}

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

#[derive(Debug)]
pub struct TimerState {
    started: bool,
    count_up: bool,
    val_cycles: u64,
    prescaler: Prescaler,
}

impl TimerState {
    pub fn new() -> TimerState {
        TimerState {
            started: false,
            count_up: false,
            val_cycles: 0,
            prescaler: Prescaler::Div1
        }
    }
}

pub fn handle_clock_update(timer_states: &TimerStates, clock_diff: usize, irq_tx: &mut irq::IrqRequests) {
    let iter_started = TimerIter::new(&timer_states).filter(|t| t.started());

    for mut timer in iter_started {
        let overflows = timer.will_overflow(clock_diff as u64);
        timer.incr_val(clock_diff as u64);
        if let Some(overflow_index) = overflows {
            // Overflow happened
            irq_tx.add(irq(overflow_index as usize))
        }
    }
}

iodevice!(TimerDevice, {
    internal_state: TimerStates;
    regs: {
        0x000 => val0: u16 {
            read_effect = |dev: &mut TimerDevice| reg_val_read(dev, 0);
            write_effect = |dev: &mut TimerDevice| reg_val_update(dev, 0);
        }
        0x002 => cnt0: u16 { write_effect = |dev: &mut TimerDevice| reg_cnt_update(dev, 0); }

        0x004 => val1: u16 {
            read_effect = |dev: &mut TimerDevice| reg_val_read(dev, 1);
            write_effect = |dev: &mut TimerDevice| reg_val_update(dev, 1);
        }
        0x006 => cnt1: u16 { write_effect = |dev: &mut TimerDevice| reg_cnt_update(dev, 1); }

        0x008 => val2: u16 {
            read_effect = |dev: &mut TimerDevice| reg_val_read(dev, 2);
            write_effect = |dev: &mut TimerDevice| reg_val_update(dev, 2);
        }
        0x00A => cnt2: u16 { write_effect = |dev: &mut TimerDevice| reg_cnt_update(dev, 2); }

        0x00C => val3: u16 {
            read_effect = |dev: &mut TimerDevice| reg_val_read(dev, 3);
            write_effect = |dev: &mut TimerDevice| reg_val_update(dev, 3);
        }
        0x00E => cnt3: u16 { write_effect = |dev: &mut TimerDevice| reg_cnt_update(dev, 3); }
    }
});<commit_msg>TIMER: implement overflow detection for Joined timers<commit_after>// TODO: How can we test this module?

use std::fmt;
use std::sync::Arc;

use parking_lot::{Mutex, MutexGuard};

use cpu::irq;
use io::regs::IoReg;

#[derive(Clone, Copy, Debug)]
pub enum Prescaler {
    Div1 = 0,
    Div64 = 1,
    Div256 = 2,
    Div1024 = 3,
}

impl Prescaler {
    fn new(val: u16) -> Prescaler {
        match val {
            0 => Prescaler::Div1,
            1 => Prescaler::Div64,
            2 => Prescaler::Div256,
            3 => Prescaler::Div1024,
            _ => unreachable!()
        }
    }
}

bfdesc!(CntReg: u16, {
    prescaler: 0 => 1,
    count_up: 2 => 2,
    irq_enable: 6 => 6,
    started: 7 => 7
});

fn get_regs(dev: &mut TimerDevice, index: usize) -> (&mut IoReg<u16>, &mut IoReg<u16>) {
    match index {
        0 => (&mut dev.val0, &mut dev.cnt0),
        1 => (&mut dev.val1, &mut dev.cnt1),
        2 => (&mut dev.val2, &mut dev.cnt2),
        3 => (&mut dev.val3, &mut dev.cnt3),
        _ => unreachable!()
    }
}

fn reg_val_update(dev: &mut TimerDevice, index: usize) {
    let val = {
        let (val, _) = get_regs(dev, index);
        val.get()
    };
    let states = &dev._internal_state;
    for mut t in TimerIter::new(states).filter(|t| t.has_index(index)) {
        t.set_val_hword(index, val)
    }
}

fn reg_val_read(dev: &mut TimerDevice, index: usize) {
    let new_val = {
        let states = &dev._internal_state;
        TimerIter::new(states).filter_map(|t| t.val_hword(index))
                              .next().unwrap()
    };
    let (val, _) = get_regs(dev, index);
    val.set_unchecked(new_val);
}

fn reg_cnt_update(dev: &mut TimerDevice, index: usize) {
    let (val, cnt) = {
        let (val, cnt) = get_regs(dev, index);
        (val.get(), cnt.get())
    };
    let mut timer = {
        let states = &dev._internal_state;
        TimerIter::new(states).filter(|t| t.has_index(index))
                              .next().unwrap()
    };
    {
        let state = match timer {
            Timer::Unit(_, ref mut state) | Timer::Joined { lowstate: ref mut state, .. } => state
        };
        state.started = bf!(cnt @ CntReg::started) == 1;
        state.count_up = bf!(cnt @ CntReg::count_up) == 1;
        state.prescaler = Prescaler::new(bf!(cnt @ CntReg::prescaler));
    }

    timer.set_val_hword(index, val);
    trace!("Setting TIMER CNT{}: {:?}", index, timer);
}

enum Timer<'a> {
    Joined { bitset: u8, lowstate: MutexGuard<'a, TimerState> },
    Unit(u8, MutexGuard<'a, TimerState>)
}

impl<'a> fmt::Debug for Timer<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Timer::Unit(num, ref state) => write!(f, "Timer::Unit({}, {:?})", num, **state),
            Timer::Joined { bitset: bs, lowstate: ref state } => {
                write!(f, "Timer::Joined {{ bitset: 0b{:b}, lowstate: {:?} }}", bs, **state)
            }
        }
    }
}

impl<'a> Timer<'a> {
    fn started(&self) -> bool {
        match *self {
            Timer::Unit(_, ref state) | Timer::Joined { lowstate: ref state, .. } => state.started
        }
    }

    fn val(&self) -> u64 {
        let state = match *self {
            Timer::Unit(_, ref state) | Timer::Joined { lowstate: ref state, .. } => state
        };
        scale(state.val_cycles, state.prescaler)
    }

    fn incr_val(&mut self, clock_diff: u64) {
        let state = match *self {
            Timer::Unit(_, ref mut state) | Timer::Joined { lowstate: ref mut state, .. } => state
        };
        state.val_cycles += clock_diff;
    }

    fn val_diff(&self, clock_diff: u64) -> u64 {
        let state = match *self {
            Timer::Unit(_, ref state) | Timer::Joined { lowstate: ref state, .. } => state
        };

        let new_val = scale(state.val_cycles + clock_diff, state.prescaler);
        new_val - self.val()
    }

    fn will_overflow_words(&self, clock_diff: u64) -> [bool; 4] {
        let diff = self.val_diff(clock_diff);
        let mut overflow_words = [false; 4];
        match *self {
            Timer::Unit(num, _) => {
                let overflows = (self.val() as u16).checked_add(diff as u16).is_none();
                if diff >> 16 != 0 || overflows {
                    overflow_words[num as usize] = true;
                }
            }
            Timer::Joined { bitset: bs, .. } => {
                let lowest_word = bs.trailing_zeros() as usize;
                let val = self.val();
                let newval = val + diff;
                for i in 0..4 {
                    let val_hword = (val >> (16*i)) as u16;
                    let newval_hword = (newval >> (16*i)) as u16;
                    if newval_hword < val_hword {
                        overflow_words[lowest_word+i] = true;
                    }
                }
            }
        }
        overflow_words
    }

    fn has_index(&self, t_index: usize) -> bool {
        let t_index = t_index as u8;

        match *self {
            Timer::Unit(num, ..) => num == t_index,
            Timer::Joined { bitset: bs, .. } => bs & (1 << t_index) != 0
        }
    }

    fn val_hword(&self, t_index: usize) -> Option<u16> {
        let t_index = t_index as u8;

        match *self {
            Timer::Unit(num, ..) if num == t_index => {
                Some(self.val() as u16)
            }
            Timer::Joined { bitset: bs, .. } if bs & (1 << t_index) != 0 => {
                Some((self.val() >> (16*t_index)) as u16)
            }
            _ => None
        }
    }

    fn set_val_hword(&mut self, t_index: usize, val: u16) {
        let t_index = t_index as u8;
        let new = match *self {
            Timer::Unit(num, ..) if num == t_index => {
                val as u64
            }
            Timer::Joined { bitset: bs, .. } if bs & (1 << t_index) != 0 => {
                self.val() & !(0xFFFF << (16*t_index)) | ((val as u64) << (16*t_index))
            }
            _ => return
        };

        let state = match *self {
            Timer::Unit(_, ref mut state) | Timer::Joined { lowstate: ref mut state, .. } => state
        };
        state.val_cycles = unscale(new, state.prescaler);
    }
}

struct TimerIter<'a> {
    next_timer: u8,
    timer_states: Vec<MutexGuard<'a, TimerState>>,
}

impl<'a> TimerIter<'a> {
    fn new<'b>(states: &'b TimerStates) -> TimerIter<'b> {
        TimerIter {
            next_timer: 3,
            timer_states: states.0.iter().map(|x| x.lock()).collect(),
        }
    }
}

impl<'a> Iterator for TimerIter<'a> {
    type Item = Timer<'a>;
    fn next(&mut self) -> Option<Self::Item> {
        if self.next_timer == 0xFF {
            return None
        }
        let mut item = self.timer_states.pop().unwrap();
        if !item.count_up {
            let ret = Timer::Unit(self.next_timer, item);
            self.next_timer -= 1;
            return Some(ret)
        }

        // assert!(bf!(cnt(0) @ CntReg::count_up) == 0);

        let mut joined_bitset = 0;
        loop {
            joined_bitset |= 1 << self.next_timer;
            self.next_timer = self.next_timer.wrapping_sub(1);
            if !item.count_up || self.next_timer == 0xFF {
                break
            }
            item = self.timer_states.pop().unwrap();
        }
        Some(Timer::Joined { bitset: joined_bitset, lowstate: item })
    }
}

fn scale(cycles: u64, prescaler: Prescaler) -> u64 {
    match prescaler {
        Prescaler::Div1 => cycles,
        Prescaler::Div64 => cycles >> 6,
        Prescaler::Div256 => cycles >> 8,
        Prescaler::Div1024 => cycles >> 10
    }
}

fn unscale(clock_ticks: u64, prescaler: Prescaler) -> u64 {
    match prescaler {
        Prescaler::Div1 => clock_ticks,
        Prescaler::Div64 => clock_ticks << 6,
        Prescaler::Div256 => clock_ticks << 8,
        Prescaler::Div1024 => clock_ticks << 10
    }
}

fn irq(t_index: usize) -> irq::IrqType {
    match t_index {
        0 => irq::IrqType::Timer0,
        1 => irq::IrqType::Timer1,
        2 => irq::IrqType::Timer2,
        3 => irq::IrqType::Timer3,
        _ => unreachable!()
    }
}

#[derive(Clone)]
pub struct TimerStates(Arc<[Mutex<TimerState>; 4]>);
impl TimerStates {
    pub fn new() -> TimerStates {
        TimerStates(Arc::new([
            Mutex::new(TimerState::new()), Mutex::new(TimerState::new()),
            Mutex::new(TimerState::new()), Mutex::new(TimerState::new()),
        ]))
    }
}

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

#[derive(Debug)]
pub struct TimerState {
    started: bool,
    count_up: bool,
    val_cycles: u64,
    prescaler: Prescaler,
}

impl TimerState {
    pub fn new() -> TimerState {
        TimerState {
            started: false,
            count_up: false,
            val_cycles: 0,
            prescaler: Prescaler::Div1
        }
    }
}

pub fn handle_clock_update(timer_states: &TimerStates, clock_diff: usize, irq_tx: &mut irq::IrqRequests) {
    let iter_started = TimerIter::new(&timer_states).filter(|t| t.started());

    for mut timer in iter_started {
        let overflows = timer.will_overflow_words(clock_diff as u64);
        timer.incr_val(clock_diff as u64);
        for (index, status) in overflows.iter().enumerate() {
            if *status {
                // Overflow happened
                irq_tx.add(irq(index))
            }
        }
    }
}

iodevice!(TimerDevice, {
    internal_state: TimerStates;
    regs: {
        0x000 => val0: u16 {
            read_effect = |dev: &mut TimerDevice| reg_val_read(dev, 0);
            write_effect = |dev: &mut TimerDevice| reg_val_update(dev, 0);
        }
        0x002 => cnt0: u16 { write_effect = |dev: &mut TimerDevice| reg_cnt_update(dev, 0); }

        0x004 => val1: u16 {
            read_effect = |dev: &mut TimerDevice| reg_val_read(dev, 1);
            write_effect = |dev: &mut TimerDevice| reg_val_update(dev, 1);
        }
        0x006 => cnt1: u16 { write_effect = |dev: &mut TimerDevice| reg_cnt_update(dev, 1); }

        0x008 => val2: u16 {
            read_effect = |dev: &mut TimerDevice| reg_val_read(dev, 2);
            write_effect = |dev: &mut TimerDevice| reg_val_update(dev, 2);
        }
        0x00A => cnt2: u16 { write_effect = |dev: &mut TimerDevice| reg_cnt_update(dev, 2); }

        0x00C => val3: u16 {
            read_effect = |dev: &mut TimerDevice| reg_val_read(dev, 3);
            write_effect = |dev: &mut TimerDevice| reg_val_update(dev, 3);
        }
        0x00E => cnt3: u16 { write_effect = |dev: &mut TimerDevice| reg_cnt_update(dev, 3); }
    }
});<|endoftext|>


<commit_before>//! impl bool {}

#[lang = "bool"]
impl bool {
    /// Returns `Some(t)` if the `bool` is `true`, or `None` otherwise.
    ///
    /// # Examples
    ///
    /// ```
    /// #![feature(bool_to_option)]
    ///
    /// assert_eq!(false.then_some(0), None);
    /// assert_eq!(true.then_some(0), Some(0));
    /// ```
    #[unstable(feature = "bool_to_option", issue = "64260")]
    #[inline]
    pub fn then_some<T>(self, t: T) -> Option<T> {
        if self { Some(t) } else { None }
    }

    /// Returns `Some(f())` if the `bool` is `true`, or `None` otherwise.
    ///
    /// # Examples
    ///
    /// ```
    /// assert_eq!(false.then(|| 0), None);
    /// assert_eq!(true.then(|| 0), Some(0));
    /// ```
    #[stable(feature = "lazy_bool_to_option", since = "1.50.0")]
    #[inline]
    pub fn then<T, F: FnOnce() -> T>(self, f: F) -> Option<T> {
        if self { Some(f()) } else { None }
    }
}
<commit_msg>Rollup merge of #87045 - jhpratt:fix-tracking-issue, r=jyn514<commit_after>//! impl bool {}

#[lang = "bool"]
impl bool {
    /// Returns `Some(t)` if the `bool` is `true`, or `None` otherwise.
    ///
    /// # Examples
    ///
    /// ```
    /// #![feature(bool_to_option)]
    ///
    /// assert_eq!(false.then_some(0), None);
    /// assert_eq!(true.then_some(0), Some(0));
    /// ```
    #[unstable(feature = "bool_to_option", issue = "80967")]
    #[inline]
    pub fn then_some<T>(self, t: T) -> Option<T> {
        if self { Some(t) } else { None }
    }

    /// Returns `Some(f())` if the `bool` is `true`, or `None` otherwise.
    ///
    /// # Examples
    ///
    /// ```
    /// assert_eq!(false.then(|| 0), None);
    /// assert_eq!(true.then(|| 0), Some(0));
    /// ```
    #[stable(feature = "lazy_bool_to_option", since = "1.50.0")]
    #[inline]
    pub fn then<T, F: FnOnce() -> T>(self, f: F) -> Option<T> {
        if self { Some(f()) } else { None }
    }
}
<|endoftext|>


<commit_before>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

use super::objects::{DocType, EsId, Admin};
use chrono;
use regex;
use hyper;
use hyper::status::StatusCode;
use rs_es::error::EsError;
use rs_es;
use rs_es::EsResponse;
use serde;
use serde_json;
use rs_es::operations::search::ScanResult;

use super::objects::{AliasOperations, AliasOperation, AliasParameter};
use rs_es::units::Duration;

// Rubber is an wrapper around elasticsearch API
pub struct Rubber {
    es_client: rs_es::Client,
    // some operation are not implemented in rs_es, we need to use a raw http client
    http_client: hyper::client::Client,
}

/// return the index associated to the given type and dataset
/// this will be an alias over another real index
fn get_main_type_and_dataset_index(doc_type: &str, dataset: &str) -> String {
    format!("munin_{}_{}", doc_type, dataset)
}

/// return the index associated to the given type
/// this will be an alias over another real index
fn get_main_type_index(doc_type: &str) -> String {
    format!("munin_{}", doc_type)
}

impl Rubber {
    // build a rubber with a connection string (http://host:port/)
    pub fn new(cnx: &str) -> Rubber {
        let re = regex::Regex::new(r"(?:https?://)?(?P<host>.+?):(?P<port>\d+)").unwrap();
        let cap = re.captures(cnx).unwrap();
        let host = cap.name("host").unwrap().as_str();
        let port = cap.name("port").unwrap().as_str().parse::<u32>().unwrap();
        info!("elastic search host {:?} port {:?}", host, port);

        Rubber {
            es_client: rs_es::Client::new(&host, port),
            http_client: hyper::client::Client::new(),
        }
    }

    pub fn get(&self, path: &str) -> Result<hyper::client::response::Response, EsError> {
        // Note: a bit duplicate on rs_es because some ES operations are not implemented
        debug!("doing a get on {}", path);
        let url = self.es_client.full_url(path);
        let result = try!(self.http_client
            .get(&url)
            .send());
        rs_es::do_req(result)
    }
    fn put(&self, path: &str, body: &str) -> Result<hyper::client::response::Response, EsError> {
        // Note: a bit duplicate on rs_es because some ES operations are not implemented
        debug!("doing a put on {} with {}", path, body);
        let url = self.es_client.full_url(path);
        let result = try!(self.http_client
            .put(&url)
            .body(body)
            .send());
        rs_es::do_req(result)
    }
    fn post(&self, path: &str, body: &str) -> Result<hyper::client::response::Response, EsError> {
        // Note: a bit duplicate on rs_es because some ES operations are not implemented
        debug!("doing a post on {} with {}", path, body);
        let url = self.es_client.full_url(path);
        let result = try!(self.http_client
            .post(&url)
            .body(body)
            .send());
        rs_es::do_req(result)
    }

    pub fn make_index(&self, doc_type: &str, dataset: &str) -> Result<String, String> {
        let current_time = chrono::UTC::now().format("%Y%m%d_%H%M%S_%f");
        let index_name = format!("munin_{}_{}_{}", doc_type, dataset, current_time);
        info!("creating index {}", index_name);
        self.create_index(&index_name.to_string()).map(|_| index_name)
    }

    fn create_index(&self, name: &String) -> Result<(), String> {
        debug!("creating index");
        // Note: in rs_es it can be done with MappingOperation but for the moment I think
        // storing the mapping in json is more convenient
        let analysis = include_str!("../../../json/settings.json");
        self.put(name, &analysis)
            .map_err(|e| {
                info!("Error while creating new index {}", name);
                e.to_string()
            })
            .and_then(|res| if res.status == StatusCode::Ok {
                Ok(())
            } else {
                Err(format!("cannot create index: {:?}", res))
            })
    }

    fn get_last_index(&self, doc_type: &str, dataset: &str) -> Result<Vec<String>, String> {
        let base_index = get_main_type_and_dataset_index(doc_type, dataset);
        self.get(&format!("{}/_aliases", base_index))
            .map_err(|e| e.to_string())
            .and_then(|res| match res.status {
                StatusCode::Ok => {
                    let value: serde_json::Value = try!(res.read_response()
                        .map_err(|e| e.to_string()));
                    Ok(value.as_object()
                        .and_then(|aliases| Some(aliases.keys().cloned().collect()))
                        .unwrap_or_else(|| {
                            info!("no previous index to delete for type {} and dataset {}",
                                  doc_type,
                                  dataset);
                            vec![]
                        }))
                }
                StatusCode::NotFound => {
                    info!("impossible to find alias {}, no last index to remove",
                          base_index);
                    Ok(vec![])
                }
                _ => Err(format!("invalid elasticsearch response: {:?}", res)),
            })
    }

    /// publish the index as the new index for this doc_type and this dataset
    /// move the index alias of the doc_type and the dataset to point to this indexes
    /// and remove the old index
    pub fn publish_index(&mut self,
                         doc_type: &str,
                         dataset: &str,
                         index: String)
                         -> Result<(), String> {
        let v = vec!["addr", "poi", "admin", "way"];
        debug!("publishing index");
        let last_indexes = try!(self.get_last_index(doc_type, dataset));

        let dataset_index = get_main_type_and_dataset_index(doc_type, dataset);
        try!(self.alias(&dataset_index, &vec![index.clone()], &last_indexes));

        let type_index = get_main_type_index(doc_type);
        try!(self.alias(&type_index, &vec![dataset_index.clone()], &last_indexes));

        if v.contains(&doc_type) {
            try!(self.alias("munin_geo_data",
                            &vec![type_index.to_string()],
                            &vec![]));
            try!(self.alias("munin", &vec!["munin_geo_data".to_string()], &vec![]));
        } else {
            try!(self.alias("munin", &vec![type_index.to_string()], &vec![]));
        }
        for i in last_indexes {
            try!(self.delete_index(&i));
        }
        Ok(())
    }

    pub fn is_existing_index(&self, name: &String) -> Result<bool, String> {
        self.get(&name)
            .map_err(|e| e.to_string())
            .map(|res| res.status == StatusCode::Ok)
    }

    /// add a list of new indexes to the alias
    /// remove a list of indexes from the alias
    fn alias(&self, alias: &str, add: &Vec<String>, remove: &Vec<String>) -> Result<(), String> {
        info!("for {}, adding alias {:?}, removing {:?}",
              alias,
              add,
              remove);
        let add_operations = add.iter().map(|x| {
            AliasOperation {
                remove: None,
                add: Some(AliasParameter {
                    index: x.clone(),
                    alias: alias.to_string(),
                }),
            }
        });
        let remove_operations = remove.iter().map(|x| {
            AliasOperation {
                add: None,
                remove: Some(AliasParameter {
                    index: x.clone(),
                    alias: alias.to_string(),
                }),
            }
        });
        let operations =
            AliasOperations { actions: add_operations.chain(remove_operations).collect() };
        let json = serde_json::to_string(&operations).unwrap();
        self.post("_aliases", &json)
            .map_err(|e| e.to_string())
            .and_then(|res| if res.status == StatusCode::Ok {
                Ok(())
            } else {
                error!("failed to change aliases for {}, es response: {:?}",
                       alias,
                       res);
                Err(format!("failed to post aliases for {}: {:?}", alias, res).to_string())
            })
    }

    pub fn delete_index(&mut self, index: &String) -> Result<(), String> {
        debug!("deleting index {}", &index);
        let res = self.es_client
            .delete_index(&index)
            .map(|res| res.acknowledged)
            .unwrap_or(false);
        if !res {
            Err(format!("Error deleting index {}", &index).into())
        } else {
            Ok(())
        }
    }

    pub fn bulk_index<T, I>(&mut self,
                            index: &String,
                            iter: I)
                            -> Result<usize, rs_es::error::EsError>
        where T: serde::Serialize + DocType + EsId,
              I: Iterator<Item = T>
    {
        use rs_es::operations::bulk::Action;
        let mut nb = 0;
        let chunk_size = 1000;
        // fold is used for creating the action and optionally set the id of the object
        let mut actions = iter.map(|v| {
            v.es_id()
                .into_iter()
                .fold(Action::index(v), |action, id| action.with_id(id))
        });
        loop {
            let chunk = actions.by_ref().take(chunk_size).collect::<Vec<_>>();
            nb += chunk.len();
            try!(self.es_client
                .bulk(&chunk)
                .with_index(&index)
                .with_doc_type(T::doc_type())
                .send());

            if chunk.len() < chunk_size {
                break;
            }
        }

        Ok(nb)
    }

    /// add all the element of 'iter' into elasticsearch
    ///
    /// To have zero downtime:
    /// first all the elements are added in a temporary index and when all has been indexed
    /// the index is published and the old index is removed
    pub fn index<T, I>(&mut self, doc_type: &str, dataset: &str, iter: I) -> Result<usize, String>
        where T: serde::Serialize + DocType + EsId,
              I: Iterator<Item = T>
    {
        // TODO better error handling
        let index = try!(self.make_index(doc_type, dataset));
        let nb_elements = try!(self.bulk_index(&index, iter).map_err(|e| e.to_string()));
        try!(self.publish_index(doc_type, dataset, index));
        Ok(nb_elements)
    }

    pub fn get_admins_from_dataset(&mut self,
                                   dataset: &str)
                                   -> Result<Vec<Admin>, rs_es::error::EsError> {
        self.get_admins_from_index(&get_main_type_and_dataset_index("admin", dataset))
    }

    pub fn get_all_admins(&mut self) -> Result<Vec<Admin>, rs_es::error::EsError> {
        self.get_admins_from_index(&get_main_type_index("admin"))
    }

    fn get_admins_from_index(&mut self, index: &str) -> Result<Vec<Admin>, rs_es::error::EsError> {
        let mut result: Vec<Admin> = vec![];
        let mut scan: ScanResult<Admin> = try!(self.es_client
            .search_query()
            .with_indexes(&[&index])
            .with_size(1000)
            .with_types(&[&"admin"])
            .scan(&Duration::minutes(1)));
        loop {
            let page = try!(scan.scroll(&mut self.es_client, &Duration::minutes(1)));
            if page.hits.hits.len() == 0 {
                break;
            }
            result.extend(page.hits.hits.into_iter().filter_map(|hit| hit.source).map(|ad| *ad));
        }
        try!(scan.close(&mut self.es_client));
        Ok(result)
    }
}

#[test]
pub fn test_valid_url() {
    Rubber::new("http://localhost:9200");
    Rubber::new("localhost:9200");
}

#[test]
#[should_panic]
pub fn test_invalid_url() {
    Rubber::new("http://bob");
}

#[test]
#[should_panic]
pub fn test_invalid_url_no_port() {
    Rubber::new("localhost");
}
<commit_msg>fmt passed<commit_after>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

use super::objects::{DocType, EsId, Admin};
use chrono;
use regex;
use hyper;
use hyper::status::StatusCode;
use rs_es::error::EsError;
use rs_es;
use rs_es::EsResponse;
use serde;
use serde_json;
use rs_es::operations::search::ScanResult;

use super::objects::{AliasOperations, AliasOperation, AliasParameter};
use rs_es::units::Duration;

// Rubber is an wrapper around elasticsearch API
pub struct Rubber {
    es_client: rs_es::Client,
    // some operation are not implemented in rs_es, we need to use a raw http client
    http_client: hyper::client::Client,
}

/// return the index associated to the given type and dataset
/// this will be an alias over another real index
fn get_main_type_and_dataset_index(doc_type: &str, dataset: &str) -> String {
    format!("munin_{}_{}", doc_type, dataset)
}

/// return the index associated to the given type
/// this will be an alias over another real index
fn get_main_type_index(doc_type: &str) -> String {
    format!("munin_{}", doc_type)
}

impl Rubber {
    // build a rubber with a connection string (http://host:port/)
    pub fn new(cnx: &str) -> Rubber {
        let re = regex::Regex::new(r"(?:https?://)?(?P<host>.+?):(?P<port>\d+)").unwrap();
        let cap = re.captures(cnx).unwrap();
        let host = cap.name("host").unwrap().as_str();
        let port = cap.name("port").unwrap().as_str().parse::<u32>().unwrap();
        info!("elastic search host {:?} port {:?}", host, port);

        Rubber {
            es_client: rs_es::Client::new(&host, port),
            http_client: hyper::client::Client::new(),
        }
    }

    pub fn get(&self, path: &str) -> Result<hyper::client::response::Response, EsError> {
        // Note: a bit duplicate on rs_es because some ES operations are not implemented
        debug!("doing a get on {}", path);
        let url = self.es_client.full_url(path);
        let result = try!(self.http_client
            .get(&url)
            .send());
        rs_es::do_req(result)
    }
    fn put(&self, path: &str, body: &str) -> Result<hyper::client::response::Response, EsError> {
        // Note: a bit duplicate on rs_es because some ES operations are not implemented
        debug!("doing a put on {} with {}", path, body);
        let url = self.es_client.full_url(path);
        let result = try!(self.http_client
            .put(&url)
            .body(body)
            .send());
        rs_es::do_req(result)
    }
    fn post(&self, path: &str, body: &str) -> Result<hyper::client::response::Response, EsError> {
        // Note: a bit duplicate on rs_es because some ES operations are not implemented
        debug!("doing a post on {} with {}", path, body);
        let url = self.es_client.full_url(path);
        let result = try!(self.http_client
            .post(&url)
            .body(body)
            .send());
        rs_es::do_req(result)
    }

    pub fn make_index(&self, doc_type: &str, dataset: &str) -> Result<String, String> {
        let current_time = chrono::UTC::now().format("%Y%m%d_%H%M%S_%f");
        let index_name = format!("munin_{}_{}_{}", doc_type, dataset, current_time);
        info!("creating index {}", index_name);
        self.create_index(&index_name.to_string()).map(|_| index_name)
    }

    fn create_index(&self, name: &String) -> Result<(), String> {
        debug!("creating index");
        // Note: in rs_es it can be done with MappingOperation but for the moment I think
        // storing the mapping in json is more convenient
        let analysis = include_str!("../../../json/settings.json");
        self.put(name, &analysis)
            .map_err(|e| {
                info!("Error while creating new index {}", name);
                e.to_string()
            })
            .and_then(|res| if res.status == StatusCode::Ok {
                Ok(())
            } else {
                Err(format!("cannot create index: {:?}", res))
            })
    }

    fn get_last_index(&self, doc_type: &str, dataset: &str) -> Result<Vec<String>, String> {
        let base_index = get_main_type_and_dataset_index(doc_type, dataset);
        self.get(&format!("{}/_aliases", base_index))
            .map_err(|e| e.to_string())
            .and_then(|res| match res.status {
                StatusCode::Ok => {
                    let value: serde_json::Value = try!(res.read_response()
                        .map_err(|e| e.to_string()));
                    Ok(value.as_object()
                        .and_then(|aliases| Some(aliases.keys().cloned().collect()))
                        .unwrap_or_else(|| {
                            info!("no previous index to delete for type {} and dataset {}",
                                  doc_type,
                                  dataset);
                            vec![]
                        }))
                }
                StatusCode::NotFound => {
                    info!("impossible to find alias {}, no last index to remove",
                          base_index);
                    Ok(vec![])
                }
                _ => Err(format!("invalid elasticsearch response: {:?}", res)),
            })
    }

    /// publish the index as the new index for this doc_type and this dataset
    /// move the index alias of the doc_type and the dataset to point to this indexes
    /// and remove the old index
    pub fn publish_index(&mut self,
                         doc_type: &str,
                         dataset: &str,
                         index: String)
                         -> Result<(), String> {
        let v = vec!["addr", "poi", "admin", "way"];
        debug!("publishing index");
        let last_indexes = try!(self.get_last_index(doc_type, dataset));

        let dataset_index = get_main_type_and_dataset_index(doc_type, dataset);
        try!(self.alias(&dataset_index, &vec![index.clone()], &last_indexes));

        let type_index = get_main_type_index(doc_type);
        try!(self.alias(&type_index, &vec![dataset_index.clone()], &last_indexes));

        if v.contains(&doc_type) {
            try!(self.alias("munin_geo_data", &vec![type_index.to_string()], &vec![]));
            try!(self.alias("munin", &vec!["munin_geo_data".to_string()], &vec![]));
        } else {
            try!(self.alias("munin", &vec![type_index.to_string()], &vec![]));
        }
        for i in last_indexes {
            try!(self.delete_index(&i));
        }
        Ok(())
    }

    pub fn is_existing_index(&self, name: &String) -> Result<bool, String> {
        self.get(&name)
            .map_err(|e| e.to_string())
            .map(|res| res.status == StatusCode::Ok)
    }

    /// add a list of new indexes to the alias
    /// remove a list of indexes from the alias
    fn alias(&self, alias: &str, add: &Vec<String>, remove: &Vec<String>) -> Result<(), String> {
        info!("for {}, adding alias {:?}, removing {:?}",
              alias,
              add,
              remove);
        let add_operations = add.iter().map(|x| {
            AliasOperation {
                remove: None,
                add: Some(AliasParameter {
                    index: x.clone(),
                    alias: alias.to_string(),
                }),
            }
        });
        let remove_operations = remove.iter().map(|x| {
            AliasOperation {
                add: None,
                remove: Some(AliasParameter {
                    index: x.clone(),
                    alias: alias.to_string(),
                }),
            }
        });
        let operations =
            AliasOperations { actions: add_operations.chain(remove_operations).collect() };
        let json = serde_json::to_string(&operations).unwrap();
        self.post("_aliases", &json)
            .map_err(|e| e.to_string())
            .and_then(|res| if res.status == StatusCode::Ok {
                Ok(())
            } else {
                error!("failed to change aliases for {}, es response: {:?}",
                       alias,
                       res);
                Err(format!("failed to post aliases for {}: {:?}", alias, res).to_string())
            })
    }

    pub fn delete_index(&mut self, index: &String) -> Result<(), String> {
        debug!("deleting index {}", &index);
        let res = self.es_client
            .delete_index(&index)
            .map(|res| res.acknowledged)
            .unwrap_or(false);
        if !res {
            Err(format!("Error deleting index {}", &index).into())
        } else {
            Ok(())
        }
    }

    pub fn bulk_index<T, I>(&mut self,
                            index: &String,
                            iter: I)
                            -> Result<usize, rs_es::error::EsError>
        where T: serde::Serialize + DocType + EsId,
              I: Iterator<Item = T>
    {
        use rs_es::operations::bulk::Action;
        let mut nb = 0;
        let chunk_size = 1000;
        // fold is used for creating the action and optionally set the id of the object
        let mut actions = iter.map(|v| {
            v.es_id()
                .into_iter()
                .fold(Action::index(v), |action, id| action.with_id(id))
        });
        loop {
            let chunk = actions.by_ref().take(chunk_size).collect::<Vec<_>>();
            nb += chunk.len();
            try!(self.es_client
                .bulk(&chunk)
                .with_index(&index)
                .with_doc_type(T::doc_type())
                .send());

            if chunk.len() < chunk_size {
                break;
            }
        }

        Ok(nb)
    }

    /// add all the element of 'iter' into elasticsearch
    ///
    /// To have zero downtime:
    /// first all the elements are added in a temporary index and when all has been indexed
    /// the index is published and the old index is removed
    pub fn index<T, I>(&mut self, doc_type: &str, dataset: &str, iter: I) -> Result<usize, String>
        where T: serde::Serialize + DocType + EsId,
              I: Iterator<Item = T>
    {
        // TODO better error handling
        let index = try!(self.make_index(doc_type, dataset));
        let nb_elements = try!(self.bulk_index(&index, iter).map_err(|e| e.to_string()));
        try!(self.publish_index(doc_type, dataset, index));
        Ok(nb_elements)
    }

    pub fn get_admins_from_dataset(&mut self,
                                   dataset: &str)
                                   -> Result<Vec<Admin>, rs_es::error::EsError> {
        self.get_admins_from_index(&get_main_type_and_dataset_index("admin", dataset))
    }

    pub fn get_all_admins(&mut self) -> Result<Vec<Admin>, rs_es::error::EsError> {
        self.get_admins_from_index(&get_main_type_index("admin"))
    }

    fn get_admins_from_index(&mut self, index: &str) -> Result<Vec<Admin>, rs_es::error::EsError> {
        let mut result: Vec<Admin> = vec![];
        let mut scan: ScanResult<Admin> = try!(self.es_client
            .search_query()
            .with_indexes(&[&index])
            .with_size(1000)
            .with_types(&[&"admin"])
            .scan(&Duration::minutes(1)));
        loop {
            let page = try!(scan.scroll(&mut self.es_client, &Duration::minutes(1)));
            if page.hits.hits.len() == 0 {
                break;
            }
            result.extend(page.hits.hits.into_iter().filter_map(|hit| hit.source).map(|ad| *ad));
        }
        try!(scan.close(&mut self.es_client));
        Ok(result)
    }
}

#[test]
pub fn test_valid_url() {
    Rubber::new("http://localhost:9200");
    Rubber::new("localhost:9200");
}

#[test]
#[should_panic]
pub fn test_invalid_url() {
    Rubber::new("http://bob");
}

#[test]
#[should_panic]
pub fn test_invalid_url_no_port() {
    Rubber::new("localhost");
}
<|endoftext|>