<commit_before>//! Various extension methods to ast Expr Nodes, which are hard to code-generate.

use crate::{
    ast::{self, child_opt, children, AstChildren, AstNode},
    SmolStr,
    SyntaxKind::*,
    SyntaxToken, T,
};

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ElseBranch {
    Block(ast::Block),
    IfExpr(ast::IfExpr),
}

impl ast::IfExpr {
    pub fn then_branch(&self) -> Option<ast::Block> {
        self.blocks().nth(0)
    }
    pub fn else_branch(&self) -> Option<ElseBranch> {
        let res = match self.blocks().nth(1) {
            Some(block) => ElseBranch::Block(block),
            None => {
                let elif: ast::IfExpr = child_opt(self)?;
                ElseBranch::IfExpr(elif)
            }
        };
        Some(res)
    }

    fn blocks(&self) -> AstChildren<ast::Block> {
        children(self)
    }
}

impl ast::RefExpr {
    pub fn is_mut(&self) -> bool {
        self.syntax().children_with_tokens().any(|n| n.kind() == T![mut])
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum PrefixOp {
    /// The `*` operator for dereferencing
    Deref,
    /// The `!` operator for logical inversion
    Not,
    /// The `-` operator for negation
    Neg,
}

impl ast::PrefixExpr {
    pub fn op_kind(&self) -> Option<PrefixOp> {
        match self.op_token()?.kind() {
            T![*] => Some(PrefixOp::Deref),
            T![!] => Some(PrefixOp::Not),
            T![-] => Some(PrefixOp::Neg),
            _ => None,
        }
    }

    pub fn op_token(&self) -> Option<SyntaxToken> {
        self.syntax().first_child_or_token()?.into_token()
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum BinOp {
    /// The `||` operator for boolean OR
    BooleanOr,
    /// The `&&` operator for boolean AND
    BooleanAnd,
    /// The `==` operator for equality testing
    EqualityTest,
    /// The `!=` operator for equality testing
    NegatedEqualityTest,
    /// The `<=` operator for lesser-equal testing
    LesserEqualTest,
    /// The `>=` operator for greater-equal testing
    GreaterEqualTest,
    /// The `<` operator for comparison
    LesserTest,
    /// The `>` operator for comparison
    GreaterTest,
    /// The `+` operator for addition
    Addition,
    /// The `*` operator for multiplication
    Multiplication,
    /// The `-` operator for subtraction
    Subtraction,
    /// The `/` operator for division
    Division,
    /// The `%` operator for remainder after division
    Remainder,
    /// The `<<` operator for left shift
    LeftShift,
    /// The `>>` operator for right shift
    RightShift,
    /// The `^` operator for bitwise XOR
    BitwiseXor,
    /// The `|` operator for bitwise OR
    BitwiseOr,
    /// The `&` operator for bitwise AND
    BitwiseAnd,
    /// The `..` operator for right-open ranges
    RangeRightOpen,
    /// The `..=` operator for right-closed ranges
    RangeRightClosed,
    /// The `=` operator for assignment
    Assignment,
    /// The `+=` operator for assignment after addition
    AddAssign,
    /// The `/=` operator for assignment after division
    DivAssign,
    /// The `*=` operator for assignment after multiplication
    MulAssign,
    /// The `%=` operator for assignment after remainders
    RemAssign,
    /// The `>>=` operator for assignment after shifting right
    ShrAssign,
    /// The `<<=` operator for assignment after shifting left
    ShlAssign,
    /// The `-=` operator for assignment after subtraction
    SubAssign,
    /// The `|=` operator for assignment after bitwise OR
    BitOrAssign,
    /// The `&=` operator for assignment after bitwise AND
    BitAndAssign,
    /// The `^=` operator for assignment after bitwise XOR
    BitXorAssign,
}

impl ast::BinExpr {
    fn op_details(&self) -> Option<(SyntaxToken, BinOp)> {
        self.syntax().children_with_tokens().filter_map(|it| it.into_token()).find_map(|c| {
            let bin_op = match c.kind() {
                T![||] => BinOp::BooleanOr,
                T![&&] => BinOp::BooleanAnd,
                T![==] => BinOp::EqualityTest,
                T![!=] => BinOp::NegatedEqualityTest,
                T![<=] => BinOp::LesserEqualTest,
                T![>=] => BinOp::GreaterEqualTest,
                T![<] => BinOp::LesserTest,
                T![>] => BinOp::GreaterTest,
                T![+] => BinOp::Addition,
                T![*] => BinOp::Multiplication,
                T![-] => BinOp::Subtraction,
                T![/] => BinOp::Division,
                T![%] => BinOp::Remainder,
                T![<<] => BinOp::LeftShift,
                T![>>] => BinOp::RightShift,
                T![^] => BinOp::BitwiseXor,
                T![|] => BinOp::BitwiseOr,
                T![&] => BinOp::BitwiseAnd,
                T![..] => BinOp::RangeRightOpen,
                T![..=] => BinOp::RangeRightClosed,
                T![=] => BinOp::Assignment,
                T![+=] => BinOp::AddAssign,
                T![/=] => BinOp::DivAssign,
                T![*=] => BinOp::MulAssign,
                T![%=] => BinOp::RemAssign,
                T![>>=] => BinOp::ShrAssign,
                T![<<=] => BinOp::ShlAssign,
                T![-=] => BinOp::SubAssign,
                T![|=] => BinOp::BitOrAssign,
                T![&=] => BinOp::BitAndAssign,
                T![^=] => BinOp::BitXorAssign,
                _ => return None,
            };
            Some((c, bin_op))
        })
    }

    pub fn op_kind(&self) -> Option<BinOp> {
        self.op_details().map(|t| t.1)
    }

    pub fn op_token(&self) -> Option<SyntaxToken> {
        self.op_details().map(|t| t.0)
    }

    pub fn lhs(&self) -> Option<ast::Expr> {
        children(self).nth(0)
    }

    pub fn rhs(&self) -> Option<ast::Expr> {
        children(self).nth(1)
    }

    pub fn sub_exprs(&self) -> (Option<ast::Expr>, Option<ast::Expr>) {
        let mut children = children(self);
        let first = children.next();
        let second = children.next();
        (first, second)
    }
}

pub enum ArrayExprKind {
    Repeat { initializer: Option<ast::Expr>, repeat: Option<ast::Expr> },
    ElementList(AstChildren<ast::Expr>),
}

impl ast::ArrayExpr {
    pub fn kind(&self) -> ArrayExprKind {
        if self.is_repeat() {
            ArrayExprKind::Repeat {
                initializer: children(self).nth(0),
                repeat: children(self).nth(1),
            }
        } else {
            ArrayExprKind::ElementList(children(self))
        }
    }

    fn is_repeat(&self) -> bool {
        self.syntax().children_with_tokens().any(|it| it.kind() == T![;])
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum LiteralKind {
    String,
    ByteString,
    Char,
    Byte,
    IntNumber { suffix: Option<SmolStr> },
    FloatNumber { suffix: Option<SmolStr> },
    Bool,
}

impl ast::Literal {
    pub fn token(&self) -> SyntaxToken {
        self.syntax()
            .children_with_tokens()
            .find(|e| e.kind() != ATTR && !e.kind().is_trivia())
            .and_then(|e| e.into_token())
            .unwrap()
    }

    pub fn kind(&self) -> LiteralKind {
        match self.token().kind() {
            INT_NUMBER => {
                let int_suffix_list = [
                    "isize", "i128", "i64", "i32", "i16", "i8", "usize", "u128", "u64", "u32",
                    "u16", "u8",
                ];

                // The lexer treats e.g. `1f64` as an integer literal. See
                // https://github.com/rust-analyzer/rust-analyzer/issues/1592
                // and the comments on the linked PR.
                let float_suffix_list = ["f32", "f64"];

                let text = self.token().text().to_string();

                let float_suffix = float_suffix_list
                    .iter()
                    .find(|&s| text.ends_with(s))
                    .map(|&suf| SmolStr::new(suf));

                if float_suffix.is_some() {
                    LiteralKind::FloatNumber { suffix: float_suffix }
                } else {
                    let suffix = int_suffix_list
                        .iter()
                        .find(|&s| text.ends_with(s))
                        .map(|&suf| SmolStr::new(suf));
                    LiteralKind::IntNumber { suffix }
                }
            }
            FLOAT_NUMBER => {
                let allowed_suffix_list = ["f64", "f32"];
                let text = self.token().text().to_string();
                let suffix = allowed_suffix_list
                    .iter()
                    .find(|&s| text.ends_with(s))
                    .map(|&suf| SmolStr::new(suf));
                LiteralKind::FloatNumber { suffix }
            }
            STRING | RAW_STRING => LiteralKind::String,
            T![true] | T![false] => LiteralKind::Bool,
            BYTE_STRING | RAW_BYTE_STRING => LiteralKind::ByteString,
            CHAR => LiteralKind::Char,
            BYTE => LiteralKind::Byte,
            _ => unreachable!(),
        }
    }
}

#[test]
fn test_literal_with_attr() {
    let parse = ast::SourceFile::parse(r#"const _: &str = { #[attr] "Hello" };"#);
    let lit = parse.tree().syntax().descendants().find_map(ast::Literal::cast).unwrap();
    assert_eq!(lit.token().text(), r#""Hello""#);
}

impl ast::NamedField {
    pub fn parent_struct_lit(&self) -> ast::StructLit {
        self.syntax().ancestors().find_map(ast::StructLit::cast).unwrap()
    }
}
<commit_msg>implement accessors for IndexExpr<commit_after>//! Various extension methods to ast Expr Nodes, which are hard to code-generate.

use crate::{
    ast::{self, child_opt, children, AstChildren, AstNode},
    SmolStr,
    SyntaxKind::*,
    SyntaxToken, T,
};

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ElseBranch {
    Block(ast::Block),
    IfExpr(ast::IfExpr),
}

impl ast::IfExpr {
    pub fn then_branch(&self) -> Option<ast::Block> {
        self.blocks().nth(0)
    }
    pub fn else_branch(&self) -> Option<ElseBranch> {
        let res = match self.blocks().nth(1) {
            Some(block) => ElseBranch::Block(block),
            None => {
                let elif: ast::IfExpr = child_opt(self)?;
                ElseBranch::IfExpr(elif)
            }
        };
        Some(res)
    }

    fn blocks(&self) -> AstChildren<ast::Block> {
        children(self)
    }
}

impl ast::RefExpr {
    pub fn is_mut(&self) -> bool {
        self.syntax().children_with_tokens().any(|n| n.kind() == T![mut])
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum PrefixOp {
    /// The `*` operator for dereferencing
    Deref,
    /// The `!` operator for logical inversion
    Not,
    /// The `-` operator for negation
    Neg,
}

impl ast::PrefixExpr {
    pub fn op_kind(&self) -> Option<PrefixOp> {
        match self.op_token()?.kind() {
            T![*] => Some(PrefixOp::Deref),
            T![!] => Some(PrefixOp::Not),
            T![-] => Some(PrefixOp::Neg),
            _ => None,
        }
    }

    pub fn op_token(&self) -> Option<SyntaxToken> {
        self.syntax().first_child_or_token()?.into_token()
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum BinOp {
    /// The `||` operator for boolean OR
    BooleanOr,
    /// The `&&` operator for boolean AND
    BooleanAnd,
    /// The `==` operator for equality testing
    EqualityTest,
    /// The `!=` operator for equality testing
    NegatedEqualityTest,
    /// The `<=` operator for lesser-equal testing
    LesserEqualTest,
    /// The `>=` operator for greater-equal testing
    GreaterEqualTest,
    /// The `<` operator for comparison
    LesserTest,
    /// The `>` operator for comparison
    GreaterTest,
    /// The `+` operator for addition
    Addition,
    /// The `*` operator for multiplication
    Multiplication,
    /// The `-` operator for subtraction
    Subtraction,
    /// The `/` operator for division
    Division,
    /// The `%` operator for remainder after division
    Remainder,
    /// The `<<` operator for left shift
    LeftShift,
    /// The `>>` operator for right shift
    RightShift,
    /// The `^` operator for bitwise XOR
    BitwiseXor,
    /// The `|` operator for bitwise OR
    BitwiseOr,
    /// The `&` operator for bitwise AND
    BitwiseAnd,
    /// The `..` operator for right-open ranges
    RangeRightOpen,
    /// The `..=` operator for right-closed ranges
    RangeRightClosed,
    /// The `=` operator for assignment
    Assignment,
    /// The `+=` operator for assignment after addition
    AddAssign,
    /// The `/=` operator for assignment after division
    DivAssign,
    /// The `*=` operator for assignment after multiplication
    MulAssign,
    /// The `%=` operator for assignment after remainders
    RemAssign,
    /// The `>>=` operator for assignment after shifting right
    ShrAssign,
    /// The `<<=` operator for assignment after shifting left
    ShlAssign,
    /// The `-=` operator for assignment after subtraction
    SubAssign,
    /// The `|=` operator for assignment after bitwise OR
    BitOrAssign,
    /// The `&=` operator for assignment after bitwise AND
    BitAndAssign,
    /// The `^=` operator for assignment after bitwise XOR
    BitXorAssign,
}

impl ast::BinExpr {
    fn op_details(&self) -> Option<(SyntaxToken, BinOp)> {
        self.syntax().children_with_tokens().filter_map(|it| it.into_token()).find_map(|c| {
            let bin_op = match c.kind() {
                T![||] => BinOp::BooleanOr,
                T![&&] => BinOp::BooleanAnd,
                T![==] => BinOp::EqualityTest,
                T![!=] => BinOp::NegatedEqualityTest,
                T![<=] => BinOp::LesserEqualTest,
                T![>=] => BinOp::GreaterEqualTest,
                T![<] => BinOp::LesserTest,
                T![>] => BinOp::GreaterTest,
                T![+] => BinOp::Addition,
                T![*] => BinOp::Multiplication,
                T![-] => BinOp::Subtraction,
                T![/] => BinOp::Division,
                T![%] => BinOp::Remainder,
                T![<<] => BinOp::LeftShift,
                T![>>] => BinOp::RightShift,
                T![^] => BinOp::BitwiseXor,
                T![|] => BinOp::BitwiseOr,
                T![&] => BinOp::BitwiseAnd,
                T![..] => BinOp::RangeRightOpen,
                T![..=] => BinOp::RangeRightClosed,
                T![=] => BinOp::Assignment,
                T![+=] => BinOp::AddAssign,
                T![/=] => BinOp::DivAssign,
                T![*=] => BinOp::MulAssign,
                T![%=] => BinOp::RemAssign,
                T![>>=] => BinOp::ShrAssign,
                T![<<=] => BinOp::ShlAssign,
                T![-=] => BinOp::SubAssign,
                T![|=] => BinOp::BitOrAssign,
                T![&=] => BinOp::BitAndAssign,
                T![^=] => BinOp::BitXorAssign,
                _ => return None,
            };
            Some((c, bin_op))
        })
    }

    pub fn op_kind(&self) -> Option<BinOp> {
        self.op_details().map(|t| t.1)
    }

    pub fn op_token(&self) -> Option<SyntaxToken> {
        self.op_details().map(|t| t.0)
    }

    pub fn lhs(&self) -> Option<ast::Expr> {
        children(self).nth(0)
    }

    pub fn rhs(&self) -> Option<ast::Expr> {
        children(self).nth(1)
    }

    pub fn sub_exprs(&self) -> (Option<ast::Expr>, Option<ast::Expr>) {
        let mut children = children(self);
        let first = children.next();
        let second = children.next();
        (first, second)
    }
}

impl ast::IndexExpr {
    pub fn base(&self) -> Option<ast::Expr> {
        children(self).nth(0)
    }
    pub fn index(&self) -> Option<ast::Expr> {
        children(self).nth(1)
    }
}

pub enum ArrayExprKind {
    Repeat { initializer: Option<ast::Expr>, repeat: Option<ast::Expr> },
    ElementList(AstChildren<ast::Expr>),
}

impl ast::ArrayExpr {
    pub fn kind(&self) -> ArrayExprKind {
        if self.is_repeat() {
            ArrayExprKind::Repeat {
                initializer: children(self).nth(0),
                repeat: children(self).nth(1),
            }
        } else {
            ArrayExprKind::ElementList(children(self))
        }
    }

    fn is_repeat(&self) -> bool {
        self.syntax().children_with_tokens().any(|it| it.kind() == T![;])
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum LiteralKind {
    String,
    ByteString,
    Char,
    Byte,
    IntNumber { suffix: Option<SmolStr> },
    FloatNumber { suffix: Option<SmolStr> },
    Bool,
}

impl ast::Literal {
    pub fn token(&self) -> SyntaxToken {
        self.syntax()
            .children_with_tokens()
            .find(|e| e.kind() != ATTR && !e.kind().is_trivia())
            .and_then(|e| e.into_token())
            .unwrap()
    }

    pub fn kind(&self) -> LiteralKind {
        match self.token().kind() {
            INT_NUMBER => {
                let int_suffix_list = [
                    "isize", "i128", "i64", "i32", "i16", "i8", "usize", "u128", "u64", "u32",
                    "u16", "u8",
                ];

                // The lexer treats e.g. `1f64` as an integer literal. See
                // https://github.com/rust-analyzer/rust-analyzer/issues/1592
                // and the comments on the linked PR.
                let float_suffix_list = ["f32", "f64"];

                let text = self.token().text().to_string();

                let float_suffix = float_suffix_list
                    .iter()
                    .find(|&s| text.ends_with(s))
                    .map(|&suf| SmolStr::new(suf));

                if float_suffix.is_some() {
                    LiteralKind::FloatNumber { suffix: float_suffix }
                } else {
                    let suffix = int_suffix_list
                        .iter()
                        .find(|&s| text.ends_with(s))
                        .map(|&suf| SmolStr::new(suf));
                    LiteralKind::IntNumber { suffix }
                }
            }
            FLOAT_NUMBER => {
                let allowed_suffix_list = ["f64", "f32"];
                let text = self.token().text().to_string();
                let suffix = allowed_suffix_list
                    .iter()
                    .find(|&s| text.ends_with(s))
                    .map(|&suf| SmolStr::new(suf));
                LiteralKind::FloatNumber { suffix }
            }
            STRING | RAW_STRING => LiteralKind::String,
            T![true] | T![false] => LiteralKind::Bool,
            BYTE_STRING | RAW_BYTE_STRING => LiteralKind::ByteString,
            CHAR => LiteralKind::Char,
            BYTE => LiteralKind::Byte,
            _ => unreachable!(),
        }
    }
}

#[test]
fn test_literal_with_attr() {
    let parse = ast::SourceFile::parse(r#"const _: &str = { #[attr] "Hello" };"#);
    let lit = parse.tree().syntax().descendants().find_map(ast::Literal::cast).unwrap();
    assert_eq!(lit.token().text(), r#""Hello""#);
}

impl ast::NamedField {
    pub fn parent_struct_lit(&self) -> ast::StructLit {
        self.syntax().ancestors().find_map(ast::StructLit::cast).unwrap()
    }
}
<|endoftext|>


<commit_before>#[cfg(all(feature = "runtime-rng", not(all(feature = "compile-time-rng", test))))]
use crate::convert::Convert;
#[cfg(feature = "specialize")]
use crate::BuildHasherExt;

#[cfg(any(
    all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "aes", not(miri)),
    all(any(target_arch = "arm", target_arch = "aarch64"), target_feature = "crypto", not(miri), feature = "stdsimd")
))]
pub use crate::aes_hash::*;

#[cfg(not(any(
    all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "aes", not(miri)),
    all(any(target_arch = "arm", target_arch = "aarch64"), target_feature = "crypto", not(miri), feature = "stdsimd")
)))]
pub use crate::fallback_hash::*;

#[cfg(all(feature = "compile-time-rng", any(not(feature = "runtime-rng"), test)))]
use const_random::const_random;
use core::any::{Any, TypeId};
use core::fmt;
use core::hash::BuildHasher;
#[cfg(feature = "specialize")]
use core::hash::Hash;
use core::hash::Hasher;

#[cfg(not(feature = "std"))]
extern crate alloc;
#[cfg(feature = "std")]
extern crate std as alloc;

use alloc::boxed::Box;
use core::sync::atomic::{AtomicUsize, Ordering};
#[cfg(not(all(target_arch = "arm", target_os = "none")))]
use once_cell::race::OnceBox;

#[cfg(any(
    all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "aes", not(miri)),
    all(any(target_arch = "arm", target_arch = "aarch64"), target_feature = "crypto", not(miri), feature = "stdsimd")
))]
use crate::aes_hash::*;
#[cfg(not(any(
    all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "aes", not(miri)),
    all(any(target_arch = "arm", target_arch = "aarch64"), target_feature = "crypto", not(miri), feature = "stdsimd")
)))]
use crate::fallback_hash::*;

