<commit_before>use tracker::Tracker;
use response::TrackerResponse;
use response::error::ErrorResponse;
use response::success::SuccessResponse;
use tracker::announce::{Action, Announce};
use tracker::scrape::Scrape;

use hyper::server::{Request, Response, Handler};
use hyper::uri::RequestUri::AbsolutePath;
use std::net::{IpAddr, SocketAddr, SocketAddrV4, SocketAddrV6};
use std::collections::HashMap;
use std::sync::Arc;
use url::form_urlencoded;
use url::{Url, UrlParser};
use std::str::FromStr;
use std::cmp;

pub struct RequestHandler {
    pub tracker: Arc<Tracker>,
}

impl RequestHandler {
    pub fn new(tracker: Arc<Tracker>) -> RequestHandler {
        RequestHandler { tracker: tracker }
    }
}

impl Handler for RequestHandler {
    fn handle(&self, req: Request, res: Response) {
        let resp = match req.uri {
            AbsolutePath(ref path) => {
                let base = Url::parse("http://localhost").unwrap();
                let url = UrlParser::new().base_url(&base).parse(path).unwrap();
                self.handle_url(&req, url)
            }
            _ => Err(ErrorResponse::BadAction),
        };
        res.send(serialize_resp(resp).as_slice()).unwrap();
    }
}

impl RequestHandler {
    fn handle_url(&self, req: &Request, url: Url) -> Result<SuccessResponse, ErrorResponse> {
        if url.path().is_none() {
            return Err(ErrorResponse::BadAction);
        }
        let path = url.path().unwrap();
        let params = url.query_pairs();

        if cfg!(feature = "private") {
            if path.len() != 2 {
                Err(ErrorResponse::BadRequest)
            } else {
                self.handle_private_req(&path[1], &path[0], params)
            }
        } else {
            if path.len() != 1 {
                Err(ErrorResponse::BadRequest)
            } else {
                self.handle_public_req(req, &path[0], params)
            }
        }
    }

    fn handle_public_req(&self,
                         req: &Request,
                         path: &String,
                         params: Option<Vec<(String, String)>>)
                         -> Result<SuccessResponse, ErrorResponse> {
        match &path[..] {
            "stats" => self.tracker.get_stats(),
            "announce" => {
                let announce = try!(request_to_announce(req, params));
                self.tracker.handle_announce(announce)
            }
            "scrape" => {
                let scrape = try!(request_to_scrape(params));
                self.tracker.handle_scrape(scrape)
            }
            _ => Err(ErrorResponse::BadAction),
        }
    }

    fn handle_private_req(&self,
                          path: &String,
                          key: &String,
                          params: Option<Vec<(String, String)>>)
                          -> Result<SuccessResponse, ErrorResponse> {
        Err(ErrorResponse::BadAction)
    }
}

fn request_to_scrape(params: Option<Vec<(String, String)>>) -> Result<Scrape, ErrorResponse> {
    if params.is_none() {
        return Err(ErrorResponse::BadRequest);
    }
    let param_vec = params.unwrap();

    let hashes = param_vec.into_iter()
                          .map(|(_, hash)| hash)
                          .collect();
    Ok(Scrape::new(hashes))
}

fn request_to_announce(req: &Request,
                       params: Option<Vec<(String, String)>>)
                       -> Result<Announce, ErrorResponse> {
    if params.is_none() {
        return Err(ErrorResponse::BadRequest);
    }
    let param_vec = params.unwrap();
    if param_vec.len() > 10 {
        return Err(ErrorResponse::BadRequest);
    }

    let mut params = HashMap::new();
    for (key, val) in param_vec {
        params.insert(key, val);
    }

    let info_hash: String = try!(get_from_params(&params, String::from("info_hash")));
    let pid: String = try!(get_from_params(&params, String::from("peer_id")));
    if info_hash.len() > 40 || pid.len() > 30 {
        return Err(ErrorResponse::BadRequest);
    }
    let ul = try!(get_from_params(&params, String::from("uploaded")));
    let dl = try!(get_from_params(&params, String::from("downloaded")));
    let left = try!(get_from_params(&params, String::from("left")));

    // IP parsing according to BEP 0007 with additional proxy forwarding check
    let port = try!(get_from_params(&params, String::from("port")));
    let (ipv4, ipv6) = get_ips(&params, req, &port);
    let action = match get_from_params::<String>(&params, String::from("event")) {
        Ok(ev_str) => {
            match &ev_str[..] {
                "started" => get_action(left),
                "stopped" => Action::Stopped,
                "completed" => Action::Completed,
                _ => get_action(left),
            }
        }
        Err(_) => get_action(left),
    };

    let numwant = cmp::min(get_from_params::<u8>(&params, String::from("numwant")).unwrap_or(25),
                           25);

    let compact = get_from_params::<u8>(&params, String::from("compact")).unwrap_or(1) != 0;

    Ok(Announce {
        info_hash: info_hash,
        peer_id: pid,
        ipv4: ipv4,
        ipv6: ipv6,
        ul: ul,
        dl: dl,
        left: left,
        action: action,
        numwant: numwant,
        compact: compact,
    })
}

fn get_ips(params: &HashMap<String, String>,
           req: &Request,
           port: &u16)
           -> (Option<SocketAddrV4>, Option<SocketAddrV6>) {
    let port = *port;
    let default_ip = match req.headers.get_raw("X-Forwarded-For") {
        Some(bytes) => {
            match String::from_utf8(bytes[0].clone()) {
                Ok(ip_str) => {
                    match ip_str.parse::<IpAddr>() {
                        Ok(ip) => SocketAddr::new(ip, port),
                        Err(_) => req.remote_addr,
                    }
                }
                Err(_) => req.remote_addr,
            }
        }
        None => req.remote_addr,
    };
    let ip = match get_from_params(&params, String::from("ip")) {
        Ok(ip) => SocketAddr::new(ip, port),
        Err(_) => default_ip,
    };

    match ip {
        SocketAddr::V4(v4) => {
            let v6 = match get_socket(&params, String::from("ipv6"), port) {
                Some(sock) => {
                    match sock {
                        SocketAddr::V6(v6) => Some(v6),
                        _ => None,
                    }
                }
                None => None,
            };
            (Some(v4), v6)
        }
        SocketAddr::V6(v6) => {
            let v4 = match get_socket(&params, String::from("ipv4"), port) {
                Some(sock) => {
                    match sock {
                        SocketAddr::V4(v4) => Some(v4),
                        _ => None,
                    }
                }
                None => None,
            };
            (v4, Some(v6))
        }
    }
}

fn get_from_params<T: FromStr>(map: &HashMap<String, String>,
                               key: String)
                               -> Result<T, ErrorResponse> {
    match map.get(&key) {
        Some(res) => {
            match res.parse::<T>() {
                Ok(val) => Ok(val),
                Err(_) => Err(ErrorResponse::BadRequest),
            }
        }
        None => Err(ErrorResponse::BadRequest),
    }
}

fn get_socket(params: &HashMap<String, String>, key: String, port: u16) -> Option<SocketAddr> {
    let ip = get_from_params(params, key.clone());
    let socket = get_from_params(params, key.clone());
    match (ip, socket) {
        (Err(_), Err(_)) => None,
        (Ok(ip), Err(_)) => Some(SocketAddr::new(ip, port)),
        (Err(_), Ok(sock)) => Some(sock),
        _ => None,
    }
}

fn get_action(left: u64) -> Action {
    if left == 0 {
        Action::Seeding
    } else {
        Action::Leeching
    }
}

fn serialize_resp(result: Result<SuccessResponse, ErrorResponse>) -> Vec<u8> {
    match result {
        Ok(resp) => resp.http_resp(),
        Err(err) => err.to_bencode(),
    }
}
<commit_msg>Some minor tweaks to tracker request handling<commit_after>use tracker::Tracker;
use response::TrackerResponse;
use response::error::ErrorResponse;
use response::success::SuccessResponse;
use tracker::announce::{Action, Announce};
use tracker::scrape::Scrape;

use hyper::server::{Request, Response, Handler};
use hyper::uri::RequestUri::AbsolutePath;
use std::net::{IpAddr, SocketAddr, SocketAddrV4, SocketAddrV6};
use std::collections::HashMap;
use std::sync::Arc;
use url::{Url, UrlParser};
use std::str::FromStr;
use std::cmp;

pub struct RequestHandler {
    pub tracker: Arc<Tracker>,
}

impl RequestHandler {
    pub fn new(tracker: Arc<Tracker>) -> RequestHandler {
        RequestHandler { tracker: tracker }
    }
}

impl Handler for RequestHandler {
    fn handle(&self, req: Request, res: Response) {
        let resp = match req.uri {
            AbsolutePath(ref path) => {
                let base = Url::parse("http://localhost").unwrap();
                let url = UrlParser::new().base_url(&base).parse(path).unwrap();
                self.handle_url(&req, url)
            }
            _ => Err(ErrorResponse::BadAction),
        };
        res.send(serialize_resp(resp).as_slice()).unwrap();
    }
}

impl RequestHandler {
    fn handle_url(&self, req: &Request, url: Url) -> Result<SuccessResponse, ErrorResponse> {
        if url.path().is_none() {
            return Err(ErrorResponse::BadAction);
        }
        let path = url.path().unwrap();
        let params = url.query_pairs();

        if cfg!(feature = "private") {
            if path.len() != 2 {
                Err(ErrorResponse::BadRequest)
            } else {
                // Do key validation
                if false {
                    Err(ErrorResponse::BadAuth)
                } else {
                    self.handle_req(req, &path[1], params)
                }
            }
        } else {
            if path.len() != 1 {
                Err(ErrorResponse::BadRequest)
            } else {
                self.handle_req(req, &path[0], params)
            }
        }
    }

    fn handle_req(&self,
                  req: &Request,
                  path: &String,
                  params: Option<Vec<(String, String)>>)
                  -> Result<SuccessResponse, ErrorResponse> {
        match &path[..] {
            "stats" => self.tracker.get_stats(),
            "announce" => {
                let announce = try!(request_to_announce(req, params));
                self.tracker.handle_announce(announce)
            }
            "scrape" => {
                let scrape = try!(request_to_scrape(params));
                self.tracker.handle_scrape(scrape)
            }
            _ => Err(ErrorResponse::BadAction),
        }
    }
}

fn request_to_scrape(params: Option<Vec<(String, String)>>) -> Result<Scrape, ErrorResponse> {
    if params.is_none() {
        return Err(ErrorResponse::BadRequest);
    }
    let param_vec = params.unwrap();

    let hashes = param_vec.into_iter()
                          .map(|(_, hash)| hash)
                          .collect();
    Ok(Scrape::new(hashes))
}

fn request_to_announce(req: &Request,
                       params: Option<Vec<(String, String)>>)
                       -> Result<Announce, ErrorResponse> {
    if params.is_none() {
        return Err(ErrorResponse::BadRequest);
    }
    let param_vec = params.unwrap();
    if param_vec.len() > 10 {
        return Err(ErrorResponse::BadRequest);
    }

    let mut params = HashMap::new();
    for (key, val) in param_vec {
        params.insert(key, val);
    }

    let info_hash: String = try!(get_from_params(&params, String::from("info_hash")));
    let pid: String = try!(get_from_params(&params, String::from("peer_id")));
    if info_hash.len() > 40 || pid.len() > 30 {
        return Err(ErrorResponse::BadRequest);
    }
    let ul = try!(get_from_params(&params, String::from("uploaded")));
    let dl = try!(get_from_params(&params, String::from("downloaded")));
    let left = try!(get_from_params(&params, String::from("left")));

    // IP parsing according to BEP 0007 with additional proxy forwarding check
    let port = try!(get_from_params(&params, String::from("port")));
    let (ipv4, ipv6) = get_ips(&params, req, &port);
    let action = match get_from_params::<String>(&params, String::from("event")) {
        Ok(ev_str) => {
            match &ev_str[..] {
                "started" => get_action(left),
                "stopped" => Action::Stopped,
                "completed" => Action::Completed,
                _ => get_action(left),
            }
        }
        Err(_) => get_action(left),
    };

    let numwant = cmp::min(get_from_params::<u8>(&params, String::from("numwant")).unwrap_or(25),
                           25);

    let compact = get_from_params::<u8>(&params, String::from("compact")).unwrap_or(1) != 0;

    Ok(Announce {
        info_hash: info_hash,
        peer_id: pid,
        ipv4: ipv4,
        ipv6: ipv6,
        ul: ul,
        dl: dl,
        left: left,
        action: action,
        numwant: numwant,
        compact: compact,
    })
}

