<commit_before>use seed::prelude::*;
use crate::Msg;

use num_traits::{ToPrimitive};

const PERFECTS_COUNT: usize = 27;

pub mod perfects_utils {
	use num_bigint::{BigUint,ToBigUint};
    use seed::prelude::*;
    use crate::Msg;

	#[derive(Clone)]
	pub struct Perfect {
		pub n: u64,
		pub p: u64,
		pub digits: u64,
		pub discovery: String,
	}

	#[derive(Clone)]
	pub struct PerfectDownload {
		pub n: u64,
		pub p: u64,
	}

	pub fn perfects() -> Vec<Perfect> {
        vec![
			/*Perfect {n: 51, p: 82589934, digits: 49724095, discovery: "2018 Laroche, Woltman, Kurowski, Blosser, et al."},
			Perfect {n: 50, p: 77232917, digits: 46498850, discovery: "2017 Jonathan Pace, George Woltman, Scott Kurowski, Aaron Blosser, et al.." },
			Perfect {n: 49, p: 74207281, digits: 44677235, discovery: "2016 Cooper, Woltman, Kurowski, Blosser et al." },
			Perfect {n: 48, p: 57885161, digits: 34850340, discovery: "2013 Cooper, Woltman, Kurowski, et al." },
			Perfect {n: 47, p: 43112609, digits: 25956377, discovery: "2008 Smith, Woltman, Kurowski, et al." },
			Perfect {n: 46, p: 42643801, digits: 42643801, discovery: "2009 Strindmo, Woltman, Kurowski, et al." },
			Perfect {n: 45, p: 37156667, digits: 22370543, discovery: "2008 Elvenich, Woltman, Kurowski, et al." },
			Perfect {n: 44, p: 32582657, digits: 19616714, discovery: "2006 Cooper, Boone, Woltman, Kurowski, et al." },
			Perfect {n: 43, p: 30402457, digits: 18304103, discovery: "2005 Cooper, Boone, Woltman, Kurowski, et al." },
			Perfect {n: 42, p: 25964951, digits: 15632458, discovery: "2005 Nowak, Woltman, Kurowski, et. al." },
			Perfect {n: 41, p: 24036583, digits: 14471465, discovery: "2004 Findley, Woltman, Kurowski, et. al." },
			Perfect {n: 40, p: 20996011, digits: 12640858, discovery: "2003 Shafer, Woltman, Kurowski, et. al." },
			Perfect {n: 39, p: 13466917, digits: 8107892, discovery: "2001 Cameron, Woltman, Kurowski, et. al." },
			Perfect {n: 38, p: 6972593, digits: 4197919, discovery: "1999 Hajratwala, Woltman, Kurowski, et. al." },
			Perfect {n: 37, p: 3021377, digits: 1819050, discovery: "1998 Clarkson, Woltman, Kurowski, et. al." },
			Perfect {n: 36, p: 2976221, digits: 1791864, discovery: "1997 Spence, Woltman, et. al." },
			Perfect {n: 35, p: 1398269, digits: 841842, discovery: "1996 Armengaud, Woltman, et. al." },
			Perfect {n: 34, p: 1257787, digits: 757263, discovery: "1996 Slowinski&Gage" },
			Perfect {n: 33, p: 859433, digits: 517430, discovery: "1994 Slowinski&Gage" },
			Perfect {n: 32, p: 756839, digits: 455663, discovery: "1992 Slowinski&Gage" },
			Perfect {n: 31, p: 216091, digits: 130100, discovery: "1985 Slowinski" },
			Perfect {n: 30, p: 132049, digits: 79502, discovery: "1983 Slowinski" },
			Perfect {n: 29, p: 110503, digits: 66530, discovery: "1988 Colquitt&Welsh" },
			Perfect {n: 28, p: 86243, digits: 51924, discovery: "1982 Slowinski" },*/
			Perfect {n: 27, p: 44497, digits: 26790, discovery: String::from("1979 Nelson&Slowinski") },
			Perfect {n: 26, p: 23209, digits: 13973, discovery: String::from("1979 Noll") },
			Perfect {n: 25, p: 21701, digits: 13066, discovery: String::from("1978 Noll&Nickel") },
			Perfect {n: 24, p: 19937, digits: 12003, discovery: String::from("1971 Tuckerman") },
			Perfect {n: 23, p: 11213, digits: 6751, discovery: String::from("1963 Gillies") },
			Perfect {n: 22, p: 9941, digits: 5985, discovery: String::from("1963 Gillies") },
			Perfect {n: 21, p: 9689, digits: 5834, discovery: String::from("1963 Gillies") },
			Perfect {n: 20, p: 4423, digits: 2663, discovery: String::from("1961 Hurwitz") },
			Perfect {n: 19, p: 4253, digits: 2561, discovery: String::from("1961 Hurwitz") },
			Perfect {n: 18, p: 3217, digits: 1937, discovery: String::from("1957 Riesel") },
			Perfect {n: 17, p: 2281, digits: 1373, discovery: String::from("1952 Robinson") },
			Perfect {n: 16, p: 2203, digits: 1327, discovery: String::from("1952 Robinson") },
			Perfect {n: 14, p: 607, digits: 366, discovery: String::from("1952 Robinson") },
			Perfect {n: 13, p: 521, digits: 314, discovery: String::from("1952 Robinson") },
			Perfect {n: 12, p: 127, digits: 77, discovery: String::from("1876 Lucas") },
			Perfect {n: 11, p: 107, digits: 65, discovery: String::from("1914 Powers") },
			Perfect {n: 10, p: 89, digits: 54, discovery: String::from("1911 Powers") },
			Perfect {n: 9, p: 61, digits: 37, discovery: String::from("1883 Pervushin") },
			Perfect {n: 8, p: 31, digits: 19, discovery: String::from("1772 Euler") },
			Perfect {n: 7, p: 19, digits: 12, discovery: String::from("1588 Cataldi") },
			Perfect {n: 6, p: 17, digits: 10, discovery: String::from("1588 Cataldi") },
			Perfect {n: 5, p: 13, digits: 8, discovery: String::from("1456 ?") },
			Perfect {n: 4, p: 7, digits: 4, discovery: String::from("?") },
			Perfect {n: 3, p: 5, digits: 3, discovery: String::from("?") },
			Perfect {n: 2, p: 3, digits: 2, discovery: String::from("?") },
			Perfect {n: 1, p: 2, digits: 1, discovery: String::from("?") },
		]
	}

	pub fn save_as_file(filename:String, filecontent:String) -> seed::dom_types::Node<Msg> {
        let href:String = vec!["data:text/plain,",&filecontent].into_iter().collect();
        a![attrs!{At::Download => &filename, At::Href => &href}, "TXT"]
    }

	pub fn generate_file(n:usize, p:usize) -> seed::dom_types::Node<Msg> {
		let download_filename:String = format!("P{}.txt",&n.to_string());
		let two:BigUint = 2.to_biguint().unwrap();
		let power:BigUint = two.clone() << (p-1-1);
        let perfect_value:BigUint = power.clone() * ((power.clone()*two.clone()) -1.to_biguint().unwrap());
        a![attrs!{At::Download => ""}, "TXT"]
    }
}

pub fn render() -> seed::dom_types::Node<Msg> {
    let mut html = vec![];
    let perfects = perfects_utils::perfects().to_owned();

    for perfect in perfects.iter() {
		let perfectp:u64 = perfect.p;
		let perfect_download = perfects_utils::PerfectDownload {n: perfect.n, p: perfect.p};
		let equation:String = format!("2{} × (2{}-1)",&(perfectp-1).to_string(),&(perfectp).to_string());

        html.push(
            tr![
                td![perfect.n.to_string()],//rank
				td![El::from_html(&equation)],//perfect number as a formula

                td![perfect.digits.to_string()],//digits in length
                td![perfect.discovery],//discovery
				td![button!["Generate",input_ev("input", move |text| Msg::GenerateDownload(perfect_download))]],

            ]
        );
    }

    div![
        h1!["The Perfect Numbers"],
        br![],
        br![],
        br![],
        table![
            attrs!{At::Class => "perfecttable text"},
            tbody![
                tr![
                    td![
                        b!["No."]
                    ],
                    td![
                        b!["Prime"]
                    ],
                    td![
                        b!["Digits"]
                    ],
                    td![
                        b!["Discovered"]
                    ],
                    td![
                        b!["Download"]
                    ]
                ],
                html
            ]
        ]
    ]
}

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

    #[test]
    fn perfect_test<'a>() {
        let mut perfects:Vec<perfects_utils::Perfect> = perfects_utils::perfects();
		perfects.reverse();
        assert_eq!(perfects[49].digits, 46498850);
        assert_eq!(perfects[39].digits, 12640858);
    }
}<commit_msg>Removed unused code<commit_after>use seed::prelude::*;
use crate::Msg;

pub mod perfects_utils {
	use num_bigint::{BigUint,ToBigUint};
    use seed::prelude::*;
    use crate::Msg;

	#[derive(Clone)]
	pub struct Perfect {
		pub n: u64,
		pub p: u64,
		pub digits: u64,
		pub discovery: String,
	}

	#[derive(Clone)]
	pub struct PerfectDownload {
		pub n: u64,
		pub p: u64,
	}

