<commit_before>//! Lints concerned with the grouping of digits with underscores in integral or
//! floating-point literal expressions.

use rustc::lint::*;
use syntax::ast::*;
use syntax_pos;
use utils::{span_help_and_lint, snippet_opt, in_external_macro};

/// **What it does:** Warns if a long integral or floating-point constant does
/// not contain underscores.
///
/// **Why is this bad?** Reading long numbers is difficult without separators.
///
/// **Known problems:** None.
///
/// **Example:**
///
/// ```rust
/// 61864918973511
/// ```
declare_lint! {
    pub UNREADABLE_LITERAL,
    Warn,
    "long integer literal without underscores"
}

/// **What it does:** Warns if an integral or floating-point constant is
/// grouped inconsistently with underscores.
///
/// **Why is this bad?** Readers may incorrectly interpret inconsistently
/// grouped digits.
///
/// **Known problems:** None.
///
/// **Example:**
///
/// ```rust
/// 618_64_9189_73_511
/// ```
declare_lint! {
    pub INCONSISTENT_DIGIT_GROUPING,
    Warn,
    "integer literals with digits grouped inconsistently"
}

/// **What it does:** Warns if the digits of an integral or floating-point
/// constant are grouped into groups that
/// are too large.
///
/// **Why is this bad?** Negatively impacts readability.
///
/// **Known problems:** None.
///
/// **Example:**
///
/// ```rust
/// 6186491_8973511
/// ```
declare_lint! {
    pub LARGE_DIGIT_GROUPS,
    Warn,
    "grouping digits into groups that are too large"
}

#[derive(Debug)]
enum Radix {
    Binary,
    Octal,
    Decimal,
    Hexadecimal,
}

impl Radix {
    /// Return a reasonable digit group size for this radix.
    pub fn suggest_grouping(&self) -> usize {
        match *self {
            Radix::Binary | Radix::Hexadecimal => 4,
            Radix::Octal | Radix::Decimal => 3,
        }
    }
}

#[derive(Debug)]
struct DigitInfo<'a> {
    /// Characters of a literal between the radix prefix and type suffix.
    pub digits: &'a str,
    /// Which radix the literal was represented in.
    pub radix: Radix,
    /// The radix prefix, if present.
    pub prefix: Option<&'a str>,
    /// The type suffix, including preceding underscore if present.
    pub suffix: Option<&'a str>,
    /// True for floating-point literals.
    pub float: bool,
}

impl<'a> DigitInfo<'a> {
    pub fn new(lit: &str, float: bool) -> DigitInfo {
        // Determine delimiter for radix prefix, if present, and radix.
        let radix = if lit.starts_with("0x") {
            Radix::Hexadecimal
        } else if lit.starts_with("0b") {
            Radix::Binary
        } else if lit.starts_with("0o") {
            Radix::Octal
        } else {
            Radix::Decimal
        };

        // Grab part of the literal after prefix, if present.
        let (prefix, sans_prefix) = if let Radix::Decimal = radix {
            (None, lit)
        } else {
            let (p, s) = lit.split_at(2);
            (Some(p), s)
        };

        let mut last_d = '\0';
        for (d_idx, d) in sans_prefix.char_indices() {
            if !float && (d == 'i' || d == 'u') || float && d == 'f' {
                let suffix_start = if last_d == '_' { d_idx - 1 } else { d_idx };
                let (digits, suffix) = sans_prefix.split_at(suffix_start);
                return DigitInfo {
                    digits: digits,
                    radix: radix,
                    prefix: prefix,
                    suffix: Some(suffix),
                    float: float,
                };
            }
            last_d = d
        }

        // No suffix found
        DigitInfo {
            digits: sans_prefix,
            radix: radix,
            prefix: prefix,
            suffix: None,
            float: float,
        }
    }
}

enum WarningType {
    UnreadableLiteral,
    InconsistentDigitGrouping,
    LargeDigitGroups,
}

impl WarningType {
    pub fn display(&self, grouping_hint: &str, cx: &EarlyContext, span: &syntax_pos::Span) {
        match *self {
            WarningType::UnreadableLiteral => {
                span_help_and_lint(
                    cx,
                    UNREADABLE_LITERAL,
                    *span,
                    "long literal lacking separators",
                    &format!("consider: {}", grouping_hint),
                )
            },
            WarningType::LargeDigitGroups => {
                span_help_and_lint(
                    cx,
                    LARGE_DIGIT_GROUPS,
                    *span,
                    "digit groups should be smaller",
                    &format!("consider: {}", grouping_hint),
                )
            },
            WarningType::InconsistentDigitGrouping => {
                span_help_and_lint(
                    cx,
                    INCONSISTENT_DIGIT_GROUPING,
                    *span,
                    "digits grouped inconsistently by underscores",
                    &format!("consider: {}", grouping_hint),
                )
            },
        };
    }
}

#[derive(Copy, Clone)]
pub struct LiteralDigitGrouping;

impl LintPass for LiteralDigitGrouping {
    fn get_lints(&self) -> LintArray {
        lint_array!(UNREADABLE_LITERAL, INCONSISTENT_DIGIT_GROUPING, LARGE_DIGIT_GROUPS)
    }
}

impl EarlyLintPass for LiteralDigitGrouping {
    fn check_expr(&mut self, cx: &EarlyContext, expr: &Expr) {
        if in_external_macro(cx, expr.span) {
            return;
        }

        if let ExprKind::Lit(ref lit) = expr.node {
            self.check_lit(cx, lit)
        }
    }
}

impl LiteralDigitGrouping {
    fn check_lit(&self, cx: &EarlyContext, lit: &Lit) {
        // Lint integral literals.
        if_let_chain! {[
            let LitKind::Int(..) = lit.node,
            let Some(src) = snippet_opt(cx, lit.span),
            let Some(firstch) = src.chars().next(),
            char::to_digit(firstch, 10).is_some()
        ], {
            let digits = LiteralDigitGrouping::get_digits(&src, false);

            LiteralDigitGrouping::do_lint(digits, cx, &lit.span);
        }}

        // Lint floating-point literals.
        if_let_chain! {[
            let LitKind::Float(..) = lit.node,
            let Some(src) = snippet_opt(cx, lit.span),
            let Some(firstch) = src.chars().next(),
            char::to_digit(firstch, 10).is_some()
        ], {
            let digits: Vec<&str> = LiteralDigitGrouping::get_digits(&src, true)
                .split_terminator('.')
                .collect();

            // Lint integral and fractional parts separately, and then check consistency of digit
            // groups if both pass.
            if let Some(integral_group_size) = LiteralDigitGrouping::do_lint(digits[0], cx, &lit.span) {
                if digits.len() > 1 {
                    // Lint the fractional part of literal just like integral part, but reversed.
                    let fractional_part = &digits[1].chars().rev().collect::<String>();
                    if let Some(fractional_group_size) = LiteralDigitGrouping::do_lint(fractional_part, cx, &lit.span) {
                        let consistent = match (integral_group_size, fractional_group_size) {
                            // No groups on either side of decimal point - good to go.
                            (0, 0) => true,
                            // Integral part has grouped digits, fractional part does not.
                            (_, 0) => digits[1].len() <= integral_group_size,
                            // Fractional part has grouped digits, integral part does not.
                            (0, _) => digits[0].len() <= fractional_group_size,
                            // Both parts have grouped digits. Groups should be the same size.
                            (_, _) => integral_group_size == fractional_group_size,
                        };

                        if !consistent {
                            span_help_and_lint(cx, INCONSISTENT_DIGIT_GROUPING, lit.span,
                                           "digits grouped inconsistently by underscores",
                                           "consider making each group three or four digits");
                        }
                    }
                }
            }
        }}
    }

    /// Returns the digits of an integral or floating-point literal.
    fn get_digits(lit: &str, float: bool) -> &str {
        // Determine delimiter for radix prefix, if present.
        let mb_r = if lit.starts_with("0x") {
            Some('x')
        } else if lit.starts_with("0b") {
            Some('b')
        } else if lit.starts_with("0o") {
            Some('o')
        } else {
            None
        };

        let has_suffix = !float && (lit.contains('i') || lit.contains('u')) || float && lit.contains('f');

        // Grab part of literal between the radix prefix and type suffix.
        let mut digits = if let Some(r) = mb_r {
            lit.split(|c| c == 'i' || c == 'u' || c == r || (float && c == 'f')).nth(1).unwrap()
        } else {
            lit.split(|c| c == 'i' || c == 'u' || (float && c == 'f')).next().unwrap()
        };

        // If there was an underscore before type suffix, drop it.
        if has_suffix && digits.chars().last().unwrap() == '_' {
            digits = digits.split_at(digits.len() - 1).0;
        }

        digits
    }

    /// Performs lint on `digits` (no decimal point) and returns the group size. `None` is
    /// returned when emitting a warning.
    fn do_lint(digits: &str, cx: &EarlyContext, span: &syntax_pos::Span) -> Option<usize> {
        // Grab underscore indices with respect to the units digit.
        let underscore_positions: Vec<usize> = digits.chars().rev().enumerate()
            .filter_map(|(idx, digit)|
                if digit == '_' {
                    Some(idx)
                } else {
                    None
                })
            .collect();

        if underscore_positions.is_empty() {
            // Check if literal needs underscores.
            if digits.len() > 4 {
                span_help_and_lint(cx, UNREADABLE_LITERAL, *span,
                                   "long literal lacking separators",
                                   "consider using underscores to make literal more readable");
                return None;
            } else {
                return Some(0);
            }
        } else {
            // Check consistency and the sizes of the groups.
            let group_size = underscore_positions[0];
            let consistent = underscore_positions
                .windows(2)
                .all(|ps| ps[1] - ps[0] == group_size + 1)
                // number of digits to the left of the last group cannot be bigger than group size.
                && (digits.len() - underscore_positions.last().unwrap() <= group_size + 1);

            if !consistent {
                span_help_and_lint(cx, INCONSISTENT_DIGIT_GROUPING, *span,
                                   "digits grouped inconsistently by underscores",
                                   "consider making each group three or four digits");
                return None;
            } else if group_size > 4 {
                span_help_and_lint(cx, LARGE_DIGIT_GROUPS, *span,
                                   "digit groups should be smaller",
                                   "consider using groups of three or four digits");
                return None;
            }
            return Some(group_size);
        }
    }
}
<commit_msg>Implement method for digit grouping suggestions<commit_after>//! Lints concerned with the grouping of digits with underscores in integral or
//! floating-point literal expressions.

use rustc::lint::*;
use syntax::ast::*;
use syntax_pos;
use utils::{span_help_and_lint, snippet_opt, in_external_macro};