fn get_ips(params: &HashMap<String, String>,
           req: &Request,
           port: &u16)
           -> (Option<SocketAddrV4>, Option<SocketAddrV6>) {
    let port = *port;
    let default_ip = match req.headers.get_raw("X-Forwarded-For") {
        Some(bytes) => {
            match String::from_utf8(bytes[0].clone()) {
                Ok(ip_str) => {
                    match ip_str.parse::<IpAddr>() {
                        Ok(ip) => SocketAddr::new(ip, port),
                        Err(_) => req.remote_addr,
                    }
                }
                Err(_) => req.remote_addr,
            }
        }
        None => req.remote_addr,
    };
    let ip = match get_from_params(&params, String::from("ip")) {
        Ok(ip) => SocketAddr::new(ip, port),
        Err(_) => default_ip,
    };

    match ip {
        SocketAddr::V4(v4) => {
            let v6 = match get_socket(&params, String::from("ipv6"), port) {
                Some(sock) => {
                    match sock {
                        SocketAddr::V6(v6) => Some(v6),
                        _ => None,
                    }
                }
                None => None,
            };
            (Some(v4), v6)
        }
        SocketAddr::V6(v6) => {
            let v4 = match get_socket(&params, String::from("ipv4"), port) {
                Some(sock) => {
                    match sock {
                        SocketAddr::V4(v4) => Some(v4),
                        _ => None,
                    }
                }
                None => None,
            };
            (v4, Some(v6))
        }
    }
}

fn get_from_params<T: FromStr>(map: &HashMap<String, String>,
                               key: String)
                               -> Result<T, ErrorResponse> {
    match map.get(&key) {
        Some(res) => {
            match res.parse::<T>() {
                Ok(val) => Ok(val),
                Err(_) => Err(ErrorResponse::BadRequest),
            }
        }
        None => Err(ErrorResponse::BadRequest),
    }
}

fn get_socket(params: &HashMap<String, String>, key: String, port: u16) -> Option<SocketAddr> {
    let ip = get_from_params(params, key.clone());
    let socket = get_from_params(params, key.clone());
    match (ip, socket) {
        (Err(_), Err(_)) => None,
        (Ok(ip), Err(_)) => Some(SocketAddr::new(ip, port)),
        (Err(_), Ok(sock)) => Some(sock),
        _ => None,
    }
}

fn get_action(left: u64) -> Action {
    if left == 0 {
        Action::Seeding
    } else {
        Action::Leeching
    }
}

fn serialize_resp(result: Result<SuccessResponse, ErrorResponse>) -> Vec<u8> {
    match result {
        Ok(resp) => resp.http_resp(),
        Err(err) => err.to_bencode(),
    }
}
<|endoftext|>


<commit_before>use int::Int;

macro_rules! div {
    ($intrinsic:ident: $ty:ty, $uty:ty) => {
        div!($intrinsic: $ty, $uty, $ty, |i| {i});
    };
    ($intrinsic:ident: $ty:ty, $uty:ty, $tyret:ty, $conv:expr) => {
        /// Returns `a / b`
        #[cfg_attr(not(test), no_mangle)]
        pub extern "C" fn $intrinsic(a: $ty, b: $ty) -> $tyret {
            let s_a = a >> (<$ty>::bits() - 1);
            let s_b = b >> (<$ty>::bits() - 1);
            let a = (a ^ s_a) - s_a;
            let b = (b ^ s_b) - s_b;
            let s = s_a ^ s_b;

            let r = udiv!(a as $uty, b as $uty);
            ($conv)((r as $ty ^ s) - s)
        }
    }
}

macro_rules! mod_ {
    ($intrinsic:ident: $ty:ty, $uty:ty) => {
        mod_!($intrinsic: $ty, $uty, $ty, |i| {i});
    };
    ($intrinsic:ident: $ty:ty, $uty:ty, $tyret:ty, $conv:expr) => {
        /// Returns `a % b`
        #[cfg_attr(not(test), no_mangle)]
        pub extern "C" fn $intrinsic(a: $ty, b: $ty) -> $tyret {
            let s = b >> (<$ty>::bits() - 1);
            let b = (b ^ s) - s;
            let s = a >> (<$ty>::bits() - 1);
            let a = (a ^ s) - s;

            let r = urem!(a as $uty, b as $uty);
            ($conv)((r as $ty ^ s) - s)
        }
    }
}

macro_rules! divmod {
    ($abi:tt, $intrinsic:ident, $div:ident: $ty:ty) => {
        /// Returns `a / b` and sets `*rem = n % d`
        #[cfg_attr(not(test), no_mangle)]
        pub extern $abi fn $intrinsic(a: $ty, b: $ty, rem: &mut $ty) -> $ty {
            #[cfg(all(feature = "c", any(target_arch = "x86")))]
            extern {
                fn $div(a: $ty, b: $ty) -> $ty;
            }

            let r = match () {
                #[cfg(not(all(feature = "c", any(target_arch = "x86"))))]
                () => $div(a, b),
                #[cfg(all(feature = "c", any(target_arch = "x86")))]
                () => unsafe { $div(a, b) },
            };
            // NOTE won't overflow because it's using the result from the
            // previous division
            *rem = a - r.wrapping_mul(b);
            r
        }
    }
}

#[cfg(not(all(feature = "c", target_arch = "arm", not(target_os = "ios"), not(thumbv6m))))]
div!(__divsi3: i32, u32);

#[cfg(not(all(feature = "c", target_arch = "x86")))]
div!(__divdi3: i64, u64);

#[cfg(not(all(windows, target_pointer_width="64")))]
div!(__divti3: i128, u128);

#[cfg(all(windows, target_pointer_width="64"))]
div!(__divti3: i128, u128, ::U64x2, ::sconv);

#[cfg(not(all(feature = "c", target_arch = "arm", not(target_os = "ios"))))]
mod_!(__modsi3: i32, u32);

#[cfg(not(all(feature = "c", target_arch = "x86")))]
mod_!(__moddi3: i64, u64);

#[cfg(not(all(windows, target_pointer_width="64")))]
mod_!(__modti3: i128, u128);

#[cfg(all(windows, target_pointer_width="64"))]
mod_!(__modti3: i128, u128, ::U64x2, ::sconv);

#[cfg(not(all(feature = "c", target_arch = "arm", not(target_os = "ios"))))]
divmod!("C", __divmodsi4, __divsi3: i32);

#[cfg(target_arch = "arm")]
divmod!("aapcs", __divmoddi4, __divdi3: i64);

#[cfg(not(target_arch = "arm"))]
divmod!("C", __divmoddi4, __divdi3: i64);

#[cfg(test)]
mod tests {
    use qc::{U32, U64};

    check! {
        fn __divdi3(f: extern "C" fn(i64, i64) -> i64, n: U64, d: U64) -> Option<i64> {
            let (n, d) = (n.0 as i64, d.0 as i64);
            if d == 0 {
                None
            } else {
                Some(f(n, d))
            }
        }

        fn __moddi3(f: extern "C" fn(i64, i64) -> i64, n: U64, d: U64) -> Option<i64> {
            let (n, d) = (n.0 as i64, d.0 as i64);
            if d == 0 {
                None
            } else {
                Some(f(n, d))
            }
        }

        #[cfg(target_arch = "arm")]
        fn __divmoddi4(f: extern "aapcs" fn(i64, i64, &mut i64) -> i64,
                       n: U64,
                       d: U64) -> Option<(i64, i64)> {
            let (n, d) = (n.0 as i64, d.0 as i64);
            if d == 0 {
                None
            } else {
                let mut r = 0;
                let q = f(n, d, &mut r);
                Some((q, r))
            }
        }

        #[cfg(not(target_arch = "arm"))]
        fn __divmoddi4(f: extern "C" fn(i64, i64, &mut i64) -> i64,
                       n: U64,
                       d: U64) -> Option<(i64, i64)> {
            let (n, d) = (n.0 as i64, d.0 as i64);
            if d == 0 {
                None
            } else {
                let mut r = 0;
                let q = f(n, d, &mut r);
                Some((q, r))
            }
        }

        fn __divsi3(f: extern "C" fn(i32, i32) -> i32,
                    n: U32,
                    d: U32) -> Option<i32> {
            let (n, d) = (n.0 as i32, d.0 as i32);
            if d == 0 {
                None
            } else {
                Some(f(n, d))
            }
        }

        fn __modsi3(f: extern "C" fn(i32, i32) -> i32,
                    n: U32,
                    d: U32) -> Option<i32> {
            let (n, d) = (n.0 as i32, d.0 as i32);
            if d == 0 {
                None
            } else {
                Some(f(n, d))
            }
        }

        fn __divmodsi4(f: extern "C" fn(i32, i32, &mut i32) -> i32,
                       n: U32,
                       d: U32) -> Option<(i32, i32)> {
            let (n, d) = (n.0 as i32, d.0 as i32);
            if d == 0 {
                None
            } else {
                let mut r = 0;
                let q = f(n, d, &mut r);
                Some((q, r))
            }
        }
    }
}

#[cfg(test)]
#[cfg(all(not(windows),
          not(target_arch = "mips64"),
          not(target_arch = "mips64el"),
          target_pointer_width="64"))]
mod tests_i128 {
    use qc::U128;
    check! {

        fn __divti3(f: extern "C" fn(i128, i128) -> i128, n: U128, d: U128) -> Option<i128> {
            let (n, d) = (n.0 as i128, d.0 as i128);
            if d == 0 {
                None
            } else {
                Some(f(n, d))
            }
        }

        fn __modti3(f: extern "C" fn(i128, i128) -> i128, n: U128, d: U128) -> Option<i128> {
            let (n, d) = (n.0 as i128, d.0 as i128);
            if d == 0 {
                None
            } else {
                Some(f(n, d))
            }
        }
    }
}
<commit_msg>Auto merge of #149 - rust-lang-nursery:gh148, r=alexcrichton<commit_after>use int::Int;

macro_rules! div {
    ($intrinsic:ident: $ty:ty, $uty:ty) => {
        div!($intrinsic: $ty, $uty, $ty, |i| {i});
    };
    ($intrinsic:ident: $ty:ty, $uty:ty, $tyret:ty, $conv:expr) => {
        /// Returns `a / b`
        #[cfg_attr(not(test), no_mangle)]
        pub extern "C" fn $intrinsic(a: $ty, b: $ty) -> $tyret {
            let s_a = a >> (<$ty>::bits() - 1);
            let s_b = b >> (<$ty>::bits() - 1);
            // NOTE it's OK to overflow here because of the `as $uty` cast below
            // This whole operation is computing the absolute value of the inputs
            // So some overflow will happen when dealing with e.g. `i64::MIN`
            // where the absolute value is `(-i64::MIN) as u64`
            let a = (a ^ s_a).wrapping_sub(s_a);
            let b = (b ^ s_b).wrapping_sub(s_b);
            let s = s_a ^ s_b;

            let r = udiv!(a as $uty, b as $uty);
            ($conv)((r as $ty ^ s) - s)
        }
    }
}

macro_rules! mod_ {
    ($intrinsic:ident: $ty:ty, $uty:ty) => {
        mod_!($intrinsic: $ty, $uty, $ty, |i| {i});
    };
    ($intrinsic:ident: $ty:ty, $uty:ty, $tyret:ty, $conv:expr) => {
        /// Returns `a % b`
        #[cfg_attr(not(test), no_mangle)]
        pub extern "C" fn $intrinsic(a: $ty, b: $ty) -> $tyret {
            let s = b >> (<$ty>::bits() - 1);
            let b = (b ^ s) - s;
            let s = a >> (<$ty>::bits() - 1);
            let a = (a ^ s) - s;

            let r = urem!(a as $uty, b as $uty);
            ($conv)((r as $ty ^ s) - s)
        }
    }
}

macro_rules! divmod {
    ($abi:tt, $intrinsic:ident, $div:ident: $ty:ty) => {
        /// Returns `a / b` and sets `*rem = n % d`
        #[cfg_attr(not(test), no_mangle)]
        pub extern $abi fn $intrinsic(a: $ty, b: $ty, rem: &mut $ty) -> $ty {
            #[cfg(all(feature = "c", any(target_arch = "x86")))]
            extern {
                fn $div(a: $ty, b: $ty) -> $ty;
            }

            let r = match () {
                #[cfg(not(all(feature = "c", any(target_arch = "x86"))))]
                () => $div(a, b),
                #[cfg(all(feature = "c", any(target_arch = "x86")))]
                () => unsafe { $div(a, b) },
            };
            // NOTE won't overflow because it's using the result from the
            // previous division
            *rem = a - r.wrapping_mul(b);
            r
        }
    }
}

#[cfg(not(all(feature = "c", target_arch = "arm", not(target_os = "ios"), not(thumbv6m))))]
div!(__divsi3: i32, u32);