	pub fn perfects() -> Vec<Perfect> {
        vec![
			/*Perfect {n: 51, p: 82589934, digits: 49724095, discovery: "2018 Laroche, Woltman, Kurowski, Blosser, et al."},
			Perfect {n: 50, p: 77232917, digits: 46498850, discovery: "2017 Jonathan Pace, George Woltman, Scott Kurowski, Aaron Blosser, et al.." },
			Perfect {n: 49, p: 74207281, digits: 44677235, discovery: "2016 Cooper, Woltman, Kurowski, Blosser et al." },
			Perfect {n: 48, p: 57885161, digits: 34850340, discovery: "2013 Cooper, Woltman, Kurowski, et al." },
			Perfect {n: 47, p: 43112609, digits: 25956377, discovery: "2008 Smith, Woltman, Kurowski, et al." },
			Perfect {n: 46, p: 42643801, digits: 42643801, discovery: "2009 Strindmo, Woltman, Kurowski, et al." },
			Perfect {n: 45, p: 37156667, digits: 22370543, discovery: "2008 Elvenich, Woltman, Kurowski, et al." },
			Perfect {n: 44, p: 32582657, digits: 19616714, discovery: "2006 Cooper, Boone, Woltman, Kurowski, et al." },
			Perfect {n: 43, p: 30402457, digits: 18304103, discovery: "2005 Cooper, Boone, Woltman, Kurowski, et al." },
			Perfect {n: 42, p: 25964951, digits: 15632458, discovery: "2005 Nowak, Woltman, Kurowski, et. al." },
			Perfect {n: 41, p: 24036583, digits: 14471465, discovery: "2004 Findley, Woltman, Kurowski, et. al." },
			Perfect {n: 40, p: 20996011, digits: 12640858, discovery: "2003 Shafer, Woltman, Kurowski, et. al." },
			Perfect {n: 39, p: 13466917, digits: 8107892, discovery: "2001 Cameron, Woltman, Kurowski, et. al." },
			Perfect {n: 38, p: 6972593, digits: 4197919, discovery: "1999 Hajratwala, Woltman, Kurowski, et. al." },
			Perfect {n: 37, p: 3021377, digits: 1819050, discovery: "1998 Clarkson, Woltman, Kurowski, et. al." },
			Perfect {n: 36, p: 2976221, digits: 1791864, discovery: "1997 Spence, Woltman, et. al." },
			Perfect {n: 35, p: 1398269, digits: 841842, discovery: "1996 Armengaud, Woltman, et. al." },
			Perfect {n: 34, p: 1257787, digits: 757263, discovery: "1996 Slowinski&Gage" },
			Perfect {n: 33, p: 859433, digits: 517430, discovery: "1994 Slowinski&Gage" },
			Perfect {n: 32, p: 756839, digits: 455663, discovery: "1992 Slowinski&Gage" },
			Perfect {n: 31, p: 216091, digits: 130100, discovery: "1985 Slowinski" },
			Perfect {n: 30, p: 132049, digits: 79502, discovery: "1983 Slowinski" },
			Perfect {n: 29, p: 110503, digits: 66530, discovery: "1988 Colquitt&Welsh" },
			Perfect {n: 28, p: 86243, digits: 51924, discovery: "1982 Slowinski" },*/
			Perfect {n: 27, p: 44497, digits: 26790, discovery: String::from("1979 Nelson&Slowinski") },
			Perfect {n: 26, p: 23209, digits: 13973, discovery: String::from("1979 Noll") },
			Perfect {n: 25, p: 21701, digits: 13066, discovery: String::from("1978 Noll&Nickel") },
			Perfect {n: 24, p: 19937, digits: 12003, discovery: String::from("1971 Tuckerman") },
			Perfect {n: 23, p: 11213, digits: 6751, discovery: String::from("1963 Gillies") },
			Perfect {n: 22, p: 9941, digits: 5985, discovery: String::from("1963 Gillies") },
			Perfect {n: 21, p: 9689, digits: 5834, discovery: String::from("1963 Gillies") },
			Perfect {n: 20, p: 4423, digits: 2663, discovery: String::from("1961 Hurwitz") },
			Perfect {n: 19, p: 4253, digits: 2561, discovery: String::from("1961 Hurwitz") },
			Perfect {n: 18, p: 3217, digits: 1937, discovery: String::from("1957 Riesel") },
			Perfect {n: 17, p: 2281, digits: 1373, discovery: String::from("1952 Robinson") },
			Perfect {n: 16, p: 2203, digits: 1327, discovery: String::from("1952 Robinson") },
			Perfect {n: 14, p: 607, digits: 366, discovery: String::from("1952 Robinson") },
			Perfect {n: 13, p: 521, digits: 314, discovery: String::from("1952 Robinson") },
			Perfect {n: 12, p: 127, digits: 77, discovery: String::from("1876 Lucas") },
			Perfect {n: 11, p: 107, digits: 65, discovery: String::from("1914 Powers") },
			Perfect {n: 10, p: 89, digits: 54, discovery: String::from("1911 Powers") },
			Perfect {n: 9, p: 61, digits: 37, discovery: String::from("1883 Pervushin") },
			Perfect {n: 8, p: 31, digits: 19, discovery: String::from("1772 Euler") },
			Perfect {n: 7, p: 19, digits: 12, discovery: String::from("1588 Cataldi") },
			Perfect {n: 6, p: 17, digits: 10, discovery: String::from("1588 Cataldi") },
			Perfect {n: 5, p: 13, digits: 8, discovery: String::from("1456 ?") },
			Perfect {n: 4, p: 7, digits: 4, discovery: String::from("?") },
			Perfect {n: 3, p: 5, digits: 3, discovery: String::from("?") },
			Perfect {n: 2, p: 3, digits: 2, discovery: String::from("?") },
			Perfect {n: 1, p: 2, digits: 1, discovery: String::from("?") },
		]
	}

	pub fn save_as_file(filename:String, filecontent:String) -> seed::dom_types::Node<Msg> {
        let href:String = vec!["data:text/plain,",&filecontent].into_iter().collect();
        a![attrs!{At::Download => &filename, At::Href => &href}, "TXT"]
    }

	pub fn generate_file(n:usize, p:usize) -> seed::dom_types::Node<Msg> {
		let download_filename:String = format!("P{}.txt",&n.to_string());
		let two:BigUint = 2.to_biguint().unwrap();
		let power:BigUint = two.clone() << (p-1-1);
        let perfect_value:BigUint = power.clone() * ((power.clone()*two.clone()) -1.to_biguint().unwrap());
        a![attrs!{At::Download => ""}, "TXT"]
    }
}

pub fn render() -> seed::dom_types::Node<Msg> {
    let mut html = vec![];
    let perfects = perfects_utils::perfects().to_owned();

    for perfect in perfects.iter() {
		let perfectp:u64 = perfect.p;
		let perfect_download = perfects_utils::PerfectDownload {n: perfect.n, p: perfect.p};
		let equation:String = format!("2{} × (2{}-1)",&(perfectp-1).to_string(),&(perfectp).to_string());

        html.push(
            tr![
                td![perfect.n.to_string()],//rank
				td![El::from_html(&equation)],//perfect number as a formula

                td![perfect.digits.to_string()],//digits in length
                td![perfect.discovery],//discovery
				td![button!["Generate",input_ev("input", move |_| Msg::GenerateDownload(perfect_download))]],

            ]
        );
    }

    div![
        h1!["The Perfect Numbers"],
        br![],
        br![],
        br![],
        table![
            attrs!{At::Class => "perfecttable text"},
            tbody![
                tr![
                    td![
                        b!["No."]
                    ],
                    td![
                        b!["Prime"]
                    ],
                    td![
                        b!["Digits"]
                    ],
                    td![
                        b!["Discovered"]
                    ],
                    td![
                        b!["Download"]
                    ]
                ],
                html
            ]
        ]
    ]
}

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

    #[test]
    fn perfect_test<'a>() {
        let mut perfects:Vec<perfects_utils::Perfect> = perfects_utils::perfects();
		perfects.reverse();
        assert_eq!(perfects[49].digits, 46498850);
        assert_eq!(perfects[39].digits, 12640858);
    }
}<|endoftext|>


<commit_before>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

//! CWT parser tests.
//!
//! These are hung off to the side to avoid cluttering the main cwt.rs.

use ::ring::error::Unspecified;

use crate::cert;
use crate::cert::Algo;
use crate::cert::Cert;
use crate::cert::CertFormat;
use crate::cert::PublicKeyParams;
use crate::crypto;
use crate::crypto::ring;
use crate::crypto::sig;
use crate::crypto::sig::Sign as _;

const UINT: u8 = 0;
const NINT: u8 = 1;
const UTF8: u8 = 3;
const BYTES: u8 = 2;
const ARRAY: u8 = 4;
const MAP: u8 = 5;

/// A `Ciphers` built on top of `ring`.
struct RingCiphers {
    verifier: Option<Box<dyn sig::Verify<Error = Unspecified>>>,
}
impl RingCiphers {
    fn new() -> Self {
        Self { verifier: None }
    }
}
impl cert::Ciphers for RingCiphers {
    type Error = Unspecified;
    fn verifier<'a>(
        &'a mut self,
        algo: Algo,
        key: &PublicKeyParams,
    ) -> Option<&'a mut dyn sig::Verify<Error = Unspecified>> {
        match (key, algo) {
            (
                PublicKeyParams::Rsa { modulus, exponent },
                Algo::RsaPkcs1Sha256,
            ) => {
                use crate::crypto::rsa::Builder as _;

                let key = ring::rsa::PublicKey::new(
                    (*modulus).into(),
                    (*exponent).into(),
                )
                .unwrap();
                let rsa = ring::rsa::Builder::new();
                self.verifier = Some(Box::new(rsa.new_verifier(key).unwrap()));
                match &mut self.verifier {
                    Some(x) => Some(&mut **x),
                    None => None,
                }
            }
        }
    }
}