/// **What it does:** Warns if a long integral or floating-point constant does
/// not contain underscores.
///
/// **Why is this bad?** Reading long numbers is difficult without separators.
///
/// **Known problems:** None.
///
/// **Example:**
///
/// ```rust
/// 61864918973511
/// ```
declare_lint! {
    pub UNREADABLE_LITERAL,
    Warn,
    "long integer literal without underscores"
}

/// **What it does:** Warns if an integral or floating-point constant is
/// grouped inconsistently with underscores.
///
/// **Why is this bad?** Readers may incorrectly interpret inconsistently
/// grouped digits.
///
/// **Known problems:** None.
///
/// **Example:**
///
/// ```rust
/// 618_64_9189_73_511
/// ```
declare_lint! {
    pub INCONSISTENT_DIGIT_GROUPING,
    Warn,
    "integer literals with digits grouped inconsistently"
}

/// **What it does:** Warns if the digits of an integral or floating-point
/// constant are grouped into groups that
/// are too large.
///
/// **Why is this bad?** Negatively impacts readability.
///
/// **Known problems:** None.
///
/// **Example:**
///
/// ```rust
/// 6186491_8973511
/// ```
declare_lint! {
    pub LARGE_DIGIT_GROUPS,
    Warn,
    "grouping digits into groups that are too large"
}

#[derive(Debug)]
enum Radix {
    Binary,
    Octal,
    Decimal,
    Hexadecimal,
}

impl Radix {
    /// Return a reasonable digit group size for this radix.
    pub fn suggest_grouping(&self) -> usize {
        match *self {
            Radix::Binary | Radix::Hexadecimal => 4,
            Radix::Octal | Radix::Decimal => 3,
        }
    }
}

#[derive(Debug)]
struct DigitInfo<'a> {
    /// Characters of a literal between the radix prefix and type suffix.
    pub digits: &'a str,
    /// Which radix the literal was represented in.
    pub radix: Radix,
    /// The radix prefix, if present.
    pub prefix: Option<&'a str>,
    /// The type suffix, including preceding underscore if present.
    pub suffix: Option<&'a str>,
    /// True for floating-point literals.
    pub float: bool,
}

impl<'a> DigitInfo<'a> {
    pub fn new(lit: &str, float: bool) -> DigitInfo {
        // Determine delimiter for radix prefix, if present, and radix.
        let radix = if lit.starts_with("0x") {
            Radix::Hexadecimal
        } else if lit.starts_with("0b") {
            Radix::Binary
        } else if lit.starts_with("0o") {
            Radix::Octal
        } else {
            Radix::Decimal
        };

        // Grab part of the literal after prefix, if present.
        let (prefix, sans_prefix) = if let Radix::Decimal = radix {
            (None, lit)
        } else {
            let (p, s) = lit.split_at(2);
            (Some(p), s)
        };

        let mut last_d = '\0';
        for (d_idx, d) in sans_prefix.char_indices() {
            if !float && (d == 'i' || d == 'u') || float && d == 'f' {
                let suffix_start = if last_d == '_' { d_idx - 1 } else { d_idx };
                let (digits, suffix) = sans_prefix.split_at(suffix_start);
                return DigitInfo {
                    digits: digits,
                    radix: radix,
                    prefix: prefix,
                    suffix: Some(suffix),
                    float: float,
                };
            }
            last_d = d
        }

        // No suffix found
        DigitInfo {
            digits: sans_prefix,
            radix: radix,
            prefix: prefix,
            suffix: None,
            float: float,
        }
    }

    /// Returns digits grouped in a sensible way.
    fn grouping_hint(&self) -> String {
        let group_size = self.radix.suggest_grouping();
        if self.digits.contains('.') {
            let mut parts = self.digits.split(".");
            let int_part_hint = parts
                .next()
                .unwrap()
                .chars()
                .rev()
                .filter(|&c| c != '_')
                .collect::<Vec<_>>()
                .chunks(group_size)
                .map(|chunk| chunk.into_iter().rev().collect())
                .rev()
                .collect::<Vec<String>>()
                .join("_");
            let frac_part_hint = parts
                .next()
                .unwrap()
                .chars()
                .filter(|&c| c != '_')
                .collect::<Vec<_>>()
                .chunks(group_size)
                .map(|chunk| chunk.into_iter().collect())
                .collect::<Vec<String>>()
                .join("_");
            format!("{}.{}{}", int_part_hint, frac_part_hint, self.suffix.unwrap_or(""))
        } else {
            let hint = self.digits
                .chars()
                .rev()
                .filter(|&c| c != '_')
                .collect::<Vec<_>>()
                .chunks(group_size)
                .map(|chunk| chunk.into_iter().rev().collect())
                .rev()
                .collect::<Vec<String>>()
                .join("_");
            format!("{}{}{}", self.prefix.unwrap_or(""), hint, self.suffix.unwrap_or(""))
        }
    }
}

enum WarningType {
    UnreadableLiteral,
    InconsistentDigitGrouping,
    LargeDigitGroups,
}

impl WarningType {
    pub fn display(&self, grouping_hint: &str, cx: &EarlyContext, span: &syntax_pos::Span) {
        match *self {
            WarningType::UnreadableLiteral => {
                span_help_and_lint(
                    cx,
                    UNREADABLE_LITERAL,
                    *span,
                    "long literal lacking separators",
                    &format!("consider: {}", grouping_hint),
                )
            },
            WarningType::LargeDigitGroups => {
                span_help_and_lint(
                    cx,
                    LARGE_DIGIT_GROUPS,
                    *span,
                    "digit groups should be smaller",
                    &format!("consider: {}", grouping_hint),
                )
            },
            WarningType::InconsistentDigitGrouping => {
                span_help_and_lint(
                    cx,
                    INCONSISTENT_DIGIT_GROUPING,
                    *span,
                    "digits grouped inconsistently by underscores",
                    &format!("consider: {}", grouping_hint),
                )
            },
        };
    }
}

#[derive(Copy, Clone)]
pub struct LiteralDigitGrouping;

impl LintPass for LiteralDigitGrouping {
    fn get_lints(&self) -> LintArray {
        lint_array!(UNREADABLE_LITERAL, INCONSISTENT_DIGIT_GROUPING, LARGE_DIGIT_GROUPS)
    }
}

impl EarlyLintPass for LiteralDigitGrouping {
    fn check_expr(&mut self, cx: &EarlyContext, expr: &Expr) {
        if in_external_macro(cx, expr.span) {
            return;
        }

        if let ExprKind::Lit(ref lit) = expr.node {
            self.check_lit(cx, lit)
        }
    }
}

impl LiteralDigitGrouping {
    fn check_lit(&self, cx: &EarlyContext, lit: &Lit) {
        // Lint integral literals.
        if_let_chain! {[
            let LitKind::Int(..) = lit.node,
            let Some(src) = snippet_opt(cx, lit.span),
            let Some(firstch) = src.chars().next(),
            char::to_digit(firstch, 10).is_some()
        ], {
            let digits = LiteralDigitGrouping::get_digits(&src, false);

            LiteralDigitGrouping::do_lint(digits, cx, &lit.span);
        }}

        // Lint floating-point literals.
        if_let_chain! {[
            let LitKind::Float(..) = lit.node,
            let Some(src) = snippet_opt(cx, lit.span),
            let Some(firstch) = src.chars().next(),
            char::to_digit(firstch, 10).is_some()
        ], {
            let digits: Vec<&str> = LiteralDigitGrouping::get_digits(&src, true)
                .split_terminator('.')
                .collect();

            // Lint integral and fractional parts separately, and then check consistency of digit
            // groups if both pass.
            if let Some(integral_group_size) = LiteralDigitGrouping::do_lint(digits[0], cx, &lit.span) {
                if digits.len() > 1 {
                    // Lint the fractional part of literal just like integral part, but reversed.
                    let fractional_part = &digits[1].chars().rev().collect::<String>();
                    if let Some(fractional_group_size) = LiteralDigitGrouping::do_lint(fractional_part, cx, &lit.span) {
                        let consistent = match (integral_group_size, fractional_group_size) {
                            // No groups on either side of decimal point - good to go.
                            (0, 0) => true,
                            // Integral part has grouped digits, fractional part does not.
                            (_, 0) => digits[1].len() <= integral_group_size,
                            // Fractional part has grouped digits, integral part does not.
                            (0, _) => digits[0].len() <= fractional_group_size,
                            // Both parts have grouped digits. Groups should be the same size.
                            (_, _) => integral_group_size == fractional_group_size,
                        };

                        if !consistent {
                            span_help_and_lint(cx, INCONSISTENT_DIGIT_GROUPING, lit.span,
                                           "digits grouped inconsistently by underscores",
                                           "consider making each group three or four digits");
                        }
                    }
                }
            }
        }}
    }

    /// Returns the digits of an integral or floating-point literal.
    fn get_digits(lit: &str, float: bool) -> &str {
        // Determine delimiter for radix prefix, if present.
        let mb_r = if lit.starts_with("0x") {
            Some('x')
        } else if lit.starts_with("0b") {
            Some('b')
        } else if lit.starts_with("0o") {
            Some('o')
        } else {
            None
        };

        let has_suffix = !float && (lit.contains('i') || lit.contains('u')) || float && lit.contains('f');

        // Grab part of literal between the radix prefix and type suffix.
        let mut digits = if let Some(r) = mb_r {
            lit.split(|c| c == 'i' || c == 'u' || c == r || (float && c == 'f')).nth(1).unwrap()
        } else {
            lit.split(|c| c == 'i' || c == 'u' || (float && c == 'f')).next().unwrap()
        };

        // If there was an underscore before type suffix, drop it.
        if has_suffix && digits.chars().last().unwrap() == '_' {
            digits = digits.split_at(digits.len() - 1).0;
        }

        digits
    }