#[cfg(not(all(feature = "c", target_arch = "x86")))]
div!(__divdi3: i64, u64);

#[cfg(not(all(windows, target_pointer_width="64")))]
div!(__divti3: i128, u128);

#[cfg(all(windows, target_pointer_width="64"))]
div!(__divti3: i128, u128, ::U64x2, ::sconv);

#[cfg(not(all(feature = "c", target_arch = "arm", not(target_os = "ios"))))]
mod_!(__modsi3: i32, u32);

#[cfg(not(all(feature = "c", target_arch = "x86")))]
mod_!(__moddi3: i64, u64);

#[cfg(not(all(windows, target_pointer_width="64")))]
mod_!(__modti3: i128, u128);

#[cfg(all(windows, target_pointer_width="64"))]
mod_!(__modti3: i128, u128, ::U64x2, ::sconv);

#[cfg(not(all(feature = "c", target_arch = "arm", not(target_os = "ios"))))]
divmod!("C", __divmodsi4, __divsi3: i32);

#[cfg(target_arch = "arm")]
divmod!("aapcs", __divmoddi4, __divdi3: i64);

#[cfg(not(target_arch = "arm"))]
divmod!("C", __divmoddi4, __divdi3: i64);

#[cfg(test)]
mod tests {
    use qc::{U32, U64};

    check! {
        fn __divdi3(f: extern "C" fn(i64, i64) -> i64, n: U64, d: U64) -> Option<i64> {
            let (n, d) = (n.0 as i64, d.0 as i64);
            if d == 0 {
                None
            } else {
                Some(f(n, d))
            }
        }

        fn __moddi3(f: extern "C" fn(i64, i64) -> i64, n: U64, d: U64) -> Option<i64> {
            let (n, d) = (n.0 as i64, d.0 as i64);
            if d == 0 {
                None
            } else {
                Some(f(n, d))
            }
        }

        #[cfg(target_arch = "arm")]
        fn __divmoddi4(f: extern "aapcs" fn(i64, i64, &mut i64) -> i64,
                       n: U64,
                       d: U64) -> Option<(i64, i64)> {
            let (n, d) = (n.0 as i64, d.0 as i64);
            if d == 0 {
                None
            } else {
                let mut r = 0;
                let q = f(n, d, &mut r);
                Some((q, r))
            }
        }

        #[cfg(not(target_arch = "arm"))]
        fn __divmoddi4(f: extern "C" fn(i64, i64, &mut i64) -> i64,
                       n: U64,
                       d: U64) -> Option<(i64, i64)> {
            let (n, d) = (n.0 as i64, d.0 as i64);
            if d == 0 {
                None
            } else {
                let mut r = 0;
                let q = f(n, d, &mut r);
                Some((q, r))
            }
        }

        fn __divsi3(f: extern "C" fn(i32, i32) -> i32,
                    n: U32,
                    d: U32) -> Option<i32> {
            let (n, d) = (n.0 as i32, d.0 as i32);
            if d == 0 {
                None
            } else {
                Some(f(n, d))
            }
        }

        fn __modsi3(f: extern "C" fn(i32, i32) -> i32,
                    n: U32,
                    d: U32) -> Option<i32> {
            let (n, d) = (n.0 as i32, d.0 as i32);
            if d == 0 {
                None
            } else {
                Some(f(n, d))
            }
        }

        fn __divmodsi4(f: extern "C" fn(i32, i32, &mut i32) -> i32,
                       n: U32,
                       d: U32) -> Option<(i32, i32)> {
            let (n, d) = (n.0 as i32, d.0 as i32);
            if d == 0 {
                None
            } else {
                let mut r = 0;
                let q = f(n, d, &mut r);
                Some((q, r))
            }
        }
    }
}

#[cfg(test)]
#[cfg(all(not(windows),
          not(target_arch = "mips64"),
          not(target_arch = "mips64el"),
          target_pointer_width="64"))]