#[test]
fn self_signed() {
    let protected = raw_cbor!(MAP [
        // Algorithm.
        UINT:1  NINT:256,
        // Criticalities.
        UINT:2  ARRAY[UINT:1, UINT:2],
    ]);
    let payload = raw_cbor!(MAP [
        // Issuer.
        UINT:1  UTF8 {"my cool ca"},
        // Subject; self-signed.
        UINT:2  UTF8 {"my cool ca"},

        // DICE SPKI.
        NINT:4670551    MAP [
            // Key type, RSA.
            UINT:1  UINT:3,
            // Algorithm,i RSA.
            UINT:3  NINT:256,
            // Modulus.
            NINT:0  BYTES {h"
                a6be96e18c5c826a748b887b194ed3f6
                daaa8897cdbaee69e5a980b9dab5af7e
                5f4fed0236d78400837b28229fd4ee24
                1c0d6f3178651265a2762aa1f449dcec
                2a8b679058f9c97a3ba93a7c2bb6d62d
                3baaa70a47991ee6f0daefc55892035f
                a4f4d1c493a73bf61018c6cf86bb43aa
                151f9f0bfa5f474a7fa2bba6a1c77e17
                e160cee2387586e43ff78efd95327d22
                0fc12b0e717089b18cf667ce672d6351
                d187d9330b8777cbe2fb28fccba0b2ad
                c560c9d03e2cba50d624736545c587cc
                0f1d3ef3d7ec69f4108fff270455155b
                d9a3d72c41a397915bad3c9aba753e7b
                f91614d25642cd968816108647e2f7a8
                580e4c05ff0a7f233beba82ca8569a53
            "},
            // Exponent.
            NINT:1  BYTES {h"010001"},
        ],
        // DICE KeyUsage
        NINT:4670552    BYTES {[0b0010_0000]},
    ]);

    let tbs = raw_cbor!(ARRAY [
        UTF8 {"Signature1"},
        BYTES {protected},
        BYTES {},
        BYTES {payload},
    ]);

    let (_, mut signer) = crypto::testdata::rsa();
    let mut sig = vec![0; signer.sig_bytes()];
    signer.sign(&[&tbs], &mut sig).unwrap();

    let data = raw_cbor! {
        BYTES {protected}
        MAP []
        BYTES {payload}
        BYTES {sig}
    };

    let cert = Cert::parse(
        &data,
        CertFormat::OpenDiceCwt,
        None,
        &mut RingCiphers::new(),
    )
    .unwrap();

    assert_eq!(cert.subject(), cert.issuer());
    assert!(cert.supports_cert_signing());
}
<commit_msg>[cert] Remove redundant RingCiphers definition<commit_after>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

//! CWT parser tests.
//!
//! These are hung off to the side to avoid cluttering the main cwt.rs.

use crate::cert::Cert;
use crate::cert::CertFormat;
use crate::cert::RingCiphers;
use crate::crypto;
use crate::crypto::sig::Sign as _;

const UINT: u8 = 0;
const NINT: u8 = 1;
const UTF8: u8 = 3;
const BYTES: u8 = 2;
const ARRAY: u8 = 4;
const MAP: u8 = 5;

#[test]
fn self_signed() {
    let protected = raw_cbor!(MAP [
        // Algorithm.
        UINT:1  NINT:256,
        // Criticalities.
        UINT:2  ARRAY[UINT:1, UINT:2],
    ]);
    let payload = raw_cbor!(MAP [
        // Issuer.
        UINT:1  UTF8 {"my cool ca"},
        // Subject; self-signed.
        UINT:2  UTF8 {"my cool ca"},

        // DICE SPKI.
        NINT:4670551    MAP [
            // Key type, RSA.
            UINT:1  UINT:3,
            // Algorithm,i RSA.
            UINT:3  NINT:256,
            // Modulus.
            NINT:0  BYTES {h"
                a6be96e18c5c826a748b887b194ed3f6
                daaa8897cdbaee69e5a980b9dab5af7e
                5f4fed0236d78400837b28229fd4ee24
                1c0d6f3178651265a2762aa1f449dcec
                2a8b679058f9c97a3ba93a7c2bb6d62d
                3baaa70a47991ee6f0daefc55892035f
                a4f4d1c493a73bf61018c6cf86bb43aa
                151f9f0bfa5f474a7fa2bba6a1c77e17
                e160cee2387586e43ff78efd95327d22
                0fc12b0e717089b18cf667ce672d6351
                d187d9330b8777cbe2fb28fccba0b2ad
                c560c9d03e2cba50d624736545c587cc
                0f1d3ef3d7ec69f4108fff270455155b
                d9a3d72c41a397915bad3c9aba753e7b
                f91614d25642cd968816108647e2f7a8
                580e4c05ff0a7f233beba82ca8569a53
            "},
            // Exponent.
            NINT:1  BYTES {h"010001"},
        ],
        // DICE KeyUsage
        NINT:4670552    BYTES {[0b0010_0000]},
    ]);

    let tbs = raw_cbor!(ARRAY [
        UTF8 {"Signature1"},
        BYTES {protected},
        BYTES {},
        BYTES {payload},
    ]);

    let (_, mut signer) = crypto::testdata::rsa();
    let mut sig = vec![0; signer.sig_bytes()];
    signer.sign(&[&tbs], &mut sig).unwrap();

    let data = raw_cbor! {
        BYTES {protected}
        MAP []
        BYTES {payload}
        BYTES {sig}
    };

    let cert = Cert::parse(
        &data,
        CertFormat::OpenDiceCwt,
        None,
        &mut RingCiphers::new(),
    )
    .unwrap();

    assert_eq!(cert.subject(), cert.issuer());
    assert!(cert.supports_cert_signing());
}
<|endoftext|>


<commit_before>// Copyright 2014 The Rustastic SMTP Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Tools for reading/writing from SMTP clients to SMTP servers and vice-versa.

use std::io::{BufRead, Read, Write, ErrorKind};
use std::io::Result as IoResult;
use std::io::Error as IoError;
use std::vec::Vec;
#[cfg(test)]
use std::fs::File;
use std::ops::{RangeFrom, IndexMut};
#[cfg(test)]
use std::error::Error;
#[cfg(test)]
use std::fs::OpenOptions;
#[cfg(test)]
use super::{MIN_ALLOWED_LINE_SIZE};
#[cfg(test)]
use std::iter::{FromIterator, repeat};

pub static LINE_TOO_LONG: &'static str = "line too long";
pub static DATA_TOO_LONG: &'static str = "message too long";

#[test]
fn test_static_vars() {
    // Already tested in the limits test further down.
}

/// A stream that reads lines of input.
///
/// # Example
/// ```no_run
/// use std::io::TcpStream;
/// use rsmtp::common::stream::InputStream;
/// use rsmtp::common::{
///     MIN_ALLOWED_LINE_SIZE,
/// };
///
/// let mut smtp = InputStream::new(
///     TcpStream::connect("127.0.0.1:25").unwrap(),
///     MIN_ALLOWED_LINE_SIZE,
///     false
/// );
///
/// println!("{}", smtp.read_line().unwrap());
/// ```
pub struct InputStream<S> {
    /// Underlying stream
    stream: S,
    /// Must be at least 1001 per RFC 5321, 1000 chars + 1 for transparency
    /// mechanism.
    max_line_size: usize,
    /// Buffer to make reading more efficient and allow pipelining
    buf: Vec<u8>,
    /// If `true`, will print debug messages of input and output to the console.
    debug: bool,
    /// The position of the `<CRLF>` found at the previous `read_line`.
    last_crlf: Option<usize>
}

// The state of the `<CRLF>` search inside a buffer. See below.
enum CRLFState {
    // We are looking for `<CR>`.
    Cr,
    // We are looking for `<LF>`.
    Lf
}

// Find the position of the first `<CRLF>` in a buffer.
fn position_crlf(buf: &[u8]) -> Option<usize> {
    let mut state = CRLFState::Cr;
    let mut index = 0;

    for byte in buf.iter() {
        match state {
            CRLFState::Cr => {
                if byte == &13 {
                    state = CRLFState::Lf;
                }
            },
            CRLFState::Lf => {
                if byte == &10 {
                    // Subtract 1 to account for the \r, seen previously.
                    return Some(index - 1);
                }
            },
        }
        index += 1;
    }

    None
}

impl<S: Read> InputStream<S> {
    /// Create a new `InputStream` from another stream.
    pub fn new(inner: S, max_line_size: usize, debug: bool) -> InputStream<S> {
        InputStream {
            stream: inner,
            max_line_size: max_line_size,
            // TODO: make line reading work even with a buffer smaller than the maximum line size.
            // Currently, this will not work because we only fill the buffer once per line, assuming
            // that the buffer is large enough.
            buf: Vec::with_capacity(max_line_size),
            debug: debug,
            last_crlf: None
        }
    }

    /// Remove the previous line from the buffer when reading a new line.
    pub fn move_buf(&mut self) {
        // Remove the last line, since we've used it already by now.
        match self.last_crlf {
            Some(p) => {
                // TODO: This could probably be optimised by shifting bytes instead
                // of re-allocating.
                self.buf = self.buf[p + 2 ..].to_vec();
                self.buf.reserve(self.max_line_size);
            },
            _ => {}
        }

        self.last_crlf = None;
    }

    /// Fill the buffer to its limit.
    fn fill_buf(&mut self) -> IoResult<usize> {
        let len = self.buf.len();
        let cap = self.buf.capacity();

        unsafe { self.buf.set_len(cap) };

        // Read as much data as the buffer can hold without re-allocation.
        match self.stream.read(self.buf.index_mut(RangeFrom {
            start: len
        })) {
            Ok(num_bytes) => {
                unsafe { self.buf.set_len(len + num_bytes) };
                Ok(num_bytes)
            },
            Err(err) => Err(err)
        }
    }

    /// Read an SMTP command. Ends with `<CRLF>`.
    pub fn read_line(&mut self) -> IoResult<&[u8]> {
        // Remove the previous line from the buffer before reading a new one.
        self.move_buf();

        let read_line = match position_crlf(self.buf.as_ref()) {
            // First, let's check if the buffer already contains a line. This
            // reduces the number of syscalls.
            Some(last_crlf) => {
                let s = &self.buf[.. last_crlf];
                self.last_crlf = Some(last_crlf);
                Ok(s)
            },
            // If we don't have a line in the buffer, we'll read more input
            // and try again.
            None => {
                match self.fill_buf() {
                    Ok(num_bytes_read) => {
                        match position_crlf(self.buf.as_ref()) {
                            Some(last_crlf) => {
                                let s = &self.buf[.. last_crlf];
                                self.last_crlf = Some(last_crlf);
                                Ok(s)
                            },
                            None => {
                                // If we didn't find a line, it means we had
                                // no `<CRLF>` in the buffer, which means that
                                // the line is too long.
                                Err(IoError::new(ErrorKind::InvalidInput, LINE_TOO_LONG))
                            }
                        }
                    },
                    Err(err) => {
                        Err(err)
                    }
                }
            }
        };

        // If we read a line, we'll say so in the console, if debug mode is on.
        match read_line {
            Ok(bytes) => {
                if self.debug {
                    println!("rsmtp: imsg: {}", String::from_utf8_lossy(bytes.as_ref()));
                }
            },
            _ => {}
        }

        read_line
    }
}

/// A stream that writes lines of output.
pub struct OutputStream<S> {
    /// Underlying stream
    stream: S,
    /// If `true`, will print debug messages of input and output to the console.
    debug: bool,
}

impl<S: Write> OutputStream<S> {
    /// Create a new `InputStream` from another stream.
    pub fn new(inner: S, debug: bool) -> OutputStream<S> {
        OutputStream {
            stream: inner,
            debug: debug,
        }
    }

    /// Write a line ended with `<CRLF>`.
    pub fn write_line(&mut self, s: &str) -> IoResult<()> {
        if self.debug {
            println!("rsmtp: omsg: {}", s);
        }
        // We use `format!()` instead of 2 calls to `write_str()` to reduce
        // the amount of syscalls and to send the string as a single packet.
        // I'm not sure if this is the right way to go though. If you think
        // this is wrong, please open a issue on Github.
        write!(&mut self.stream, "{}\r\n", s)
    }
}

#[test]
fn test_new() {
    // This method is already tested via `test_read_line()`.
}

#[test]
fn test_write_line() {
    // Use a block so the file gets closed at the end of it.
    {
        let mut file_write: File;
        let mut stream: OutputStream<File>;

        file_write = OpenOptions::new()
            .truncate(true).write(true)
            .open("tests/stream/write_line")
            .unwrap();
        stream = OutputStream::new(file_write, false);
        stream.write_line("HelloWorld").unwrap();
        stream.write_line("ByeBye").unwrap();
    }
    let mut file_read: File;
    let mut expected = String::new();

    file_read = OpenOptions::new()
        .read(true)
        .open("tests/stream/write_line")
        .unwrap();
    file_read.read_to_string(&mut expected).unwrap();
    assert_eq!("HelloWorld\r\nByeBye\r\n", expected.as_str());
}

#[test]
fn test_limits() {
    let mut file: File;
    let mut stream: InputStream<File>;

    file = OpenOptions::new().read(true).open("tests/stream/1line1").unwrap();
    stream = InputStream::new(file, 3, false);
    match stream.read_line() {
        Ok(_) => panic!(),
        Err(err) => {
            assert_eq!("line too long", err.description());
            assert_eq!(ErrorKind::InvalidInput, err.kind());
        }
    }
}

#[test]
fn test_read_line() {
    let mut file: File;
    let mut stream: InputStream<File>;
    let mut expected: String;

    file = OpenOptions::new().read(true).open("tests/stream/0line1").unwrap();
    stream = InputStream::new(file, MIN_ALLOWED_LINE_SIZE, false);
    assert!(!stream.read_line().is_ok());

    file = OpenOptions::new().read(true).open("tests/stream/0line2").unwrap();
    stream = InputStream::new(file, MIN_ALLOWED_LINE_SIZE, false);
    assert!(!stream.read_line().is_ok());

    file = OpenOptions::new().read(true).open("tests/stream/0line3").unwrap();
    stream = InputStream::new(file, MIN_ALLOWED_LINE_SIZE, false);
    assert!(!stream.read_line().is_ok());

    file = OpenOptions::new().read(true).open("tests/stream/1line1").unwrap();
    stream = InputStream::new(file, MIN_ALLOWED_LINE_SIZE, false);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned().as_ref(), "hello world!");
    assert!(!stream.read_line().is_ok());

    file = OpenOptions::new().read(true).open("tests/stream/1line2").unwrap();
    stream = InputStream::new(file, MIN_ALLOWED_LINE_SIZE, false);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned().as_ref(), "hello world!");
    assert!(!stream.read_line().is_ok());

    file = OpenOptions::new().read(true).open("tests/stream/2lines1").unwrap();
    stream = InputStream::new(file, MIN_ALLOWED_LINE_SIZE, false);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned().as_ref(), "hello world!");
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned().as_ref(), "bye bye world!");
    assert!(!stream.read_line().is_ok());

    expected =  String::from_iter(repeat('x').take(62));
    file = OpenOptions::new().read(true).open("tests/stream/xlines1").unwrap();
    stream = InputStream::new(file, MIN_ALLOWED_LINE_SIZE, false);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned(), expected);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned(), expected);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned(), expected);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned(), expected);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned(), expected);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned(), expected);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned(), expected);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned(), expected);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned(), expected);
    assert!(!stream.read_line().is_ok());
}
<commit_msg>reset size of read buffer if there was an error during reading<commit_after>// Copyright 2014 The Rustastic SMTP Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Tools for reading/writing from SMTP clients to SMTP servers and vice-versa.