#[cfg(not(all(target_arch = "arm", target_os = "none")))]
static RAND_SOURCE: OnceBox<Box<dyn RandomSource + Send + Sync>> = OnceBox::new();

/// A supplier of Randomness used for different hashers.
/// See [RandomState.set_random_source].
pub trait RandomSource {

    fn get_fixed_seeds(&self) -> &'static [[u64; 4]; 2];

    fn gen_hasher_seed(&self) -> usize;

}

pub(crate) const PI: [u64; 4] = [
    0x243f_6a88_85a3_08d3,
    0x1319_8a2e_0370_7344,
    0xa409_3822_299f_31d0,
    0x082e_fa98_ec4e_6c89,
];

pub(crate) const PI2: [u64; 4] = [
    0x4528_21e6_38d0_1377,
    0xbe54_66cf_34e9_0c6c,
    0xc0ac_29b7_c97c_50dd,
    0x3f84_d5b5_b547_0917,
];

struct DefaultRandomSource {
    counter: AtomicUsize,
}

impl DefaultRandomSource {
    fn new() -> DefaultRandomSource {
        DefaultRandomSource {
            counter: AtomicUsize::new(&PI as *const _ as usize),
        }
    }

    const fn default() -> DefaultRandomSource {
        DefaultRandomSource {
            counter: AtomicUsize::new(PI[3] as usize),
        }
    }
}

impl RandomSource for DefaultRandomSource {

    #[cfg(all(feature = "runtime-rng", not(all(feature = "compile-time-rng", test))))]
    fn get_fixed_seeds(&self) -> &'static [[u64; 4]; 2] {
        static SEEDS: OnceBox<[[u64; 4]; 2]> = OnceBox::new();

        SEEDS.get_or_init(|| {
            let mut result: [u8; 64] = [0; 64];
            getrandom::getrandom(&mut result).expect("getrandom::getrandom() failed.");
            Box::new(result.convert())
        })
    }

    #[cfg(all(feature = "compile-time-rng", any(not(feature = "runtime-rng"), test)))]
    fn get_fixed_seeds(&self) -> &'static [[u64; 4]; 2] {
        const RAND: [[u64; 4]; 2] = [
            [
                const_random!(u64),
                const_random!(u64),
                const_random!(u64),
                const_random!(u64),
            ], [
                const_random!(u64),
                const_random!(u64),
                const_random!(u64),
                const_random!(u64),
            ]
        ];
        &RAND
    }

    #[cfg(all(not(feature = "runtime-rng"), not(feature = "compile-time-rng")))]
    fn get_fixed_seeds(&self) -> &'static [[u64; 4]; 2] {
        &[PI, PI2]
    }

    #[cfg(not(all(target_arch = "arm", target_os = "none")))]
    fn gen_hasher_seed(&self) -> usize {
        let stack = self as *const _ as usize;
        self.counter.fetch_add(stack, Ordering::Relaxed)
    }

    #[cfg(all(target_arch = "arm", target_os = "none"))]
    fn gen_hasher_seed(&self) -> usize {
        let stack = self as *const _ as usize;
        let previous = self.counter.load(Ordering::Relaxed);
        let new = previous.wrapping_add(stack);
        self.counter.store(new, Ordering::Relaxed);
        new
    }
}