mod tests_i128 {
    use qc::U128;
    check! {

        fn __divti3(f: extern "C" fn(i128, i128) -> i128, n: U128, d: U128) -> Option<i128> {
            let (n, d) = (n.0 as i128, d.0 as i128);
            if d == 0 {
                None
            } else {
                Some(f(n, d))
            }
        }

        fn __modti3(f: extern "C" fn(i128, i128) -> i128, n: U128, d: U128) -> Option<i128> {
            let (n, d) = (n.0 as i128, d.0 as i128);
            if d == 0 {
                None
            } else {
                Some(f(n, d))
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::f64;
use rand::Rng;
use rand::distributions::{Sample, IndependentSample};
use statistics::*;
use distribution::{Univariate, Discrete, Distribution};
use {Result, StatsError};

/// Implements the [Categorical](https://en.wikipedia.org/wiki/Categorical_distribution)
/// distribution, also known as the generalized Bernoulli or discrete distribution
///
/// # Examples
///
/// ```
/// use statrs::distribution::{Categorical, Discrete};
/// use statrs::statistics::Mode;
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct Categorical {
    norm_pmf: Vec<f64>,
    cdf: Vec<f64>,
}

impl Categorical {
    pub fn new(prob_mass: &[f64]) -> Result<Categorical> {
        if !is_valid_prob_mass(prob_mass) {
            Err(StatsError::BadParams)
        } else {
            // extract un-normalized cdf
            let mut cdf = vec![0.0; prob_mass.len()];
            cdf[0] = prob_mass[0];
            for i in 1..prob_mass.len() {
                unsafe {
                    let val = cdf.get_unchecked(i - 1) + prob_mass.get_unchecked(i);
                    let elem = cdf.get_unchecked_mut(i);
                    *elem = val;
                }
            }

            // extract normalized probability mass
            let sum = cdf[cdf.len() - 1];
            let mut norm_pmf = vec![0.0; prob_mass.len()];
            for i in 0..prob_mass.len() {
                unsafe {
                    let elem = norm_pmf.get_unchecked_mut(i);
                    *elem = prob_mass.get_unchecked(i) / sum;
                }
            }
            Ok(Categorical {
                norm_pmf: norm_pmf,
                cdf: cdf,
            })
        }
    }

    fn cdf_max(&self) -> f64 {
        *unsafe { self.cdf.get_unchecked(self.cdf.len() - 1) }
    }
}

impl Sample<f64> for Categorical {
    /// Generate a random sample from a categorical
    /// distribution using `r` as the source of randomness.
    /// Refer [here](#method.sample-1) for implementation details
    fn sample<R: Rng>(&mut self, r: &mut R) -> f64 {
        super::Distribution::sample(self, r)
    }
}

impl IndependentSample<f64> for Categorical {
    /// Generate a random independent sample from a categorical
    /// distribution using `r` as the source of randomness.
    /// Refer [here](#method.sample-1) for implementation details
    fn ind_sample<R: Rng>(&self, r: &mut R) -> f64 {
        super::Distribution::sample(self, r)
    }
}

impl Distribution<f64> for Categorical {
    /// Generate a random sample from the categorical distribution
    /// using `r` as the source of randomness
    ///
    /// # Examples
    ///
    /// ```
    /// # extern crate rand;
    /// # extern crate statrs;
    /// use rand::StdRng;
    /// use statrs::distribution::{Categorical, Distribution};
    ///
    /// # fn main() {
    /// let mut r = rand::StdRng::new().unwrap();
    /// let n = Categorical::new(&[1.0, 2.0, 3.0]).unwrap();
    /// print!("{}", n.sample::<StdRng>(&mut r));
    /// # }
    /// ```
    fn sample<R: Rng>(&self, r: &mut R) -> f64 {
        let draw = r.next_f64() * self.cdf_max();
        let mut idx = 0;

        if draw == 0.0 {
            // skip zero-probability categories
            let mut el = unsafe { self.cdf.get_unchecked(idx) };
            while *el == 0.0 {
                // don't need bounds checking because we do not allow
                // creating Categorical distributions with all 0.0 probs
                idx += 1;
                el = unsafe { self.cdf.get_unchecked(idx) }
            }
        }
        let mut el = unsafe { self.cdf.get_unchecked(idx) };
        while draw > *el {
            idx += 1;
            el = unsafe { self.cdf.get_unchecked(idx) };
        }
        return idx as f64;
    }
}

impl Univariate<u64, f64> for Categorical {
    /// Calculates the cumulative distribution function for the categorical
    /// distribution at `x`
    ///
    /// # Panics
    ///
    /// If `x < 0.0` or `x > k` where `k` is the number of categories
    /// (i.e. the length of the `prob_mass` slice passed to the constructor)
    ///
    /// # Formula
    ///
    /// ```ignore
    /// sum(p_j) from 0..x
    /// ```
    ///
    /// where `p_j` is the probability mass for the `j`th category
    fn cdf(&self, x: f64) -> f64 {
        assert!(x >= 0.0 && x <= self.cdf.len() as f64,
                format!("{}",
                        StatsError::ArgIntervalIncl("x", 0.0, self.cdf.len() as f64)));
        if x == self.cdf.len() as f64 {
            1.0
        } else {
            unsafe { self.cdf.get_unchecked(x as usize) / self.cdf_max() }
        }
    }
}

impl Min<u64> for Categorical {
    /// Returns the minimum value in the domain of the
    /// categorical distribution representable by a 64-bit
    /// integer
    ///
    /// # Formula
    ///
    /// ```ignore
    /// 0
    /// ```
    fn min(&self) -> u64 {
        0
    }
}

impl Max<u64> for Categorical {
    /// Returns the maximum value in the domain of the
    /// categorical distribution representable by a 64-bit
    /// integer
    ///
    /// # Formula
    ///
    /// ```ignore
    /// n
    /// ```
    fn max(&self) -> u64 {
        self.cdf.len() as u64 - 1
    }
}

impl Mean<f64> for Categorical {
    /// Returns the mean of the categorical distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// sum(j * p_j) for j in 0..k-1
    /// ```
    ///
    /// where `p_j` is the `j`th probability mass and `k` is the number
    /// of categories
    fn mean(&self) -> f64 {
        self.norm_pmf.iter().enumerate().fold(0.0, |acc, (idx, &val)| acc + idx as f64 * val)
    }
}

impl Variance<f64> for Categorical {
    /// Returns the variance of the categorical distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// sum(p_j * (j - μ)^2) for j in 0..k-1
    /// ```
    ///
    /// where `p_j` is the `j`th probability mass, `k` is the number
    /// of categories, and `μ` is the mean
    fn variance(&self) -> f64 {
        let mu = self.mean();
        self.norm_pmf.iter().enumerate().fold(0.0, |acc, (idx, &val)| {
            let r = idx as f64 - mu;
            acc + r * r * val
        })
    }

    /// Returns the standard deviation of the categorical distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// sqrt(sum(p_j * (j - μ)^2)) for j in 0..k-1
    /// ```
    ///
    /// where `p_j` is the `j`th probability mass, `k` is the number
    /// of categories, and `μ` is the mean
    fn std_dev(&self) -> f64 {
        self.variance().sqrt()
    }
}

// determines if `p` is a valid probability mass array
// for the Categorical distribution
fn is_valid_prob_mass(p: &[f64]) -> bool {
    !p.iter().any(|&x| x < 0.0 || x.is_nan()) && !p.iter().all(|&x| x == 0.0)
}

#[test]
fn test_is_valid_prob_mass() {
    let invalid = [1.0, f64::NAN, 3.0];
    assert!(!is_valid_prob_mass(&invalid));
    let invalid2 = [-2.0, 5.0, 1.0, 6.2];
    assert!(!is_valid_prob_mass(&invalid2));
    let invalid3 = [0.0, 0.0, 0.0];
    assert!(!is_valid_prob_mass(&invalid3));
    let invalid4: [f64; 0] = [];
    assert!(!is_valid_prob_mass(&invalid4));
    let valid = [5.2, 0.00001, 1e-15, 1000000.12];
    assert!(is_valid_prob_mass(&valid));
}

#[cfg_attr(rustfmt, rustfmt_skip)]
#[cfg(test)]
mod test {
    use std::fmt::Debug;
    use statistics::*;
    use distribution::{Univariate, Categorical};

    fn try_create(prob_mass: &[f64]) -> Categorical {
        let n = Categorical::new(prob_mass);
        assert!(n.is_ok());
        n.unwrap()
    }

    fn create_case(prob_mass: &[f64]) {
        try_create(prob_mass);
    }

    fn bad_create_case(prob_mass: &[f64]) {
        let n = Categorical::new(prob_mass);
        assert!(n.is_err());
    }

    fn test_case<T, F>(prob_mass: &[f64], expected: T, eval: F)
        where T: PartialEq + Debug,
              F: Fn(Categorical) -> T
    {
        let n = try_create(prob_mass);
        let x = eval(n);
        assert_eq!(expected, x);
    }

    #[test]
    fn test_create() {
        create_case(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
    }

    #[test]
    fn test_bad_create() {
        bad_create_case(&[-1.0, 1.0]);
        bad_create_case(&[0.0, 0.0]);
    }

    #[test]
    fn test_mean() {
        test_case(&[0.0, 0.25, 0.5, 0.25], 2.0, |x| x.mean());
        test_case(&[0.0, 1.0, 2.0, 1.0], 2.0, |x| x.mean());
        test_case(&[0.0, 0.5, 0.5], 1.5, |x| x.mean());
        test_case(&[0.75, 0.25], 0.25, |x| x.mean());
        test_case(&[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], 5.0, |x| x.mean());
    }

    #[test]
    fn test_variance() {
        test_case(&[0.0, 0.25, 0.5, 0.25], 0.5, |x| x.variance());
        test_case(&[0.0, 1.0, 2.0, 1.0], 0.5, |x| x.variance());
        test_case(&[0.0, 0.5, 0.5], 0.25, |x| x.variance());
        test_case(&[0.75, 0.25], 0.1875, |x| x.variance());
        test_case(&[1.0, 0.0, 1.0], 1.0, |x| x.variance());
    }

    #[test]
    fn test_std_dev() {
        test_case(&[0.0, 0.25, 0.5, 0.25], 0.70710678118654752440084436210485, |x| x.std_dev());
        test_case(&[0.0, 1.0, 2.0, 1.0], 0.70710678118654752440084436210485, |x| x.std_dev());
        test_case(&[0.0, 0.5, 0.5], 0.5, |x| x.std_dev());
        test_case(&[0.75, 0.25], 0.43301270189221932338186158537647, |x| x.std_dev());
        test_case(&[1.0, 0.0, 1.0], 1.0, |x| x.std_dev());
    }

    #[test]
    fn test_min_max() {
        test_case(&[4.0, 2.5, 2.5, 1.0], 0, |x| x.min());
        test_case(&[4.0, 2.5, 2.5, 1.0], 3, |x| x.max());
    }

    #[test]
    fn test_cdf() {
        test_case(&[0.0, 3.0, 1.0, 1.0], 3.0 / 5.0, |x| x.cdf(1.5));
        test_case(&[1.0, 1.0, 1.0, 1.0], 0.25, |x| x.cdf(0.0));
        test_case(&[4.0, 2.5, 2.5, 1.0], 0.4, |x| x.cdf(0.8));
        test_case(&[4.0, 2.5, 2.5, 1.0], 1.0, |x| x.cdf(3.2));
        test_case(&[4.0, 2.5, 2.5, 1.0], 1.0, |x| x.cdf(4.0));
    }

    #[test]
    #[should_panic]
    fn test_cdf_input_low() {
        test_case(&[4.0, 2.5, 2.5, 1.0], 1.0, |x| x.cdf(-1.0));
    }

    #[test]
    #[should_panic]
    fn test_cdf_input_high() {
        test_case(&[4.0, 2.5, 2.5, 1.0], 1.0, |x| x.cdf(4.5));
    }
}<commit_msg>feat: impl Discrete for categorical<commit_after>use std::f64;
use rand::Rng;
use rand::distributions::{Sample, IndependentSample};
use statistics::*;
use distribution::{Univariate, Discrete, Distribution};
use {Result, StatsError};

/// Implements the [Categorical](https://en.wikipedia.org/wiki/Categorical_distribution)
/// distribution, also known as the generalized Bernoulli or discrete distribution
///
/// # Examples
///
/// ```
/// use statrs::distribution::{Categorical, Discrete};
/// use statrs::statistics::Mode;
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct Categorical {
    norm_pmf: Vec<f64>,
    cdf: Vec<f64>,
}

impl Categorical {
    pub fn new(prob_mass: &[f64]) -> Result<Categorical> {
        if !is_valid_prob_mass(prob_mass) {
            Err(StatsError::BadParams)
        } else {
            // extract un-normalized cdf
            let mut cdf = vec![0.0; prob_mass.len()];
            cdf[0] = prob_mass[0];
            for i in 1..prob_mass.len() {
                unsafe {
                    let val = cdf.get_unchecked(i - 1) + prob_mass.get_unchecked(i);
                    let elem = cdf.get_unchecked_mut(i);
                    *elem = val;
                }
            }

            // extract normalized probability mass
            let sum = cdf[cdf.len() - 1];
            let mut norm_pmf = vec![0.0; prob_mass.len()];
            for i in 0..prob_mass.len() {
                unsafe {
                    let elem = norm_pmf.get_unchecked_mut(i);
                    *elem = prob_mass.get_unchecked(i) / sum;
                }
            }
            Ok(Categorical {
                norm_pmf: norm_pmf,
                cdf: cdf,
            })
        }
    }

    fn cdf_max(&self) -> f64 {
        *unsafe { self.cdf.get_unchecked(self.cdf.len() - 1) }
    }
}

impl Sample<f64> for Categorical {
    /// Generate a random sample from a categorical
    /// distribution using `r` as the source of randomness.
    /// Refer [here](#method.sample-1) for implementation details
    fn sample<R: Rng>(&mut self, r: &mut R) -> f64 {
        super::Distribution::sample(self, r)
    }
}

impl IndependentSample<f64> for Categorical {
    /// Generate a random independent sample from a categorical
    /// distribution using `r` as the source of randomness.
    /// Refer [here](#method.sample-1) for implementation details
    fn ind_sample<R: Rng>(&self, r: &mut R) -> f64 {
        super::Distribution::sample(self, r)
    }
}

impl Distribution<f64> for Categorical {
    /// Generate a random sample from the categorical distribution
    /// using `r` as the source of randomness
    ///
    /// # Examples
    ///
    /// ```
    /// # extern crate rand;
    /// # extern crate statrs;
    /// use rand::StdRng;
    /// use statrs::distribution::{Categorical, Distribution};
    ///
    /// # fn main() {
    /// let mut r = rand::StdRng::new().unwrap();
    /// let n = Categorical::new(&[1.0, 2.0, 3.0]).unwrap();
    /// print!("{}", n.sample::<StdRng>(&mut r));
    /// # }
    /// ```
    fn sample<R: Rng>(&self, r: &mut R) -> f64 {
        let draw = r.next_f64() * self.cdf_max();
        let mut idx = 0;

        if draw == 0.0 {
            // skip zero-probability categories
            let mut el = unsafe { self.cdf.get_unchecked(idx) };
            while *el == 0.0 {
                // don't need bounds checking because we do not allow
                // creating Categorical distributions with all 0.0 probs
                idx += 1;
                el = unsafe { self.cdf.get_unchecked(idx) }
            }
        }
        let mut el = unsafe { self.cdf.get_unchecked(idx) };
        while draw > *el {
            idx += 1;
            el = unsafe { self.cdf.get_unchecked(idx) };
        }
        return idx as f64;
    }
}

impl Univariate<u64, f64> for Categorical {
    /// Calculates the cumulative distribution function for the categorical
    /// distribution at `x`
    ///
    /// # Panics
    ///
    /// If `x < 0.0` or `x > k` where `k` is the number of categories
    /// (i.e. the length of the `prob_mass` slice passed to the constructor)
    ///
    /// # Formula
    ///
    /// ```ignore
    /// sum(p_j) from 0..x
    /// ```
    ///
    /// where `p_j` is the probability mass for the `j`th category
    fn cdf(&self, x: f64) -> f64 {
        assert!(x >= 0.0 && x <= self.cdf.len() as f64,
                format!("{}",
                        StatsError::ArgIntervalIncl("x", 0.0, self.cdf.len() as f64)));
        if x == self.cdf.len() as f64 {
            1.0
        } else {
            unsafe { self.cdf.get_unchecked(x as usize) / self.cdf_max() }
        }
    }
}

impl Min<u64> for Categorical {
    /// Returns the minimum value in the domain of the
    /// categorical distribution representable by a 64-bit
    /// integer
    ///
    /// # Formula
    ///
    /// ```ignore
    /// 0
    /// ```
    fn min(&self) -> u64 {
        0
    }
}

impl Max<u64> for Categorical {
    /// Returns the maximum value in the domain of the
    /// categorical distribution representable by a 64-bit
    /// integer
    ///
    /// # Formula
    ///
    /// ```ignore
    /// n
    /// ```
    fn max(&self) -> u64 {
        self.cdf.len() as u64 - 1
    }
}

impl Mean<f64> for Categorical {
    /// Returns the mean of the categorical distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// sum(j * p_j) for j in 0..k-1
    /// ```
    ///
    /// where `p_j` is the `j`th probability mass and `k` is the number
    /// of categories
    fn mean(&self) -> f64 {
        self.norm_pmf.iter().enumerate().fold(0.0, |acc, (idx, &val)| acc + idx as f64 * val)
    }
}

impl Variance<f64> for Categorical {
    /// Returns the variance of the categorical distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// sum(p_j * (j - μ)^2) for j in 0..k-1
    /// ```
    ///
    /// where `p_j` is the `j`th probability mass, `k` is the number
    /// of categories, and `μ` is the mean
    fn variance(&self) -> f64 {
        let mu = self.mean();
        self.norm_pmf.iter().enumerate().fold(0.0, |acc, (idx, &val)| {
            let r = idx as f64 - mu;
            acc + r * r * val
        })
    }

    /// Returns the standard deviation of the categorical distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// sqrt(sum(p_j * (j - μ)^2)) for j in 0..k-1
    /// ```
    ///
    /// where `p_j` is the `j`th probability mass, `k` is the number
    /// of categories, and `μ` is the mean
    fn std_dev(&self) -> f64 {
        self.variance().sqrt()
    }
}

impl Discrete<u64, f64> for Categorical {
    /// Calculates the probability mass function for the categorical
    /// distribution at `x`
    ///
    /// # Panics
    ///
    /// If `x >= k` where `k` is the number of categories
    ///
    /// # Formula
    ///
    /// ```ignore
    /// p_x
    /// ```
    fn pmf(&self, x: u64) -> f64 {
        assert!(x < self.norm_pmf.len() as u64, format!("{}", StatsError::ArgLtArg("x", "k")));
        unsafe { *self.norm_pmf.get_unchecked(x as usize) }
    }

    /// Calculates the log probability mass function for the categorical
    /// distribution at `x`
    fn ln_pmf(&self, x: u64) -> f64 {
        self.pmf(x).ln()
    }
}

// determines if `p` is a valid probability mass array
// for the Categorical distribution
fn is_valid_prob_mass(p: &[f64]) -> bool {
    !p.iter().any(|&x| x < 0.0 || x.is_nan()) && !p.iter().all(|&x| x == 0.0)
}

#[test]
fn test_is_valid_prob_mass() {
    let invalid = [1.0, f64::NAN, 3.0];
    assert!(!is_valid_prob_mass(&invalid));
    let invalid2 = [-2.0, 5.0, 1.0, 6.2];
    assert!(!is_valid_prob_mass(&invalid2));
    let invalid3 = [0.0, 0.0, 0.0];
    assert!(!is_valid_prob_mass(&invalid3));
    let invalid4: [f64; 0] = [];
    assert!(!is_valid_prob_mass(&invalid4));
    let valid = [5.2, 0.00001, 1e-15, 1000000.12];
    assert!(is_valid_prob_mass(&valid));
}

#[cfg_attr(rustfmt, rustfmt_skip)]
#[cfg(test)]
mod test {
    use std::fmt::Debug;
    use statistics::*;
    use distribution::{Univariate, Discrete, Categorical};

    fn try_create(prob_mass: &[f64]) -> Categorical {
        let n = Categorical::new(prob_mass);
        assert!(n.is_ok());
        n.unwrap()
    }

    fn create_case(prob_mass: &[f64]) {
        try_create(prob_mass);
    }

    fn bad_create_case(prob_mass: &[f64]) {
        let n = Categorical::new(prob_mass);
        assert!(n.is_err());
    }

    fn get_value<T, F>(prob_mass: &[f64], eval: F) -> T
        where T: PartialEq + Debug,
              F: Fn(Categorical) -> T
    {
        let n = try_create(prob_mass);
        eval(n)
    }

    fn test_case<T, F>(prob_mass: &[f64], expected: T, eval: F)
        where T: PartialEq + Debug,
              F: Fn(Categorical) -> T
    {
        let x = get_value(prob_mass, eval);
        assert_eq!(expected, x);
    }

    #[test]
    fn test_create() {
        create_case(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
    }

    #[test]
    fn test_bad_create() {
        bad_create_case(&[-1.0, 1.0]);
        bad_create_case(&[0.0, 0.0]);
    }

    #[test]
    fn test_mean() {
        test_case(&[0.0, 0.25, 0.5, 0.25], 2.0, |x| x.mean());
        test_case(&[0.0, 1.0, 2.0, 1.0], 2.0, |x| x.mean());
        test_case(&[0.0, 0.5, 0.5], 1.5, |x| x.mean());
        test_case(&[0.75, 0.25], 0.25, |x| x.mean());
        test_case(&[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], 5.0, |x| x.mean());
    }

    #[test]
    fn test_variance() {
        test_case(&[0.0, 0.25, 0.5, 0.25], 0.5, |x| x.variance());
        test_case(&[0.0, 1.0, 2.0, 1.0], 0.5, |x| x.variance());
        test_case(&[0.0, 0.5, 0.5], 0.25, |x| x.variance());
        test_case(&[0.75, 0.25], 0.1875, |x| x.variance());
        test_case(&[1.0, 0.0, 1.0], 1.0, |x| x.variance());
    }

    #[test]
    fn test_std_dev() {
        test_case(&[0.0, 0.25, 0.5, 0.25], 0.70710678118654752440084436210485, |x| x.std_dev());
        test_case(&[0.0, 1.0, 2.0, 1.0], 0.70710678118654752440084436210485, |x| x.std_dev());
        test_case(&[0.0, 0.5, 0.5], 0.5, |x| x.std_dev());
        test_case(&[0.75, 0.25], 0.43301270189221932338186158537647, |x| x.std_dev());
        test_case(&[1.0, 0.0, 1.0], 1.0, |x| x.std_dev());
    }

    #[test]
    fn test_min_max() {
        test_case(&[4.0, 2.5, 2.5, 1.0], 0, |x| x.min());
        test_case(&[4.0, 2.5, 2.5, 1.0], 3, |x| x.max());
    }

    #[test]
    fn test_pmf() {
        test_case(&[0.0, 0.25, 0.5, 0.25], 0.0, |x| x.pmf(0));
        test_case(&[0.0, 0.25, 0.5, 0.25], 0.25, |x| x.pmf(1));
        test_case(&[0.0, 0.25, 0.5, 0.25], 0.25, |x| x.pmf(3));
    }

    #[test]
    #[should_panic]
    fn test_pmf_x_too_high() {
        get_value(&[4.0, 2.5, 2.5, 1.0], |x| x.pmf(4));
    }

    #[test]
    fn test_ln_pmf() {
        test_case(&[0.0, 0.25, 0.5, 0.25], 0f64.ln(), |x| x.ln_pmf(0));
        test_case(&[0.0, 0.25, 0.5, 0.25], 0.25f64.ln(), |x| x.ln_pmf(1));
        test_case(&[0.0, 0.25, 0.5, 0.25], 0.25f64.ln(), |x| x.ln_pmf(3));
    }

    #[test]
    #[should_panic]
    fn test_ln_pmf_x_too_high() {
        get_value(&[4.0, 2.5, 2.5, 1.0], |x| x.ln_pmf(4));
    }

    #[test]
    fn test_cdf() {
        test_case(&[0.0, 3.0, 1.0, 1.0], 3.0 / 5.0, |x| x.cdf(1.5));
        test_case(&[1.0, 1.0, 1.0, 1.0], 0.25, |x| x.cdf(0.0));
        test_case(&[4.0, 2.5, 2.5, 1.0], 0.4, |x| x.cdf(0.8));
        test_case(&[4.0, 2.5, 2.5, 1.0], 1.0, |x| x.cdf(3.2));
        test_case(&[4.0, 2.5, 2.5, 1.0], 1.0, |x| x.cdf(4.0));
    }

    #[test]
    #[should_panic]
    fn test_cdf_input_low() {
        get_value(&[4.0, 2.5, 2.5, 1.0], |x| x.cdf(-1.0));
    }

    #[test]
    #[should_panic]
    fn test_cdf_input_high() {
        get_value(&[4.0, 2.5, 2.5, 1.0], |x| x.cdf(4.5));
    }
}<|endoftext|>


<commit_before>//! Plugin which issues vault tokens to services.

use docker_compose::v2 as dc;
use serde_yaml;
use std::collections::BTreeMap;
use std::env;
use std::fmt::Debug;
use std::fs;
use std::path::PathBuf;
#[cfg(test)]
use std::cell::{Ref, RefCell};
use vault;
use vault::client::VaultDuration;

use plugins;
use plugins::transform::Operation;
use plugins::transform::{Plugin as PluginTransform, PluginNew};
use project::Project;
use util::Error;

#[cfg(feature = "serde_macros")]
include!(concat!("vault_config.in.rs"));
#[cfg(feature = "serde_codegen")]
include!(concat!(env!("OUT_DIR"), "/plugins/transform/vault_config.rs"));

/// The "environment" in which to interpret a configuration file.  We don't
/// want to use the OS environment variables, but rather a fake environment
/// with a few carefully selected values.
#[derive(Debug)]
struct ConfigEnvironment<'a> {
    /// The context for our transformation, including the project, pod,
    /// etc.
    ctx: &'a plugins::Context<'a>,
    /// The name of the current service.
    service: &'a str,
}

impl<'a> dc::Environment for ConfigEnvironment<'a> {
    fn var(&self, key: &str) -> Result<String, env::VarError> {
        let result = match key {
            "PROJECT" => Ok(self.ctx.project.name()),
            "OVERRIDE" => Ok(self.ctx.ovr.name()),
            "POD" => Ok(self.ctx.pod.name()),
            "SERVICE" => Ok(self.service),
            _ => Err(env::VarError::NotPresent),
        };
        result.map(|s| s.to_owned())
    }
}

/// An abstract interface to Vault's token-generation capabilities.  We use
/// this to mock vault during tests.
trait GenerateToken: Debug {
    /// Get a `VAULT_ADDR` value to use along with this token.
    fn addr(&self) -> &str;
    /// Generate a token with the specified parameters.
    fn generate_token(&self,
                      display_name: &str,
                      policies: &[&str],
                      ttl: VaultDuration)
                      -> Result<String, Error>;
}

/// A fake interface to vault for testing purposes.
#[derive(Debug)]
#[cfg(test)]
struct MockVault {
    /// The tokens we were asked to generate.  We store these in a RefCell
    /// so that we can have "interior" mutability, because we don't want
    /// `generate_token` to be `&mut self` in the general case.
    calls: RefCell<Vec<(String, Vec<String>, VaultDuration)>>,
}

#[cfg(test)]
impl MockVault {
    /// Create a new MockVault.
    fn new() -> MockVault {
        MockVault { calls: RefCell::new(vec![]) }
    }

    /// Return the calls that were made to our MockVault.
    fn calls(&self) -> Ref<Vec<(String, Vec<String>, VaultDuration)>> {
        self.calls.borrow()
    }
}

#[cfg(test)]
impl GenerateToken for MockVault {
    fn addr(&self) -> &str {
        "http://example.com:8200/"
    }

    fn generate_token(&self,
                      display_name: &str,
                      policies: &[&str],
                      ttl: VaultDuration)
                      -> Result<String, Error> {
        let saved_policies = policies.iter()
            .cloned()
            .map(|p| p.to_owned())
            .collect();
        self.calls.borrow_mut().push((display_name.to_owned(), saved_policies, ttl));
        Ok("fake_token".to_owned())
    }
}

/// An interface to an actual vault server.
#[derive(Debug)]
struct Vault {
    /// The address of our vault server.
    addr: String,
    /// The master token that we'll use to issue new tokens.
    token: String,
}

impl Vault {
    /// Create a new vault client.
    fn new() -> Result<Vault, Error> {
        let addr = try!(env::var("VAULT_ADDR"));
        let token = try!(env::var("VAULT_MASTER_TOKEN"));
        Ok(Vault {
            addr: addr,
            token: token,
        })
    }
}

impl GenerateToken for Vault {
    fn addr(&self) -> &str {
        &self.addr
    }

    fn generate_token(&self,
                      display_name: &str,
                      policies: &[&str],
                      ttl: VaultDuration)
                      -> Result<String, Error> {
        // We can't store `client` in `self`, because it has some obnoxious
        // lifetime parameters.  So we'll just recreate it.  This is
        // probably not the worst idea, because it uses `hyper` for HTTP,
        // and `hyper` HTTP connections used to have expiration issues that
        // were tricky for clients to deal with correctly.
        let client = try!(vault::Client::new(&self.addr, &self.token));
        unimplemented!()
    }
}

/// Issues `VAULT_TOKEN` values to services.
#[derive(Debug)]
pub struct Plugin {
    /// Our `config/vault.yml` YAML file, parsed and read into memory.
    config: Config,
    /// Our source of tokens.
    generator: Box<GenerateToken>,
}

impl Plugin {
    /// Get the path to this plugin's config file.
    fn config_path(project: &Project) -> PathBuf {
        project.root_dir().join("config/vault.yml")
    }

    /// Create a new plugin, specifying an alternate source for tokens.
    fn new_with_generator<G>(project: &Project, generator: G) -> Result<Plugin, Error>
        where G: GenerateToken + 'static
    {
        let f = try!(fs::File::open(&Self::config_path(project)));
        let config = try!(serde_yaml::from_reader(f));
        Ok(Plugin {
            config: config,
            generator: Box::new(generator),
        })
    }
}

impl PluginTransform for Plugin {
    fn name(&self) -> &'static str {
        "vault"
    }

    fn transform(&self,
                 _op: Operation,
                 ctx: &plugins::Context,
                 file: &mut dc::File)
                 -> Result<(), Error> {

        for (name, mut service) in &mut file.services {
            // Set up a ConfigEnvironment that we can use to perform
            // interpolations of values like `$SERVICE` in our config file.
            let env = ConfigEnvironment {
                ctx: ctx,
                service: name,
            };

            // Insert our VAULT_ADDR value into the generated files.
            service.environment
                .insert("VAULT_ADDR".to_owned(), self.generator.addr().to_owned());

            // Generate a VAULT_TOKEN.
            let token = try!(self.generator
                .generate_token("TODO", &[], VaultDuration::hours(1)));
            service.environment.insert("VAULT_TOKEN".to_owned(), token);

            // Add in any extra environment variables.
            for (var, raw_val) in &self.config.extra_environment {
                let mut val = raw_val.to_owned();
                let new_val = try!(val.interpolate_env(&env)).to_owned();
                service.environment.insert(var.to_owned(), new_val);
            }
        }
        Ok(())
    }
}

impl PluginNew for Plugin {
    /// Should we enable this plugin for this project?
    fn should_enable_for(project: &Project) -> Result<bool, Error> {
        let path = Self::config_path(project);
        Ok(path.exists())
    }

    /// Create a new plugin.
    fn new(project: &Project) -> Result<Self, Error> {
        Self::new_with_generator(project, try!(Vault::new()))
    }
}

#[test]
fn interpolates_policies() {
    use env_logger;
    let _ = env_logger::init();
    let proj = Project::from_example("rails_hello").unwrap();
    let ovr = proj.ovr("production").unwrap();

    let vault = MockVault::new();
    let plugin = Plugin::new_with_generator(&proj, vault).unwrap();

    let frontend = proj.pod("frontend").unwrap();
    let ctx = plugins::Context::new(&proj, ovr, frontend);
    let mut file = frontend.merged_file(ovr).unwrap();
    plugin.transform(Operation::Output, &ctx, &mut file).unwrap();
    let web = file.services.get("web").unwrap();
    assert_eq!(web.environment.get("VAULT_ADDR").expect("has VAULT_ADDR"),
               "http://example.com:8200/");
    assert_eq!(web.environment.get("VAULT_TOKEN").expect("has VAULT_TOKEN"),
               "fake_token");
    assert_eq!(web.environment.get("VAULT_ENV").expect("has VAULT_ENV"),
               "production");

    // TODO: Check policies used to issue tokens.
}
<commit_msg>vault: Generate new tokens using vault<commit_after>//! Plugin which issues vault tokens to services.

use docker_compose::v2 as dc;
use serde_yaml;
use std::collections::BTreeMap;
use std::env;
use std::fmt::Debug;
use std::fs;
use std::path::PathBuf;
#[cfg(test)]
use std::cell::{Ref, RefCell};
use vault;
use vault::client::VaultDuration;

use plugins;
use plugins::transform::Operation;
use plugins::transform::{Plugin as PluginTransform, PluginNew};
use project::Project;
use util::Error;

#[cfg(feature = "serde_macros")]
include!(concat!("vault_config.in.rs"));
#[cfg(feature = "serde_codegen")]
include!(concat!(env!("OUT_DIR"), "/plugins/transform/vault_config.rs"));

/// The "environment" in which to interpret a configuration file.  We don't
/// want to use the OS environment variables, but rather a fake environment
/// with a few carefully selected values.
#[derive(Debug)]
struct ConfigEnvironment<'a> {
    /// The context for our transformation, including the project, pod,
    /// etc.
    ctx: &'a plugins::Context<'a>,
    /// The name of the current service.
    service: &'a str,
}

impl<'a> dc::Environment for ConfigEnvironment<'a> {
    fn var(&self, key: &str) -> Result<String, env::VarError> {
        let result = match key {
            "PROJECT" => Ok(self.ctx.project.name()),
            "OVERRIDE" => Ok(self.ctx.ovr.name()),
            "POD" => Ok(self.ctx.pod.name()),
            "SERVICE" => Ok(self.service),
            _ => Err(env::VarError::NotPresent),
        };
        result.map(|s| s.to_owned())
    }
}

/// An abstract interface to Vault's token-generation capabilities.  We use
/// this to mock vault during tests.
trait GenerateToken: Debug {
    /// Get a `VAULT_ADDR` value to use along with this token.
    fn addr(&self) -> &str;
    /// Generate a token with the specified parameters.
    fn generate_token(&self,
                      display_name: &str,
                      policies: &[&str],
                      ttl: VaultDuration)
                      -> Result<String, Error>;
}

/// A fake interface to vault for testing purposes.
#[derive(Debug)]
#[cfg(test)]
struct MockVault {
    /// The tokens we were asked to generate.  We store these in a RefCell
    /// so that we can have "interior" mutability, because we don't want
    /// `generate_token` to be `&mut self` in the general case.
    calls: RefCell<Vec<(String, Vec<String>, VaultDuration)>>,
}

#[cfg(test)]
impl MockVault {
    /// Create a new MockVault.
    fn new() -> MockVault {
        MockVault { calls: RefCell::new(vec![]) }
    }

    /// Return the calls that were made to our MockVault.
    fn calls(&self) -> Ref<Vec<(String, Vec<String>, VaultDuration)>> {
        self.calls.borrow()
    }
}

#[cfg(test)]
impl GenerateToken for MockVault {
    fn addr(&self) -> &str {
        "http://example.com:8200/"
    }

    fn generate_token(&self,
                      display_name: &str,
                      policies: &[&str],
                      ttl: VaultDuration)
                      -> Result<String, Error> {
        let saved_policies = policies.iter()
            .cloned()
            .map(|p| p.to_owned())
            .collect();
        self.calls.borrow_mut().push((display_name.to_owned(), saved_policies, ttl));
        Ok("fake_token".to_owned())
    }
}

/// An interface to an actual vault server.
#[derive(Debug)]
struct Vault {
    /// The address of our vault server.
    addr: String,
    /// The master token that we'll use to issue new tokens.
    token: String,
}

impl Vault {
    /// Create a new vault client.
    fn new() -> Result<Vault, Error> {
        let addr = try!(env::var("VAULT_ADDR"));
        let token = try!(env::var("VAULT_MASTER_TOKEN"));
        Ok(Vault {
            addr: addr,
            token: token,
        })
    }
}

impl GenerateToken for Vault {
    fn addr(&self) -> &str {
        &self.addr
    }

    fn generate_token(&self,
                      display_name: &str,
                      policies: &[&str],
                      ttl: VaultDuration)
                      -> Result<String, Error> {
        // We can't store `client` in `self`, because it has some obnoxious
        // lifetime parameters.  So we'll just recreate it.  This is
        // probably not the worst idea, because it uses `hyper` for HTTP,
        // and `hyper` HTTP connections used to have expiration issues that
        // were tricky for clients to deal with correctly.
        let client = try!(vault::Client::new(&self.addr, &self.token));
        let opts = vault::client::TokenOptions::default()
            .display_name(display_name)
            .renewable(false)
            .ttl(ttl)
            .policies(policies.to_owned());
        let auth = try!(client.create_token(&opts));
        Ok(auth.client_token)
    }
}

/// Issues `VAULT_TOKEN` values to services.
#[derive(Debug)]
pub struct Plugin {
    /// Our `config/vault.yml` YAML file, parsed and read into memory.
    config: Config,
    /// Our source of tokens.
    generator: Box<GenerateToken>,
}

impl Plugin {
    /// Get the path to this plugin's config file.
    fn config_path(project: &Project) -> PathBuf {
        project.root_dir().join("config/vault.yml")
    }

    /// Create a new plugin, specifying an alternate source for tokens.
    fn new_with_generator<G>(project: &Project, generator: G) -> Result<Plugin, Error>
        where G: GenerateToken + 'static
    {
        let f = try!(fs::File::open(&Self::config_path(project)));
        let config = try!(serde_yaml::from_reader(f));
        Ok(Plugin {
            config: config,
            generator: Box::new(generator),
        })
    }
}

impl PluginTransform for Plugin {
    fn name(&self) -> &'static str {
        "vault"
    }

    fn transform(&self,
                 _op: Operation,
                 ctx: &plugins::Context,
                 file: &mut dc::File)
                 -> Result<(), Error> {

        for (name, mut service) in &mut file.services {
            // Set up a ConfigEnvironment that we can use to perform
            // interpolations of values like `$SERVICE` in our config file.
            let env = ConfigEnvironment {
                ctx: ctx,
                service: name,
            };

            // Insert our VAULT_ADDR value into the generated files.
            service.environment
                .insert("VAULT_ADDR".to_owned(), self.generator.addr().to_owned());

            // Generate a VAULT_TOKEN.
            let token = try!(self.generator
                .generate_token("TODO", &[], VaultDuration::hours(1)));
            service.environment.insert("VAULT_TOKEN".to_owned(), token);

            // Add in any extra environment variables.
            for (var, raw_val) in &self.config.extra_environment {
                let mut val = raw_val.to_owned();
                let new_val = try!(val.interpolate_env(&env)).to_owned();
                service.environment.insert(var.to_owned(), new_val);
            }
        }
        Ok(())
    }
}

impl PluginNew for Plugin {
    /// Should we enable this plugin for this project?
    fn should_enable_for(project: &Project) -> Result<bool, Error> {
        let path = Self::config_path(project);
        Ok(path.exists())
    }

    /// Create a new plugin.
    fn new(project: &Project) -> Result<Self, Error> {
        Self::new_with_generator(project, try!(Vault::new()))
    }
}

#[test]
fn interpolates_policies() {
    use env_logger;
    let _ = env_logger::init();
    let proj = Project::from_example("rails_hello").unwrap();
    let ovr = proj.ovr("production").unwrap();

    let vault = MockVault::new();
    let plugin = Plugin::new_with_generator(&proj, vault).unwrap();

    let frontend = proj.pod("frontend").unwrap();
    let ctx = plugins::Context::new(&proj, ovr, frontend);
    let mut file = frontend.merged_file(ovr).unwrap();
    plugin.transform(Operation::Output, &ctx, &mut file).unwrap();
    let web = file.services.get("web").unwrap();
    assert_eq!(web.environment.get("VAULT_ADDR").expect("has VAULT_ADDR"),
               "http://example.com:8200/");
    assert_eq!(web.environment.get("VAULT_TOKEN").expect("has VAULT_TOKEN"),
               "fake_token");
    assert_eq!(web.environment.get("VAULT_ENV").expect("has VAULT_ENV"),
               "production");

    // TODO: Check policies used to issue tokens.
}
<|endoftext|>


<commit_before>use reexport::*;
use rustc::lint::*;
use rustc::hir::def::Def;
use rustc::hir::*;
use rustc::hir::intravisit::{Visitor, walk_ty, walk_ty_param_bound, walk_fn_decl, walk_generics};
use std::collections::{HashSet, HashMap};
use syntax::codemap::Span;
use utils::{in_external_macro, span_lint};

/// **What it does:** This lint checks for lifetime annotations which can be removed by relying on lifetime elision.
///
/// **Why is this bad?** The additional lifetimes make the code look more complicated, while there is nothing out of the ordinary going on. Removing them leads to more readable code.
///
/// **Known problems:** Potential false negatives: we bail out if the function has a `where` clause where lifetimes are mentioned.
///
/// **Example:** `fn in_and_out<'a>(x: &'a u8, y: u8) -> &'a u8 { x }`
declare_lint! {
    pub NEEDLESS_LIFETIMES,
    Warn,
    "using explicit lifetimes for references in function arguments when elision rules \
     would allow omitting them"
}

/// **What it does:** This lint checks for lifetimes in generics that are never used anywhere else.
///
/// **Why is this bad?** The additional lifetimes make the code look more complicated, while there is nothing out of the ordinary going on. Removing them leads to more readable code.
///
/// **Known problems:** None
///
/// **Example:** `fn unused_lifetime<'a>(x: u8) { .. }`
declare_lint! {
    pub UNUSED_LIFETIMES,
    Warn,
    "unused lifetimes in function definitions"
}

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

impl LintPass for LifetimePass {
    fn get_lints(&self) -> LintArray {
        lint_array!(NEEDLESS_LIFETIMES, UNUSED_LIFETIMES)
    }
}

impl LateLintPass for LifetimePass {
    fn check_item(&mut self, cx: &LateContext, item: &Item) {
        if let ItemFn(ref decl, _, _, _, ref generics, _) = item.node {
            check_fn_inner(cx, decl, generics, item.span);
        }
    }

    fn check_impl_item(&mut self, cx: &LateContext, item: &ImplItem) {
        if let ImplItemKind::Method(ref sig, _) = item.node {
            check_fn_inner(cx, &sig.decl, &sig.generics, item.span);
        }
    }

    fn check_trait_item(&mut self, cx: &LateContext, item: &TraitItem) {
        if let MethodTraitItem(ref sig, _) = item.node {
            check_fn_inner(cx, &sig.decl, &sig.generics, item.span);
        }
    }
}

/// The lifetime of a &-reference.
#[derive(PartialEq, Eq, Hash, Debug)]
enum RefLt {
    Unnamed,
    Static,
    Named(Name),
}

fn bound_lifetimes(bound: &TyParamBound) -> Option<HirVec<&Lifetime>> {
    if let TraitTyParamBound(ref trait_ref, _) = *bound {
        let lt = trait_ref.trait_ref
                          .path
                          .segments
                          .last()
                          .expect("a path must have at least one segment")
                          .parameters
                          .lifetimes();

        Some(lt)
    } else {
        None
    }
}

fn check_fn_inner(cx: &LateContext, decl: &FnDecl, generics: &Generics, span: Span) {
    if in_external_macro(cx, span) || has_where_lifetimes(cx, &generics.where_clause) {
        return;
    }

    let bounds_lts = generics.ty_params
                             .iter()
                             .flat_map(|ref typ| typ.bounds.iter().filter_map(bound_lifetimes).flat_map(|lts| lts));

    if could_use_elision(cx, decl, &generics.lifetimes, bounds_lts) {
        span_lint(cx,
                  NEEDLESS_LIFETIMES,
                  span,
                  "explicit lifetimes given in parameter types where they could be elided");
    }
    report_extra_lifetimes(cx, decl, generics);
}

fn could_use_elision<'a, T: Iterator<Item = &'a Lifetime>>(cx: &LateContext, func: &FnDecl,
                                                           named_lts: &[LifetimeDef], bounds_lts: T)
                                                           -> bool {
    // There are two scenarios where elision works:
    // * no output references, all input references have different LT
    // * output references, exactly one input reference with same LT
    // All lifetimes must be unnamed, 'static or defined without bounds on the
    // level of the current item.

    // check named LTs
    let allowed_lts = allowed_lts_from(named_lts);

    // these will collect all the lifetimes for references in arg/return types
    let mut input_visitor = RefVisitor::new(cx);
    let mut output_visitor = RefVisitor::new(cx);

    // extract lifetimes in input argument types
    for arg in &func.inputs {
        input_visitor.visit_ty(&arg.ty);
    }
    // extract lifetimes in output type
    if let Return(ref ty) = func.output {
        output_visitor.visit_ty(ty);
    }

    let input_lts = lts_from_bounds(input_visitor.into_vec(), bounds_lts);
    let output_lts = output_visitor.into_vec();

    // check for lifetimes from higher scopes
    for lt in input_lts.iter().chain(output_lts.iter()) {
        if !allowed_lts.contains(lt) {
            return false;
        }
    }

    // no input lifetimes? easy case!
    if input_lts.is_empty() {
        false
    } else if output_lts.is_empty() {
        // no output lifetimes, check distinctness of input lifetimes

        // only unnamed and static, ok
        if input_lts.iter().all(|lt| *lt == RefLt::Unnamed || *lt == RefLt::Static) {
            return false;
        }
        // we have no output reference, so we only need all distinct lifetimes
        input_lts.len() == unique_lifetimes(&input_lts)
    } else {
        // we have output references, so we need one input reference,
        // and all output lifetimes must be the same
        if unique_lifetimes(&output_lts) > 1 {
            return false;
        }
        if input_lts.len() == 1 {
            match (&input_lts[0], &output_lts[0]) {
                (&RefLt::Named(n1), &RefLt::Named(n2)) if n1 == n2 => true,
                (&RefLt::Named(_), &RefLt::Unnamed) => true,
                _ => false, // already elided, different named lifetimes
                // or something static going on
            }
        } else {
            false
        }
    }
}

fn allowed_lts_from(named_lts: &[LifetimeDef]) -> HashSet<RefLt> {
    let mut allowed_lts = HashSet::new();
    for lt in named_lts {
        if lt.bounds.is_empty() {
            allowed_lts.insert(RefLt::Named(lt.lifetime.name));
        }
    }
    allowed_lts.insert(RefLt::Unnamed);
    allowed_lts.insert(RefLt::Static);
    allowed_lts
}

fn lts_from_bounds<'a, T: Iterator<Item = &'a Lifetime>>(mut vec: Vec<RefLt>, bounds_lts: T) -> Vec<RefLt> {
    for lt in bounds_lts {
        if lt.name.as_str() != "'static" {
            vec.push(RefLt::Named(lt.name));
        }
    }

    vec
}

/// Number of unique lifetimes in the given vector.
fn unique_lifetimes(lts: &[RefLt]) -> usize {
    lts.iter().collect::<HashSet<_>>().len()
}

/// A visitor usable for `rustc_front::visit::walk_ty()`.
struct RefVisitor<'v, 't: 'v> {
    cx: &'v LateContext<'v, 't>,
    lts: Vec<RefLt>,
}

impl<'v, 't> RefVisitor<'v, 't> {
    fn new(cx: &'v LateContext<'v, 't>) -> RefVisitor<'v, 't> {
        RefVisitor {
            cx: cx,
            lts: Vec::new(),
        }
    }