    /// Performs lint on `digits` (no decimal point) and returns the group size. `None` is
    /// returned when emitting a warning.
    fn do_lint(digits: &str, cx: &EarlyContext, span: &syntax_pos::Span) -> Option<usize> {
        // Grab underscore indices with respect to the units digit.
        let underscore_positions: Vec<usize> = digits.chars().rev().enumerate()
            .filter_map(|(idx, digit)|
                if digit == '_' {
                    Some(idx)
                } else {
                    None
                })
            .collect();

        if underscore_positions.is_empty() {
            // Check if literal needs underscores.
            if digits.len() > 4 {
                span_help_and_lint(cx, UNREADABLE_LITERAL, *span,
                                   "long literal lacking separators",
                                   "consider using underscores to make literal more readable");
                return None;
            } else {
                return Some(0);
            }
        } else {
            // Check consistency and the sizes of the groups.
            let group_size = underscore_positions[0];
            let consistent = underscore_positions
                .windows(2)
                .all(|ps| ps[1] - ps[0] == group_size + 1)
                // number of digits to the left of the last group cannot be bigger than group size.
                && (digits.len() - underscore_positions.last().unwrap() <= group_size + 1);

            if !consistent {
                span_help_and_lint(cx, INCONSISTENT_DIGIT_GROUPING, *span,
                                   "digits grouped inconsistently by underscores",
                                   "consider making each group three or four digits");
                return None;
            } else if group_size > 4 {
                span_help_and_lint(cx, LARGE_DIGIT_GROUPS, *span,
                                   "digit groups should be smaller",
                                   "consider using groups of three or four digits");
                return None;
            }
            return Some(group_size);
        }
    }
}
<|endoftext|>


<commit_before>//! Lints concerned with the grouping of digits with underscores in integral or
//! floating-point literal expressions.

use rustc::lint::*;
use syntax::ast::*;
use syntax_pos;
use utils::{span_help_and_lint, snippet_opt, in_external_macro};

/// **What it does:** Warns if a long integral or floating-point constant does
/// not contain underscores.
///
/// **Why is this bad?** Reading long numbers is difficult without separators.
///
/// **Known problems:** None.
///
/// **Example:**
///
/// ```rust
/// 61864918973511
/// ```
declare_lint! {
    pub UNREADABLE_LITERAL,
    Warn,
    "long integer literal without underscores"
}

/// **What it does:** Warns if an integral or floating-point constant is
/// grouped inconsistently with underscores.
///
/// **Why is this bad?** Readers may incorrectly interpret inconsistently
/// grouped digits.
///
/// **Known problems:** None.
///
/// **Example:**
///
/// ```rust
/// 618_64_9189_73_511
/// ```
declare_lint! {
    pub INCONSISTENT_DIGIT_GROUPING,
    Warn,
    "integer literals with digits grouped inconsistently"
}

/// **What it does:** Warns if the digits of an integral or floating-point
/// constant are grouped into groups that
/// are too large.
///
/// **Why is this bad?** Negatively impacts readability.
///
/// **Known problems:** None.
///
/// **Example:**
///
/// ```rust
/// 6186491_8973511
/// ```
declare_lint! {
    pub LARGE_DIGIT_GROUPS,
    Warn,
    "grouping digits into groups that are too large"
}

#[derive(Copy, Clone)]
pub struct LiteralDigitGrouping;

impl LintPass for LiteralDigitGrouping {
    fn get_lints(&self) -> LintArray {
        lint_array!(UNREADABLE_LITERAL, INCONSISTENT_DIGIT_GROUPING, LARGE_DIGIT_GROUPS)
    }
}

impl EarlyLintPass for LiteralDigitGrouping {
    fn check_expr(&mut self, cx: &EarlyContext, expr: &Expr) {
        if in_external_macro(cx, expr.span) {
            return;
        }

        if let ExprKind::Lit(ref lit) = expr.node {
            self.check_lit(cx, lit)
        }
    }
}

impl LiteralDigitGrouping {
    fn check_lit(&self, cx: &EarlyContext, lit: &Lit) {
        // Lint integral literals.
        if_let_chain! {[
            let LitKind::Int(..) = lit.node,
            let Some(src) = snippet_opt(cx, lit.span),
            let Some(firstch) = src.chars().next(),
            char::to_digit(firstch, 10).is_some()
        ], {
            let digits = LiteralDigitGrouping::get_digits(&src, false);

            LiteralDigitGrouping::do_lint(digits, cx, &lit.span);
        }}

        // Lint floating-point literals.
        if_let_chain! {[
            let LitKind::Float(..) = lit.node,
            let Some(src) = snippet_opt(cx, lit.span),
            let Some(firstch) = src.chars().next(),
            char::to_digit(firstch, 10).is_some()
        ], {
            let digits: Vec<&str> = LiteralDigitGrouping::get_digits(&src, true)
                .split_terminator('.')
                .collect();

            // Lint integral and fractional parts separately, and then check consistency of digit
            // groups if both pass.
            if let Some(integral_group_size) = LiteralDigitGrouping::do_lint(digits[0], cx, &lit.span) {
                if digits.len() > 1 {
                    // Lint the fractional part of literal just like integral part, but reversed.
                    let fractional_part = &digits[1].chars().rev().collect::<String>();
                    if let Some(fractional_group_size) = LiteralDigitGrouping::do_lint(fractional_part, cx, &lit.span) {
                        let consistent = match (integral_group_size, fractional_group_size) {
                            // No groups on either side of decimal point - good to go.
                            (0, 0) => true,
                            // Integral part has grouped digits, fractional part does not.
                            (_, 0) => digits[1].len() <= integral_group_size,
                            // Fractional part has grouped digits, integral part does not.
                            (0, _) => digits[0].len() <= fractional_group_size,
                            // Both parts have grouped digits. Groups should be the same size.
                            (_, _) => integral_group_size == fractional_group_size,
                        };

                        if !consistent {
                            span_help_and_lint(cx, INCONSISTENT_DIGIT_GROUPING, lit.span,
                                           "digits grouped inconsistently by underscores",
                                           "consider making each group three or four digits");
                        }
                    }
                }
            }
        }}
    }

    /// Returns the digits of an integral or floating-point literal.
    fn get_digits(lit: &str, float: bool) -> &str {
        // Determine delimiter for radix prefix, if present.
        let mb_r = if lit.starts_with("0x") {
            Some('x')
        } else if lit.starts_with("0b") {
            Some('b')
        } else if lit.starts_with("0o") {
            Some('o')
        } else {
            None
        };

        let has_suffix = !float && (lit.contains('i') || lit.contains('u')) || float && lit.contains('f');

        // Grab part of literal between the radix prefix and type suffix.
        let mut digits = if let Some(r) = mb_r {
            lit.split(|c| c == 'i' || c == 'u' || c == r || (float && c == 'f')).nth(1).unwrap()
        } else {
            lit.split(|c| c == 'i' || c == 'u' || (float && c == 'f')).next().unwrap()
        };

        // If there was an underscore before type suffix, drop it.
        if has_suffix && digits.chars().last().unwrap() == '_' {
            digits = digits.split_at(digits.len() - 1).0;
        }

        digits
    }

    /// Performs lint on `digits` (no decimal point) and returns the group size. `None` is
    /// returned when emitting a warning.
    fn do_lint(digits: &str, cx: &EarlyContext, span: &syntax_pos::Span) -> Option<usize> {
        // Grab underscore indices with respect to the units digit.
        let underscore_positions: Vec<usize> = digits.chars().rev().enumerate()
            .filter_map(|(idx, digit)|
                if digit == '_' {
                    Some(idx)
                } else {
                    None
                })
            .collect();

        if underscore_positions.is_empty() {
            // Check if literal needs underscores.
            if digits.len() > 4 {
                span_help_and_lint(cx, UNREADABLE_LITERAL, *span,
                                   "long literal lacking separators",
                                   "consider using underscores to make literal more readable");
                return None;
            } else {
                return Some(0);
            }
        } else {
            // Check consistency and the sizes of the groups.
            let group_size = underscore_positions[0];
            let consistent = underscore_positions
                .windows(2)
                .all(|ps| ps[1] - ps[0] == group_size + 1)
                // number of digits to the left of the last group cannot be bigger than group size.
                && (digits.len() - underscore_positions.last().unwrap() <= group_size + 1);

            if !consistent {
                span_help_and_lint(cx, INCONSISTENT_DIGIT_GROUPING, *span,
                                   "digits grouped inconsistently by underscores",
                                   "consider making each group three or four digits");
                return None;
            } else if group_size > 4 {
                span_help_and_lint(cx, LARGE_DIGIT_GROUPS, *span,
                                   "digit groups should be smaller",
                                   "consider using groups of three or four digits");
                return None;
            }
            return Some(group_size);
        }
    }
}
<commit_msg>Add some structs for cleaning up lint logic later.<commit_after>//! Lints concerned with the grouping of digits with underscores in integral or
//! floating-point literal expressions.

use rustc::lint::*;
use syntax::ast::*;
use syntax_pos;
use utils::{span_help_and_lint, snippet_opt, in_external_macro};

/// **What it does:** Warns if a long integral or floating-point constant does
/// not contain underscores.
///
/// **Why is this bad?** Reading long numbers is difficult without separators.
///
/// **Known problems:** None.
///
/// **Example:**
///
/// ```rust
/// 61864918973511
/// ```
declare_lint! {
    pub UNREADABLE_LITERAL,
    Warn,
    "long integer literal without underscores"
}

/// **What it does:** Warns if an integral or floating-point constant is
/// grouped inconsistently with underscores.
///
/// **Why is this bad?** Readers may incorrectly interpret inconsistently
/// grouped digits.
///
/// **Known problems:** None.
///
/// **Example:**
///
/// ```rust
/// 618_64_9189_73_511
/// ```
declare_lint! {
    pub INCONSISTENT_DIGIT_GROUPING,
    Warn,
    "integer literals with digits grouped inconsistently"
}

/// **What it does:** Warns if the digits of an integral or floating-point
/// constant are grouped into groups that
/// are too large.
///
/// **Why is this bad?** Negatively impacts readability.
///
/// **Known problems:** None.
///
/// **Example:**
///
/// ```rust
/// 6186491_8973511
/// ```
declare_lint! {
    pub LARGE_DIGIT_GROUPS,
    Warn,
    "grouping digits into groups that are too large"
}

#[derive(Debug)]
enum Radix {
    Binary,
    Octal,
    Decimal,
    Hexadecimal,
}

impl Radix {
    /// Return a reasonable digit group size for this radix.
    pub fn suggest_grouping(&self) -> usize {
        match *self {
            Radix::Binary | Radix::Hexadecimal => 4,
            Radix::Octal | Radix::Decimal => 3,
        }
    }
}

#[derive(Debug)]
struct DigitInfo<'a> {
    /// Characters of a literal between the radix prefix and type suffix.
    pub digits: &'a str,
    /// Which radix the literal was represented in.
    pub radix: Radix,
    /// The radix prefix, if present.
    pub prefix: Option<&'a str>,
    /// The type suffix, including preceding underscore if present.
    pub suffix: Option<&'a str>,
    /// True for floating-point literals.
    pub float: bool,
}