use std::io::{BufRead, Read, Write, ErrorKind};
use std::io::Result as IoResult;
use std::io::Error as IoError;
use std::vec::Vec;
#[cfg(test)]
use std::fs::File;
use std::ops::{RangeFrom, IndexMut};
#[cfg(test)]
use std::error::Error;
#[cfg(test)]
use std::fs::OpenOptions;
#[cfg(test)]
use super::{MIN_ALLOWED_LINE_SIZE};
#[cfg(test)]
use std::iter::{FromIterator, repeat};

pub static LINE_TOO_LONG: &'static str = "line too long";
pub static DATA_TOO_LONG: &'static str = "message too long";

#[test]
fn test_static_vars() {
    // Already tested in the limits test further down.
}

/// A stream that reads lines of input.
///
/// # Example
/// ```no_run
/// use std::io::TcpStream;
/// use rsmtp::common::stream::InputStream;
/// use rsmtp::common::{
///     MIN_ALLOWED_LINE_SIZE,
/// };
///
/// let mut smtp = InputStream::new(
///     TcpStream::connect("127.0.0.1:25").unwrap(),
///     MIN_ALLOWED_LINE_SIZE,
///     false
/// );
///
/// println!("{}", smtp.read_line().unwrap());
/// ```
pub struct InputStream<S> {
    /// Underlying stream
    stream: S,
    /// Must be at least 1001 per RFC 5321, 1000 chars + 1 for transparency
    /// mechanism.
    max_line_size: usize,
    /// Buffer to make reading more efficient and allow pipelining
    buf: Vec<u8>,
    /// If `true`, will print debug messages of input and output to the console.
    debug: bool,
    /// The position of the `<CRLF>` found at the previous `read_line`.
    last_crlf: Option<usize>
}

// The state of the `<CRLF>` search inside a buffer. See below.
enum CRLFState {
    // We are looking for `<CR>`.
    Cr,
    // We are looking for `<LF>`.
    Lf
}

// Find the position of the first `<CRLF>` in a buffer.
fn position_crlf(buf: &[u8]) -> Option<usize> {
    let mut state = CRLFState::Cr;
    let mut index = 0;

    for byte in buf.iter() {
        match state {
            CRLFState::Cr => {
                if byte == &13 {
                    state = CRLFState::Lf;
                }
            },
            CRLFState::Lf => {
                if byte == &10 {
                    // Subtract 1 to account for the \r, seen previously.
                    return Some(index - 1);
                }
            },
        }
        index += 1;
    }

    None
}

impl<S: Read> InputStream<S> {
    /// Create a new `InputStream` from another stream.
    pub fn new(inner: S, max_line_size: usize, debug: bool) -> InputStream<S> {
        InputStream {
            stream: inner,
            max_line_size: max_line_size,
            // TODO: make line reading work even with a buffer smaller than the maximum line size.
            // Currently, this will not work because we only fill the buffer once per line, assuming
            // that the buffer is large enough.
            buf: Vec::with_capacity(max_line_size),
            debug: debug,
            last_crlf: None
        }
    }

    /// Remove the previous line from the buffer when reading a new line.
    pub fn move_buf(&mut self) {
        // Remove the last line, since we've used it already by now.
        match self.last_crlf {
            Some(p) => {
                // TODO: This could probably be optimised by shifting bytes instead
                // of re-allocating.
                self.buf = self.buf[p + 2 ..].to_vec();
                self.buf.reserve(self.max_line_size);
            },
            _ => {}
        }

        self.last_crlf = None;
    }

    /// Fill the buffer to its limit.
    fn fill_buf(&mut self) -> IoResult<usize> {
        let len = self.buf.len();
        let cap = self.buf.capacity();

        unsafe { self.buf.set_len(cap) };

        // Read as much data as the buffer can hold without re-allocation.
        match self.stream.read(self.buf.index_mut(RangeFrom {
            start: len
        })) {
            Ok(num_bytes) => {
                unsafe { self.buf.set_len(len + num_bytes) };
                Ok(num_bytes)
            },
            Err(err) => {
                unsafe { self.buf.set_len(len) };
                Err(err)
            }
        }
    }

    /// Read an SMTP command. Ends with `<CRLF>`.
    pub fn read_line(&mut self) -> IoResult<&[u8]> {
        // Remove the previous line from the buffer before reading a new one.
        self.move_buf();

        let read_line = match position_crlf(self.buf.as_ref()) {
            // First, let's check if the buffer already contains a line. This
            // reduces the number of syscalls.
            Some(last_crlf) => {
                let s = &self.buf[.. last_crlf];
                self.last_crlf = Some(last_crlf);
                Ok(s)
            },
            // If we don't have a line in the buffer, we'll read more input
            // and try again.
            None => {
                match self.fill_buf() {
                    Ok(num_bytes_read) => {
                        match position_crlf(self.buf.as_ref()) {
                            Some(last_crlf) => {
                                let s = &self.buf[.. last_crlf];
                                self.last_crlf = Some(last_crlf);
                                Ok(s)
                            },
                            None => {
                                // If we didn't find a line, it means we had
                                // no `<CRLF>` in the buffer, which means that
                                // the line is too long.
                                Err(IoError::new(ErrorKind::InvalidInput, LINE_TOO_LONG))
                            }
                        }
                    },
                    Err(err) => {
                        Err(err)
                    }
                }
            }
        };

        // If we read a line, we'll say so in the console, if debug mode is on.
        match read_line {
            Ok(bytes) => {
                if self.debug {
                    println!("rsmtp: imsg: {}", String::from_utf8_lossy(bytes.as_ref()));
                }
            },
            _ => {}
        }

        read_line
    }
}

/// A stream that writes lines of output.
pub struct OutputStream<S> {
    /// Underlying stream
    stream: S,
    /// If `true`, will print debug messages of input and output to the console.
    debug: bool,
}

impl<S: Write> OutputStream<S> {
    /// Create a new `InputStream` from another stream.
    pub fn new(inner: S, debug: bool) -> OutputStream<S> {
        OutputStream {
            stream: inner,
            debug: debug,
        }
    }

    /// Write a line ended with `<CRLF>`.
    pub fn write_line(&mut self, s: &str) -> IoResult<()> {
        if self.debug {
            println!("rsmtp: omsg: {}", s);
        }
        // We use `format!()` instead of 2 calls to `write_str()` to reduce
        // the amount of syscalls and to send the string as a single packet.
        // I'm not sure if this is the right way to go though. If you think
        // this is wrong, please open a issue on Github.
        write!(&mut self.stream, "{}\r\n", s)
    }
}

#[test]
fn test_new() {
    // This method is already tested via `test_read_line()`.
}

#[test]
fn test_write_line() {
    // Use a block so the file gets closed at the end of it.
    {
        let mut file_write: File;
        let mut stream: OutputStream<File>;

        file_write = OpenOptions::new()
            .truncate(true).write(true)
            .open("tests/stream/write_line")
            .unwrap();
        stream = OutputStream::new(file_write, false);
        stream.write_line("HelloWorld").unwrap();
        stream.write_line("ByeBye").unwrap();
    }
    let mut file_read: File;
    let mut expected = String::new();

    file_read = OpenOptions::new()
        .read(true)
        .open("tests/stream/write_line")
        .unwrap();
    file_read.read_to_string(&mut expected).unwrap();
    assert_eq!("HelloWorld\r\nByeBye\r\n", expected.as_str());
}

#[test]
fn test_limits() {
    let mut file: File;
    let mut stream: InputStream<File>;

    file = OpenOptions::new().read(true).open("tests/stream/1line1").unwrap();
    stream = InputStream::new(file, 3, false);
    match stream.read_line() {
        Ok(_) => panic!(),
        Err(err) => {
            assert_eq!("line too long", err.description());
            assert_eq!(ErrorKind::InvalidInput, err.kind());
        }
    }
}

#[test]
fn test_read_line() {
    let mut file: File;
    let mut stream: InputStream<File>;
    let mut expected: String;

    file = OpenOptions::new().read(true).open("tests/stream/0line1").unwrap();
    stream = InputStream::new(file, MIN_ALLOWED_LINE_SIZE, false);
    assert!(!stream.read_line().is_ok());

    file = OpenOptions::new().read(true).open("tests/stream/0line2").unwrap();
    stream = InputStream::new(file, MIN_ALLOWED_LINE_SIZE, false);
    assert!(!stream.read_line().is_ok());

    file = OpenOptions::new().read(true).open("tests/stream/0line3").unwrap();
    stream = InputStream::new(file, MIN_ALLOWED_LINE_SIZE, false);
    assert!(!stream.read_line().is_ok());

    file = OpenOptions::new().read(true).open("tests/stream/1line1").unwrap();
    stream = InputStream::new(file, MIN_ALLOWED_LINE_SIZE, false);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned().as_ref(), "hello world!");
    assert!(!stream.read_line().is_ok());

    file = OpenOptions::new().read(true).open("tests/stream/1line2").unwrap();
    stream = InputStream::new(file, MIN_ALLOWED_LINE_SIZE, false);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned().as_ref(), "hello world!");
    assert!(!stream.read_line().is_ok());

    file = OpenOptions::new().read(true).open("tests/stream/2lines1").unwrap();
    stream = InputStream::new(file, MIN_ALLOWED_LINE_SIZE, false);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned().as_ref(), "hello world!");
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned().as_ref(), "bye bye world!");
    assert!(!stream.read_line().is_ok());

    expected =  String::from_iter(repeat('x').take(62));
    file = OpenOptions::new().read(true).open("tests/stream/xlines1").unwrap();
    stream = InputStream::new(file, MIN_ALLOWED_LINE_SIZE, false);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned(), expected);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned(), expected);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned(), expected);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned(), expected);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned(), expected);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned(), expected);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned(), expected);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned(), expected);
    assert_eq!(String::from_utf8_lossy(stream.read_line().unwrap().as_ref()).to_owned(), expected);
    assert!(!stream.read_line().is_ok());
}
<|endoftext|>


<commit_before>use crate::stream_ext::Fuse;
use crate::Stream;
use tokio::time::{sleep, Instant, Sleep};

use core::future::Future;
use core::pin::Pin;
use core::task::{Context, Poll};
use pin_project_lite::pin_project;
use std::time::Duration;

pin_project! {
    /// Stream returned by the [`chunks_timeout`](super::StreamExt::chunks_timeout) method.
    #[must_use = "streams do nothing unless polled"]
    #[derive(Debug)]
    pub struct ChunksTimeout<S: Stream> {
        #[pin]
        stream: Fuse<S>,
        #[pin]
        deadline: Sleep,
        duration: Duration,
        items: Vec<S::Item>,
        cap: usize, // https://github.com/rust-lang/futures-rs/issues/1475
    }
}

impl<S: Stream> ChunksTimeout<S> {
    pub(super) fn new(stream: S, max_size: usize, duration: Duration) -> Self {
        ChunksTimeout {
            stream: Fuse::new(stream),
            deadline: sleep(duration),
            duration,
            items: Vec::with_capacity(max_size),
            cap: max_size,
        }
    }
}

impl<S: Stream> Stream for ChunksTimeout<S> {
    type Item = Vec<S::Item>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let mut me = self.as_mut().project();
        loop {
            match me.stream.as_mut().poll_next(cx) {
                Poll::Pending => break,
                Poll::Ready(Some(item)) => {
                    if me.items.is_empty() {
                        me.deadline.as_mut().reset(Instant::now() + *me.duration);
                        me.items.reserve_exact(*me.cap);
                    }
                    me.items.push(item);
                    if me.items.len() >= *me.cap {
                        return Poll::Ready(Some(std::mem::take(me.items)));
                    }
                }
                Poll::Ready(None) => {
                    // Returning Some here is only correct because we fuse the inner stream.
                    let last = if me.items.is_empty() {
                        None
                    } else {
                        Some(std::mem::take(me.items))
                    };

                    return Poll::Ready(last);
                }
            }
        }