/// Provides a [Hasher] factory. This is typically used (e.g. by [HashMap]) to create
/// [AHasher]s in order to hash the keys of the map. See `build_hasher` below.
///
/// [build_hasher]: ahash::
/// [Hasher]: std::hash::Hasher
/// [BuildHasher]: std::hash::BuildHasher
/// [HashMap]: std::collections::HashMap
#[derive(Clone)]
pub struct RandomState {
    pub(crate) k0: u64,
    pub(crate) k1: u64,
    pub(crate) k2: u64,
    pub(crate) k3: u64,
}

impl fmt::Debug for RandomState {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.pad("RandomState { .. }")
    }
}

impl RandomState {

    /// Provides an optional way to manually supply a source of randomness for Hasher keys.
    ///
    /// The provided [RandomSource] will be used to be used as a source of randomness by [RandomState] to generate new states.
    /// If this method is not invoked the standard source of randomness is used as described in the Readme.
    ///
    /// The source of randomness can only be set once, and must be set before the first RandomState is created.
    /// If the source has already been specified `Err` is returned with a `bool` indicating if the set failed because
    /// method was previously invoked (true) or if the default source is already being used (false).
    #[cfg(not(all(target_arch = "arm", target_os = "none")))]
    pub fn set_random_source(source: impl RandomSource + Send + Sync + 'static) -> Result<(), bool> {
        RAND_SOURCE.set(Box::new(Box::new(source))).map_err(|s| s.as_ref().type_id() != TypeId::of::<&DefaultRandomSource>())
    }

    #[inline]
    #[cfg(not(all(target_arch = "arm", target_os = "none")))]
    fn get_src() -> &'static dyn RandomSource {
        RAND_SOURCE.get_or_init(|| Box::new(Box::new(DefaultRandomSource::new()))).as_ref()
    }

    #[inline]
    #[cfg(all(target_arch = "arm", target_os = "none"))]
    fn get_src() -> &'static dyn RandomSource {
        static RAND_SOURCE: DefaultRandomSource = DefaultRandomSource::default();
        &RAND_SOURCE
    }

    /// Use randomly generated keys
    #[inline]
    pub fn new() -> RandomState {
        let src = Self::get_src();
        let fixed = src.get_fixed_seeds();
        Self::from_keys(&fixed[0], &fixed[1], src.gen_hasher_seed())
    }

    /// Allows for supplying seeds, but each time it is called the resulting state will be different.
    /// This is done using a static counter, so it can safely be used with a fixed keys.
    #[inline]
    pub fn generate_with(k0: u64, k1: u64, k2: u64, k3: u64) -> RandomState {
        let src = Self::get_src();
        let fixed = src.get_fixed_seeds();
        RandomState::from_keys(&fixed[0], &[k0, k1, k2, k3], src.gen_hasher_seed())
    }

    fn from_keys(a: &[u64; 4], b: &[u64; 4], c: usize) -> RandomState {
        let &[k0, k1, k2, k3] = a;
        let mut hasher = AHasher::from_random_state(&RandomState { k0, k1, k2, k3 });
        hasher.write_usize(c);
        let mix = |k: u64| {
            let mut h = hasher.clone();
            h.write_u64(k);
            h.finish()
        };
        RandomState {
            k0: mix(b[0]),
            k1: mix(b[1]),
            k2: mix(b[2]),
            k3: mix(b[3]),
        }
    }

    /// Internal. Used by Default.
    #[inline]
    pub(crate) fn with_fixed_keys() -> RandomState {
        let [k0, k1, k2, k3] = Self::get_src().get_fixed_seeds()[0];
        RandomState { k0, k1, k2, k3 }
    }

    /// Allows for explicitly setting a seed to used.
    ///
    /// Note: This method does not require the provided seed to be strong.
    #[inline]
    pub fn with_seed(key: usize) -> RandomState {
        let fixed = Self::get_src().get_fixed_seeds();
        RandomState::from_keys(&fixed[0], &fixed[1], key)
    }

    /// Allows for explicitly setting the seeds to used.
    ///
    /// Note: This method is robust against 0s being passed for one or more of the parameters
    /// or the same value being passed for more than one parameter.
    #[inline]
    pub const fn with_seeds(k0: u64, k1: u64, k2: u64, k3: u64) -> RandomState {
        RandomState { k0: k0 ^ PI2[0], k1: k1 ^ PI2[1], k2: k2 ^ PI2[2], k3: k3 ^ PI2[3] }
    }
}

impl Default for RandomState {
    #[inline]
    fn default() -> Self {
        Self::new()
    }
}

impl BuildHasher for RandomState {
    type Hasher = AHasher;

    /// Constructs a new [AHasher] with keys based on this [RandomState] object.
    /// This means that two different [RandomState]s will will generate
    /// [AHasher]s that will return different hashcodes, but [Hasher]s created from the same [BuildHasher]
    /// will generate the same hashes for the same input data.
    ///
    /// # Examples
    ///
    /// ```
    /// use ahash::{AHasher, RandomState};
    /// use std::hash::{Hasher, BuildHasher};
    ///
    /// let build_hasher = RandomState::new();
    /// let mut hasher_1 = build_hasher.build_hasher();
    /// let mut hasher_2 = build_hasher.build_hasher();
    ///
    /// hasher_1.write_u32(1234);
    /// hasher_2.write_u32(1234);
    ///
    /// assert_eq!(hasher_1.finish(), hasher_2.finish());
    ///
    /// let other_build_hasher = RandomState::new();
    /// let mut different_hasher = other_build_hasher.build_hasher();
    /// different_hasher.write_u32(1234);
    /// assert_ne!(different_hasher.finish(), hasher_1.finish());
    /// ```
    /// [Hasher]: std::hash::Hasher
    /// [BuildHasher]: std::hash::BuildHasher
    /// [HashMap]: std::collections::HashMap
    #[inline]
    fn build_hasher(&self) -> AHasher {
        AHasher::from_random_state(self)
    }
}

#[cfg(feature = "specialize")]
impl BuildHasherExt for RandomState {
    #[inline]
    fn hash_as_u64<T: Hash + ?Sized>(&self, value: &T) -> u64 {
        let mut hasher = AHasherU64 {
            buffer: self.k0,
            pad: self.k1,
        };
        value.hash(&mut hasher);
        hasher.finish()
    }

    #[inline]
    fn hash_as_fixed_length<T: Hash + ?Sized>(&self, value: &T) -> u64 {
        let mut hasher = AHasherFixed(self.build_hasher());
        value.hash(&mut hasher);
        hasher.finish()
    }

    #[inline]
    fn hash_as_str<T: Hash + ?Sized>(&self, value: &T) -> u64 {
        let mut hasher = AHasherStr(self.build_hasher());
        value.hash(&mut hasher);
        hasher.finish()
    }
}

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

    #[test]
    fn test_unique() {
        let a = RandomState::new();
        let b = RandomState::new();
        assert_ne!(a.build_hasher().finish(), b.build_hasher().finish());
    }

    #[cfg(all(feature = "runtime-rng", not(all(feature = "compile-time-rng", test))))]
    #[test]
    fn test_not_pi() {
        assert_ne!(PI, RandomState::get_src().get_fixed_seeds()[0]);
    }

    #[cfg(all(feature = "compile-time-rng", any(not(feature = "runtime-rng"), test)))]
    #[test]
    fn test_not_pi_const() {
        assert_ne!(PI, RandomState::get_src().get_fixed_seeds()[0]);
    }

    #[cfg(all(not(feature = "runtime-rng"), not(feature = "compile-time-rng")))]
    #[test]
    fn test_pi() {
        assert_eq!(PI, RandomState::get_src().get_fixed_seeds()[0]);
    }

    #[test]
    fn test_with_seeds_const() {
        const _CONST_RANDOM_STATE: RandomState = RandomState::with_seeds(17, 19, 21, 23);
    }
}
<commit_msg>More pre-mixing to guard against partially weak keys<commit_after>#[cfg(all(feature = "runtime-rng", not(all(feature = "compile-time-rng", test))))]
use crate::convert::Convert;
#[cfg(feature = "specialize")]
use crate::BuildHasherExt;

#[cfg(any(
    all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "aes", not(miri)),
    all(any(target_arch = "arm", target_arch = "aarch64"), target_feature = "crypto", not(miri), feature = "stdsimd")
))]
pub use crate::aes_hash::*;

#[cfg(not(any(
    all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "aes", not(miri)),
    all(any(target_arch = "arm", target_arch = "aarch64"), target_feature = "crypto", not(miri), feature = "stdsimd")
)))]
pub use crate::fallback_hash::*;

#[cfg(all(feature = "compile-time-rng", any(not(feature = "runtime-rng"), test)))]
use const_random::const_random;
use core::any::{Any, TypeId};
use core::fmt;
use core::hash::BuildHasher;
#[cfg(feature = "specialize")]
use core::hash::Hash;
use core::hash::Hasher;

#[cfg(not(feature = "std"))]
extern crate alloc;
#[cfg(feature = "std")]
extern crate std as alloc;

use alloc::boxed::Box;
use core::sync::atomic::{AtomicUsize, Ordering};
#[cfg(not(all(target_arch = "arm", target_os = "none")))]
use once_cell::race::OnceBox;

#[cfg(any(
    all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "aes", not(miri)),
    all(any(target_arch = "arm", target_arch = "aarch64"), target_feature = "crypto", not(miri), feature = "stdsimd")
))]
use crate::aes_hash::*;
#[cfg(not(any(
    all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "aes", not(miri)),
    all(any(target_arch = "arm", target_arch = "aarch64"), target_feature = "crypto", not(miri), feature = "stdsimd")
)))]
use crate::fallback_hash::*;

#[cfg(not(all(target_arch = "arm", target_os = "none")))]
static RAND_SOURCE: OnceBox<Box<dyn RandomSource + Send + Sync>> = OnceBox::new();

/// A supplier of Randomness used for different hashers.
/// See [RandomState.set_random_source].
pub trait RandomSource {

    fn get_fixed_seeds(&self) -> &'static [[u64; 4]; 2];

    fn gen_hasher_seed(&self) -> usize;

}

pub(crate) const PI: [u64; 4] = [
    0x243f_6a88_85a3_08d3,
    0x1319_8a2e_0370_7344,
    0xa409_3822_299f_31d0,
    0x082e_fa98_ec4e_6c89,
];

pub(crate) const PI2: [u64; 4] = [
    0x4528_21e6_38d0_1377,
    0xbe54_66cf_34e9_0c6c,
    0xc0ac_29b7_c97c_50dd,
    0x3f84_d5b5_b547_0917,
];

struct DefaultRandomSource {
    counter: AtomicUsize,
}

impl DefaultRandomSource {
    fn new() -> DefaultRandomSource {
        DefaultRandomSource {
            counter: AtomicUsize::new(&PI as *const _ as usize),
        }
    }

    const fn default() -> DefaultRandomSource {
        DefaultRandomSource {
            counter: AtomicUsize::new(PI[3] as usize),
        }
    }
}

impl RandomSource for DefaultRandomSource {