    fn record(&mut self, lifetime: &Option<Lifetime>) {
        if let Some(ref lt) = *lifetime {
            if lt.name.as_str() == "'static" {
                self.lts.push(RefLt::Static);
            } else {
                self.lts.push(RefLt::Named(lt.name));
            }
        } else {
            self.lts.push(RefLt::Unnamed);
        }
    }

    fn into_vec(self) -> Vec<RefLt> {
        self.lts
    }

    fn collect_anonymous_lifetimes(&mut self, path: &Path, ty: &Ty) {
        let last_path_segment = path.segments.last().map(|s| &s.parameters);
        if let Some(&AngleBracketedParameters(ref params)) = last_path_segment {
            if params.lifetimes.is_empty() {
                if let Some(def) = self.cx.tcx.def_map.borrow().get(&ty.id).map(|r| r.full_def()) {
                    match def {
                        Def::TyAlias(def_id) |
                        Def::Struct(def_id) => {
                            let type_scheme = self.cx.tcx.lookup_item_type(def_id);
                            for _ in type_scheme.generics.regions.as_slice() {
                                self.record(&None);
                            }
                        }
                        Def::Trait(def_id) => {
                            let trait_def = self.cx.tcx.trait_defs.borrow()[&def_id];
                            for _ in &trait_def.generics.regions {
                                self.record(&None);
                            }
                        }
                        _ => (),
                    }
                }
            }
        }
    }
}

impl<'v, 't> Visitor<'v> for RefVisitor<'v, 't> {
    // for lifetimes as parameters of generics
    fn visit_lifetime(&mut self, lifetime: &'v Lifetime) {
        self.record(&Some(*lifetime));
    }