impl<'a> DigitInfo<'a> {
    pub fn new(lit: &str, float: bool) -> DigitInfo {
        // Determine delimiter for radix prefix, if present, and radix.
        let radix = if lit.starts_with("0x") {
            Radix::Hexadecimal
        } else if lit.starts_with("0b") {
            Radix::Binary
        } else if lit.starts_with("0o") {
            Radix::Octal
        } else {
            Radix::Decimal
        };

        // Grab part of the literal after prefix, if present.
        let (prefix, sans_prefix) = if let Radix::Decimal = radix {
            (None, lit)
        } else {
            let (p, s) = lit.split_at(2);
            (Some(p), s)
        };

        let mut last_d = '\0';
        for (d_idx, d) in sans_prefix.char_indices() {
            if !float && (d == 'i' || d == 'u') || float && d == 'f' {
                let suffix_start = if last_d == '_' { d_idx - 1 } else { d_idx };
                let (digits, suffix) = sans_prefix.split_at(suffix_start);
                return DigitInfo {
                    digits: digits,
                    radix: radix,
                    prefix: prefix,
                    suffix: Some(suffix),
                    float: float,
                };
            }
            last_d = d
        }

        // No suffix found
        DigitInfo {
            digits: sans_prefix,
            radix: radix,
            prefix: prefix,
            suffix: None,
            float: float,
        }
    }
}

enum WarningType {
    UnreadableLiteral,
    InconsistentDigitGrouping,
    LargeDigitGroups,
}

impl WarningType {
    pub fn display(&self, grouping_hint: &str, cx: &EarlyContext, span: &syntax_pos::Span) {
        match *self {
            WarningType::UnreadableLiteral => {
                span_help_and_lint(
                    cx,
                    UNREADABLE_LITERAL,
                    *span,
                    "long literal lacking separators",
                    &format!("consider: {}", grouping_hint),
                )
            },
            WarningType::LargeDigitGroups => {
                span_help_and_lint(
                    cx,
                    LARGE_DIGIT_GROUPS,
                    *span,
                    "digit groups should be smaller",
                    &format!("consider: {}", grouping_hint),
                )
            },
            WarningType::InconsistentDigitGrouping => {
                span_help_and_lint(
                    cx,
                    INCONSISTENT_DIGIT_GROUPING,
                    *span,
                    "digits grouped inconsistently by underscores",
                    &format!("consider: {}", grouping_hint),
                )
            },
        };
    }
}

#[derive(Copy, Clone)]
pub struct LiteralDigitGrouping;

impl LintPass for LiteralDigitGrouping {
    fn get_lints(&self) -> LintArray {
        lint_array!(UNREADABLE_LITERAL, INCONSISTENT_DIGIT_GROUPING, LARGE_DIGIT_GROUPS)
    }
}

impl EarlyLintPass for LiteralDigitGrouping {
    fn check_expr(&mut self, cx: &EarlyContext, expr: &Expr) {
        if in_external_macro(cx, expr.span) {
            return;
        }

        if let ExprKind::Lit(ref lit) = expr.node {
            self.check_lit(cx, lit)
        }
    }
}

impl LiteralDigitGrouping {
    fn check_lit(&self, cx: &EarlyContext, lit: &Lit) {
        // Lint integral literals.
        if_let_chain! {[
            let LitKind::Int(..) = lit.node,
            let Some(src) = snippet_opt(cx, lit.span),
            let Some(firstch) = src.chars().next(),
            char::to_digit(firstch, 10).is_some()
        ], {
            let digits = LiteralDigitGrouping::get_digits(&src, false);

            LiteralDigitGrouping::do_lint(digits, cx, &lit.span);
        }}

        // Lint floating-point literals.
        if_let_chain! {[
            let LitKind::Float(..) = lit.node,
            let Some(src) = snippet_opt(cx, lit.span),
            let Some(firstch) = src.chars().next(),
            char::to_digit(firstch, 10).is_some()
        ], {
            let digits: Vec<&str> = LiteralDigitGrouping::get_digits(&src, true)
                .split_terminator('.')
                .collect();

            // Lint integral and fractional parts separately, and then check consistency of digit
            // groups if both pass.
            if let Some(integral_group_size) = LiteralDigitGrouping::do_lint(digits[0], cx, &lit.span) {
                if digits.len() > 1 {
                    // Lint the fractional part of literal just like integral part, but reversed.
                    let fractional_part = &digits[1].chars().rev().collect::<String>();
                    if let Some(fractional_group_size) = LiteralDigitGrouping::do_lint(fractional_part, cx, &lit.span) {
                        let consistent = match (integral_group_size, fractional_group_size) {
                            // No groups on either side of decimal point - good to go.
                            (0, 0) => true,
                            // Integral part has grouped digits, fractional part does not.
                            (_, 0) => digits[1].len() <= integral_group_size,
                            // Fractional part has grouped digits, integral part does not.
                            (0, _) => digits[0].len() <= fractional_group_size,
                            // Both parts have grouped digits. Groups should be the same size.
                            (_, _) => integral_group_size == fractional_group_size,
                        };

                        if !consistent {
                            span_help_and_lint(cx, INCONSISTENT_DIGIT_GROUPING, lit.span,
                                           "digits grouped inconsistently by underscores",
                                           "consider making each group three or four digits");
                        }
                    }
                }
            }
        }}
    }

    /// Returns the digits of an integral or floating-point literal.
    fn get_digits(lit: &str, float: bool) -> &str {
        // Determine delimiter for radix prefix, if present.
        let mb_r = if lit.starts_with("0x") {
            Some('x')
        } else if lit.starts_with("0b") {
            Some('b')
        } else if lit.starts_with("0o") {
            Some('o')
        } else {
            None
        };

        let has_suffix = !float && (lit.contains('i') || lit.contains('u')) || float && lit.contains('f');

        // Grab part of literal between the radix prefix and type suffix.
        let mut digits = if let Some(r) = mb_r {
            lit.split(|c| c == 'i' || c == 'u' || c == r || (float && c == 'f')).nth(1).unwrap()
        } else {
            lit.split(|c| c == 'i' || c == 'u' || (float && c == 'f')).next().unwrap()
        };

        // If there was an underscore before type suffix, drop it.
        if has_suffix && digits.chars().last().unwrap() == '_' {
            digits = digits.split_at(digits.len() - 1).0;
        }

        digits
    }

    /// Performs lint on `digits` (no decimal point) and returns the group size. `None` is
    /// returned when emitting a warning.
    fn do_lint(digits: &str, cx: &EarlyContext, span: &syntax_pos::Span) -> Option<usize> {
        // Grab underscore indices with respect to the units digit.
        let underscore_positions: Vec<usize> = digits.chars().rev().enumerate()
            .filter_map(|(idx, digit)|
                if digit == '_' {
                    Some(idx)
                } else {
                    None
                })
            .collect();

        if underscore_positions.is_empty() {
            // Check if literal needs underscores.
            if digits.len() > 4 {
                span_help_and_lint(cx, UNREADABLE_LITERAL, *span,
                                   "long literal lacking separators",
                                   "consider using underscores to make literal more readable");
                return None;
            } else {
                return Some(0);
            }
        } else {
            // Check consistency and the sizes of the groups.
            let group_size = underscore_positions[0];
            let consistent = underscore_positions
                .windows(2)
                .all(|ps| ps[1] - ps[0] == group_size + 1)
                // number of digits to the left of the last group cannot be bigger than group size.
                && (digits.len() - underscore_positions.last().unwrap() <= group_size + 1);

            if !consistent {
                span_help_and_lint(cx, INCONSISTENT_DIGIT_GROUPING, *span,
                                   "digits grouped inconsistently by underscores",
                                   "consider making each group three or four digits");
                return None;
            } else if group_size > 4 {
                span_help_and_lint(cx, LARGE_DIGIT_GROUPS, *span,
                                   "digit groups should be smaller",
                                   "consider using groups of three or four digits");
                return None;
            }
            return Some(group_size);
        }
    }
}
<|endoftext|>


<commit_before>#[cfg(not(windows))]
#[test]
fn build_ffi_test() {
    use std::process::Command;

    let output = Command::new("make")
        .arg("-C")
        .arg("examples")
        .arg("check")
        .output()
        .expect("failed to execute process");

    println!("status: {}", output.status);
    println!("--- stdout ---");
    println!("{}", String::from_utf8_lossy(&output.stdout));
    println!("-- stderr ---");
    println!("{}", String::from_utf8_lossy(&output.stderr));
    assert!(output.status.success());
}
<commit_msg>Run `make clean` as part of the build_ffi test.<commit_after>#[cfg(not(windows))]
#[test]
fn build_ffi_test() {
    use std::process::Command;

    let output = Command::new("make")
        .arg("-C")
        .arg("examples")
        .arg("clean")
        .arg("check")
        .output()
        .expect("failed to execute process");

    println!("status: {}", output.status);
    println!("--- stdout ---");
    println!("{}", String::from_utf8_lossy(&output.stdout));
    println!("-- stderr ---");
    println!("{}", String::from_utf8_lossy(&output.stderr));
    assert!(output.status.success());
}
<|endoftext|>


<commit_before>use distribution::Distribution;
use random::Source;

/// A categorical distribution.
#[derive(Clone)]
pub struct Categorical {
    k: usize,
    p: Vec<f64>,
}

impl Categorical {
    /// Create a categorical distribution with success probability `p`.
    ///
    /// It should hold that `p[i] >= 0`, `p[i] <= 1`, and `sum(p) == 1`.
    #[inline]
    pub fn new(p: &[f64]) -> Categorical {
        should!(is_probability_vector(p), {
            const EPSILON: f64 = 1e-12;
            let mut in_unit = true;
            let mut sum = 0.0;
            for &p in p.iter() {
                if p < 0.0 || p > 1.0 {
                    in_unit = false;
                    break;
                }
                sum += p;
            }
            in_unit && (sum - 1.0).abs() <= EPSILON
        });
        Categorical { k: p.len(), p: p.to_vec() }
    }

    /// Return the number of categories.
    #[inline(always)]
    pub fn k(&self) -> usize { self.k }

    /// Return the event probabilities.
    #[inline(always)]
    pub fn p(&self) -> &[f64] { &self.p }
}

impl Distribution for Categorical {
    type Value = usize;

    fn mean(&self) -> f64 {
        self.p.iter().enumerate().fold(0.0, |sum, (i, p)| sum + i as f64 * p)
    }

    fn var(&self) -> f64 {
        let mu = self.mean();
        self.p.iter().enumerate().fold(0.0, |sum, (i, p)| sum + (i as f64 - mu).powi(2) * p)
    }