    #[cfg(all(feature = "runtime-rng", not(all(feature = "compile-time-rng", test))))]
    fn get_fixed_seeds(&self) -> &'static [[u64; 4]; 2] {
        static SEEDS: OnceBox<[[u64; 4]; 2]> = OnceBox::new();

        SEEDS.get_or_init(|| {
            let mut result: [u8; 64] = [0; 64];
            getrandom::getrandom(&mut result).expect("getrandom::getrandom() failed.");
            Box::new(result.convert())
        })
    }

    #[cfg(all(feature = "compile-time-rng", any(not(feature = "runtime-rng"), test)))]
    fn get_fixed_seeds(&self) -> &'static [[u64; 4]; 2] {
        const RAND: [[u64; 4]; 2] = [
            [
                const_random!(u64),
                const_random!(u64),
                const_random!(u64),
                const_random!(u64),
            ], [
                const_random!(u64),
                const_random!(u64),
                const_random!(u64),
                const_random!(u64),
            ]
        ];
        &RAND
    }

    #[cfg(all(not(feature = "runtime-rng"), not(feature = "compile-time-rng")))]
    fn get_fixed_seeds(&self) -> &'static [[u64; 4]; 2] {
        &[PI, PI2]
    }

    #[cfg(not(all(target_arch = "arm", target_os = "none")))]
    fn gen_hasher_seed(&self) -> usize {
        let stack = self as *const _ as usize;
        self.counter.fetch_add(stack, Ordering::Relaxed)
    }

    #[cfg(all(target_arch = "arm", target_os = "none"))]
    fn gen_hasher_seed(&self) -> usize {
        let stack = self as *const _ as usize;
        let previous = self.counter.load(Ordering::Relaxed);
        let new = previous.wrapping_add(stack);
        self.counter.store(new, Ordering::Relaxed);
        new
    }
}

/// Provides a [Hasher] factory. This is typically used (e.g. by [HashMap]) to create
/// [AHasher]s in order to hash the keys of the map. See `build_hasher` below.
///
/// [build_hasher]: ahash::
/// [Hasher]: std::hash::Hasher
/// [BuildHasher]: std::hash::BuildHasher
/// [HashMap]: std::collections::HashMap
#[derive(Clone)]
pub struct RandomState {
    pub(crate) k0: u64,
    pub(crate) k1: u64,
    pub(crate) k2: u64,
    pub(crate) k3: u64,
}

impl fmt::Debug for RandomState {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.pad("RandomState { .. }")
    }
}

impl RandomState {

    /// Provides an optional way to manually supply a source of randomness for Hasher keys.
    ///
    /// The provided [RandomSource] will be used to be used as a source of randomness by [RandomState] to generate new states.
    /// If this method is not invoked the standard source of randomness is used as described in the Readme.
    ///
    /// The source of randomness can only be set once, and must be set before the first RandomState is created.
    /// If the source has already been specified `Err` is returned with a `bool` indicating if the set failed because
    /// method was previously invoked (true) or if the default source is already being used (false).
    #[cfg(not(all(target_arch = "arm", target_os = "none")))]
    pub fn set_random_source(source: impl RandomSource + Send + Sync + 'static) -> Result<(), bool> {
        RAND_SOURCE.set(Box::new(Box::new(source))).map_err(|s| s.as_ref().type_id() != TypeId::of::<&DefaultRandomSource>())
    }

    #[inline]
    #[cfg(not(all(target_arch = "arm", target_os = "none")))]
    fn get_src() -> &'static dyn RandomSource {
        RAND_SOURCE.get_or_init(|| Box::new(Box::new(DefaultRandomSource::new()))).as_ref()
    }

    #[inline]
    #[cfg(all(target_arch = "arm", target_os = "none"))]
    fn get_src() -> &'static dyn RandomSource {
        static RAND_SOURCE: DefaultRandomSource = DefaultRandomSource::default();
        &RAND_SOURCE
    }

    /// Use randomly generated keys
    #[inline]
    pub fn new() -> RandomState {
        let src = Self::get_src();
        let fixed = src.get_fixed_seeds();
        Self::from_keys(&fixed[0], &fixed[1], src.gen_hasher_seed())
    }

    /// Allows for supplying seeds, but each time it is called the resulting state will be different.
    /// This is done using a static counter, so it can safely be used with a fixed keys.
    #[inline]
    pub fn generate_with(k0: u64, k1: u64, k2: u64, k3: u64) -> RandomState {
        let src = Self::get_src();
        let fixed = src.get_fixed_seeds();
        RandomState::from_keys(&fixed[0], &[k0, k1, k2, k3], src.gen_hasher_seed())
    }

    fn from_keys(a: &[u64; 4], b: &[u64; 4], c: usize) -> RandomState {
        let &[k0, k1, k2, k3] = a;
        let mut hasher = AHasher::from_random_state(&RandomState { k0, k1, k2, k3 });
        hasher.write_usize(c);
        let mix = |l: u64, r: u64| {
            let mut h = hasher.clone();
            h.write_u64(l);
            h.write_u64(r);
            h.finish()
        };
        RandomState {
            k0: mix(b[0], b[2]),
            k1: mix(b[1], b[3]),
            k2: mix(b[2], b[1]),
            k3: mix(b[3], b[0]),
        }
    }

    /// Internal. Used by Default.
    #[inline]
    pub(crate) fn with_fixed_keys() -> RandomState {
        let [k0, k1, k2, k3] = Self::get_src().get_fixed_seeds()[0];
        RandomState { k0, k1, k2, k3 }
    }

    /// Allows for explicitly setting a seed to used.
    ///
    /// Note: This method does not require the provided seed to be strong.
    #[inline]
    pub fn with_seed(key: usize) -> RandomState {
        let fixed = Self::get_src().get_fixed_seeds();
        RandomState::from_keys(&fixed[0], &fixed[1], key)
    }

    /// Allows for explicitly setting the seeds to used.
    ///
    /// Note: This method is robust against 0s being passed for one or more of the parameters
    /// or the same value being passed for more than one parameter.
    #[inline]
    pub const fn with_seeds(k0: u64, k1: u64, k2: u64, k3: u64) -> RandomState {
        RandomState { k0: k0 ^ PI2[0], k1: k1 ^ PI2[1], k2: k2 ^ PI2[2], k3: k3 ^ PI2[3] }
    }
}

impl Default for RandomState {
    #[inline]
    fn default() -> Self {
        Self::new()
    }
}

impl BuildHasher for RandomState {
    type Hasher = AHasher;

    /// Constructs a new [AHasher] with keys based on this [RandomState] object.
    /// This means that two different [RandomState]s will will generate
    /// [AHasher]s that will return different hashcodes, but [Hasher]s created from the same [BuildHasher]
    /// will generate the same hashes for the same input data.
    ///
    /// # Examples
    ///
    /// ```
    /// use ahash::{AHasher, RandomState};
    /// use std::hash::{Hasher, BuildHasher};
    ///
    /// let build_hasher = RandomState::new();
    /// let mut hasher_1 = build_hasher.build_hasher();
    /// let mut hasher_2 = build_hasher.build_hasher();
    ///
    /// hasher_1.write_u32(1234);
    /// hasher_2.write_u32(1234);
    ///
    /// assert_eq!(hasher_1.finish(), hasher_2.finish());
    ///
    /// let other_build_hasher = RandomState::new();
    /// let mut different_hasher = other_build_hasher.build_hasher();
    /// different_hasher.write_u32(1234);
    /// assert_ne!(different_hasher.finish(), hasher_1.finish());
    /// ```
    /// [Hasher]: std::hash::Hasher
    /// [BuildHasher]: std::hash::BuildHasher
    /// [HashMap]: std::collections::HashMap
    #[inline]
    fn build_hasher(&self) -> AHasher {
        AHasher::from_random_state(self)
    }
}

#[cfg(feature = "specialize")]
impl BuildHasherExt for RandomState {
    #[inline]
    fn hash_as_u64<T: Hash + ?Sized>(&self, value: &T) -> u64 {
        let mut hasher = AHasherU64 {
            buffer: self.k0,
            pad: self.k1,
        };
        value.hash(&mut hasher);
        hasher.finish()
    }

    #[inline]
    fn hash_as_fixed_length<T: Hash + ?Sized>(&self, value: &T) -> u64 {
        let mut hasher = AHasherFixed(self.build_hasher());
        value.hash(&mut hasher);
        hasher.finish()
    }

    #[inline]
    fn hash_as_str<T: Hash + ?Sized>(&self, value: &T) -> u64 {
        let mut hasher = AHasherStr(self.build_hasher());
        value.hash(&mut hasher);
        hasher.finish()
    }
}

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

    #[test]
    fn test_unique() {
        let a = RandomState::new();
        let b = RandomState::new();
        assert_ne!(a.build_hasher().finish(), b.build_hasher().finish());
    }

    #[cfg(all(feature = "runtime-rng", not(all(feature = "compile-time-rng", test))))]
    #[test]
    fn test_not_pi() {
        assert_ne!(PI, RandomState::get_src().get_fixed_seeds()[0]);
    }

    #[cfg(all(feature = "compile-time-rng", any(not(feature = "runtime-rng"), test)))]
    #[test]
    fn test_not_pi_const() {
        assert_ne!(PI, RandomState::get_src().get_fixed_seeds()[0]);
    }

    #[cfg(all(not(feature = "runtime-rng"), not(feature = "compile-time-rng")))]
    #[test]
    fn test_pi() {
        assert_eq!(PI, RandomState::get_src().get_fixed_seeds()[0]);
    }

    #[test]
    fn test_with_seeds_const() {
        const _CONST_RANDOM_STATE: RandomState = RandomState::with_seeds(17, 19, 21, 23);
    }
}
<|endoftext|>


<commit_before>use crate::{decode_config_slice, Config, DecodeError};
use std::io::Read;
use std::{cmp, fmt, io};

// This should be large, but it has to fit on the stack.
pub(crate) const BUF_SIZE: usize = 1024;

// 4 bytes of base64 data encode 3 bytes of raw data (modulo padding).
const BASE64_CHUNK_SIZE: usize = 4;
const DECODED_CHUNK_SIZE: usize = 3;

/// A `Read` implementation that decodes base64 data read from an underlying reader.
///
/// # Examples
///
/// ```
/// use std::io::Read;
/// use std::io::Cursor;
///
/// // use a cursor as the simplest possible `Read` -- in real code this is probably a file, etc.
/// let mut wrapped_reader = Cursor::new(b"YXNkZg==");
/// let mut decoder = base64::read::DecoderReader::new(
///     &mut wrapped_reader, base64::STANDARD);
///
/// // handle errors as you normally would
/// let mut result = Vec::new();
/// decoder.read_to_end(&mut result).unwrap();
///
/// assert_eq!(b"asdf", &result[..]);
///
/// ```
pub struct DecoderReader<'a, R: 'a + io::Read> {
    config: Config,
    /// Where b64 data is read from
    r: &'a mut R,

    // Holds b64 data read from the delegate reader.
    b64_buffer: [u8; BUF_SIZE],
    // The start of the pending buffered data in b64_buffer.
    b64_offset: usize,
    // The amount of buffered b64 data.
    b64_len: usize,
    // Since the caller may provide us with a buffer of size 1 or 2 that's too small to copy a
    // decoded chunk in to, we have to be able to hang on to a few decoded bytes.
    // Technically we only need to hold 2 bytes but then we'd need a separate temporary buffer to
    // decode 3 bytes into and then juggle copying one byte into the provided read buf and the rest
    // into here, which seems like a lot of complexity for 1 extra byte of storage.
    decoded_buffer: [u8; 3],
    // index of start of decoded data
    decoded_offset: usize,
    // length of decoded data
    decoded_len: usize,
    // used to provide accurate offsets in errors
    total_b64_decoded: usize,
}