        if !me.items.is_empty() {
            ready!(me.deadline.poll(cx));
            return Poll::Ready(Some(std::mem::take(me.items)));
        }

        Poll::Pending
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let chunk_len = if self.items.is_empty() { 0 } else { 1 };
        let (lower, upper) = self.stream.size_hint();
        let lower = (lower / self.cap).saturating_add(chunk_len);
        let upper = upper.and_then(|x| x.checked_add(chunk_len));
        (lower, upper)
    }
}
<commit_msg>stream: fix panic in `ChunksTimeout::new` (#5036)<commit_after>use crate::stream_ext::Fuse;
use crate::Stream;
use tokio::time::{sleep, Sleep};

use core::future::Future;
use core::pin::Pin;
use core::task::{Context, Poll};
use pin_project_lite::pin_project;
use std::time::Duration;

pin_project! {
    /// Stream returned by the [`chunks_timeout`](super::StreamExt::chunks_timeout) method.
    #[must_use = "streams do nothing unless polled"]
    #[derive(Debug)]
    pub struct ChunksTimeout<S: Stream> {
        #[pin]
        stream: Fuse<S>,
        #[pin]
        deadline: Option<Sleep>,
        duration: Duration,
        items: Vec<S::Item>,
        cap: usize, // https://github.com/rust-lang/futures-rs/issues/1475
    }
}

impl<S: Stream> ChunksTimeout<S> {
    pub(super) fn new(stream: S, max_size: usize, duration: Duration) -> Self {
        ChunksTimeout {
            stream: Fuse::new(stream),
            deadline: None,
            duration,
            items: Vec::with_capacity(max_size),
            cap: max_size,
        }
    }
}

impl<S: Stream> Stream for ChunksTimeout<S> {
    type Item = Vec<S::Item>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let mut me = self.as_mut().project();
        loop {
            match me.stream.as_mut().poll_next(cx) {
                Poll::Pending => break,
                Poll::Ready(Some(item)) => {
                    if me.items.is_empty() {
                        me.deadline.set(Some(sleep(*me.duration)));
                        me.items.reserve_exact(*me.cap);
                    }
                    me.items.push(item);
                    if me.items.len() >= *me.cap {
                        return Poll::Ready(Some(std::mem::take(me.items)));
                    }
                }
                Poll::Ready(None) => {
                    // Returning Some here is only correct because we fuse the inner stream.
                    let last = if me.items.is_empty() {
                        None
                    } else {
                        Some(std::mem::take(me.items))
                    };

                    return Poll::Ready(last);
                }
            }
        }

        if !me.items.is_empty() {
            if let Some(deadline) = me.deadline.as_pin_mut() {
                ready!(deadline.poll(cx));
            }
            return Poll::Ready(Some(std::mem::take(me.items)));
        }

        Poll::Pending
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let chunk_len = if self.items.is_empty() { 0 } else { 1 };
        let (lower, upper) = self.stream.size_hint();
        let lower = (lower / self.cap).saturating_add(chunk_len);
        let upper = upper.and_then(|x| x.checked_add(chunk_len));
        (lower, upper)
    }
}
<|endoftext|>


<commit_before>#[allow(unused_imports)]
use log::Event::*;
use std::error::Error;
use std::fmt;
use registry;

/// Custom error type for the rayon thread pool configuration.
#[derive(Debug,PartialEq)]
pub enum InitError {
    /// Error if number of threads is set to zero.
    NumberOfThreadsZero,

    /// Error if the gloal thread pool is initialized multiple times
    /// and the configuration is not equal for all configurations.
    GlobalPoolAlreadyInitialized,
}

impl fmt::Display for InitError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            InitError::NumberOfThreadsZero => {
                write!(f,
                       "The number of threads was set to zero but must be greater than zero.")
            }
            InitError::GlobalPoolAlreadyInitialized => {
                write!(f,
                       "The gobal thread pool has already been initialized with a different \
                        configuration. Only one valid configuration is allowed.")
            }
        }
    }
}

impl Error for InitError {
    fn description(&self) -> &str {
        match *self {
            InitError::NumberOfThreadsZero => "number of threads set to zero",
            InitError::GlobalPoolAlreadyInitialized => {
                "global thread pool has already been initialized"
            }
        }
    }
}

/// Contains the rayon thread pool configuration.
pub struct Configuration {
    /// The number of threads in the rayon thread pool. Must not be zero.
    num_threads: Option<usize>,
    get_thread_name: Option<Box<FnMut(usize) -> String>>,
}

impl Configuration {
    /// Creates and return a valid rayon thread pool configuration, but does not initialize it.
    pub fn new() -> Configuration {
        Configuration {
            num_threads: None,
            get_thread_name: None,
        }
    }

    /// Get the number of threads that will be used for the thread
    /// pool. See `set_num_threads` for more information.
    pub fn num_threads(&self) -> Option<usize> {
        self.num_threads
    }

    /// Get the base thread name for the given index.
    pub fn thread_name(&mut self, index: usize) -> Option<String> {
        self.get_thread_name.as_mut().map(|c| c(index))
    }

    /// Set a closure which takes the thread index and return
    /// the threads name.
    pub fn set_thread_name<F>(mut self, closure: F) -> Self
    where F: FnMut(usize) -> String + 'static {
        self.get_thread_name = Some(Box::new(closure));
        self
    }

    /// Set the number of threads to be used in the rayon threadpool.
    /// The argument `num_threads` must not be zero. If you do not
    /// call this function, rayon will select a suitable default
    /// (currently, the default is one thread per logical CPU).
    pub fn set_num_threads(mut self, num_threads: usize) -> Configuration {
        self.num_threads = Some(num_threads);
        self
    }

    /// Checks whether the configuration is valid.
    pub fn validate(&self) -> Result<(), InitError> {
        if let Some(value) = self.num_threads {
            if value == 0 {
                return Err(InitError::NumberOfThreadsZero);
            }
        }

        Ok(())
    }
}

/// Initializes the global thread pool. This initialization is
/// **optional**.  If you do not call this function, the thread pool
/// will be automatically initialized with the default
/// configuration. In fact, calling `initialize` is not recommended,
/// except for in two scenarios:
///
/// - You wish to change the default configuration.
/// - You are running a benchmark, in which case initializing may
///   yield slightly more consistent results, since the worker threads
///   will already be ready to go even in the first iteration.  But
///   this cost is minimal.
///
/// Initialization of the global thread pool happens exactly
/// once. Once started, the configuration cannot be
/// changed. Therefore, if you call `initialize` a second time, it
/// will simply check that the global thread pool already has the
/// configuration you requested, rather than making changes.
///
/// An `Ok` result indicates that the thread pool is running with the
/// given configuration. Otherwise, a suitable error is returned.
pub fn initialize(config: Configuration) -> Result<(), InitError> {
    try!(config.validate());

    let num_threads = config.num_threads;

    let registry = registry::get_registry_with_config(config);

    if let Some(value) = num_threads {
        if value != registry.num_threads() {
            return Err(InitError::GlobalPoolAlreadyInitialized);
        }
    }

    registry.wait_until_primed();

    Ok(())
}

/// This is a debugging API not really intended for end users. It will
/// dump some performance statistics out using `println`.
pub fn dump_stats() {
    dump_stats!();
}

<commit_msg>Update docs for thread naming functions.<commit_after>#[allow(unused_imports)]
use log::Event::*;
use std::error::Error;
use std::fmt;
use registry;

/// Custom error type for the rayon thread pool configuration.
#[derive(Debug,PartialEq)]
pub enum InitError {
    /// Error if number of threads is set to zero.
    NumberOfThreadsZero,

    /// Error if the gloal thread pool is initialized multiple times
    /// and the configuration is not equal for all configurations.
    GlobalPoolAlreadyInitialized,
}

impl fmt::Display for InitError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            InitError::NumberOfThreadsZero => {
                write!(f,
                       "The number of threads was set to zero but must be greater than zero.")
            }
            InitError::GlobalPoolAlreadyInitialized => {
                write!(f,
                       "The gobal thread pool has already been initialized with a different \
                        configuration. Only one valid configuration is allowed.")
            }
        }
    }
}

impl Error for InitError {
    fn description(&self) -> &str {
        match *self {
            InitError::NumberOfThreadsZero => "number of threads set to zero",
            InitError::GlobalPoolAlreadyInitialized => {
                "global thread pool has already been initialized"
            }
        }
    }
}

/// Contains the rayon thread pool configuration.
pub struct Configuration {
    /// The number of threads in the rayon thread pool. Must not be zero.
    num_threads: Option<usize>,
    get_thread_name: Option<Box<FnMut(usize) -> String>>,
}

impl Configuration {
    /// Creates and return a valid rayon thread pool configuration, but does not initialize it.
    pub fn new() -> Configuration {
        Configuration {
            num_threads: None,
            get_thread_name: None,
        }
    }

    /// Get the number of threads that will be used for the thread
    /// pool. See `set_num_threads` for more information.
    pub fn num_threads(&self) -> Option<usize> {
        self.num_threads
    }

    /// Get the thread name for the thread with the given index.
    pub fn thread_name(&mut self, index: usize) -> Option<String> {
        self.get_thread_name.as_mut().map(|c| c(index))
    }

    /// Set a closure which takes a thread index and returns
    /// the thread's name.
    pub fn set_thread_name<F>(mut self, closure: F) -> Self
    where F: FnMut(usize) -> String + 'static {
        self.get_thread_name = Some(Box::new(closure));
        self
    }

    /// Set the number of threads to be used in the rayon threadpool.
    /// The argument `num_threads` must not be zero. If you do not
    /// call this function, rayon will select a suitable default
    /// (currently, the default is one thread per logical CPU).
    pub fn set_num_threads(mut self, num_threads: usize) -> Configuration {
        self.num_threads = Some(num_threads);
        self
    }

    /// Checks whether the configuration is valid.
    pub fn validate(&self) -> Result<(), InitError> {
        if let Some(value) = self.num_threads {
            if value == 0 {
                return Err(InitError::NumberOfThreadsZero);
            }
        }

        Ok(())
    }
}

/// Initializes the global thread pool. This initialization is
/// **optional**.  If you do not call this function, the thread pool
/// will be automatically initialized with the default
/// configuration. In fact, calling `initialize` is not recommended,
/// except for in two scenarios:
///
/// - You wish to change the default configuration.
/// - You are running a benchmark, in which case initializing may
///   yield slightly more consistent results, since the worker threads
///   will already be ready to go even in the first iteration.  But
///   this cost is minimal.
///
/// Initialization of the global thread pool happens exactly
/// once. Once started, the configuration cannot be
/// changed. Therefore, if you call `initialize` a second time, it
/// will simply check that the global thread pool already has the
/// configuration you requested, rather than making changes.
///
/// An `Ok` result indicates that the thread pool is running with the
/// given configuration. Otherwise, a suitable error is returned.
pub fn initialize(config: Configuration) -> Result<(), InitError> {
    try!(config.validate());

    let num_threads = config.num_threads;

    let registry = registry::get_registry_with_config(config);

    if let Some(value) = num_threads {
        if value != registry.num_threads() {
            return Err(InitError::GlobalPoolAlreadyInitialized);
        }
    }

    registry.wait_until_primed();

    Ok(())
}

/// This is a debugging API not really intended for end users. It will
/// dump some performance statistics out using `println`.
pub fn dump_stats() {
    dump_stats!();
}

<|endoftext|>


<commit_before>extern crate sfunc;

use std::rand::Rng;
use std::rand::distributions::{Normal, IndependentSample};

use super::Distribution;