    fn skewness(&self) -> f64 {
        let mu = self.mean();
        let sigma2 = self.var();
        let skew = self.p.iter().enumerate()
                         .fold(0.0, |sum, (i, p)| sum + (i as f64 - mu).powi(3) * p);
        skew / (sigma2 * sigma2.sqrt())
    }

    fn kurtosis(&self) -> f64 {
        let mu = self.mean();
        let sigma2 = self.var();
        let kurt = self.p.iter().enumerate()
                         .fold(0.0, |sum, (i, p)| sum + (i as f64 - mu).powi(4) * p);
        kurt / sigma2.powi(2) - 3.0
    }

    fn median(&self) -> f64 {
        if self.p[0] > 0.5 {
            return 0.0;
        } else if self.p[0] == 0.5 {
            return 0.5;
        }
        let mut sum = 0.0;
        for i in 0..self.k {
            sum += self.p[i];
            if sum == 0.5 {
                return (2 * i - 1) as f64 / 2.0;
            } else if sum > 0.5 {
                return i as f64;
            }
        }
        unreachable!()
    }

    fn modes(&self) -> Vec<usize> {
        let mut modes = Vec::new();
        let mut max = 0.0;
        for (i, &p) in self.p.iter().enumerate() {
            if p == max {
                modes.push(i);
            }
            if p > max {
                max = p;
                modes = vec![i];
            }
        }
        modes
    }

    #[inline]
    fn entropy(&self) -> f64 {
        -self.p.iter().fold(0.0, |sum, p| sum + p * p.ln())
    }

    #[inline]
    fn cdf(&self, x: usize) -> f64 {
        if x >= self.k - 1 {
            1.0
        } else {
            self.p.iter().take(x + 1).fold(0.0, |a, b| a + b)
        }
    }

    fn inv_cdf(&self, p: f64) -> usize {
        should!(0.0 <= p && p <= 1.0);
        if p == 0.0 {
            return self.p.iter().position(|&p| p > 0.0).unwrap();
        }
        let mut sum = 0.0;
        for i in 0..self.k {
            sum += self.p[i];
            if sum >= p || sum == 1.0 {
                return i;
            }
        }
        self.k - 1
    }

    #[inline]
    fn pdf(&self, x: usize) -> f64 {
        self.p[x]
    }

    #[inline]
    fn sample<S>(&self, source: &mut S) -> usize where S: Source {
        self.inv_cdf(source.read::<f64>())
    }
}

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

    macro_rules! new(
        (equal $k:expr) => { Categorical::new(&[1.0 / $k as f64; $k]) };
        ($p:expr) => { Categorical::new(&$p); }
    );

    #[test]
    fn mean() {
        assert_eq!(new!(equal 3).mean(), 1.0);
        assert_eq!(new!([0.3, 0.3, 0.4]).mean(), 1.1);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).mean(), 1.5);
    }

    #[test]
    fn var() {
        assert_eq!(new!(equal 3).var(), 2.0 / 3.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).var(), 11.0 / 12.0);
    }

    #[test]
    fn sd() {
        assert_eq!(new!(equal 2).sd(), 0.5);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).sd(), 0.9574271077563381);
    }

    #[test]
    fn skewness() {
        assert_eq!(new!(equal 6).skewness(), 0.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).skewness(), 0.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).skewness(), -0.6);
    }

    #[test]
    fn kurtosis() {
        assert_eq!(new!(equal 2).kurtosis(), -2.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).kurtosis(), -0.7999999999999998);
    }

    #[test]
    fn median() {
        assert_eq!(new!([0.6, 0.2, 0.2]).median(), 0.0);
        assert_eq!(new!(equal 2).median(), 0.5);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).median(), 2.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).median(), 0.5);
    }

    #[test]
    fn modes() {
        assert_eq!(new!([0.6, 0.2, 0.2]).modes(), vec![0]);
        assert_eq!(new!(equal 2).modes(), vec![0, 1]);
        assert_eq!(new!(equal 3).modes(), vec![0, 1, 2]);
        assert_eq!(new!([0.4, 0.2, 0.4]).modes(), vec![0, 2]);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).modes(), vec![1, 2]);
    }

    #[test]
    fn entropy() {
        use std::f64::consts::LN_2;
        assert_eq!(new!(equal 2).entropy(), LN_2);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).entropy(), 1.2798542258336676);
    }

    #[test]
    fn pdf() {
        let p = [0.0, 0.75, 0.25, 0.0];
        let d1 = new!(p);
        assert_eq!(&(0..4).map(|x| d1.pdf(x)).collect::<Vec<_>>(), &p.to_vec());

        let d2 = new!(equal 3);
        assert_eq!(&(0..3).map(|x| d2.pdf(x)).collect::<Vec<_>>(), &vec![1.0 / 3.0; 3])
    }

    #[test]
    fn cdf() {
        let d1 = new!([0.0, 0.75, 0.25, 0.0]);
        assert_eq!(&(0..4).map(|x| d1.cdf(x)).collect::<Vec<_>>(),
                   &vec![0.0, 0.75, 1.0, 1.0]);

        let d2 = new!(equal 3);
        assert_eq!(&(0..3).map(|x| d2.cdf(x)).collect::<Vec<_>>(),
                   &vec![1.0 / 3.0, 2.0 / 3.0, 1.0]);
    }

    #[test]
    fn inv_cdf() {
        let d1 = new!([0.0, 0.75, 0.25, 0.0]);
        let p1 = vec![0.0, 0.75, 0.7500001, 1.0];
        assert_eq!(&p1.iter().map(|&p| d1.inv_cdf(p)).collect::<Vec<_>>(), &vec![1, 1, 2, 2]);

        let d2 = new!(equal 3);
        let p2 = vec![0.0, 0.5, 0.75, 1.0];
        assert_eq!(&p2.iter().map(|&p| d2.inv_cdf(p)).collect::<Vec<_>>(), &vec![0, 1, 2, 2]);

    }

    #[test]
    fn sample() {
        let mut source = random::default();
        let sum = Independent(&new!([0.0, 0.5, 0.5]), &mut source).take(100).fold(0, |a, b| a + b);
        assert!(100 <= sum && sum <= 200);
    }
}
<commit_msg>Cosmetic adjustments<commit_after>use distribution::Distribution;
use random::Source;

/// A categorical distribution.
#[derive(Clone)]
pub struct Categorical {
    k: usize,
    p: Vec<f64>,
}

impl Categorical {
    /// Create a categorical distribution with success probability `p`.
    ///
    /// It should hold that `p[i] >= 0`, `p[i] <= 1`, and `sum(p) == 1`.
    #[inline]
    pub fn new(p: &[f64]) -> Categorical {
        should!(is_probability_vector(p), {
            const EPSILON: f64 = 1e-12;
            let mut in_unit = true;
            let mut sum = 0.0;
            for &p in p.iter() {
                if p < 0.0 || p > 1.0 {
                    in_unit = false;
                    break;
                }
                sum += p;
            }
            in_unit && (sum - 1.0).abs() <= EPSILON
        });
        Categorical { k: p.len(), p: p.to_vec() }
    }

    /// Return the number of categories.
    #[inline(always)]
    pub fn k(&self) -> usize { self.k }

    /// Return the event probabilities.
    #[inline(always)]
    pub fn p(&self) -> &[f64] { &self.p }
}

impl Distribution for Categorical {
    type Value = usize;

    fn mean(&self) -> f64 {
        self.p.iter().enumerate().fold(0.0, |sum, (i, p)| sum + i as f64 * p)
    }

    fn var(&self) -> f64 {
        let mean = self.mean();
        self.p.iter().enumerate().fold(0.0, |sum, (i, p)| sum + (i as f64 - mean).powi(2) * p)
    }

    fn skewness(&self) -> f64 {
        let (mean, var) = (self.mean(), self.var());
        let skew = self.p.iter().enumerate()
                                .fold(0.0, |sum, (i, p)| sum + (i as f64 - mean).powi(3) * p);
        skew / (var * var.sqrt())
    }

    fn kurtosis(&self) -> f64 {
        let (mean, var) = (self.mean(), self.var());
        let kurt = self.p.iter().enumerate()
                                .fold(0.0, |sum, (i, p)| sum + (i as f64 - mean).powi(4) * p);
        kurt / var.powi(2) - 3.0
    }

    fn median(&self) -> f64 {
        if self.p[0] > 0.5 {
            return 0.0;
        } else if self.p[0] == 0.5 {
            return 0.5;
        }
        let mut sum = 0.0;
        for i in 0..self.k {
            sum += self.p[i];
            if sum == 0.5 {
                return (2 * i - 1) as f64 / 2.0;
            } else if sum > 0.5 {
                return i as f64;
            }
        }
        unreachable!()
    }

    fn modes(&self) -> Vec<usize> {
        let mut modes = Vec::new();
        let mut max = 0.0;
        for (i, &p) in self.p.iter().enumerate() {
            if p == max {
                modes.push(i);
            }
            if p > max {
                max = p;
                modes = vec![i];
            }
        }
        modes
    }

    #[inline]
    fn entropy(&self) -> f64 {
        -self.p.iter().fold(0.0, |sum, p| sum + p * p.ln())
    }

    #[inline]
    fn cdf(&self, x: usize) -> f64 {
        if x >= self.k - 1 {
            1.0
        } else {
            self.p.iter().take(x + 1).fold(0.0, |a, b| a + b)
        }
    }

    fn inv_cdf(&self, p: f64) -> usize {
        should!(0.0 <= p && p <= 1.0);
        if p == 0.0 {
            return self.p.iter().position(|&p| p > 0.0).unwrap();
        }
        let mut sum = 0.0;
        for i in 0..self.k {
            sum += self.p[i];
            if sum >= p || sum == 1.0 {
                return i;
            }
        }
        self.k - 1
    }

    #[inline]
    fn pdf(&self, x: usize) -> f64 {
        self.p[x]
    }