impl<'a, R: io::Read> fmt::Debug for DecoderReader<'a, R> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("DecoderReader")
            .field("config", &self.config)
            .field("b64_offset", &self.b64_offset)
            .field("b64_len", &self.b64_len)
            .field("decoded_buffer", &self.decoded_buffer)
            .field("decoded_offset", &self.decoded_offset)
            .field("decoded_len", &self.decoded_len)
            .field("total_b64_decoded", &self.total_b64_decoded)
            .finish()
    }
}

impl<'a, R: io::Read> DecoderReader<'a, R> {
    /// Create a new decoder that will read from the provided reader `r`.
    pub fn new(r: &'a mut R, config: Config) -> Self {
        DecoderReader {
            config,
            r,
            b64_buffer: [0; BUF_SIZE],
            b64_offset: 0,
            b64_len: 0,
            decoded_buffer: [0; DECODED_CHUNK_SIZE],
            decoded_offset: 0,
            decoded_len: 0,
            total_b64_decoded: 0,
        }
    }

    /// Write as much as possible of the decoded buffer into the target buffer.
    /// Must only be called when there is something to write and space to write into.
    /// Returns a Result with the number of (decoded) bytes copied.
    fn flush_decoded_buf(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        debug_assert!(self.decoded_len > 0);
        debug_assert!(buf.len() > 0);

        let copy_len = cmp::min(self.decoded_len, buf.len());
        debug_assert!(copy_len > 0);
        debug_assert!(copy_len <= self.decoded_len);

        buf[..copy_len].copy_from_slice(
            &self.decoded_buffer[self.decoded_offset..self.decoded_offset + copy_len],
        );

        self.decoded_offset += copy_len;
        self.decoded_len -= copy_len;

        debug_assert!(self.decoded_len < DECODED_CHUNK_SIZE);

        Ok(copy_len)
    }

    /// Read into the remaining space in the buffer after the current contents.
    /// Must only be called when there is space to read into in the buffer.
    /// Returns the number of bytes read.
    fn read_from_delegate(&mut self) -> io::Result<usize> {
        debug_assert!(self.b64_offset + self.b64_len < BUF_SIZE);

        let read = self
            .r
            .read(&mut self.b64_buffer[self.b64_offset + self.b64_len..])?;
        self.b64_len += read;

        debug_assert!(self.b64_offset + self.b64_len <= BUF_SIZE);

        return Ok(read);
    }

    /// Decode the requested number of bytes from the b64 buffer into the provided buffer. It's the
    /// caller's responsibility to choose the number of b64 bytes to decode correctly.
    ///
    /// Returns a Result with the number of decoded bytes written.
    fn decode_to_buf(&mut self, num_bytes: usize, buf: &mut [u8]) -> io::Result<usize> {
        debug_assert!(self.b64_len >= num_bytes);
        debug_assert!(self.b64_offset + self.b64_len <= BUF_SIZE);
        debug_assert!(buf.len() > 0);

        let decoded = decode_config_slice(
            &self.b64_buffer[self.b64_offset..self.b64_offset + num_bytes],
            self.config,
            &mut buf[..],
        )
        .map_err(|e| match e {
            DecodeError::InvalidByte(offset, byte) => {
                DecodeError::InvalidByte(self.total_b64_decoded + offset, byte)
            }
            DecodeError::InvalidLength => DecodeError::InvalidLength,
            DecodeError::InvalidLastSymbol(offset, byte) => {
                DecodeError::InvalidLastSymbol(self.total_b64_decoded + offset, byte)
            }
        })
        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;

        self.total_b64_decoded += num_bytes;
        self.b64_len -= num_bytes;
        self.b64_offset += num_bytes;

        debug_assert!(self.b64_offset + self.b64_len <= BUF_SIZE);

        Ok(decoded)
    }
}

impl<'a, R: Read> Read for DecoderReader<'a, R> {
    /// Decode input from the wrapped reader.
    ///
    /// Under non-error circumstances, this returns `Ok` with the value being the number of bytes
    /// written in `buf`.
    ///
    /// Where possible, this function buffers base64 to minimize the number of read() calls to the
    /// delegate reader.
    ///
    /// # Errors
    ///
    /// Any errors emitted by the delegate reader are returned. Decoding errors due to invalid
    /// base64 are also possible, and will have `io::ErrorKind::InvalidData`.
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if buf.len() == 0 {
            return Ok(0);
        }

        // offset == BUF_SIZE when we copied it all last time
        debug_assert!(self.b64_offset <= BUF_SIZE);
        debug_assert!(self.b64_offset + self.b64_len <= BUF_SIZE);
        debug_assert!(if self.b64_offset == BUF_SIZE {
            self.b64_len == 0
        } else {
            self.b64_len <= BUF_SIZE
        });

        debug_assert!(if self.decoded_len == 0 {
            // can be = when we were able to copy the complete chunk
            self.decoded_offset <= DECODED_CHUNK_SIZE
        } else {
            self.decoded_offset < DECODED_CHUNK_SIZE
        });

        // We shouldn't ever decode into here when we can't immediately write at least one byte into
        // the provided buf, so the effective length should only be 3 momentarily between when we
        // decode and when we copy into the target buffer.
        debug_assert!(self.decoded_len < DECODED_CHUNK_SIZE);
        debug_assert!(self.decoded_len + self.decoded_offset <= DECODED_CHUNK_SIZE);

        if self.decoded_len > 0 {
            // we have a few leftover decoded bytes; flush that rather than pull in more b64
            self.flush_decoded_buf(buf)
        } else {
            let mut at_eof = false;
            while self.b64_len < BASE64_CHUNK_SIZE {
                // scoot the few bytes we have to the beginning
                self.b64_buffer
                    .copy_within(self.b64_offset..self.b64_offset + self.b64_len, 0);
                self.b64_offset = 0;

                // then fill in more data
                let read = self.read_from_delegate()?;
                if read == 0 {
                    // we never pass in an empty buf, so 0 => we've hit EOF
                    at_eof = true;
                    break;
                }
            }

            if self.b64_len == 0 {
                // we must be at EOF, and we have no data left to decode
                return Ok(0);
            };

            debug_assert!(if at_eof {
                // if we are at eof, we may not have a complete chunk
                self.b64_len > 0
            } else {
                // otherwise, we must have at least one chunk
                self.b64_len >= BASE64_CHUNK_SIZE
            });

            if buf.len() < DECODED_CHUNK_SIZE {
                // caller requested an annoyingly short read
                debug_assert_eq!(0, self.decoded_len);

                // have to write to a tmp buf first to avoid double mutable borrow
                let mut decoded_chunk = [0_u8; DECODED_CHUNK_SIZE];
                // if we are at eof, could have less than BASE64_CHUNK_SIZE
                let to_decode = cmp::min(self.b64_len, BASE64_CHUNK_SIZE);

                let decoded = self.decode_to_buf(to_decode, &mut decoded_chunk[..])?;
                self.decoded_buffer[..decoded].copy_from_slice(&decoded_chunk[..decoded]);

                self.decoded_offset = 0;
                self.decoded_len = decoded;

                // can be less than 3 on last block due to padding
                debug_assert!(decoded <= 3);

                self.flush_decoded_buf(buf)
            } else {
                let bytes_that_can_decode_into_buf = (buf.len() / DECODED_CHUNK_SIZE)
                    .checked_mul(BASE64_CHUNK_SIZE)
                    .expect("too many chunks");
                debug_assert!(bytes_that_can_decode_into_buf >= BASE64_CHUNK_SIZE);

                let bytes_available_to_decode = if at_eof {
                    self.b64_len
                } else {
                    // only use complete chunks
                    self.b64_len - self.b64_len % 4
                };

                let actual_decode_len =
                    cmp::min(bytes_that_can_decode_into_buf, bytes_available_to_decode);
                self.decode_to_buf(actual_decode_len, buf)
            }
        }
    }
}
<commit_msg>Fix 1.34.0 compatibility<commit_after>use crate::{decode_config_slice, Config, DecodeError};
use std::io::Read;
use std::{cmp, fmt, io};

// This should be large, but it has to fit on the stack.
pub(crate) const BUF_SIZE: usize = 1024;

// 4 bytes of base64 data encode 3 bytes of raw data (modulo padding).
const BASE64_CHUNK_SIZE: usize = 4;
const DECODED_CHUNK_SIZE: usize = 3;

/// A `Read` implementation that decodes base64 data read from an underlying reader.
///
/// # Examples
///
/// ```
/// use std::io::Read;
/// use std::io::Cursor;
///
/// // use a cursor as the simplest possible `Read` -- in real code this is probably a file, etc.
/// let mut wrapped_reader = Cursor::new(b"YXNkZg==");
/// let mut decoder = base64::read::DecoderReader::new(
///     &mut wrapped_reader, base64::STANDARD);
///
/// // handle errors as you normally would
/// let mut result = Vec::new();
/// decoder.read_to_end(&mut result).unwrap();
///
/// assert_eq!(b"asdf", &result[..]);
///
/// ```
pub struct DecoderReader<'a, R: 'a + io::Read> {
    config: Config,
    /// Where b64 data is read from
    r: &'a mut R,

    // Holds b64 data read from the delegate reader.
    b64_buffer: [u8; BUF_SIZE],
    // The start of the pending buffered data in b64_buffer.
    b64_offset: usize,
    // The amount of buffered b64 data.
    b64_len: usize,
    // Since the caller may provide us with a buffer of size 1 or 2 that's too small to copy a
    // decoded chunk in to, we have to be able to hang on to a few decoded bytes.
    // Technically we only need to hold 2 bytes but then we'd need a separate temporary buffer to
    // decode 3 bytes into and then juggle copying one byte into the provided read buf and the rest
    // into here, which seems like a lot of complexity for 1 extra byte of storage.
    decoded_buffer: [u8; 3],
    // index of start of decoded data
    decoded_offset: usize,
    // length of decoded data
    decoded_len: usize,
    // used to provide accurate offsets in errors
    total_b64_decoded: usize,
}

impl<'a, R: io::Read> fmt::Debug for DecoderReader<'a, R> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("DecoderReader")
            .field("config", &self.config)
            .field("b64_offset", &self.b64_offset)
            .field("b64_len", &self.b64_len)
            .field("decoded_buffer", &self.decoded_buffer)
            .field("decoded_offset", &self.decoded_offset)
            .field("decoded_len", &self.decoded_len)
            .field("total_b64_decoded", &self.total_b64_decoded)
            .finish()
    }
}

impl<'a, R: io::Read> DecoderReader<'a, R> {
    /// Create a new decoder that will read from the provided reader `r`.
    pub fn new(r: &'a mut R, config: Config) -> Self {
        DecoderReader {
            config,
            r,
            b64_buffer: [0; BUF_SIZE],
            b64_offset: 0,
            b64_len: 0,
            decoded_buffer: [0; DECODED_CHUNK_SIZE],
            decoded_offset: 0,
            decoded_len: 0,
            total_b64_decoded: 0,
        }
    }

    /// Write as much as possible of the decoded buffer into the target buffer.
    /// Must only be called when there is something to write and space to write into.
    /// Returns a Result with the number of (decoded) bytes copied.
    fn flush_decoded_buf(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        debug_assert!(self.decoded_len > 0);
        debug_assert!(buf.len() > 0);

        let copy_len = cmp::min(self.decoded_len, buf.len());
        debug_assert!(copy_len > 0);
        debug_assert!(copy_len <= self.decoded_len);

        buf[..copy_len].copy_from_slice(
            &self.decoded_buffer[self.decoded_offset..self.decoded_offset + copy_len],
        );

        self.decoded_offset += copy_len;
        self.decoded_len -= copy_len;

        debug_assert!(self.decoded_len < DECODED_CHUNK_SIZE);

        Ok(copy_len)
    }