/// A Gaussian distribution.
pub struct Gaussian {
    /// The mean value.
    pub mu: f64,
    /// The standard deviation.
    pub sigma: f64,

    normal: Normal,
}

impl Gaussian {
    /// Creates a Gaussian distribution with the mean value `mu` and standard
    /// deviation `sigma`.
    #[inline]
    pub fn new(mu: f64, sigma: f64) -> Gaussian {
        Gaussian {
            mu: mu,
            sigma: sigma,
            normal: Normal::new(mu, sigma),
        }
    }
}

impl Distribution<f64> for Gaussian {
    fn cdf(&self, x: f64) -> f64 {
        use self::sfunc::erf;
        (1.0 + erf((x - self.mu) / (self.sigma * Float::sqrt2()))) / 2.0
    }

    /// Compute the inverse of the cummulative distribution function at
    /// probability `p`.
    ///
    /// The code is based on a [C implementation][1] by John Burkardt.
    ///
    /// [1]: http://people.sc.fsu.edu/~jburkardt/c_src/asa241/asa241.html
    fn inv_cdf(&self, p: f64) -> f64 {
        #[inline(always)]
        fn ln(x: f64) -> f64 { x.ln() }
        #[inline(always)]
        fn sqrt(x: f64) -> f64 { x.sqrt() }

        const CONST1: f64 = 0.180625;
        const CONST2: f64 = 1.6;
        const SPLIT1: f64 = 0.425;
        const SPLIT2: f64 = 5.0;
        const A: [f64, ..8] = [3.3871328727963666080,
                               1.3314166789178437745e+2,
                               1.9715909503065514427e+3,
                               1.3731693765509461125e+4,
                               4.5921953931549871457e+4,
                               6.7265770927008700853e+4,
                               3.3430575583588128105e+4,
                               2.5090809287301226727e+3];
        const B: [f64, ..8] = [1.0000000000000000000,
                               4.2313330701600911252e+1,
                               6.8718700749205790830e+2,
                               5.3941960214247511077e+3,
                               2.1213794301586595867e+4,
                               3.9307895800092710610e+4,
                               2.8729085735721942674e+4,
                               5.2264952788528545610e+3];
        const C: [f64, ..8] = [1.42343711074968357734,
                               4.63033784615654529590,
                               5.76949722146069140550,
                               3.64784832476320460504,
                               1.27045825245236838258,
                               2.41780725177450611770e-1,
                               2.27238449892691845833e-2,
                               7.74545014278341407640e-4];
        const D: [f64, ..8] = [1.00000000000000000000,
                               2.05319162663775882187,
                               1.67638483018380384940,
                               6.89767334985100004550e-1,
                               1.48103976427480074590e-1,
                               1.51986665636164571966e-2,
                               5.47593808499534494600e-4,
                               1.05075007164441684324e-9];
        const E: [f64, ..8] = [6.65790464350110377720,
                               5.46378491116411436990,
                               1.78482653991729133580,
                               2.96560571828504891230e-1,
                               2.65321895265761230930e-2,
                               1.24266094738807843860e-3,
                               2.71155556874348757815e-5,
                               2.01033439929228813265e-7];
        const F: [f64, ..8] = [1.00000000000000000000,
                               5.99832206555887937690e-1,
                               1.36929880922735805310e-1,
                               1.48753612908506148525e-2,
                               7.86869131145613259100e-4,
                               1.84631831751005468180e-5,
                               1.42151175831644588870e-7,
                               2.04426310338993978564e-15];

        #[inline]
        fn poly(c: &[f64], x: f64) -> f64 {
            c.iter().rev().fold(0.0, |y, &c| y * x + c)
        }

        if p <= 0.0 {
            return Float::neg_infinity();
        }
        if 1.0 <= p {
            return Float::infinity();
        }

        let q = p - 0.5;

        if ::std::num::abs(q) <= SPLIT1 {
            let x = CONST1 - q * q;
            return self.mu + self.sigma * q * poly(&A, x) / poly(&B, x);
        }

        let mut x = if q < 0.0 { p } else { 1.0 - p };

        x = sqrt(-ln(x));

        if x <= SPLIT2 {
            x -= CONST2;
            x = poly(&C, x) / poly(&D, x);
        } else {
            x -= SPLIT2;
            x = poly(&E, x) / poly(&F, x);
        }

        if q < 0.0 {
            x = -x;
        }

        self.mu + self.sigma * x
    }

    #[inline]
    fn sample<R: Rng>(&self, rng: &mut R) -> f64 {
        self.normal.ind_sample(rng)
    }
}

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

    macro_rules! assert_almost_eq(
        ($x:expr, $y:expr) => ({
            let e: f64 = ::std::f64::EPSILON.sqrt();
            let x: Vec<f64> = $x;
            let y: Vec<f64> = $y;
            for i in range(0u, x.len()) {
                if x[i].is_finite() {
                    assert!(::std::num::abs(x[i] - y[i]) < e,
                            "expected {:e} ~ {:e}", x[i], y[i]);
                } else {
                    assert_eq!(x[i], y[i]);
                }
            }
        });
    )

    #[test]
    fn cdf() {
        let dist = Gaussian::new(1.0, 2.0);

        let x = vec![-4.0, -3.5, -3.0, -2.5, -2.0, -1.5, -1.0, -0.5, 0.0, 0.5,
                     1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0];
        let p = vec![6.209665325776139e-03, 1.222447265504470e-02,
                     2.275013194817922e-02, 4.005915686381709e-02,
                     6.680720126885809e-02, 1.056497736668553e-01,
                     1.586552539314571e-01, 2.266273523768682e-01,
                     3.085375387259869e-01, 4.012936743170763e-01,
                     5.000000000000000e-01, 5.987063256829237e-01,
                     6.914624612740131e-01, 7.733726476231317e-01,
                     8.413447460685429e-01, 8.943502263331446e-01,
                     9.331927987311419e-01];

        assert_almost_eq!(x.iter().map(|&x| dist.cdf(x)).collect(), p);
    }

    #[test]
    fn inv_cdf() {
        let dist = Gaussian::new(-1.0, 0.25);

        let p = vec![0.00, 0.05, 0.10, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40,
                     0.45, 0.50, 0.55, 0.60, 0.65, 0.70, 0.75, 0.80, 0.85,
                     0.90, 0.95, 1.00];
        let x = vec![Float::neg_infinity(), -1.411213406737868e+00,
                     -1.320387891386150e+00, -1.259108347373447e+00,
                     -1.210405308393228e+00, -1.168622437549020e+00,
                     -1.131100128177010e+00, -1.096330116601892e+00,
                     -1.063336775783950e+00, -1.031415336713768e+00,
                     -1.000000000000000e+00, -9.685846632862315e-01,
                     -9.366632242160501e-01, -9.036698833981082e-01,
                     -8.688998718229899e-01, -8.313775624509796e-01,
                     -7.895946916067714e-01, -7.408916526265525e-01,
                     -6.796121086138498e-01, -5.887865932621319e-01,
                     Float::infinity()];

        assert_almost_eq!(p.iter().map(|&p| dist.inv_cdf(p)).collect(), x);
    }
}

#[cfg(test)]
mod bench {
    extern crate test;

    use super::super::{Distribution, Sampler, Uniform};
    use super::Gaussian;

    #[bench]
    fn cdf(bench: &mut test::Bencher) {
        let mut rng = ::std::rand::task_rng();
        let dist = Gaussian::new(0.0, 1.0);
        let x = Sampler(&dist, &mut rng).take(1000).collect::<Vec<_>>();

        bench.iter(|| {
            test::black_box(x.iter().map(|&x| dist.cdf(x)).collect::<Vec<_>>())
        })
    }

    #[bench]
    fn inv_cdf(bench: &mut test::Bencher) {
        let mut rng = ::std::rand::task_rng();
        let dist = Gaussian::new(0.0, 1.0);
        let p = Sampler(&Uniform::new(0.0, 1.0), &mut rng).take(1000).collect::<Vec<_>>();

        bench.iter(|| {
            test::black_box(p.iter().map(|&p| dist.inv_cdf(p)).collect::<Vec<_>>())
        })
    }
}
<commit_msg>Improved Gaussian#inv_cdf<commit_after>extern crate sfunc;

use std::rand::Rng;
use std::rand::distributions::{Normal, IndependentSample};

use super::Distribution;

/// A Gaussian distribution.
pub struct Gaussian {
    /// The mean value.
    pub mu: f64,
    /// The standard deviation.
    pub sigma: f64,

    normal: Normal,
}

impl Gaussian {
    /// Creates a Gaussian distribution with the mean value `mu` and standard
    /// deviation `sigma`.
    #[inline]
    pub fn new(mu: f64, sigma: f64) -> Gaussian {
        Gaussian {
            mu: mu,
            sigma: sigma,
            normal: Normal::new(mu, sigma),
        }
    }
}

impl Distribution<f64> for Gaussian {
    #[inline]
    fn cdf(&self, x: f64) -> f64 {
        use self::sfunc::erf;
        (1.0 + erf((x - self.mu) / (self.sigma * Float::sqrt2()))) / 2.0
    }

    /// Compute the inverse of the cummulative distribution function at
    /// probability `p`.
    ///
    /// The code is based on a [C implementation][1] by John Burkardt.
    ///
    /// [1]: http://people.sc.fsu.edu/~jburkardt/c_src/asa241/asa241.html
    fn inv_cdf(&self, p: f64) -> f64 {
        #[inline(always)]
        fn ln(x: f64) -> f64 { x.ln() }
        #[inline(always)]
        fn sqrt(x: f64) -> f64 { x.sqrt() }

        const CONST1: f64 = 0.180625;
        const CONST2: f64 = 1.6;
        const SPLIT1: f64 = 0.425;
        const SPLIT2: f64 = 5.0;
        const A: [f64, ..8] = [3.3871328727963666080,
                               1.3314166789178437745e+2,
                               1.9715909503065514427e+3,
                               1.3731693765509461125e+4,
                               4.5921953931549871457e+4,
                               6.7265770927008700853e+4,
                               3.3430575583588128105e+4,
                               2.5090809287301226727e+3];
        const B: [f64, ..8] = [1.0000000000000000000,
                               4.2313330701600911252e+1,
                               6.8718700749205790830e+2,
                               5.3941960214247511077e+3,
                               2.1213794301586595867e+4,
                               3.9307895800092710610e+4,
                               2.8729085735721942674e+4,
                               5.2264952788528545610e+3];
        const C: [f64, ..8] = [1.42343711074968357734,
                               4.63033784615654529590,
                               5.76949722146069140550,
                               3.64784832476320460504,
                               1.27045825245236838258,
                               2.41780725177450611770e-1,
                               2.27238449892691845833e-2,
                               7.74545014278341407640e-4];
        const D: [f64, ..8] = [1.00000000000000000000,
                               2.05319162663775882187,
                               1.67638483018380384940,
                               6.89767334985100004550e-1,
                               1.48103976427480074590e-1,
                               1.51986665636164571966e-2,
                               5.47593808499534494600e-4,
                               1.05075007164441684324e-9];
        const E: [f64, ..8] = [6.65790464350110377720,
                               5.46378491116411436990,
                               1.78482653991729133580,
                               2.96560571828504891230e-1,
                               2.65321895265761230930e-2,
                               1.24266094738807843860e-3,
                               2.71155556874348757815e-5,
                               2.01033439929228813265e-7];
        const F: [f64, ..8] = [1.00000000000000000000,
                               5.99832206555887937690e-1,
                               1.36929880922735805310e-1,
                               1.48753612908506148525e-2,
                               7.86869131145613259100e-4,
                               1.84631831751005468180e-5,
                               1.42151175831644588870e-7,
                               2.04426310338993978564e-15];

        #[inline(always)]
        fn poly(c: &[f64], x: f64) -> f64 {
            c[0] + x * (c[1] + x * (c[2] + x * (c[3] + x * (c[4] + x * (c[5] + x * (c[6] + x * c[7]))))))
        }

        if p <= 0.0 {
            return Float::neg_infinity();
        }
        if 1.0 <= p {
            return Float::infinity();
        }

        let q = p - 0.5;

        if (if q < 0.0 { -q } else { q }) <= SPLIT1 {
            let x = CONST1 - q * q;
            return self.mu + self.sigma * q * poly(&A, x) / poly(&B, x);
        }

        let mut x = if q < 0.0 { p } else { 1.0 - p };

        x = sqrt(-ln(x));

        if x <= SPLIT2 {
            x -= CONST2;
            x = poly(&C, x) / poly(&D, x);
        } else {
            x -= SPLIT2;
            x = poly(&E, x) / poly(&F, x);
        }

        self.mu + self.sigma * if q < 0.0 { -x } else { x }
    }