    fn visit_ty(&mut self, ty: &'v Ty) {
        match ty.node {
            TyRptr(None, _) => {
                self.record(&None);
            }
            TyPath(_, ref path) => {
                self.collect_anonymous_lifetimes(path, ty);
            }
            _ => (),
        }
        walk_ty(self, ty);
    }
}

/// Are any lifetimes mentioned in the `where` clause? If yes, we don't try to
/// reason about elision.
fn has_where_lifetimes(cx: &LateContext, where_clause: &WhereClause) -> bool {
    for predicate in &where_clause.predicates {
        match *predicate {
            WherePredicate::RegionPredicate(..) => return true,
            WherePredicate::BoundPredicate(ref pred) => {
                // a predicate like F: Trait or F: for<'a> Trait<'a>
                let mut visitor = RefVisitor::new(cx);
                // walk the type F, it may not contain LT refs
                walk_ty(&mut visitor, &pred.bounded_ty);
                if !visitor.lts.is_empty() {
                    return true;
                }
                // if the bounds define new lifetimes, they are fine to occur
                let allowed_lts = allowed_lts_from(&pred.bound_lifetimes);
                // now walk the bounds
                for bound in pred.bounds.iter() {
                    walk_ty_param_bound(&mut visitor, bound);
                }
                // and check that all lifetimes are allowed
                for lt in visitor.into_vec() {
                    if !allowed_lts.contains(&lt) {
                        return true;
                    }
                }
            }
            WherePredicate::EqPredicate(ref pred) => {
                let mut visitor = RefVisitor::new(cx);
                walk_ty(&mut visitor, &pred.ty);
                if !visitor.lts.is_empty() {
                    return true;
                }
            }
        }
    }
    false
}