    #[inline]
    fn sample<S>(&self, source: &mut S) -> usize where S: Source {
        self.inv_cdf(source.read::<f64>())
    }
}

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

    macro_rules! new(
        (equal $k:expr) => { Categorical::new(&[1.0 / $k as f64; $k]) };
        ($p:expr) => { Categorical::new(&$p); }
    );

    #[test]
    fn mean() {
        assert_eq!(new!(equal 3).mean(), 1.0);
        assert_eq!(new!([0.3, 0.3, 0.4]).mean(), 1.1);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).mean(), 1.5);
    }

    #[test]
    fn var() {
        assert_eq!(new!(equal 3).var(), 2.0 / 3.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).var(), 11.0 / 12.0);
    }

    #[test]
    fn sd() {
        assert_eq!(new!(equal 2).sd(), 0.5);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).sd(), 0.9574271077563381);
    }

    #[test]
    fn skewness() {
        assert_eq!(new!(equal 6).skewness(), 0.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).skewness(), 0.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).skewness(), -0.6);
    }

    #[test]
    fn kurtosis() {
        assert_eq!(new!(equal 2).kurtosis(), -2.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).kurtosis(), -0.7999999999999998);
    }

    #[test]
    fn median() {
        assert_eq!(new!([0.6, 0.2, 0.2]).median(), 0.0);
        assert_eq!(new!(equal 2).median(), 0.5);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).median(), 2.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).median(), 0.5);
    }

    #[test]
    fn modes() {
        assert_eq!(new!([0.6, 0.2, 0.2]).modes(), vec![0]);
        assert_eq!(new!(equal 2).modes(), vec![0, 1]);
        assert_eq!(new!(equal 3).modes(), vec![0, 1, 2]);
        assert_eq!(new!([0.4, 0.2, 0.4]).modes(), vec![0, 2]);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).modes(), vec![1, 2]);
    }

    #[test]
    fn entropy() {
        use std::f64::consts::LN_2;
        assert_eq!(new!(equal 2).entropy(), LN_2);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).entropy(), 1.2798542258336676);
    }

    #[test]
    fn pdf() {
        let p = [0.0, 0.75, 0.25, 0.0];
        let d1 = new!(p);
        assert_eq!(&(0..4).map(|x| d1.pdf(x)).collect::<Vec<_>>(), &p.to_vec());

        let d2 = new!(equal 3);
        assert_eq!(&(0..3).map(|x| d2.pdf(x)).collect::<Vec<_>>(), &vec![1.0 / 3.0; 3])
    }

    #[test]
    fn cdf() {
        let d1 = new!([0.0, 0.75, 0.25, 0.0]);
        assert_eq!(&(0..4).map(|x| d1.cdf(x)).collect::<Vec<_>>(),
                   &vec![0.0, 0.75, 1.0, 1.0]);

        let d2 = new!(equal 3);
        assert_eq!(&(0..3).map(|x| d2.cdf(x)).collect::<Vec<_>>(),
                   &vec![1.0 / 3.0, 2.0 / 3.0, 1.0]);
    }

    #[test]
    fn inv_cdf() {
        let d1 = new!([0.0, 0.75, 0.25, 0.0]);
        let p1 = vec![0.0, 0.75, 0.7500001, 1.0];
        assert_eq!(&p1.iter().map(|&p| d1.inv_cdf(p)).collect::<Vec<_>>(), &vec![1, 1, 2, 2]);

        let d2 = new!(equal 3);
        let p2 = vec![0.0, 0.5, 0.75, 1.0];
        assert_eq!(&p2.iter().map(|&p| d2.inv_cdf(p)).collect::<Vec<_>>(), &vec![0, 1, 2, 2]);

    }

    #[test]
    fn sample() {
        let mut source = random::default();
        let sum = Independent(&new!([0.0, 0.5, 0.5]), &mut source).take(100).fold(0, |a, b| a + b);
        assert!(100 <= sum && sum <= 200);
    }
}
<|endoftext|>


<commit_before>//!
//! A demonstration of constructing and using a blocking stream.
//!
//! Audio from the default input device is passed directly to the default output device in a duplex
//! stream, so beware of feedback!
//!

extern crate portaudio;

use portaudio::pa;
use std::error::Error;
use std::mem::replace;

const SAMPLE_RATE: f64 = 44_100.0;
const CHANNELS: u32 = 2;
const FRAMES: u32 = 256;

fn main() {

    println!("PortAudio version : {}", pa::get_version());
    println!("PortAudio version text : {}", pa::get_version_text());

    match pa::initialize() {
        Ok(()) => println!("Successfully initialized PortAudio"),
        Err(err) => println!("An error occurred while initializing PortAudio: {}", err.description()),
    }

    println!("PortAudio host count : {}", pa::host::get_api_count() as isize);

    let default_host = pa::host::get_default_api();
    println!("PortAudio default host : {}", default_host as isize);

    match pa::host::get_api_info(default_host) {
        None => println!("Couldn't retrieve api info for the default host."),
        Some(info) => println!("PortAudio host name : {}", info.name),
    }

    let type_id = pa::host::api_type_id_to_host_api_index(pa::HostApiTypeId::CoreAudio) as isize;
    println!("PortAudio type id : {}", type_id);

    let def_input = pa::device::get_default_input();
    let input_info = match pa::device::get_info(def_input) {
        Ok(info) => info,
        Err(err) => panic!("An error occurred while retrieving input info: {}", err.description()),
    };
    println!("Default input device info :");
    println!("\tversion : {}", input_info.struct_version);
    println!("\tname : {}", input_info.name);
    println!("\tmax input channels : {}", input_info.max_input_channels);
    println!("\tmax output channels : {}", input_info.max_output_channels);
    println!("\tdefault sample rate : {}", input_info.default_sample_rate);

    // Construct the input stream parameters.
    let input_stream_params = pa::StreamParameters {
        device : def_input,
        channel_count : CHANNELS as i32,
        sample_format : pa::SampleFormat::Float32,
        suggested_latency : input_info.default_low_input_latency
    };

    let def_output = pa::device::get_default_output();
    let output_info = match pa::device::get_info(def_output) {
        Ok(info) => info,
        Err(err) => panic!("An error occurred while retrieving output info: {}", err.description()),
    };

    println!("Default output device name : {}", output_info.name);

    // Construct the output stream parameters.
    let output_stream_params = pa::StreamParameters {
        device : def_output,
        channel_count : CHANNELS as i32,
        sample_format : pa::SampleFormat::Float32,
        suggested_latency : output_info.default_low_output_latency
    };

    let mut stream : pa::Stream<f32, f32> = pa::Stream::new();

    match stream.open(Some(&input_stream_params),
                      Some(&output_stream_params),
                      SAMPLE_RATE,
                      FRAMES,
                      pa::StreamFlags::ClipOff,
                      None) {
        Ok(()) => println!("Successfully opened the stream."),
        Err(err) => println!("An error occurred while opening the stream: {}", err.description()),
    }

    match stream.start() {
        Ok(()) => println!("Successfully started the stream."),
        Err(err) => println!("An error occurred while starting the stream: {}", err.description()),
    }

    // We'll use this function to wait for read/write availability.
    fn wait_for_stream<F: Fn() -> Result<pa::StreamAvailable, pa::error::Error>>(f: F, name: &str)
        -> u32
    {
        'waiting_for_stream: loop {
            match f() {
                Ok(available) => match available {
                    pa::StreamAvailable::Frames(frames) => return frames as u32,
                    pa::StreamAvailable::InputOverflowed => println!("Input stream has overflowed"),
                    pa::StreamAvailable::OutputUnderflowed => println!("Output stream has underflowed"),
                },
                Err(err) => panic!("An error occurred while waiting for the {} stream: {}", name, err.description()),
            }
        }
    };

    // We'll use this buffer to transfer samples from the input stream to the output stream.
    let mut buffer = Vec::with_capacity((FRAMES * CHANNELS) as usize);

    // Now start the main read/write loop! In this example, we pass the input buffer directly to
    // the output buffer, so watch out for feedback.
    'stream: loop {

        // How many frames are available on the input stream?
        let in_frames = wait_for_stream(|| stream.get_stream_read_available(), "Read");

        // If there are frames available, let's take them and add them to our buffer.
        if in_frames > 0 {
            match stream.read(in_frames) {
                Ok(input_samples) => {
                    buffer.extend(input_samples.into_iter());
                    println!("Read {:?} frames from the input stream.", in_frames);
                },
                Err(err) => {
                    println!("An error occurred while reading from the input stream: {}", err.description());
                    break 'stream
                },
            }
        }

        // How many frames are available for writing on the output stream?
        let out_frames = wait_for_stream(|| stream.get_stream_write_available(), "Write");

        // How many frames do we have so far?
        let buffer_frames = (buffer.len() / CHANNELS as usize) as u32;

        // If there are frames available for writing and we have some to write, then write!
        if out_frames > 0 && buffer_frames > 0 {
            // If we have more than enough frames for writing, take them from the start of the buffer.
            let (write_buffer, write_frames) = if buffer_frames >= out_frames {
                let out_samples = (out_frames * CHANNELS as u32) as usize;
                let remaining_buffer = buffer[out_samples..].iter().map(|&sample| sample).collect();
                let write_buffer = replace(&mut buffer, remaining_buffer);
                (write_buffer, out_frames)
            }
            // Otherwise if we have less, just take what we can for now.
            else {
                let write_buffer = replace(&mut buffer, Vec::with_capacity((FRAMES * CHANNELS) as usize));
                (write_buffer, buffer_frames)
            };
            match stream.write(write_buffer, write_frames) {
                Ok(_) => println!("Wrote {:?} frames to the output stream.", out_frames),
                Err(err) => {
                    println!("An error occurred while writing to the output stream: {}", err.description());
                    break 'stream
                },
            }
        }

    }

    match stream.close() {
        Ok(()) => println!("Successfully closed the stream."),
        Err(err) => println!("An error occurred while closing the stream: {}", err.description()),
    }

    println!("");

    match pa::terminate() {
        Ok(()) => println!("Successfully terminated PortAudio."),
        Err(err) => println!("An error occurred while terminating PortAudio: {}", err.description()),
    }

}
<commit_msg>Added missing truncation of buffer after taking samples<commit_after>//!
//! A demonstration of constructing and using a blocking stream.
//!
//! Audio from the default input device is passed directly to the default output device in a duplex
//! stream, so beware of feedback!
//!

extern crate portaudio;

use portaudio::pa;
use std::error::Error;
use std::mem::replace;

const SAMPLE_RATE: f64 = 44_100.0;
const CHANNELS: u32 = 2;
const FRAMES: u32 = 256;