    /// Read into the remaining space in the buffer after the current contents.
    /// Must only be called when there is space to read into in the buffer.
    /// Returns the number of bytes read.
    fn read_from_delegate(&mut self) -> io::Result<usize> {
        debug_assert!(self.b64_offset + self.b64_len < BUF_SIZE);

        let read = self
            .r
            .read(&mut self.b64_buffer[self.b64_offset + self.b64_len..])?;
        self.b64_len += read;

        debug_assert!(self.b64_offset + self.b64_len <= BUF_SIZE);

        return Ok(read);
    }

    /// Decode the requested number of bytes from the b64 buffer into the provided buffer. It's the
    /// caller's responsibility to choose the number of b64 bytes to decode correctly.
    ///
    /// Returns a Result with the number of decoded bytes written.
    fn decode_to_buf(&mut self, num_bytes: usize, buf: &mut [u8]) -> io::Result<usize> {
        debug_assert!(self.b64_len >= num_bytes);
        debug_assert!(self.b64_offset + self.b64_len <= BUF_SIZE);
        debug_assert!(buf.len() > 0);

        let decoded = decode_config_slice(
            &self.b64_buffer[self.b64_offset..self.b64_offset + num_bytes],
            self.config,
            &mut buf[..],
        )
        .map_err(|e| match e {
            DecodeError::InvalidByte(offset, byte) => {
                DecodeError::InvalidByte(self.total_b64_decoded + offset, byte)
            }
            DecodeError::InvalidLength => DecodeError::InvalidLength,
            DecodeError::InvalidLastSymbol(offset, byte) => {
                DecodeError::InvalidLastSymbol(self.total_b64_decoded + offset, byte)
            }
        })
        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;

        self.total_b64_decoded += num_bytes;
        self.b64_len -= num_bytes;
        self.b64_offset += num_bytes;

        debug_assert!(self.b64_offset + self.b64_len <= BUF_SIZE);

        Ok(decoded)
    }
}

impl<'a, R: Read> Read for DecoderReader<'a, R> {
    /// Decode input from the wrapped reader.
    ///
    /// Under non-error circumstances, this returns `Ok` with the value being the number of bytes
    /// written in `buf`.
    ///
    /// Where possible, this function buffers base64 to minimize the number of read() calls to the
    /// delegate reader.
    ///
    /// # Errors
    ///
    /// Any errors emitted by the delegate reader are returned. Decoding errors due to invalid
    /// base64 are also possible, and will have `io::ErrorKind::InvalidData`.
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if buf.len() == 0 {
            return Ok(0);
        }

        // offset == BUF_SIZE when we copied it all last time
        debug_assert!(self.b64_offset <= BUF_SIZE);
        debug_assert!(self.b64_offset + self.b64_len <= BUF_SIZE);
        debug_assert!(if self.b64_offset == BUF_SIZE {
            self.b64_len == 0
        } else {
            self.b64_len <= BUF_SIZE
        });

        debug_assert!(if self.decoded_len == 0 {
            // can be = when we were able to copy the complete chunk
            self.decoded_offset <= DECODED_CHUNK_SIZE
        } else {
            self.decoded_offset < DECODED_CHUNK_SIZE
        });

        // We shouldn't ever decode into here when we can't immediately write at least one byte into
        // the provided buf, so the effective length should only be 3 momentarily between when we
        // decode and when we copy into the target buffer.
        debug_assert!(self.decoded_len < DECODED_CHUNK_SIZE);
        debug_assert!(self.decoded_len + self.decoded_offset <= DECODED_CHUNK_SIZE);

        if self.decoded_len > 0 {
            // we have a few leftover decoded bytes; flush that rather than pull in more b64
            self.flush_decoded_buf(buf)
        } else {
            let mut at_eof = false;
            while self.b64_len < BASE64_CHUNK_SIZE {
                // Work around lack of copy_within, which is only present in 1.37
                // Copy any bytes we have to the start of the buffer.
                // We know we have < 1 chunk, so we can use a tiny tmp buffer.
                let mut memmove_buf = [0_u8; BASE64_CHUNK_SIZE];
                memmove_buf[..self.b64_len].copy_from_slice(
                    &self.b64_buffer[self.b64_offset..self.b64_offset + self.b64_len],
                );
                self.b64_buffer[0..self.b64_len].copy_from_slice(&memmove_buf[..self.b64_len]);
                self.b64_offset = 0;

                // then fill in more data
                let read = self.read_from_delegate()?;
                if read == 0 {
                    // we never pass in an empty buf, so 0 => we've hit EOF
                    at_eof = true;
                    break;
                }
            }

            if self.b64_len == 0 {
                // we must be at EOF, and we have no data left to decode
                return Ok(0);
            };

            debug_assert!(if at_eof {
                // if we are at eof, we may not have a complete chunk
                self.b64_len > 0
            } else {
                // otherwise, we must have at least one chunk
                self.b64_len >= BASE64_CHUNK_SIZE
            });

            if buf.len() < DECODED_CHUNK_SIZE {
                // caller requested an annoyingly short read
                debug_assert_eq!(0, self.decoded_len);

                // have to write to a tmp buf first to avoid double mutable borrow
                let mut decoded_chunk = [0_u8; DECODED_CHUNK_SIZE];
                // if we are at eof, could have less than BASE64_CHUNK_SIZE
                let to_decode = cmp::min(self.b64_len, BASE64_CHUNK_SIZE);

                let decoded = self.decode_to_buf(to_decode, &mut decoded_chunk[..])?;
                self.decoded_buffer[..decoded].copy_from_slice(&decoded_chunk[..decoded]);

                self.decoded_offset = 0;
                self.decoded_len = decoded;

                // can be less than 3 on last block due to padding
                debug_assert!(decoded <= 3);

                self.flush_decoded_buf(buf)
            } else {
                let bytes_that_can_decode_into_buf = (buf.len() / DECODED_CHUNK_SIZE)
                    .checked_mul(BASE64_CHUNK_SIZE)
                    .expect("too many chunks");
                debug_assert!(bytes_that_can_decode_into_buf >= BASE64_CHUNK_SIZE);

                let bytes_available_to_decode = if at_eof {
                    self.b64_len
                } else {
                    // only use complete chunks
                    self.b64_len - self.b64_len % 4
                };

                let actual_decode_len =
                    cmp::min(bytes_that_can_decode_into_buf, bytes_available_to_decode);
                self.decode_to_buf(actual_decode_len, buf)
            }
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (C) 2021, Cloudflare, Inc.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright notice,
//       this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

//! Proportional Rate Reduction
//!
//! This implementation is based on the following RFC:
//!
//! <https://datatracker.ietf.org/doc/html/rfc6937>

use std::cmp;

#[derive(Default, Debug)]
pub struct PRR {
    // Total bytes delivered during recovery.
    prr_delivered: usize,

    // FlightSize at the start of recovery.
    recoverfs: usize,

    // Total bytes sent during recovery.
    prr_out: usize,

    // Total additional bytes can be sent for retransmit during recovery.
    pub snd_cnt: usize,
}

impl PRR {
    pub fn on_packet_sent(&mut self, sent_bytes: usize) {
        self.prr_out += sent_bytes;

        self.snd_cnt = self.snd_cnt.saturating_sub(sent_bytes);
    }

    pub fn congestion_event(&mut self, bytes_in_flight: usize) {
        self.prr_delivered = 0;

        self.recoverfs = bytes_in_flight;

        self.prr_out = 0;

        self.snd_cnt = 0;
    }

    pub fn on_packet_acked(
        &mut self, delivered_data: usize, pipe: usize, ssthresh: usize,
        max_datagram_size: usize,
    ) {
        self.prr_delivered += delivered_data;

        self.snd_cnt = if pipe > ssthresh {
            // Proportional Rate Reduction.
            ((self.prr_delivered * ssthresh + self.recoverfs - 1) /
                self.recoverfs)
                .saturating_sub(self.prr_out)
        } else {
            // PRR-SSRB.
            let limit = cmp::max(
                self.prr_delivered.saturating_sub(self.prr_out),
                delivered_data,
            ) + max_datagram_size;

            // Attempt to catch up, as permitted by limit
            cmp::min(ssthresh - pipe, limit)
        };

        // snd_cnt should be a positive number.
        self.snd_cnt = cmp::max(self.snd_cnt, 0);
    }
}

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

    #[test]
    fn congestion_event() {
        let mut prr = PRR::default();
        let bytes_in_flight = 1000;

        prr.congestion_event(bytes_in_flight);

        assert_eq!(prr.recoverfs, bytes_in_flight);
        assert_eq!(prr.snd_cnt, 0);
    }

    #[test]
    fn on_packet_sent() {
        let mut prr = PRR::default();
        let bytes_in_flight = 1000;
        let bytes_sent = 500;

        prr.congestion_event(bytes_in_flight);

        prr.on_packet_sent(bytes_sent);

        assert_eq!(prr.prr_out, bytes_sent);
        assert_eq!(prr.snd_cnt, 0);
    }

    #[test]
    fn on_packet_acked_prr() {
        let mut prr = PRR::default();
        let max_datagram_size = 1000;
        let bytes_in_flight = max_datagram_size * 10;
        let ssthresh = bytes_in_flight / 2;
        let acked = 1000;

        prr.congestion_event(bytes_in_flight);

        // pipe > ssthresh uses PRR algorithm.
        let pipe = bytes_in_flight;

        prr.on_packet_acked(acked, pipe, ssthresh, max_datagram_size);

        assert_eq!(prr.snd_cnt, 500);

        let snd_cnt = prr.snd_cnt;

        // send one more allowed by snd_cnt
        prr.on_packet_sent(snd_cnt);

        prr.on_packet_acked(acked, pipe, ssthresh, max_datagram_size);

        assert_eq!(prr.snd_cnt, 500);
    }

    #[test]
    fn on_packet_acked_prr_ssrb() {
        let mut prr = PRR::default();
        let max_datagram_size = 1000;
        let bytes_in_flight = max_datagram_size * 10;
        let ssthresh = bytes_in_flight / 2;
        let acked = 1000;

        prr.congestion_event(bytes_in_flight);

        // pipe <= ssthresh uses PRR-SSRB algorithm.
        let pipe = max_datagram_size;

        prr.on_packet_acked(acked, pipe, ssthresh, max_datagram_size);

        assert_eq!(prr.snd_cnt, 2000);

        let snd_cnt = prr.snd_cnt;

        // send one more allowed by snd_cnt
        prr.on_packet_sent(snd_cnt);

        prr.on_packet_acked(acked, pipe, ssthresh, max_datagram_size);

        assert_eq!(prr.snd_cnt, 2000);
    }
}
<commit_msg>recovery: fix more PRR calculation issues<commit_after>// Copyright (C) 2021, Cloudflare, Inc.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright notice,
//       this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

//! Proportional Rate Reduction
//!
//! This implementation is based on the following RFC:
//!
//! <https://datatracker.ietf.org/doc/html/rfc6937>

use std::cmp;

#[derive(Default, Debug)]
pub struct PRR {
    // Total bytes delivered during recovery.
    prr_delivered: usize,

    // FlightSize at the start of recovery.
    recoverfs: usize,

    // Total bytes sent during recovery.
    prr_out: usize,

    // Total additional bytes can be sent for retransmit during recovery.
    pub snd_cnt: usize,
}

impl PRR {
    pub fn on_packet_sent(&mut self, sent_bytes: usize) {
        self.prr_out += sent_bytes;

        self.snd_cnt = self.snd_cnt.saturating_sub(sent_bytes);
    }