struct LifetimeChecker(HashMap<Name, Span>);

impl<'v> Visitor<'v> for LifetimeChecker {
    // for lifetimes as parameters of generics
    fn visit_lifetime(&mut self, lifetime: &'v Lifetime) {
        self.0.remove(&lifetime.name);
    }

    fn visit_lifetime_def(&mut self, _: &'v LifetimeDef) {
        // don't actually visit `<'a>` or `<'a: 'b>`
        // we've already visited the `'a` declarations and
        // don't want to spuriously remove them
        // `'b` in `'a: 'b` is useless unless used elsewhere in
        // a non-lifetime bound
    }
}

fn report_extra_lifetimes(cx: &LateContext, func: &FnDecl, generics: &Generics) {
    let hs = generics.lifetimes
                     .iter()
                     .map(|lt| (lt.lifetime.name, lt.lifetime.span))
                     .collect();
    let mut checker = LifetimeChecker(hs);

    walk_generics(&mut checker, generics);
    walk_fn_decl(&mut checker, func);

    for &v in checker.0.values() {
        span_lint(cx, UNUSED_LIFETIMES, v, "this lifetime isn't used in the function definition");
    }
}
<commit_msg>fallout<commit_after>use reexport::*;
use rustc::lint::*;
use rustc::hir::def::Def;
use rustc::hir::*;
use rustc::hir::intravisit::{Visitor, walk_ty, walk_ty_param_bound, walk_fn_decl, walk_generics};
use std::collections::{HashSet, HashMap};
use syntax::codemap::Span;
use utils::{in_external_macro, span_lint};

/// **What it does:** This lint checks for lifetime annotations which can be removed by relying on lifetime elision.
///
/// **Why is this bad?** The additional lifetimes make the code look more complicated, while there is nothing out of the ordinary going on. Removing them leads to more readable code.
///
/// **Known problems:** Potential false negatives: we bail out if the function has a `where` clause where lifetimes are mentioned.
///
/// **Example:** `fn in_and_out<'a>(x: &'a u8, y: u8) -> &'a u8 { x }`
declare_lint! {
    pub NEEDLESS_LIFETIMES,
    Warn,
    "using explicit lifetimes for references in function arguments when elision rules \
     would allow omitting them"
}

/// **What it does:** This lint checks for lifetimes in generics that are never used anywhere else.
///
/// **Why is this bad?** The additional lifetimes make the code look more complicated, while there is nothing out of the ordinary going on. Removing them leads to more readable code.
///
/// **Known problems:** None
///
/// **Example:** `fn unused_lifetime<'a>(x: u8) { .. }`
declare_lint! {
    pub UNUSED_LIFETIMES,
    Warn,
    "unused lifetimes in function definitions"
}

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

impl LintPass for LifetimePass {
    fn get_lints(&self) -> LintArray {
        lint_array!(NEEDLESS_LIFETIMES, UNUSED_LIFETIMES)
    }
}

impl LateLintPass for LifetimePass {
    fn check_item(&mut self, cx: &LateContext, item: &Item) {
        if let ItemFn(ref decl, _, _, _, ref generics, _) = item.node {
            check_fn_inner(cx, decl, generics, item.span);
        }
    }

    fn check_impl_item(&mut self, cx: &LateContext, item: &ImplItem) {
        if let ImplItemKind::Method(ref sig, _) = item.node {
            check_fn_inner(cx, &sig.decl, &sig.generics, item.span);
        }
    }

    fn check_trait_item(&mut self, cx: &LateContext, item: &TraitItem) {
        if let MethodTraitItem(ref sig, _) = item.node {
            check_fn_inner(cx, &sig.decl, &sig.generics, item.span);
        }
    }
}

/// The lifetime of a &-reference.
#[derive(PartialEq, Eq, Hash, Debug)]
enum RefLt {
    Unnamed,
    Static,
    Named(Name),
}

fn bound_lifetimes(bound: &TyParamBound) -> HirVec<&Lifetime> {
    if let TraitTyParamBound(ref trait_ref, _) = *bound {
        trait_ref.trait_ref
                 .path
                 .segments
                 .last()
                 .expect("a path must have at least one segment")
                 .parameters
                 .lifetimes()
    } else {
        HirVec::new()
    }
}

fn check_fn_inner(cx: &LateContext, decl: &FnDecl, generics: &Generics, span: Span) {
    if in_external_macro(cx, span) || has_where_lifetimes(cx, &generics.where_clause) {
        return;
    }

    let bounds_lts = generics.ty_params
                             .iter()
                             .flat_map(|typ| typ.bounds.iter().flat_map(bound_lifetimes));

    if could_use_elision(cx, decl, &generics.lifetimes, bounds_lts) {
        span_lint(cx,
                  NEEDLESS_LIFETIMES,
                  span,
                  "explicit lifetimes given in parameter types where they could be elided");
    }
    report_extra_lifetimes(cx, decl, generics);
}

fn could_use_elision<'a, T: Iterator<Item = &'a Lifetime>>(cx: &LateContext, func: &FnDecl,
                                                           named_lts: &[LifetimeDef], bounds_lts: T)
                                                           -> bool {
    // There are two scenarios where elision works:
    // * no output references, all input references have different LT
    // * output references, exactly one input reference with same LT
    // All lifetimes must be unnamed, 'static or defined without bounds on the
    // level of the current item.

    // check named LTs
    let allowed_lts = allowed_lts_from(named_lts);

    // these will collect all the lifetimes for references in arg/return types
    let mut input_visitor = RefVisitor::new(cx);
    let mut output_visitor = RefVisitor::new(cx);

    // extract lifetimes in input argument types
    for arg in &func.inputs {
        input_visitor.visit_ty(&arg.ty);
    }
    // extract lifetimes in output type
    if let Return(ref ty) = func.output {
        output_visitor.visit_ty(ty);
    }

    let input_lts = lts_from_bounds(input_visitor.into_vec(), bounds_lts);
    let output_lts = output_visitor.into_vec();

    // check for lifetimes from higher scopes
    for lt in input_lts.iter().chain(output_lts.iter()) {
        if !allowed_lts.contains(lt) {
            return false;
        }
    }

    // no input lifetimes? easy case!
    if input_lts.is_empty() {
        false
    } else if output_lts.is_empty() {
        // no output lifetimes, check distinctness of input lifetimes

        // only unnamed and static, ok
        if input_lts.iter().all(|lt| *lt == RefLt::Unnamed || *lt == RefLt::Static) {
            return false;
        }
        // we have no output reference, so we only need all distinct lifetimes
        input_lts.len() == unique_lifetimes(&input_lts)
    } else {
        // we have output references, so we need one input reference,
        // and all output lifetimes must be the same
        if unique_lifetimes(&output_lts) > 1 {
            return false;
        }
        if input_lts.len() == 1 {
            match (&input_lts[0], &output_lts[0]) {
                (&RefLt::Named(n1), &RefLt::Named(n2)) if n1 == n2 => true,
                (&RefLt::Named(_), &RefLt::Unnamed) => true,
                _ => false, // already elided, different named lifetimes
                // or something static going on
            }
        } else {
            false
        }
    }
}