fn main() {

    println!("PortAudio version : {}", pa::get_version());
    println!("PortAudio version text : {}", pa::get_version_text());

    match pa::initialize() {
        Ok(()) => println!("Successfully initialized PortAudio"),
        Err(err) => println!("An error occurred while initializing PortAudio: {}", err.description()),
    }

    println!("PortAudio host count : {}", pa::host::get_api_count() as isize);

    let default_host = pa::host::get_default_api();
    println!("PortAudio default host : {}", default_host as isize);

    match pa::host::get_api_info(default_host) {
        None => println!("Couldn't retrieve api info for the default host."),
        Some(info) => println!("PortAudio host name : {}", info.name),
    }

    let type_id = pa::host::api_type_id_to_host_api_index(pa::HostApiTypeId::CoreAudio) as isize;
    println!("PortAudio type id : {}", type_id);

    let def_input = pa::device::get_default_input();
    let input_info = match pa::device::get_info(def_input) {
        Ok(info) => info,
        Err(err) => panic!("An error occurred while retrieving input info: {}", err.description()),
    };
    println!("Default input device info :");
    println!("\tversion : {}", input_info.struct_version);
    println!("\tname : {}", input_info.name);
    println!("\tmax input channels : {}", input_info.max_input_channels);
    println!("\tmax output channels : {}", input_info.max_output_channels);
    println!("\tdefault sample rate : {}", input_info.default_sample_rate);

    // Construct the input stream parameters.
    let input_stream_params = pa::StreamParameters {
        device : def_input,
        channel_count : CHANNELS as i32,
        sample_format : pa::SampleFormat::Float32,
        suggested_latency : input_info.default_low_input_latency
    };

    let def_output = pa::device::get_default_output();
    let output_info = match pa::device::get_info(def_output) {
        Ok(info) => info,
        Err(err) => panic!("An error occurred while retrieving output info: {}", err.description()),
    };

    println!("Default output device name : {}", output_info.name);

    // Construct the output stream parameters.
    let output_stream_params = pa::StreamParameters {
        device : def_output,
        channel_count : CHANNELS as i32,
        sample_format : pa::SampleFormat::Float32,
        suggested_latency : output_info.default_low_output_latency
    };

    let mut stream : pa::Stream<f32, f32> = pa::Stream::new();

    match stream.open(Some(&input_stream_params),
                      Some(&output_stream_params),
                      SAMPLE_RATE,
                      FRAMES,
                      pa::StreamFlags::ClipOff,
                      None) {
        Ok(()) => println!("Successfully opened the stream."),
        Err(err) => println!("An error occurred while opening the stream: {}", err.description()),
    }

    match stream.start() {
        Ok(()) => println!("Successfully started the stream."),
        Err(err) => println!("An error occurred while starting the stream: {}", err.description()),
    }

    // We'll use this function to wait for read/write availability.
    fn wait_for_stream<F: Fn() -> Result<pa::StreamAvailable, pa::error::Error>>(f: F, name: &str)
        -> u32
    {
        'waiting_for_stream: loop {
            match f() {
                Ok(available) => match available {
                    pa::StreamAvailable::Frames(frames) => return frames as u32,
                    pa::StreamAvailable::InputOverflowed => println!("Input stream has overflowed"),
                    pa::StreamAvailable::OutputUnderflowed => println!("Output stream has underflowed"),
                },
                Err(err) => panic!("An error occurred while waiting for the {} stream: {}", name, err.description()),
            }
        }
    };

    // We'll use this buffer to transfer samples from the input stream to the output stream.
    let mut buffer = Vec::with_capacity((FRAMES * CHANNELS) as usize);

    // Now start the main read/write loop! In this example, we pass the input buffer directly to
    // the output buffer, so watch out for feedback.
    'stream: loop {

        // How many frames are available on the input stream?
        let in_frames = wait_for_stream(|| stream.get_stream_read_available(), "Read");

        // If there are frames available, let's take them and add them to our buffer.
        if in_frames > 0 {
            match stream.read(in_frames) {
                Ok(input_samples) => {
                    buffer.extend(input_samples.into_iter());
                    println!("Read {:?} frames from the input stream.", in_frames);
                },
                Err(err) => {
                    println!("An error occurred while reading from the input stream: {}", err.description());
                    break 'stream
                },
            }
        }

        // How many frames are available for writing on the output stream?
        let out_frames = wait_for_stream(|| stream.get_stream_write_available(), "Write");

        // How many frames do we have so far?
        let buffer_frames = (buffer.len() / CHANNELS as usize) as u32;

        // If there are frames available for writing and we have some to write, then write!
        if out_frames > 0 && buffer_frames > 0 {
            // If we have more than enough frames for writing, take them from the start of the buffer.
            let (write_buffer, write_frames) = if buffer_frames >= out_frames {
                let out_samples = (out_frames * CHANNELS as u32) as usize;
                let remaining_buffer = buffer[out_samples..].iter().map(|&sample| sample).collect();
                buffer.truncate(out_samples);
                let write_buffer = replace(&mut buffer, remaining_buffer);
                (write_buffer, out_frames)
            }
            // Otherwise if we have less, just take what we can for now.
            else {
                let write_buffer = replace(&mut buffer, Vec::with_capacity((FRAMES * CHANNELS) as usize));
                (write_buffer, buffer_frames)
            };
            match stream.write(write_buffer, write_frames) {
                Ok(_) => println!("Wrote {:?} frames to the output stream.", out_frames),
                Err(err) => {
                    println!("An error occurred while writing to the output stream: {}", err.description());
                    break 'stream
                },
            }
        }

    }

    match stream.close() {
        Ok(()) => println!("Successfully closed the stream."),
        Err(err) => println!("An error occurred while closing the stream: {}", err.description()),
    }

    println!("");

    match pa::terminate() {
        Ok(()) => println!("Successfully terminated PortAudio."),
        Err(err) => println!("An error occurred while terminating PortAudio: {}", err.description()),
    }

}
<|endoftext|>


<commit_before>extern crate janeiro;

#[macro_use]
extern crate log;
extern crate env_logger;

use std::str;

use janeiro::{Rio, Transport, Protocol, Reason};

struct HttpClientProtocol;
impl HttpClientProtocol {
    fn new() -> HttpClientProtocol {
        HttpClientProtocol
    }
}

impl Protocol for HttpClientProtocol {

    fn connection_made(&self, transport: &mut Transport) {
        debug!("Connection made");
        let data = b"GET /\r\nHost: www.gandi.net\r\nAccept: text/html\r\n\r\n";
        transport.write(data);
    }

    #[allow(unused_variables)]
    fn data_received(&self, data: &[u8], transport: &mut Transport) {
        let s_data = str::from_utf8(data).unwrap().trim();

        println!("==========================================================");
        println!("{}", s_data);
        println!("==========================================================");
        //transport.hang_up();

    }

    fn connection_lost(&self, reason: Reason) {
        match reason {
            Reason::ConnectionLost => info!("Connection closed by peer"),
            Reason::HangUp => info!("Hang hup"),
            Reason::ConnectionError => println!("Connection error"),
        }
    }
}

fn main() {
    env_logger::init().unwrap();
    info!("Start the client");
    let mut rio = Rio::new();
    let protocol = HttpClientProtocol::new();
    rio.connect("217.70.184.1:80", Box::new(protocol));
    info!("Start running the loop");
    rio.run_forever();
}
<commit_msg>Fix http_cli<commit_after>extern crate janeiro;

#[macro_use]
extern crate log;
extern crate env_logger;

use std::str;

use janeiro::{Rio, Transport, Protocol, Reason};

struct HttpClientProtocol;
impl HttpClientProtocol {
    fn new() -> HttpClientProtocol {
        HttpClientProtocol
    }
}

impl Protocol for HttpClientProtocol {

    fn connection_made(&mut self, transport: &mut Transport) {
        debug!("Connection made");
        let data = b"GET /\r\nHost: www.gandi.net\r\nAccept: text/html\r\n\r\n";
        transport.write(data);
    }

    #[allow(unused_variables)]
    fn data_received(&mut self, data: &[u8], transport: &mut Transport) {
        let s_data = str::from_utf8(data).unwrap().trim();

        println!("==========================================================");
        println!("{}", s_data);
        println!("==========================================================");
        //transport.hang_up();

    }

    fn connection_lost(&mut self, reason: Reason) {
        match reason {
            Reason::ConnectionLost => info!("Connection closed by peer"),
            Reason::HangUp => info!("Hang hup"),
            Reason::ConnectionError => println!("Connection error"),
        }
    }
}

fn main() {
    env_logger::init().unwrap();
    info!("Start the client");
    let mut rio = Rio::new();
    let protocol = HttpClientProtocol::new();
    let _ = rio.connect("217.70.184.1:80", Box::new(protocol));
    info!("Start running the loop");
    rio.run_forever();
}
<|endoftext|>


<commit_before>//! Shadowsocks TCP server

use std::{
    io::{self, ErrorKind},
    net::SocketAddr,
    sync::Arc,
    time::Duration,
};

use futures::future::{self, Either};
use log::{debug, error, info, trace, warn};
use shadowsocks::{
    crypto::v1::CipherKind,
    net::{AcceptOpts, TcpStream as OutboundTcpStream},
    relay::{
        socks5::{Address, Error as Socks5Error},
        tcprelay::{
            utils::{copy_from_encrypted, copy_to_encrypted},
            ProxyServerStream,
        },
    },
    ProxyListener,
    ServerConfig,
};
use tokio::{net::TcpStream as TokioTcpStream, time};

use crate::net::{utils::ignore_until_end, MonProxyStream};

use super::context::ServiceContext;

pub struct TcpServer {
    context: Arc<ServiceContext>,
    accept_opts: AcceptOpts,
}

impl TcpServer {
    pub fn new(context: Arc<ServiceContext>, accept_opts: AcceptOpts) -> TcpServer {
        TcpServer { context, accept_opts }
    }

    pub async fn run(self, svr_cfg: &ServerConfig) -> io::Result<()> {
        let listener = ProxyListener::bind_with_opts(self.context.context(), svr_cfg, self.accept_opts).await?;

        info!(
            "shadowsocks tcp server listening on {}, inbound address {}",
            listener.local_addr().expect("listener.local_addr"),
            svr_cfg.addr()
        );

        loop {
            let flow_stat = self.context.flow_stat();

            let (local_stream, peer_addr) =
                match listener.accept_map(|s| MonProxyStream::from_stream(s, flow_stat)).await {
                    Ok(s) => s,
                    Err(err) => {
                        error!("tcp server accept failed with error: {}", err);
                        time::sleep(Duration::from_secs(1)).await;
                        continue;
                    }
                };

            let client = TcpServerClient {
                context: self.context.clone(),
                method: svr_cfg.method(),
                peer_addr,
                stream: local_stream,
                timeout: svr_cfg.timeout(),
            };

            tokio::spawn(async move {
                if let Err(err) = client.serve().await {
                    debug!("tcp server stream aborted with error: {}", err);
                }
            });
        }
    }
}

struct TcpServerClient {
    context: Arc<ServiceContext>,
    method: CipherKind,
    peer_addr: SocketAddr,
    stream: ProxyServerStream<MonProxyStream<TokioTcpStream>>,
    timeout: Option<Duration>,
}