    pub fn congestion_event(&mut self, bytes_in_flight: usize) {
        self.prr_delivered = 0;

        self.recoverfs = bytes_in_flight;

        self.prr_out = 0;

        self.snd_cnt = 0;
    }

    pub fn on_packet_acked(
        &mut self, delivered_data: usize, pipe: usize, ssthresh: usize,
        max_datagram_size: usize,
    ) {
        self.prr_delivered += delivered_data;

        self.snd_cnt = if pipe > ssthresh {
            // Proportional Rate Reduction.
            if self.recoverfs > 0 {
                ((self.prr_delivered * ssthresh + self.recoverfs - 1) /
                    self.recoverfs)
                    .saturating_sub(self.prr_out)
            } else {
                0
            }
        } else {
            // PRR-SSRB.
            let limit = cmp::max(
                self.prr_delivered.saturating_sub(self.prr_out),
                delivered_data,
            ) + max_datagram_size;

            // Attempt to catch up, as permitted by limit
            cmp::min(ssthresh - pipe, limit)
        };

        // snd_cnt should be a positive number.
        self.snd_cnt = cmp::max(self.snd_cnt, 0);
    }
}

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

    #[test]
    fn congestion_event() {
        let mut prr = PRR::default();
        let bytes_in_flight = 1000;

        prr.congestion_event(bytes_in_flight);

        assert_eq!(prr.recoverfs, bytes_in_flight);
        assert_eq!(prr.snd_cnt, 0);
    }

    #[test]
    fn on_packet_sent() {
        let mut prr = PRR::default();
        let bytes_in_flight = 1000;
        let bytes_sent = 500;

        prr.congestion_event(bytes_in_flight);

        prr.on_packet_sent(bytes_sent);

        assert_eq!(prr.prr_out, bytes_sent);
        assert_eq!(prr.snd_cnt, 0);
    }

    #[test]
    fn on_packet_acked_prr() {
        let mut prr = PRR::default();
        let max_datagram_size = 1000;
        let bytes_in_flight = max_datagram_size * 10;
        let ssthresh = bytes_in_flight / 2;
        let acked = 1000;

        prr.congestion_event(bytes_in_flight);

        // pipe > ssthresh uses PRR algorithm.
        let pipe = bytes_in_flight;

        prr.on_packet_acked(acked, pipe, ssthresh, max_datagram_size);

        assert_eq!(prr.snd_cnt, 500);

        let snd_cnt = prr.snd_cnt;

        // send one more allowed by snd_cnt
        prr.on_packet_sent(snd_cnt);

        prr.on_packet_acked(acked, pipe, ssthresh, max_datagram_size);

        assert_eq!(prr.snd_cnt, 500);
    }

    #[test]
    fn on_packet_acked_prr_overflow() {
        let mut prr = PRR::default();
        let max_datagram_size = 1000;
        let bytes_in_flight = max_datagram_size * 10;
        let ssthresh = bytes_in_flight / 2;
        let acked = 1000;

        prr.congestion_event(bytes_in_flight);

        prr.on_packet_sent(max_datagram_size);

        // pipe > ssthresh uses PRR algorithm.
        let pipe = bytes_in_flight + max_datagram_size;

        prr.on_packet_acked(acked, pipe, ssthresh, max_datagram_size);

        assert_eq!(prr.snd_cnt, 0);
    }

    #[test]
    fn on_packet_acked_prr_zero_in_flight() {
        let mut prr = PRR::default();
        let max_datagram_size = 1000;
        let bytes_in_flight = 0;
        let ssthresh = 3000;
        let acked = 1000;

        prr.congestion_event(bytes_in_flight);

        // pipe > ssthresh uses PRR algorithm.
        let pipe = ssthresh + 1000;

        prr.on_packet_acked(acked, pipe, ssthresh, max_datagram_size);

        assert_eq!(prr.snd_cnt, 0);
    }

    #[test]
    fn on_packet_acked_prr_ssrb() {
        let mut prr = PRR::default();
        let max_datagram_size = 1000;
        let bytes_in_flight = max_datagram_size * 10;
        let ssthresh = bytes_in_flight / 2;
        let acked = 1000;

        prr.congestion_event(bytes_in_flight);

        // pipe <= ssthresh uses PRR-SSRB algorithm.
        let pipe = max_datagram_size;

        prr.on_packet_acked(acked, pipe, ssthresh, max_datagram_size);

        assert_eq!(prr.snd_cnt, 2000);

        let snd_cnt = prr.snd_cnt;

        // send one more allowed by snd_cnt
        prr.on_packet_sent(snd_cnt);

        prr.on_packet_acked(acked, pipe, ssthresh, max_datagram_size);

        assert_eq!(prr.snd_cnt, 2000);
    }

    #[test]
    fn on_packet_acked_prr_ssrb_overflow() {
        let mut prr = PRR::default();
        let max_datagram_size = 1000;
        let bytes_in_flight = max_datagram_size * 10;
        let ssthresh = bytes_in_flight / 2;
        let acked = 500;

        prr.congestion_event(bytes_in_flight);

        // pipe <= ssthresh uses PRR-SSRB algorithm.
        let pipe = max_datagram_size;

        prr.on_packet_sent(max_datagram_size);

        prr.on_packet_acked(acked, pipe, ssthresh, max_datagram_size);

        assert_eq!(prr.snd_cnt, 1500);
    }
}
<|endoftext|>


<commit_before>pub fn add(num_l: &str, num_r: &str) -> String {
    let mut sum: Vec<u8> = (String::from(num_l) + num_r).into();

    sum.sort();
    sum.reverse();

    compress(&String::from_utf8(sum).unwrap())
}

pub fn compress(num: &str) -> String {
    match num {
        "IIIII" => "V".to_string(),
        "IIII" => "IV".to_string(),
        _ => num.to_string(),
    }
}

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

    #[test]
    fn add_i_i() {
        assert_eq!("II", add("I", "I"));
    }

    #[test]
    fn add_i_ii() {
        assert_eq!("III", add("I", "II"));
    }

    #[test]
    fn add_ii_iii() {
        assert_eq!("V", add("II", "III"));

    }

    #[test]
    fn add_v_i() {
        assert_eq!("VI", add("V", "I"));
    }

    #[test]
    fn add_i_v() {
        assert_eq!("VI", add("I", "V"));
    }

    #[test]
    fn compress_iiiii() {
        assert_eq!("V", compress("IIIII"));
    }

    #[test]
    fn compress_iiii() {
        assert_eq!("IV", compress("IIII"));
    }
}
<commit_msg>compress -> normalize in rust<commit_after>pub fn add(num_l: &str, num_r: &str) -> String {
    let mut sum: Vec<u8> = (String::from(num_l) + num_r).into();

    sum.sort();
    sum.reverse();

    normalize(&String::from_utf8(sum).unwrap())
}

pub fn normalize(num: &str) -> String {
    match num {
        "IIIII" => "V".to_string(),
        "IIII" => "IV".to_string(),
        _ => num.to_string(),
    }
}

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

    #[test]
    fn add_i_i() {
        assert_eq!("II", add("I", "I"));
    }

    #[test]
    fn add_i_ii() {
        assert_eq!("III", add("I", "II"));
    }

    #[test]
    fn add_ii_iii() {
        assert_eq!("V", add("II", "III"));

    }

    #[test]
    fn add_v_i() {
        assert_eq!("VI", add("V", "I"));
    }

    #[test]
    fn add_i_v() {
        assert_eq!("VI", add("I", "V"));
    }

    #[test]
    fn normalize_iiiii() {
        assert_eq!("V", normalize("IIIII"));
    }

    #[test]
    fn normalize_iiii() {
        assert_eq!("IV", normalize("IIII"));
    }
}
<|endoftext|>


<commit_before>extern crate eventual;
extern crate robots;

use eventual::Async;

use std::any::Any;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{channel, Sender};
use std::time::Duration;

use robots::actors::{Actor, ActorSystem, ActorCell, ActorContext, ActorRef, Props};

#[derive(Debug, PartialEq)]
enum Res {
    Ok,
    Err,
}

#[derive(Copy, Clone)]
enum InternalStateMessage {
    Set(u32),
    Get,
    Panic,
}

struct InternalState {
    last: Mutex<u32>,
    sender: Arc<Mutex<Sender<Res>>>,
}

impl Actor for InternalState {
    fn receive(&self, message: Box<Any>, context: ActorCell) {
        if let Ok(message) = Box::<Any>::downcast::<InternalStateMessage>(message) {
            match *message {
                InternalStateMessage::Get => {
                    context.tell(context.sender(), *self.last.lock().unwrap())
                }
                InternalStateMessage::Set(message) => {
                    // Here mixing the test actir for the two tests might seem a bit weird,
                    // but we would get two very similar actors otherwise.
                    let mut last = self.last.lock().unwrap();
                    if message <= *last {
                        let _ = self.sender.lock().unwrap().send(Res::Err);
                    } else {
                        *last = message;
                    }
                    if *last == 1000 {
                        let _ = self.sender.lock().unwrap().send(Res::Ok);
                    }
                }
                InternalStateMessage::Panic => panic!(""),
            }
        }
    }
}

impl InternalState {
    fn new(sender: Arc<Mutex<Sender<Res>>>) -> InternalState {
        InternalState {
            last: Mutex::new(0),
            sender: sender,
        }
    }
}

#[test]
fn read_messages_in_order() {
    let actor_system = ActorSystem::new("test".to_owned());
    actor_system.spawn_threads(9);

    let (tx, rx) = channel();
    let tx = Arc::new(Mutex::new(tx));

    let props = Props::new(Arc::new(InternalState::new), tx);
    let actor_ref_1 = actor_system.actor_of(props.clone(), "sender".to_owned());
    let actor_ref_2 = actor_system.actor_of(props.clone(), "receiver".to_owned());

    for i in 1..1001 {
        actor_ref_1.tell_to(actor_ref_2.clone(), InternalStateMessage::Set(i as u32));
    }

    let res = rx.recv();
    assert_eq!(Ok(Res::Ok), res);

    actor_system.shutdown();
}

#[test]
fn recover_from_panic() {
    let actor_system = ActorSystem::new("test".to_owned());
    actor_system.spawn_threads(1);

    let (tx, _rx) = channel();
    let tx = Arc::new(Mutex::new(tx));

    let props = Props::new(Arc::new(InternalState::new), tx);
    let requester = actor_system.actor_of(props.clone(), "sender".to_owned());
    let answerer = actor_system.actor_of(props.clone(), "receiver".to_owned());

    requester.tell_to(answerer.clone(), InternalStateMessage::Set(10));
    let res = answerer.ask(InternalStateMessage::Get).await().unwrap();
    let res = Box::<Any>::downcast::<u32>(res).unwrap();
    assert_eq!(10u32, *res);

    requester.tell_to(answerer.clone(), InternalStateMessage::Panic);
    let res = answerer.ask(InternalStateMessage::Get).await().unwrap();
    let res = Box::<Any>::downcast::<u32>(res).unwrap();
    assert_eq!(0u32, *res);

    actor_system.shutdown();
}

struct Resolver;

impl Actor for Resolver {
    fn receive(&self, message: Box<Any>, context: ActorCell) {
        if let Ok(message) = Box::<Any>::downcast::<String>(message) {
            let res = context.identify_actor(*message)
                             .await()
                             .unwrap();
            context.tell(context.sender(), res);
        }
    }
}

impl Resolver {
    fn new(_dummy: ()) -> Resolver {
        Resolver
    }
}