fn allowed_lts_from(named_lts: &[LifetimeDef]) -> HashSet<RefLt> {
    let mut allowed_lts = HashSet::new();
    for lt in named_lts {
        if lt.bounds.is_empty() {
            allowed_lts.insert(RefLt::Named(lt.lifetime.name));
        }
    }
    allowed_lts.insert(RefLt::Unnamed);
    allowed_lts.insert(RefLt::Static);
    allowed_lts
}

fn lts_from_bounds<'a, T: Iterator<Item = &'a Lifetime>>(mut vec: Vec<RefLt>, bounds_lts: T) -> Vec<RefLt> {
    for lt in bounds_lts {
        if lt.name.as_str() != "'static" {
            vec.push(RefLt::Named(lt.name));
        }
    }

    vec
}

/// Number of unique lifetimes in the given vector.
fn unique_lifetimes(lts: &[RefLt]) -> usize {
    lts.iter().collect::<HashSet<_>>().len()
}

/// A visitor usable for `rustc_front::visit::walk_ty()`.
struct RefVisitor<'v, 't: 'v> {
    cx: &'v LateContext<'v, 't>,
    lts: Vec<RefLt>,
}

impl<'v, 't> RefVisitor<'v, 't> {
    fn new(cx: &'v LateContext<'v, 't>) -> RefVisitor<'v, 't> {
        RefVisitor {
            cx: cx,
            lts: Vec::new(),
        }
    }

    fn record(&mut self, lifetime: &Option<Lifetime>) {
        if let Some(ref lt) = *lifetime {
            if lt.name.as_str() == "'static" {
                self.lts.push(RefLt::Static);
            } else {
                self.lts.push(RefLt::Named(lt.name));
            }
        } else {
            self.lts.push(RefLt::Unnamed);
        }
    }

    fn into_vec(self) -> Vec<RefLt> {
        self.lts
    }

    fn collect_anonymous_lifetimes(&mut self, path: &Path, ty: &Ty) {
        let last_path_segment = path.segments.last().map(|s| &s.parameters);
        if let Some(&AngleBracketedParameters(ref params)) = last_path_segment {
            if params.lifetimes.is_empty() {
                if let Some(def) = self.cx.tcx.def_map.borrow().get(&ty.id).map(|r| r.full_def()) {
                    match def {
                        Def::TyAlias(def_id) |
                        Def::Struct(def_id) => {
                            let type_scheme = self.cx.tcx.lookup_item_type(def_id);
                            for _ in type_scheme.generics.regions.as_slice() {
                                self.record(&None);
                            }
                        }
                        Def::Trait(def_id) => {
                            let trait_def = self.cx.tcx.trait_defs.borrow()[&def_id];
                            for _ in &trait_def.generics.regions {
                                self.record(&None);
                            }
                        }
                        _ => (),
                    }
                }
            }
        }
    }
}

impl<'v, 't> Visitor<'v> for RefVisitor<'v, 't> {
    // for lifetimes as parameters of generics
    fn visit_lifetime(&mut self, lifetime: &'v Lifetime) {
        self.record(&Some(*lifetime));
    }

    fn visit_ty(&mut self, ty: &'v Ty) {
        match ty.node {
            TyRptr(None, _) => {
                self.record(&None);
            }
            TyPath(_, ref path) => {
                self.collect_anonymous_lifetimes(path, ty);
            }
            _ => (),
        }
        walk_ty(self, ty);
    }
}

/// Are any lifetimes mentioned in the `where` clause? If yes, we don't try to
/// reason about elision.
fn has_where_lifetimes(cx: &LateContext, where_clause: &WhereClause) -> bool {
    for predicate in &where_clause.predicates {
        match *predicate {
            WherePredicate::RegionPredicate(..) => return true,
            WherePredicate::BoundPredicate(ref pred) => {
                // a predicate like F: Trait or F: for<'a> Trait<'a>
                let mut visitor = RefVisitor::new(cx);
                // walk the type F, it may not contain LT refs
                walk_ty(&mut visitor, &pred.bounded_ty);
                if !visitor.lts.is_empty() {
                    return true;
                }
                // if the bounds define new lifetimes, they are fine to occur
                let allowed_lts = allowed_lts_from(&pred.bound_lifetimes);
                // now walk the bounds
                for bound in pred.bounds.iter() {
                    walk_ty_param_bound(&mut visitor, bound);
                }
                // and check that all lifetimes are allowed
                for lt in visitor.into_vec() {
                    if !allowed_lts.contains(&lt) {
                        return true;
                    }
                }
            }
            WherePredicate::EqPredicate(ref pred) => {
                let mut visitor = RefVisitor::new(cx);
                walk_ty(&mut visitor, &pred.ty);
                if !visitor.lts.is_empty() {
                    return true;
                }
            }
        }
    }
    false
}

struct LifetimeChecker(HashMap<Name, Span>);

impl<'v> Visitor<'v> for LifetimeChecker {
    // for lifetimes as parameters of generics
    fn visit_lifetime(&mut self, lifetime: &'v Lifetime) {
        self.0.remove(&lifetime.name);
    }

    fn visit_lifetime_def(&mut self, _: &'v LifetimeDef) {
        // don't actually visit `<'a>` or `<'a: 'b>`
        // we've already visited the `'a` declarations and
        // don't want to spuriously remove them
        // `'b` in `'a: 'b` is useless unless used elsewhere in
        // a non-lifetime bound
    }
}

fn report_extra_lifetimes(cx: &LateContext, func: &FnDecl, generics: &Generics) {
    let hs = generics.lifetimes
                     .iter()
                     .map(|lt| (lt.lifetime.name, lt.lifetime.span))
                     .collect();
    let mut checker = LifetimeChecker(hs);

    walk_generics(&mut checker, generics);
    walk_fn_decl(&mut checker, func);

    for &v in checker.0.values() {
        span_lint(cx, UNUSED_LIFETIMES, v, "this lifetime isn't used in the function definition");
    }
}
<|endoftext|>


<commit_before>use std;

pub fn iter_dibits<T: Iterator<Item = u8>>(src: T) -> Bits<DibitIterator, T> {
    Bits::new(src)
}

pub fn iter_bits<T: Iterator<Item = u8>>(src: T) -> Bits<BitIterator, T> {
    Bits::new(src)
}

pub trait IterateParams {
    type IterType;

    fn bits() -> u8;
    fn wrap(x: u8) -> Self::IterType;

    fn shift() -> u8 { Self::bits() }
    fn iterations() -> u8 { 8 / Self::bits() }

    fn validate() {
        assert!(Self::bits().is_power_of_two());
    }
}

pub struct Bit(pub u8);
pub struct BitIterator;

impl IterateParams for BitIterator {
    type IterType = Bit;
    fn bits() -> u8 { 1 }
    fn wrap(x: u8) -> Bit { Bit(x) }
}

pub struct Dibit(pub u8);
pub struct DibitIterator;

impl IterateParams for DibitIterator {
    type IterType = Dibit;
    fn bits() -> u8 { 2 }
    fn wrap(x: u8) -> Dibit { Dibit(x) }
}

pub struct Bits<P, T>
    where P: IterateParams, T: Iterator<Item = u8>
{
    params: std::marker::PhantomData<P>,
    /// The source of bits.
    src: T,
    /// The current index into the current byte.
    idx: u8,
    /// The current byte in the source.
    byte: u8,
}

impl<P, T> Bits<P, T>
    where P: IterateParams, T: Iterator<Item = u8>
{
    /// Construct a new `Bits<T>`.
    fn new(src: T) -> Bits<P, T> {
        P::validate();

        Bits {
            params: std::marker::PhantomData,
            src: src,
            byte: 0,
            idx: 0,
        }
    }
}

impl<P, T> Iterator for Bits<P, T>
    where P: IterateParams, T: Iterator<Item = u8>
{
    type Item = P::IterType;

    fn next(&mut self) -> Option<Self::Item> {
        if self.idx == 0 {
            self.byte = match self.src.next() {
                Some(b) => b,
                None => return None,
            };
        }

        // Store the current byte for later.
        let byte = self.byte;

        // Strip off the MSBs for the next iteration.
        self.byte <<= P::shift();

        // Move to the next item and reset after all have been visited.
        self.idx += 1;
        self.idx %= P::iterations();

        // Yield the MSBs.
        Some(P::wrap(byte >> (8 - P::shift())))
    }
}

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

        const BITS: &'static [u8] = &[
            0b00011011,
            0b11001100,
        ];

        {
            let mut d = iter_dibits(BITS.iter().cloned());

            assert!(d.next().unwrap().0 == 0b00);
            assert!(d.next().unwrap().0 == 0b01);
            assert!(d.next().unwrap().0 == 0b10);
            assert!(d.next().unwrap().0 == 0b11);
            assert!(d.next().unwrap().0 == 0b11);
            assert!(d.next().unwrap().0 == 0b00);
            assert!(d.next().unwrap().0 == 0b11);
            assert!(d.next().unwrap().0 == 0b00);
            assert!(d.next().is_none());
        }

        {
            let mut b = iter_bits(BITS.iter().cloned());

            assert!(b.next().unwrap().0 == 0);
            assert!(b.next().unwrap().0 == 0);
            assert!(b.next().unwrap().0 == 0);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 0);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 0);
            assert!(b.next().unwrap().0 == 0);
        }
    }
}
<commit_msg>Move parameter validation to tests<commit_after>use std;

pub fn iter_dibits<T: Iterator<Item = u8>>(src: T) -> Bits<DibitIterator, T> {
    Bits::new(src)
}

pub fn iter_bits<T: Iterator<Item = u8>>(src: T) -> Bits<BitIterator, T> {
    Bits::new(src)
}

pub trait IterateParams {
    type IterType;

    fn bits() -> u8;
    fn wrap(x: u8) -> Self::IterType;

    fn shift() -> u8 { Self::bits() }
    fn iterations() -> u8 { 8 / Self::bits() }

    fn validate() {
        assert!(Self::bits().is_power_of_two());
    }
}

pub struct Bit(pub u8);
pub struct BitIterator;

impl IterateParams for BitIterator {
    type IterType = Bit;
    fn bits() -> u8 { 1 }
    fn wrap(x: u8) -> Bit { Bit(x) }
}

pub struct Dibit(pub u8);
pub struct DibitIterator;

impl IterateParams for DibitIterator {
    type IterType = Dibit;
    fn bits() -> u8 { 2 }
    fn wrap(x: u8) -> Dibit { Dibit(x) }
}

pub struct Bits<P, T>
    where P: IterateParams, T: Iterator<Item = u8>
{
    params: std::marker::PhantomData<P>,
    /// The source of bits.
    src: T,
    /// The current index into the current byte.
    idx: u8,
    /// The current byte in the source.
    byte: u8,
}

impl<P, T> Bits<P, T>
    where P: IterateParams, T: Iterator<Item = u8>
{
    /// Construct a new `Bits<T>`.
    fn new(src: T) -> Bits<P, T> {
        Bits {
            params: std::marker::PhantomData,
            src: src,
            byte: 0,
            idx: 0,
        }
    }
}

impl<P, T> Iterator for Bits<P, T>
    where P: IterateParams, T: Iterator<Item = u8>
{
    type Item = P::IterType;

    fn next(&mut self) -> Option<Self::Item> {
        if self.idx == 0 {
            self.byte = match self.src.next() {
                Some(b) => b,
                None => return None,
            };
        }

        // Store the current byte for later.
        let byte = self.byte;

        // Strip off the MSBs for the next iteration.
        self.byte <<= P::shift();

        // Move to the next item and reset after all have been visited.
        self.idx += 1;
        self.idx %= P::iterations();

        // Yield the MSBs.
        Some(P::wrap(byte >> (8 - P::shift())))
    }
}

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

    #[test]
    fn validate_params() {
        BitIterator::validate();
        DibitIterator::validate();
    }

    #[test]
    fn test_bits() {
        const BITS: &'static [u8] = &[
            0b00011011,
            0b11001100,
        ];

        {
            let mut d = iter_dibits(BITS.iter().cloned());

            assert!(d.next().unwrap().0 == 0b00);
            assert!(d.next().unwrap().0 == 0b01);
            assert!(d.next().unwrap().0 == 0b10);
            assert!(d.next().unwrap().0 == 0b11);
            assert!(d.next().unwrap().0 == 0b11);
            assert!(d.next().unwrap().0 == 0b00);
            assert!(d.next().unwrap().0 == 0b11);
            assert!(d.next().unwrap().0 == 0b00);
            assert!(d.next().is_none());
        }

        {
            let mut b = iter_bits(BITS.iter().cloned());

            assert!(b.next().unwrap().0 == 0);
            assert!(b.next().unwrap().0 == 0);
            assert!(b.next().unwrap().0 == 0);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 0);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 0);
            assert!(b.next().unwrap().0 == 0);
        }
    }
}
<|endoftext|>