    #[inline]
    fn sample<R: Rng>(&self, rng: &mut R) -> f64 {
        self.normal.ind_sample(rng)
    }
}

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

    macro_rules! assert_almost_eq(
        ($x:expr, $y:expr) => ({
            let e: f64 = ::std::f64::EPSILON.sqrt();
            let x: Vec<f64> = $x;
            let y: Vec<f64> = $y;
            for i in range(0u, x.len()) {
                if x[i].is_finite() {
                    assert!(::std::num::abs(x[i] - y[i]) < e,
                            "expected {:e} ~ {:e}", x[i], y[i]);
                } else {
                    assert_eq!(x[i], y[i]);
                }
            }
        });
    )

    #[test]
    fn cdf() {
        let dist = Gaussian::new(1.0, 2.0);

        let x = vec![-4.0, -3.5, -3.0, -2.5, -2.0, -1.5, -1.0, -0.5, 0.0, 0.5,
                     1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0];
        let p = vec![6.209665325776139e-03, 1.222447265504470e-02,
                     2.275013194817922e-02, 4.005915686381709e-02,
                     6.680720126885809e-02, 1.056497736668553e-01,
                     1.586552539314571e-01, 2.266273523768682e-01,
                     3.085375387259869e-01, 4.012936743170763e-01,
                     5.000000000000000e-01, 5.987063256829237e-01,
                     6.914624612740131e-01, 7.733726476231317e-01,
                     8.413447460685429e-01, 8.943502263331446e-01,
                     9.331927987311419e-01];

        assert_almost_eq!(x.iter().map(|&x| dist.cdf(x)).collect(), p);
    }

    #[test]
    fn inv_cdf() {
        let dist = Gaussian::new(-1.0, 0.25);

        let p = vec![0.00, 0.05, 0.10, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40,
                     0.45, 0.50, 0.55, 0.60, 0.65, 0.70, 0.75, 0.80, 0.85,
                     0.90, 0.95, 1.00];
        let x = vec![Float::neg_infinity(), -1.411213406737868e+00,
                     -1.320387891386150e+00, -1.259108347373447e+00,
                     -1.210405308393228e+00, -1.168622437549020e+00,
                     -1.131100128177010e+00, -1.096330116601892e+00,
                     -1.063336775783950e+00, -1.031415336713768e+00,
                     -1.000000000000000e+00, -9.685846632862315e-01,
                     -9.366632242160501e-01, -9.036698833981082e-01,
                     -8.688998718229899e-01, -8.313775624509796e-01,
                     -7.895946916067714e-01, -7.408916526265525e-01,
                     -6.796121086138498e-01, -5.887865932621319e-01,
                     Float::infinity()];

        assert_almost_eq!(p.iter().map(|&p| dist.inv_cdf(p)).collect(), x);
    }
}

#[cfg(test)]
mod bench {
    extern crate test;

    use super::super::{Distribution, Sampler, Uniform};
    use super::Gaussian;

    #[bench]
    fn cdf(bench: &mut test::Bencher) {
        let mut rng = ::std::rand::task_rng();
        let dist = Gaussian::new(0.0, 1.0);
        let x = Sampler(&dist, &mut rng).take(1000).collect::<Vec<_>>();

        bench.iter(|| {
            test::black_box(x.iter().map(|&x| dist.cdf(x)).collect::<Vec<_>>())
        })
    }

    #[bench]
    fn inv_cdf(bench: &mut test::Bencher) {
        let mut rng = ::std::rand::task_rng();
        let dist = Gaussian::new(0.0, 1.0);
        let p = Sampler(&Uniform::new(0.0, 1.0), &mut rng).take(1000).collect::<Vec<_>>();

        bench.iter(|| {
            test::black_box(p.iter().map(|&p| dist.inv_cdf(p)).collect::<Vec<_>>())
        })
    }
}
<|endoftext|>


<commit_before>extern crate alsa;
extern crate libc;

use std;
use std::ffi::CString;
use std::error;
use time::{Time, TimeDelta};
use std::collections::VecDeque;
use base::*;

use super::input::*;

const ACCEPTED_RATES: [usize; 6] = [44100, 48000, 88200, 96000, 176400, 192000];

const RATE_DETECTION_PERIOD_MS: i64 = 1000;
const RATE_DETECTION_MIN_PERIOD_MS: i64 = 950;

struct RateDetector {
    history: VecDeque<(Time, usize)>,
    sum: usize,
}

impl RateDetector {
    fn new() -> RateDetector {
        return RateDetector {
            history: VecDeque::new(),
            sum: 0,
        };
    }

    fn update(&mut self, samples: usize) -> Option<usize> {
        let now = Time::now();
        self.sum += samples;
        self.history.push_back((now, samples));
        while self.history.len() > 0 &&
            now - self.history[0].0 >= TimeDelta::milliseconds(RATE_DETECTION_PERIOD_MS)
        {
            self.sum -= self.history.pop_front().unwrap().1;
        }

        let period = now - self.history[0].0;
        if period < TimeDelta::milliseconds(RATE_DETECTION_MIN_PERIOD_MS) {
            return None;
        }

        let current_rate =
            (self.sum - self.history[0].1) as i64 * 1000_000_000 / period.in_nanoseconds();

        for rate in ACCEPTED_RATES.iter() {
            // Check if the current rate is within 5% of a known value
            if ((*rate as i32 - current_rate as i32).abs() as usize) < *rate / 20 {
                return Some(*rate);
            }
        }
        None
    }
}

enum State {
    Inactive,
    Active { silent_samples: usize },
}

pub struct AlsaInput {
    pcm: alsa::PCM,
    sample_rate: usize,
    period_duration: TimeDelta,
    format: SampleFormat,
    period_size: usize,
    rate_detector: Option<RateDetector>,
    state: State,
    reference_time: Time,
    pos: i64,
}

const TARGET_SAMPLE_RATE: usize = 48000;

// Deactivate input after 30 seconds of silence.
const SILENCE_PERIOD_SECONDS: usize = 30;

impl AlsaInput {
    pub fn open(
        name: &str,
        target_sample_rate: usize,
        period_duration: TimeDelta,
        auto_sample_rate: bool,
    ) -> Result<Box<AlsaInput>> {
        let pcm = try!(alsa::PCM::open(
            &*CString::new(name).unwrap(),
            alsa::Direction::Capture,
            false
        ));

        let sample_rate;
        let period_size;
        let format;
        {
            let hwp = try!(alsa::pcm::HwParams::any(&pcm));
            try!(hwp.set_channels(2));

            try!(hwp.set_rate(
                target_sample_rate as u32,
                alsa::ValueOr::Nearest
            ));

            for fmt in [
                alsa::pcm::Format::S32LE,
                alsa::pcm::Format::S243LE,
                alsa::pcm::Format::S24LE,
                alsa::pcm::Format::S16LE,
            ].iter()
            {
                match hwp.set_format(*fmt) {
                    Ok(_) => break,
                    Err(_) => (),
                }
            }

            let target_period_size =
                period_duration * target_sample_rate as i64 / TimeDelta::seconds(1);
            try!(hwp.set_period_size_near(
                target_period_size as alsa::pcm::Frames,
                alsa::ValueOr::Nearest
            ));
            try!(hwp.set_periods(10, alsa::ValueOr::Nearest));
            try!(hwp.set_rate_resample(false));
            try!(hwp.set_access(alsa::pcm::Access::RWInterleaved));
            try!(pcm.hw_params(&hwp));

            sample_rate = try!(hwp.get_rate()) as usize;
            period_size = try!(hwp.get_period_size()) as usize;
            format = match try!(hwp.get_format()) {
                alsa::pcm::Format::S16LE => SampleFormat::S16LE,
                alsa::pcm::Format::S243LE => SampleFormat::S24LE3,
                alsa::pcm::Format::S24LE => SampleFormat::S24LE4,
                alsa::pcm::Format::S32LE => SampleFormat::S32LE,
                fmt => return Err(Error::new(&format!("Unknown format: {:?}", fmt))),
            };

            println!(
                "INFO: Reading {}. Buffer size: {}x{}. {} {}",
                name,
                period_size,
                try!(hwp.get_periods()),
                sample_rate,
                format
            );
        }

        Ok(Box::new(AlsaInput {
            pcm: pcm,
            sample_rate: sample_rate,
            period_duration: period_duration,
            format: format,
            period_size: period_size,
            rate_detector: if auto_sample_rate {
                Some(RateDetector::new())
            } else {
                None
            },
            state: State::Inactive,
            reference_time: Time::now(),
            pos: 0,
        }))
    }
}

impl Input for AlsaInput {
    fn read(&mut self) -> Result<Option<Frame>> {
        let mut buf = vec![0u8; self.period_size * CHANNELS * self.format.bytes_per_sample()];
        let read_result = self.pcm.io().readi(&mut buf[..]);
        let samples = match read_result {
            Ok(0) => return Ok(None),
            Ok(r) => r,
            Err(e) => {
                if e.code() == -libc::EWOULDBLOCK {
                    return Ok(None);
                }
                println!("Recovering AlsaInput {}", e.code());
                match self.pcm.recover(e.code(), true) {
                    Ok(_) => return self.read(),
                    Err(_) => return Err(Error::new(error::Error::description(&e))),
                }
            }
        };

        assert!(samples > 0);

        if self.rate_detector.is_some() {
            match self.rate_detector.as_mut().unwrap().update(samples) {
                Some(rate) => if rate != self.sample_rate {
                    println!("INFO: rate changed {}", rate);
                    self.sample_rate = rate;
                },
                None => {
                    // Drop all data if we don't know sample rate.
                    return Ok(None);
                }
            }
        }

        let mut all_zeros = true;
        for i in 0..(samples * 2) {
            if buf[i] != 0 {
                all_zeros = false;
                break;
            }
        }

        self.state = match (&self.state, all_zeros) {
            (_, false) => State::Active { silent_samples: 0 },
            (&State::Active { silent_samples }, true) => {
                let silence = silent_samples + samples;
                if silence > SILENCE_PERIOD_SECONDS * self.sample_rate {
                    State::Inactive
                } else {
                    State::Active {
                        silent_samples: silence,
                    }
                }
            }
            (&State::Inactive, true) => State::Inactive,
        };

        match self.state {
            State::Inactive => return Ok(None),
            _ => (),
        }

        let mut timestamp = get_sample_timestamp(self.reference_time, self.sample_rate, self.pos);
        self.pos += samples as i64;

        let now = Time::now();
        if (now - timestamp).abs() >= self.period_duration * 2 {
            println!(
                "INFO: Resetting input reference time. {} ",
                (now - timestamp).abs().in_milliseconds()
            );
            self.reference_time = now;
            timestamp = now;
            self.pos = 0;
        }

        let bytes = samples * CHANNELS * self.format.bytes_per_sample();
        Ok(Some(Frame::from_buffer(
            self.format,
            self.sample_rate,
            &buf[0..bytes],
            timestamp,
        )))
    }
}