impl TcpServerClient {
    async fn serve(mut self) -> io::Result<()> {
        let target_addr = match Address::read_from(&mut self.stream).await {
            Ok(a) => a,
            Err(Socks5Error::IoError(ref err)) if err.kind() == ErrorKind::UnexpectedEof => {
                debug!(
                    "handshake failed, received EOF before a complete target Address, peer: {}",
                    self.peer_addr
                );
                return Ok(());
            }
            Err(err) => {
                // https://github.com/shadowsocks/shadowsocks-rust/issues/292
                //
                // Keep connection open.
                warn!(
                    "handshake failed, maybe wrong method or key, or under reply attacks. peer: {}, error: {}",
                    self.peer_addr, err
                );
                let _ = ignore_until_end(&mut self.stream).await;
                return Ok(());
            }
        };

        trace!(
            "accepted tcp client connection {}, establishing tunnel to {}",
            self.peer_addr,
            target_addr
        );

        if self.context.check_outbound_blocked(&target_addr).await {
            error!(
                "tcp client {} outbound {} blocked by ACL rules",
                self.peer_addr, target_addr
            );
            return Ok(());
        }

        let mut remote_stream = match self.timeout {
            Some(d) => {
                match time::timeout(
                    d,
                    OutboundTcpStream::connect_remote_with_opts(
                        self.context.context_ref(),
                        &target_addr,
                        self.context.connect_opts_ref(),
                    ),
                )
                .await
                {
                    Ok(Ok(s)) => s,
                    Ok(Err(e)) => return Err(e),
                    Err(..) => {
                        return Err(io::Error::new(
                            ErrorKind::TimedOut,
                            format!("connect {} timeout", target_addr),
                        ))
                    }
                }
            }
            None => {
                OutboundTcpStream::connect_remote_with_opts(
                    self.context.context_ref(),
                    &target_addr,
                    self.context.connect_opts_ref(),
                )
                .await?
            }
        };

        let (mut lr, mut lw) = self.stream.into_split();
        let (mut rr, mut rw) = remote_stream.split();

        let l2r = copy_to_encrypted(self.method, &mut lr, &mut rw);
        let r2l = copy_from_encrypted(self.method, &mut rr, &mut lw);

        tokio::pin!(l2r);
        tokio::pin!(r2l);

        debug!(
            "established tcp tunnel {} <-> {} with {:?}",
            self.peer_addr,
            target_addr,
            self.context.connect_opts_ref()
        );

        match future::select(l2r, r2l).await {
            Either::Left((Ok(..), ..)) => {
                trace!("tcp tunnel {} -> {} closed", self.peer_addr, target_addr);
            }
            Either::Left((Err(err), ..)) => {
                trace!(
                    "tcp tunnel {} -> {} closed with error: {}",
                    self.peer_addr,
                    target_addr,
                    err
                );
            }
            Either::Right((Ok(..), ..)) => {
                trace!("tcp tunnel {} <- {} closed", self.peer_addr, target_addr);
            }
            Either::Right((Err(err), ..)) => {
                trace!(
                    "tcp tunnel {} <- {} closed with error: {}",
                    self.peer_addr,
                    target_addr,
                    err
                );
            }
        }

        Ok(())
    }
}
<commit_msg>send syn after 500ms if TFO is enabled for server<commit_after>//! Shadowsocks TCP server

use std::{
    future::Future,
    io::{self, ErrorKind},
    net::SocketAddr,
    sync::Arc,
    time::Duration,
};

use futures::future::{self, Either};
use log::{debug, error, info, trace, warn};
use shadowsocks::{
    crypto::v1::CipherKind,
    net::{AcceptOpts, TcpStream as OutboundTcpStream},
    relay::{
        socks5::{Address, Error as Socks5Error},
        tcprelay::{
            utils::{copy_from_encrypted, copy_to_encrypted},
            ProxyServerStream,
        },
    },
    ProxyListener,
    ServerConfig,
};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::TcpStream as TokioTcpStream,
    time,
};

use crate::net::{utils::ignore_until_end, MonProxyStream};

use super::context::ServiceContext;

pub struct TcpServer {
    context: Arc<ServiceContext>,
    accept_opts: AcceptOpts,
}

impl TcpServer {
    pub fn new(context: Arc<ServiceContext>, accept_opts: AcceptOpts) -> TcpServer {
        TcpServer { context, accept_opts }
    }

    pub async fn run(self, svr_cfg: &ServerConfig) -> io::Result<()> {
        let listener = ProxyListener::bind_with_opts(self.context.context(), svr_cfg, self.accept_opts).await?;

        info!(
            "shadowsocks tcp server listening on {}, inbound address {}",
            listener.local_addr().expect("listener.local_addr"),
            svr_cfg.addr()
        );

        loop {
            let flow_stat = self.context.flow_stat();

            let (local_stream, peer_addr) =
                match listener.accept_map(|s| MonProxyStream::from_stream(s, flow_stat)).await {
                    Ok(s) => s,
                    Err(err) => {
                        error!("tcp server accept failed with error: {}", err);
                        time::sleep(Duration::from_secs(1)).await;
                        continue;
                    }
                };

            let client = TcpServerClient {
                context: self.context.clone(),
                method: svr_cfg.method(),
                peer_addr,
                stream: local_stream,
                timeout: svr_cfg.timeout(),
            };

            tokio::spawn(async move {
                if let Err(err) = client.serve().await {
                    debug!("tcp server stream aborted with error: {}", err);
                }
            });
        }
    }
}

#[inline]
async fn timeout_fut<F, R>(duration: Option<Duration>, f: F) -> io::Result<R>
where
    F: Future<Output = io::Result<R>>,
{
    match duration {
        None => f.await,
        Some(d) => match time::timeout(d, f).await {
            Ok(o) => o,
            Err(..) => Err(ErrorKind::TimedOut.into()),
        },
    }
}

struct TcpServerClient {
    context: Arc<ServiceContext>,
    method: CipherKind,
    peer_addr: SocketAddr,
    stream: ProxyServerStream<MonProxyStream<TokioTcpStream>>,
    timeout: Option<Duration>,
}

impl TcpServerClient {
    async fn serve(mut self) -> io::Result<()> {
        let target_addr = match Address::read_from(&mut self.stream).await {
            Ok(a) => a,
            Err(Socks5Error::IoError(ref err)) if err.kind() == ErrorKind::UnexpectedEof => {
                debug!(
                    "handshake failed, received EOF before a complete target Address, peer: {}",
                    self.peer_addr
                );
                return Ok(());
            }
            Err(err) => {
                // https://github.com/shadowsocks/shadowsocks-rust/issues/292
                //
                // Keep connection open.
                warn!(
                    "handshake failed, maybe wrong method or key, or under reply attacks. peer: {}, error: {}",
                    self.peer_addr, err
                );
                let _ = ignore_until_end(&mut self.stream).await;
                return Ok(());
            }
        };

        trace!(
            "accepted tcp client connection {}, establishing tunnel to {}",
            self.peer_addr,
            target_addr
        );

        if self.context.check_outbound_blocked(&target_addr).await {
            error!(
                "tcp client {} outbound {} blocked by ACL rules",
                self.peer_addr, target_addr
            );
            return Ok(());
        }

        let mut remote_stream = match timeout_fut(
            self.timeout,
            OutboundTcpStream::connect_remote_with_opts(
                self.context.context_ref(),
                &target_addr,
                self.context.connect_opts_ref(),
            ),
        )
        .await
        {
            Ok(s) => s,
            Err(err) => {
                error!(
                    "tcp tunnel {} -> {} connect failed, error: {}",
                    self.peer_addr, target_addr, err
                );
                return Err(err);
            }
        };

        // https://github.com/shadowsocks/shadowsocks-rust/issues/232
        //
        // Protocols like FTP, clients will wait for servers to send Welcome Message without sending anything.
        //
        // Wait at most 500ms, and then sends handshake packet to remote servers.
        if self.context.connect_opts_ref().tcp.fastopen {
            let mut buffer = [0u8; 8192];
            match time::timeout(Duration::from_millis(500), self.stream.read(&mut buffer)).await {
                Ok(Ok(0)) => {
                    // EOF. Just terminate right here.
                    return Ok(());
                }
                Ok(Ok(n)) => {
                    // Send the first packet.
                    timeout_fut(self.timeout, remote_stream.write_all(&buffer[..n])).await?;
                }
                Ok(Err(err)) => return Err(err),
                Err(..) => {
                    // Timeout. Send handshake to server.
                    timeout_fut(self.timeout, remote_stream.write(&[])).await?;

                    trace!(
                        "tcp tunnel {} -> {} sent TFO connect without data",
                        self.peer_addr,
                        target_addr
                    );
                }
            }
        }

        let (mut lr, mut lw) = self.stream.into_split();
        let (mut rr, mut rw) = remote_stream.split();

        let l2r = copy_to_encrypted(self.method, &mut lr, &mut rw);
        let r2l = copy_from_encrypted(self.method, &mut rr, &mut lw);

        tokio::pin!(l2r);
        tokio::pin!(r2l);

        debug!(
            "established tcp tunnel {} <-> {} with {:?}",
            self.peer_addr,
            target_addr,
            self.context.connect_opts_ref()
        );

        match future::select(l2r, r2l).await {
            Either::Left((Ok(..), ..)) => {
                trace!("tcp tunnel {} -> {} closed", self.peer_addr, target_addr);
            }
            Either::Left((Err(err), ..)) => {
                trace!(
                    "tcp tunnel {} -> {} closed with error: {}",
                    self.peer_addr,
                    target_addr,
                    err
                );
            }
            Either::Right((Ok(..), ..)) => {
                trace!("tcp tunnel {} <- {} closed", self.peer_addr, target_addr);
            }
            Either::Right((Err(err), ..)) => {
                trace!(
                    "tcp tunnel {} <- {} closed with error: {}",
                    self.peer_addr,
                    target_addr,
                    err
                );
            }
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>extern crate tinyexpr;

fn main()
{
    let r = tinyexpr::interp("2*1/sin(3.14/2)").unwrap_or_else(|e| {
        panic!("{}", e);
    });

    println!("{:?}", r);
}
<commit_msg>updated example<commit_after>extern crate tinyexpr;

fn main()
{
    // evaluate expression and fetch result
    let result = tinyexpr::interp("2*1/sin(pi/2)").unwrap_or_else(|e| {
        panic!("{}", e);
    });

    println!("{:?}", result);
}
<|endoftext|>