#[test]
fn resolve_name_real_path() {
    let actor_system = ActorSystem::new("test".to_owned());
    actor_system.spawn_threads(2);

    let props = Props::new(Arc::new(Resolver::new), ());
    let answerer = actor_system.actor_of(props.clone(), "answerer".to_owned());
    let requester = actor_system.actor_of(props.clone(), "sender".to_owned());

    // We wait to be sure that the actors will be registered to the name resolver.
    std::thread::sleep(Duration::from_millis(100));

    let res = answerer.ask("/user/sender".to_owned()).await().unwrap();
    let res = Box::<Any>::downcast::<Option<ActorRef>>(res).unwrap();
    assert_eq!(requester.path(), (*res).unwrap().path());

    actor_system.shutdown();
}

#[test]
fn resolve_name_fake_path() {
    let actor_system = ActorSystem::new("test".to_owned());
    actor_system.spawn_threads(2);

    let props = Props::new(Arc::new(Resolver::new), ());
    let answerer = actor_system.actor_of(props.clone(), "answerer".to_owned());

    // We wait to be sure that the actors will be registered to the name resolver.
    std::thread::sleep(Duration::from_millis(100));

    let res = answerer.ask("/foo/bar".to_owned()).await().unwrap();
    let res = *Box::<Any>::downcast::<Option<ActorRef>>(res).unwrap();

    match res {
        None => {}
        Some(_) => panic!("The name resolver gave an ActorRef when he should not."),
    };

    actor_system.shutdown();
}
<commit_msg>Fix typo in comment.<commit_after>extern crate eventual;
extern crate robots;

use eventual::Async;

use std::any::Any;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{channel, Sender};
use std::time::Duration;

use robots::actors::{Actor, ActorSystem, ActorCell, ActorContext, ActorRef, Props};

#[derive(Debug, PartialEq)]
enum Res {
    Ok,
    Err,
}

#[derive(Copy, Clone)]
enum InternalStateMessage {
    Set(u32),
    Get,
    Panic,
}

struct InternalState {
    last: Mutex<u32>,
    sender: Arc<Mutex<Sender<Res>>>,
}

impl Actor for InternalState {
    fn receive(&self, message: Box<Any>, context: ActorCell) {
        if let Ok(message) = Box::<Any>::downcast::<InternalStateMessage>(message) {
            match *message {
                InternalStateMessage::Get => {
                    context.tell(context.sender(), *self.last.lock().unwrap())
                }
                InternalStateMessage::Set(message) => {
                    // Here mixing the test actor for the two tests might seem a bit weird,
                    // but we would get two very similar actors otherwise.
                    let mut last = self.last.lock().unwrap();
                    if message <= *last {
                        let _ = self.sender.lock().unwrap().send(Res::Err);
                    } else {
                        *last = message;
                    }
                    if *last == 1000 {
                        let _ = self.sender.lock().unwrap().send(Res::Ok);
                    }
                }
                InternalStateMessage::Panic => panic!(""),
            }
        }
    }
}

impl InternalState {
    fn new(sender: Arc<Mutex<Sender<Res>>>) -> InternalState {
        InternalState {
            last: Mutex::new(0),
            sender: sender,
        }
    }
}

#[test]
fn read_messages_in_order() {
    let actor_system = ActorSystem::new("test".to_owned());
    actor_system.spawn_threads(9);

    let (tx, rx) = channel();
    let tx = Arc::new(Mutex::new(tx));

    let props = Props::new(Arc::new(InternalState::new), tx);
    let actor_ref_1 = actor_system.actor_of(props.clone(), "sender".to_owned());
    let actor_ref_2 = actor_system.actor_of(props.clone(), "receiver".to_owned());

    for i in 1..1001 {
        actor_ref_1.tell_to(actor_ref_2.clone(), InternalStateMessage::Set(i as u32));
    }

    let res = rx.recv();
    assert_eq!(Ok(Res::Ok), res);

    actor_system.shutdown();
}

#[test]
fn recover_from_panic() {
    let actor_system = ActorSystem::new("test".to_owned());
    actor_system.spawn_threads(1);

    let (tx, _rx) = channel();
    let tx = Arc::new(Mutex::new(tx));

    let props = Props::new(Arc::new(InternalState::new), tx);
    let requester = actor_system.actor_of(props.clone(), "sender".to_owned());
    let answerer = actor_system.actor_of(props.clone(), "receiver".to_owned());

    requester.tell_to(answerer.clone(), InternalStateMessage::Set(10));
    let res = answerer.ask(InternalStateMessage::Get).await().unwrap();
    let res = Box::<Any>::downcast::<u32>(res).unwrap();
    assert_eq!(10u32, *res);

    requester.tell_to(answerer.clone(), InternalStateMessage::Panic);
    let res = answerer.ask(InternalStateMessage::Get).await().unwrap();
    let res = Box::<Any>::downcast::<u32>(res).unwrap();
    assert_eq!(0u32, *res);

    actor_system.shutdown();
}

struct Resolver;

impl Actor for Resolver {
    fn receive(&self, message: Box<Any>, context: ActorCell) {
        if let Ok(message) = Box::<Any>::downcast::<String>(message) {
            let res = context.identify_actor(*message)
                             .await()
                             .unwrap();
            context.tell(context.sender(), res);
        }
    }
}

impl Resolver {
    fn new(_dummy: ()) -> Resolver {
        Resolver
    }
}

#[test]
fn resolve_name_real_path() {
    let actor_system = ActorSystem::new("test".to_owned());
    actor_system.spawn_threads(2);

    let props = Props::new(Arc::new(Resolver::new), ());
    let answerer = actor_system.actor_of(props.clone(), "answerer".to_owned());
    let requester = actor_system.actor_of(props.clone(), "sender".to_owned());

    // We wait to be sure that the actors will be registered to the name resolver.
    std::thread::sleep(Duration::from_millis(100));

    let res = answerer.ask("/user/sender".to_owned()).await().unwrap();
    let res = Box::<Any>::downcast::<Option<ActorRef>>(res).unwrap();
    assert_eq!(requester.path(), (*res).unwrap().path());

    actor_system.shutdown();
}

#[test]
fn resolve_name_fake_path() {
    let actor_system = ActorSystem::new("test".to_owned());
    actor_system.spawn_threads(2);

    let props = Props::new(Arc::new(Resolver::new), ());
    let answerer = actor_system.actor_of(props.clone(), "answerer".to_owned());

    // We wait to be sure that the actors will be registered to the name resolver.
    std::thread::sleep(Duration::from_millis(100));

    let res = answerer.ask("/foo/bar".to_owned()).await().unwrap();
    let res = *Box::<Any>::downcast::<Option<ActorRef>>(res).unwrap();

    match res {
        None => {}
        Some(_) => panic!("The name resolver gave an ActorRef when he should not."),
    };

    actor_system.shutdown();
}
<|endoftext|>


<commit_before>extern crate robots;

use std::sync::{Arc, Mutex};
use std::sync::mpsc::{channel, Sender};

use robots::actors::{Actor, ActorSystem, ActorCell, Arguments, ActorContext, Props};

#[derive(Debug, PartialEq)]
enum Res {
    Ok,
    Err
}

struct InternalState {
    last: Mutex<u32>,
    sender: Arc<Mutex<Sender<Res>>>,
}

impl Actor<u32> for InternalState {
    fn receive<Args: Arguments>(&self, message: u32, _context: ActorCell<Args, u32, InternalState>) {
        let mut last = self.last.lock().unwrap();
        if message <= *last {
            let _ = self.sender.lock().unwrap().send(Res::Err);
        } else {
            *last = message;
        }
        if *last == 1000 {
            let _ = self.sender.lock().unwrap().send(Res::Ok);
        }
    }
}

impl InternalState {
    fn new(sender: Arc<Mutex<Sender<Res>>>) -> InternalState {
        InternalState {
            last: Mutex::new(0),
            sender: sender,
        }
    }
}

#[test]
fn read_messages_in_order() {
    let actor_system = ActorSystem::new("test".to_owned());
    actor_system.spawn_threads(9);

    let (tx, rx) = channel();
    let tx = Arc::new(Mutex::new(tx));

    let props = Props::new(Arc::new(InternalState::new), tx);
    let actor_ref_1 = actor_system.actor_of(props.clone(), "sender".to_owned());
    let actor_ref_2 = actor_system.actor_of(props.clone(), "receiver".to_owned());

    for i in 1..1001 {
        actor_ref_1.tell_to(actor_ref_2.clone(), i as u32);
    }

    let res = rx.recv();
    assert_eq!(Ok(Res::Ok), res);

    actor_system.shutdown();
}
<commit_msg>Add a test for the ask methods.<commit_after>extern crate eventual;
extern crate robots;

use eventual::Async;

use std::sync::{Arc, Mutex};
use std::sync::mpsc::{channel, Sender};

use robots::actors::{Actor, ActorSystem, ActorCell, Arguments, ActorContext, Props};

#[derive(Debug, PartialEq)]
enum Res {
    Ok,
    Err
}

#[derive(Copy, Clone)]
enum InternalStateMessage {
    Set(u32),
    Get,
    Panic,
}

struct InternalState {
    last: Mutex<u32>,
    sender: Arc<Mutex<Sender<Res>>>,
}

impl Actor<InternalStateMessage> for InternalState {
    fn receive<Args: Arguments>(&self, message: InternalStateMessage, context: ActorCell<Args, InternalStateMessage, InternalState>) {
        match message {
            InternalStateMessage::Get => context.tell(context.sender(), *self.last.lock().unwrap()),
            InternalStateMessage::Set(message) => {
                // Here mixing the test actir for the two tests might seem a bit weird,
                // but we would get two very similar actors otherwise.
                let mut last = self.last.lock().unwrap();
                if message <= *last {
                    let _ = self.sender.lock().unwrap().send(Res::Err);
                } else {
                    *last = message;
                }
                if *last == 1000 {
                    let _ = self.sender.lock().unwrap().send(Res::Ok);
                }
            }
            InternalStateMessage::Panic => panic!(""),
        }
    }
}

impl InternalState {
    fn new(sender: Arc<Mutex<Sender<Res>>>) -> InternalState {
        InternalState {
            last: Mutex::new(0),
            sender: sender,
        }
    }
}

#[test]
fn read_messages_in_order() {
    let actor_system = ActorSystem::new("test".to_owned());
    actor_system.spawn_threads(9);

    let (tx, rx) = channel();
    let tx = Arc::new(Mutex::new(tx));

    let props = Props::new(Arc::new(InternalState::new), tx);
    let actor_ref_1 = actor_system.actor_of(props.clone(), "sender".to_owned());
    let actor_ref_2 = actor_system.actor_of(props.clone(), "receiver".to_owned());

    for i in 1..1001 {
        actor_ref_1.tell_to(actor_ref_2.clone(), InternalStateMessage::Set(i as u32));
    }

    let res = rx.recv();
    assert_eq!(Ok(Res::Ok), res);

    actor_system.shutdown();
}

#[test]
fn recover_from_panic() {
    let actor_system = ActorSystem::new("test".to_owned());
    actor_system.spawn_threads(1);

    let (tx, _rx) = channel();
    let tx = Arc::new(Mutex::new(tx));

    let props = Props::new(Arc::new(InternalState::new), tx);
    let requester = actor_system.actor_of(props.clone(), "sender".to_owned());
    let answerer = actor_system.actor_of(props.clone(), "receiver".to_owned());

    requester.tell_to(answerer.clone(), InternalStateMessage::Set(10));
    let res: u32 = requester.ask_to::<InternalStateMessage, u32, ()>(answerer.clone(), InternalStateMessage::Get).and_then(|x| Ok(x)).await().unwrap();
    assert_eq!(10u32, res);

    requester.tell_to(answerer.clone(), InternalStateMessage::Panic);
    let res: u32 = requester.ask_to::<InternalStateMessage, u32, ()>(answerer.clone(), InternalStateMessage::Get).and_then(|x| Ok(x)).await().unwrap();
    assert_eq!(0u32, res);

    actor_system.shutdown();
}
<|endoftext|>