const RETRY_PERIOD_SECS: i64 = 3;

pub struct ResilientAlsaInput {
    device_name: String,
    period_duration: TimeDelta,
    input: Option<Box<AlsaInput>>,
    last_open_attempt: Time,
}

impl ResilientAlsaInput {
    pub fn new(name: &str, period_duration: TimeDelta) -> Box<ResilientAlsaInput> {
        let device = match AlsaInput::open(name, TARGET_SAMPLE_RATE, period_duration, true) {
            Ok(device) => Some(device),
            Err(e) => {
                println!("warning: failed to open {}: {}", name, e);
                None
            }
        };
        Box::new(ResilientAlsaInput {
            device_name: String::from(name),
            period_duration: period_duration,
            input: device,
            last_open_attempt: Time::now(),
        })
    }

    fn try_reopen(&mut self) {
        let now = Time::now();
        if (now - self.last_open_attempt).in_seconds() >= RETRY_PERIOD_SECS {
            self.last_open_attempt = now;
            match AlsaInput::open(
                &self.device_name,
                TARGET_SAMPLE_RATE,
                self.period_duration,
                true,
            ) {
                Ok(input) => self.input = Some(input),
                Err(e) => println!("info: Failed to reopen {}", e),
            }
        }
    }
}

impl Input for ResilientAlsaInput {
    fn read(&mut self) -> Result<Option<Frame>> {
        if self.input.is_none() {
            self.try_reopen()
        }
        let mut reset = false;
        let result = match self.input.as_mut() {
            Some(input) => match input.read() {
                Ok(Some(frame)) => Some(frame),
                Ok(None) => None,
                Err(e) => {
                    println!(
                        "warning: input read from {} failed: {}",
                        self.device_name,
                        e
                    );
                    reset = true;
                    None
                }
            },
            None => None,
        };
        if reset {
            self.input = None;
        }
        if self.input.is_none() {
            std::thread::sleep(self.period_duration.as_duration());
        }

        Ok(result)
    }
}
<commit_msg>Input sample rate probing<commit_after>extern crate alsa;
extern crate libc;

use std;
use std::ffi::CString;
use std::error;
use time::{Time, TimeDelta};
use std::collections::VecDeque;
use base::*;

use super::input::*;

const ACCEPTED_RATES: [usize; 6] = [44100, 48000, 88200, 96000, 176400, 192000];

const RATE_DETECTION_PERIOD_MS: i64 = 1000;
const RATE_DETECTION_MIN_PERIOD_MS: i64 = 950;

struct RateDetector {
    history: VecDeque<(Time, usize)>,
    sum: usize,
}

impl RateDetector {
    fn new() -> RateDetector {
        return RateDetector {
            history: VecDeque::new(),
            sum: 0,
        };
    }

    fn update(&mut self, samples: usize) -> Option<usize> {
        let now = Time::now();
        self.sum += samples;
        self.history.push_back((now, samples));
        while self.history.len() > 0 &&
            now - self.history[0].0 >= TimeDelta::milliseconds(RATE_DETECTION_PERIOD_MS)
        {
            self.sum -= self.history.pop_front().unwrap().1;
        }

        let period = now - self.history[0].0;
        if period < TimeDelta::milliseconds(RATE_DETECTION_MIN_PERIOD_MS) {
            return None;
        }

        let current_rate =
            (self.sum - self.history[0].1) as i64 * 1000_000_000 / period.in_nanoseconds();

        for rate in ACCEPTED_RATES.iter() {
            // Check if the current rate is within 5% of a known value
            if ((*rate as i32 - current_rate as i32).abs() as usize) < *rate / 20 {
                return Some(*rate);
            }
        }
        None
    }
}

enum State {
    Inactive,
    Active { silent_samples: usize },
}

pub struct AlsaInput {
    pcm: alsa::PCM,
    sample_rate: usize,
    period_duration: TimeDelta,
    format: SampleFormat,
    period_size: usize,
    rate_detector: Option<RateDetector>,
    state: State,
    reference_time: Time,
    pos: i64,
}

// Deactivate input after 5 seconds of silence.
const SILENCE_PERIOD_SECONDS: usize = 5;

impl AlsaInput {
    pub fn open(
        name: &str,
        target_sample_rate: usize,
        period_duration: TimeDelta,
        auto_sample_rate: bool,
    ) -> Result<Box<AlsaInput>> {
        let pcm = try!(alsa::PCM::open(
            &*CString::new(name).unwrap(),
            alsa::Direction::Capture,
            false
        ));

        let sample_rate;
        let period_size;
        let format;
        {
            let hwp = try!(alsa::pcm::HwParams::any(&pcm));
            try!(hwp.set_channels(2));

            try!(hwp.set_rate(
                target_sample_rate as u32,
                alsa::ValueOr::Nearest
            ));

            for fmt in [
                alsa::pcm::Format::S32LE,
                alsa::pcm::Format::S243LE,
                alsa::pcm::Format::S24LE,
                alsa::pcm::Format::S16LE,
            ].iter()
            {
                match hwp.set_format(*fmt) {
                    Ok(_) => break,
                    Err(_) => (),
                }
            }

            let target_period_size =
                period_duration * target_sample_rate as i64 / TimeDelta::seconds(1);
            try!(hwp.set_period_size_near(
                target_period_size as alsa::pcm::Frames,
                alsa::ValueOr::Nearest
            ));
            try!(hwp.set_periods(10, alsa::ValueOr::Nearest));
            try!(hwp.set_rate_resample(false));
            try!(hwp.set_access(alsa::pcm::Access::RWInterleaved));
            try!(pcm.hw_params(&hwp));

            sample_rate = try!(hwp.get_rate()) as usize;
            period_size = try!(hwp.get_period_size()) as usize;
            format = match try!(hwp.get_format()) {
                alsa::pcm::Format::S16LE => SampleFormat::S16LE,
                alsa::pcm::Format::S243LE => SampleFormat::S24LE3,
                alsa::pcm::Format::S24LE => SampleFormat::S24LE4,
                alsa::pcm::Format::S32LE => SampleFormat::S32LE,
                fmt => return Err(Error::new(&format!("Unknown format: {:?}", fmt))),
            };

            println!(
                "INFO: Reading {}. Buffer size: {}x{}. {} {}",
                name,
                period_size,
                try!(hwp.get_periods()),
                sample_rate,
                format
            );
        }

        Ok(Box::new(AlsaInput {
            pcm: pcm,
            sample_rate: sample_rate,
            period_duration: period_duration,
            format: format,
            period_size: period_size,
            rate_detector: if auto_sample_rate {
                Some(RateDetector::new())
            } else {
                None
            },
            state: State::Inactive,
            reference_time: Time::now(),
            pos: 0,
        }))
    }
}

impl Input for AlsaInput {
    fn read(&mut self) -> Result<Option<Frame>> {
        let mut buf = vec![0u8; self.period_size * CHANNELS * self.format.bytes_per_sample()];
        let read_result = self.pcm.io().readi(&mut buf[..]);
        let samples = match read_result {
            Ok(0) => return Ok(None),
            Ok(r) => r,
            Err(e) => {
                if e.code() == -libc::EWOULDBLOCK {
                    return Ok(None);
                }
                println!("Recovering AlsaInput {}", e.code());
                match self.pcm.recover(e.code(), true) {
                    Ok(_) => return self.read(),
                    Err(_) => return Err(Error::new(error::Error::description(&e))),
                }
            }
        };

        assert!(samples > 0);

        if self.rate_detector.is_some() {
            match self.rate_detector.as_mut().unwrap().update(samples) {
                Some(rate) => if rate != self.sample_rate {
                    println!("INFO: rate changed {}", rate);
                    self.sample_rate = rate;
                },
                None => (),
            }
        }

        let bytes = samples * CHANNELS * self.format.bytes_per_sample();

        let mut all_zeros = true;
        for i in 0..bytes {
            if buf[i] != 0 {
                all_zeros = false;
                break;
            }
        }

        self.state = match (&self.state, all_zeros) {
            (_, false) => State::Active { silent_samples: 0 },
            (&State::Active { silent_samples }, true) => {
                let silence = silent_samples + samples;
                if silence > SILENCE_PERIOD_SECONDS * self.sample_rate {
                    State::Inactive
                } else {
                    State::Active {
                        silent_samples: silence,
                    }
                }
            }
            (&State::Inactive, true) => State::Inactive,
        };

        match self.state {
            State::Inactive => return Ok(None),
            _ => (),
        }

        let mut timestamp = get_sample_timestamp(self.reference_time, self.sample_rate, self.pos);
        self.pos += samples as i64;

        let now = Time::now();
        if (now - timestamp).abs() >= self.period_duration * 2 {
            println!(
                "INFO: Resetting input reference time. {} ",
                (now - timestamp).abs().in_milliseconds()
            );
            self.reference_time = now;
            timestamp = now;
            self.pos = 0;
        }

        Ok(Some(Frame::from_buffer(
            self.format,
            self.sample_rate,
            &buf[0..bytes],
            timestamp,
        )))
    }
}

const TARGET_SAMPLE_RATES: [usize; 2] = [44100, 48000];

const RETRY_PERIOD_MS: i64 = 3000;
const PROBE_TIME_MS: i64 = 500;

pub struct ResilientAlsaInput {
    device_name: String,
    period_duration: TimeDelta,
    input: Option<Box<AlsaInput>>,
    last_open_attempt: Time,
    last_active: Time,
    next_sample_rate_to_probe: usize,
}

impl ResilientAlsaInput {
    pub fn new(name: &str, period_duration: TimeDelta) -> Box<ResilientAlsaInput> {
        Box::new(ResilientAlsaInput {
            device_name: String::from(name),
            period_duration: period_duration,
            input: None,
            last_open_attempt: Time::now(),
            last_active: Time::now(),
            next_sample_rate_to_probe: 0,
        })
    }

    fn try_reopen(&mut self) {
        let now = Time::now();
        if now - self.last_open_attempt >= TimeDelta::milliseconds(RETRY_PERIOD_MS) {
            self.last_open_attempt = now;
            match AlsaInput::open(
                &self.device_name,
                TARGET_SAMPLE_RATES[self.next_sample_rate_to_probe],
                self.period_duration,
                true,
            ) {
                Ok(input) => {
                    self.input = Some(input);
                    self.last_active = now;
                }
                Err(_) => (),
            };
            self.next_sample_rate_to_probe =
                (self.next_sample_rate_to_probe + 1) % TARGET_SAMPLE_RATES.len();
        }
    }
}

impl Input for ResilientAlsaInput {
    fn read(&mut self) -> Result<Option<Frame>> {
        if self.input.is_none() {
            self.try_reopen()
        }
        let mut reset = false;
        let result = match self.input.as_mut() {
            Some(input) => match input.read() {
                Ok(Some(frame)) => {
                    self.last_active = Time::now();
                    Some(frame)
                }
                Ok(None) => {
                    if Time::now() - self.last_active > TimeDelta::milliseconds(PROBE_TIME_MS) {
                        reset = true;
                    }
                    None
                }
                Err(e) => {
                    println!(
                        "warning: input read from {} failed: {}",
                        self.device_name,
                        e
                    );
                    reset = true;
                    None
                }
            },
            None => None,
        };
        if reset {
            self.input = None;
        }
        if self.input.is_none() {
            std::thread::sleep(self.period_duration.as_duration());
        }

        Ok(result)
    }
}
<|endoftext|>